Skip to content

Commit 2424ecc

Browse files
committed
No Consumer Reader
1 parent d980334 commit 2424ecc

File tree

5 files changed

+114
-9
lines changed

5 files changed

+114
-9
lines changed

Diff for: tests/topics/test_topic_reader.py

+54
Original file line numberDiff line numberDiff line change
@@ -251,3 +251,57 @@ async def wait(fut):
251251

252252
await reader0.close()
253253
await reader1.close()
254+
255+
256+
@pytest.mark.asyncio
257+
class TestTopicNoConsumerReaderAsyncIO:
258+
async def test_reader_with_sync_lambda(self, driver, topic_with_messages):
259+
reader = driver.topic_client.no_consumer_reader(topic_with_messages, [0], lambda x: None)
260+
await reader.wait_message()
261+
msg = await reader.receive_message()
262+
263+
assert msg is not None
264+
assert msg.seqno
265+
266+
await reader.close()
267+
268+
async def test_reader_with_async_lambda(self, driver, topic_with_messages):
269+
async def get_start_offset_lambda(partition_id: int) -> int:
270+
return 1
271+
272+
reader = driver.topic_client.no_consumer_reader(topic_with_messages, [0], get_start_offset_lambda)
273+
await reader.wait_message()
274+
msg = await reader.receive_message()
275+
276+
assert msg is not None
277+
assert msg.seqno == 2
278+
279+
await reader.close()
280+
281+
async def test_offsets_updated_after_reconnect(self, driver, topic_with_messages):
282+
current_offset = 0
283+
284+
def get_start_offset_lambda(partition_id: int) -> int:
285+
nonlocal current_offset
286+
return current_offset
287+
288+
reader = driver.topic_client.no_consumer_reader(topic_with_messages, [0], get_start_offset_lambda)
289+
await reader.wait_message()
290+
msg = await reader.receive_message()
291+
292+
assert msg is not None
293+
assert msg.seqno == current_offset + 1
294+
295+
current_offset += 2
296+
297+
reader._reconnector._stream_reader._set_first_error(ydb.Unavailable("some retriable error"))
298+
299+
await asyncio.sleep(0)
300+
301+
await reader.wait_message()
302+
msg = await reader.receive_message()
303+
304+
assert msg is not None
305+
assert msg.seqno == current_offset + 1
306+
307+
await reader.close()

Diff for: ydb/_grpc/grpcwrapper/ydb_topic.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -439,12 +439,13 @@ def from_proto(
439439
@dataclass
440440
class InitRequest(IToProto):
441441
topics_read_settings: List["StreamReadMessage.InitRequest.TopicReadSettings"]
442-
consumer: str
442+
consumer: Optional[str]
443443
auto_partitioning_support: bool
444444

445445
def to_proto(self) -> ydb_topic_pb2.StreamReadMessage.InitRequest:
446446
res = ydb_topic_pb2.StreamReadMessage.InitRequest()
447-
res.consumer = self.consumer
447+
if self.consumer is not None:
448+
res.consumer = self.consumer
448449
for settings in self.topics_read_settings:
449450
res.topics_read_settings.append(settings.to_proto())
450451
res.auto_partitioning_support = self.auto_partitioning_support

Diff for: ydb/_topic_reader/topic_reader.py

+7-3
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import datetime
44
from dataclasses import dataclass
55
from typing import (
6+
Awaitable,
67
Union,
78
Optional,
89
List,
@@ -42,7 +43,7 @@ def _to_topic_read_settings(self) -> StreamReadMessage.InitRequest.TopicReadSett
4243

4344
@dataclass
4445
class PublicReaderSettings:
45-
consumer: str
46+
consumer: Optional[str]
4647
topic: TopicSelectorTypes
4748
buffer_size_bytes: int = 50 * 1024 * 1024
4849
auto_partitioning_support: bool = True
@@ -54,12 +55,15 @@ class PublicReaderSettings:
5455
decoder_executor: Optional[concurrent.futures.Executor] = None
5556
update_token_interval: Union[int, float] = 3600
5657

58+
partition_ids: Optional[List[int]] = None
59+
get_start_offset_lambda: Optional[Union[Callable[[int], int], Callable[[int], Awaitable[int]]]] = None
60+
5761
def __post_init__(self):
5862
# check possible create init message
5963
_ = self._init_message()
6064

6165
def _init_message(self) -> StreamReadMessage.InitRequest:
62-
if not isinstance(self.consumer, str):
66+
if self.consumer is not None and not isinstance(self.consumer, str):
6367
raise TypeError("Unsupported type for customer field: '%s'" % type(self.consumer))
6468

6569
if isinstance(self.topic, list):
@@ -69,7 +73,7 @@ def _init_message(self) -> StreamReadMessage.InitRequest:
6973

7074
for index, selector in enumerate(selectors):
7175
if isinstance(selector, str):
72-
selectors[index] = PublicTopicSelector(path=selector)
76+
selectors[index] = PublicTopicSelector(path=selector, partitions=self.partition_ids)
7377
elif isinstance(selector, PublicTopicSelector):
7478
pass
7579
else:

Diff for: ydb/_topic_reader/topic_reader_asyncio.py

+22-3
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,14 @@ async def close(self, flush: bool = True):
182182
await self._reconnector.close(flush)
183183

184184

185+
class PublicAsyncIONoConsumerReader(PublicAsyncIOReader):
186+
def commit(self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]):
187+
raise issues.Error("Commit operations are not supported for topic reader without consumer.")
188+
189+
async def commit_with_ack(self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]):
190+
raise issues.Error("Commit operations are not supported for topic reader without consumer.")
191+
192+
185193
class ReaderReconnector:
186194
_static_reader_reconnector_counter = AtomicCounter()
187195

@@ -393,6 +401,7 @@ class ReaderStream:
393401
_update_token_interval: Union[int, float]
394402
_update_token_event: asyncio.Event
395403
_get_token_function: Callable[[], str]
404+
_settings: topic_reader.PublicReaderSettings
396405

397406
def __init__(
398407
self,
@@ -425,6 +434,8 @@ def __init__(
425434
self._get_token_function = get_token_function
426435
self._update_token_event = asyncio.Event()
427436

437+
self._settings = settings
438+
428439
@staticmethod
429440
async def create(
430441
reader_reconnector_id: int,
@@ -615,7 +626,7 @@ async def _read_messages_loop(self):
615626
message.server_message,
616627
StreamReadMessage.StartPartitionSessionRequest,
617628
):
618-
self._on_start_partition_session(message.server_message)
629+
await self._on_start_partition_session(message.server_message)
619630

620631
elif isinstance(
621632
message.server_message,
@@ -660,7 +671,7 @@ async def _update_token(self, token: str):
660671
finally:
661672
self._update_token_event.clear()
662673

663-
def _on_start_partition_session(self, message: StreamReadMessage.StartPartitionSessionRequest):
674+
async def _on_start_partition_session(self, message: StreamReadMessage.StartPartitionSessionRequest):
664675
try:
665676
if message.partition_session.partition_session_id in self._partition_sessions:
666677
raise TopicReaderError(
@@ -676,11 +687,19 @@ def _on_start_partition_session(self, message: StreamReadMessage.StartPartitionS
676687
reader_reconnector_id=self._reader_reconnector_id,
677688
reader_stream_id=self._id,
678689
)
690+
691+
read_offset = None
692+
callee = self._settings.get_start_offset_lambda
693+
if callee is not None:
694+
read_offset = callee(message.partition_session.partition_id)
695+
if asyncio.iscoroutinefunction(callee):
696+
read_offset = await read_offset
697+
679698
self._stream.write(
680699
StreamReadMessage.FromClient(
681700
client_message=StreamReadMessage.StartPartitionSessionResponse(
682701
partition_session_id=message.partition_session.partition_session_id,
683-
read_offset=None,
702+
read_offset=read_offset,
684703
commit_offset=None,
685704
)
686705
),

Diff for: ydb/topic.py

+28-1
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
"TopicError",
1515
"TopicMeteringMode",
1616
"TopicReader",
17+
"TopicNoConsumerReaderAsyncIO",
1718
"TopicReaderAsyncIO",
1819
"TopicReaderBatch",
1920
"TopicReaderMessage",
@@ -36,7 +37,7 @@
3637
import datetime
3738
from dataclasses import dataclass
3839
import logging
39-
from typing import List, Union, Mapping, Optional, Dict, Callable
40+
from typing import Awaitable, List, Union, Mapping, Optional, Dict, Callable
4041

4142
from . import aio, Credentials, _apis, issues
4243

@@ -56,6 +57,7 @@
5657

5758
from ._topic_reader.topic_reader_asyncio import (
5859
PublicAsyncIOReader as TopicReaderAsyncIO,
60+
PublicAsyncIONoConsumerReader as TopicNoConsumerReaderAsyncIO,
5961
PublicTopicReaderPartitionExpiredError as TopicReaderPartitionExpiredError,
6062
PublicTopicReaderUnexpectedCodecError as TopicReaderUnexpectedCodecError,
6163
)
@@ -261,6 +263,31 @@ def reader(
261263

262264
return TopicReaderAsyncIO(self._driver, settings, _parent=self)
263265

266+
def no_consumer_reader(
267+
self,
268+
topic: Union[str, TopicReaderSelector, List[Union[str, TopicReaderSelector]]],
269+
partition_ids: List[int],
270+
get_start_offset_lambda: Union[Callable[[int], int], Callable[[int], Awaitable[int]]],
271+
buffer_size_bytes: int = 50 * 1024 * 1024,
272+
# decoders: map[codec_code] func(encoded_bytes)->decoded_bytes
273+
# the func will be called from multiply threads in parallel
274+
decoders: Union[Mapping[int, Callable[[bytes], bytes]], None] = None,
275+
# custom decoder executor for call builtin and custom decoders. If None - use shared executor pool.
276+
# if max_worker in the executor is 1 - then decoders will be called from the thread without parallel
277+
decoder_executor: Optional[concurrent.futures.Executor] = None,
278+
auto_partitioning_support: Optional[bool] = True, # Auto partitioning feature flag. Default - True.
279+
) -> TopicNoConsumerReaderAsyncIO:
280+
if not decoder_executor:
281+
decoder_executor = self._executor
282+
283+
args = locals().copy()
284+
del args["self"]
285+
args["consumer"] = None
286+
287+
settings = TopicReaderSettings(**args)
288+
289+
return TopicNoConsumerReaderAsyncIO(self._driver, settings, _parent=self)
290+
264291
def writer(
265292
self,
266293
topic,

0 commit comments

Comments
 (0)