add sync reader

Author: rekby

tests/topics/test_topic_reader.py

@@ -2,10 +2,19 @@
 
 
 @pytest.mark.asyncio
-class TestTopicWriterAsyncIO:
+class TestTopicReaderAsyncIO:
     async def test_read_message(
         self, driver, topic_path, topic_with_messages, topic_consumer
     ):
         reader = driver.topic_client.topic_reader(topic_consumer, topic_path)
 
         assert await reader.receive_batch() is not None
+
+
+class TestTopicReaderSync:
+    def test_read_message(
+        self, driver_sync, topic_path, topic_with_messages, topic_consumer
+    ):
+        reader = driver_sync.topic_client.topic_reader(topic_consumer, topic_path)
+
+        assert reader.receive_batch() is not None

ydb/_topic_common/common.py

@@ -2,6 +2,7 @@
 import concurrent.futures
 import threading
 import typing
+from typing import Optional
 
 from .. import operation, issues
 from .._grpc.grpcwrapper.common_utils import IFromProtoWithProtoType
@@ -27,6 +28,7 @@ def wrapper(rpc_state, response_pb, driver=None):
 
     return wrapper
 
+
 _shared_event_loop_lock = threading.Lock()
 _shared_event_loop = None  # type: Optional[asyncio.AbstractEventLoop]
 

ydb/_topic_reader/topic_reader.py

@@ -1,16 +1,13 @@
-import concurrent.futures
 import enum
 import datetime
 from dataclasses import dataclass
 from typing import (
     Union,
     Optional,
     List,
-    Iterable,
 )
 
 from ..table import RetrySettings
-from .datatypes import ICommittable, PublicBatch, PublicMessage
 from .._topic_common.common import TokenGetterFuncType
 from .._grpc.grpcwrapper.ydb_topic import StreamReadMessage, OffsetsRange
 

ydb/_topic_reader/topic_reader_sync.py

@@ -1,12 +1,72 @@
+import asyncio
+import concurrent.futures
+import typing
+from typing import List, Union, Iterable, Optional, Coroutine
+
+from ydb._grpc.grpcwrapper.common_utils import SupportedDriverType
+from ydb._topic_common.common import _get_shared_event_loop
+from ydb._topic_reader.datatypes import PublicMessage, PublicBatch, ICommittable
+from ydb._topic_reader.topic_reader import (
+    PublicReaderSettings,
+    SessionStat,
+    CommitResult,
+)
+from ydb._topic_reader.topic_reader_asyncio import (
+    PublicAsyncIOReader,
+    TopicReaderClosedError,
+)
+
 
 class TopicReaderSync:
+    _loop: asyncio.AbstractEventLoop
+    _async_reader: PublicAsyncIOReader
+    _closed: bool
+
+    def __init__(
+        self,
+        driver: SupportedDriverType,
+        settings: PublicReaderSettings,
+        *,
+        eventloop: Optional[asyncio.AbstractEventLoop] = None,
+    ):
+        self._closed = False
+
+        if eventloop:
+            self._loop = eventloop
+        else:
+            self._loop = _get_shared_event_loop()
+
+        async def create_reader():
+            return PublicAsyncIOReader(driver, settings)
+
+        self._async_reader = asyncio.run_coroutine_threadsafe(
+            create_reader(), self._loop
+        ).result()
+
+    def __del__(self):
+        self.close()
+
+    def _call(self, coro):
+        if self._closed:
+            raise TopicReaderClosedError()
+
+        return asyncio.run_coroutine_threadsafe(coro, self._loop)
+
+    def _call_sync(self, coro: Coroutine, timeout):
+        f = self._call(coro)
+        try:
+            return f.result(timeout)
+        except TimeoutError:
+            f.cancel()
+            raise
+
     def async_sessions_stat(self) -> concurrent.futures.Future:
         """
         Receive stat from the server, return feature.
         """
         raise NotImplementedError()
 
-    async def sessions_stat(self) -> List["SessionStat"]:
+    async def sessions_stat(self) -> List[SessionStat]:
         """
         Receive stat from the server
 
@@ -67,8 +127,8 @@ def batches(
     def receive_batch(
         self,
         *,
-        max_messages: Union[int, None] = None,
-        max_bytes: Union[int, None],
+        max_messages: typing.Union[int, None] = None,
+        max_bytes: typing.Union[int, None] = None,
         timeout: Union[float, None] = None,
     ) -> Union[PublicBatch, None]:
         """
@@ -78,7 +138,12 @@ def receive_batch(
         if no new message in timeout seconds (default - infinite): raise TimeoutError()
         if timeout <= 0 - it will fast non block method, get messages from internal buffer only.
         """
-        raise NotImplementedError()
+        return self._call_sync(
+            self._async_reader.receive_batch(
+                max_messages=max_messages, max_bytes=max_bytes
+            ),
+            timeout,
+        )
 
     def commit(self, mess: ICommittable):
         """
@@ -87,11 +152,11 @@ def commit(self, mess: ICommittable):
         For the method no way check the commit result
         (for example if lost connection - commits will not re-send and committed messages will receive again)
         """
-        raise NotImplementedError()
+        self._call_sync(self._async_reader.commit(mess), None)
 
     def commit_with_ack(
         self, mess: ICommittable
-    ) -> Union["CommitResult", List["CommitResult"]]:
+    ) -> Union[CommitResult, List[CommitResult]]:
         """
         write commit message to a buffer and wait ack from the server.
 
@@ -101,7 +166,7 @@ def commit_with_ack(
 
     def async_commit_with_ack(
         self, mess: ICommittable
-    ) -> Union["CommitResult", List["CommitResult"]]:
+    ) -> Union[CommitResult, List[CommitResult]]:
         """
         write commit message to a buffer and return Future for wait result.
         """
@@ -120,4 +185,8 @@ def flush(self):
         raise NotImplementedError()
 
     def close(self):
-        raise NotImplementedError()
+        if self._closed:
+            return
+        self._closed = True
+
+        self._call_sync(self._async_reader.close())

ydb/_topic_writer/topic_writer_sync.py

@@ -17,6 +17,7 @@
 from .topic_writer_asyncio import WriterAsyncIO
 from .._topic_common.common import _get_shared_event_loop
 
+
 class WriterSync:
     _loop: asyncio.AbstractEventLoop
     _async_writer: WriterAsyncIO
@@ -27,7 +28,7 @@ def __init__(
         driver: SupportedDriverType,
         settings: PublicWriterSettings,
         *,
-        eventloop: asyncio.AbstractEventLoop = None,
+        eventloop: Optional[asyncio.AbstractEventLoop] = None,
     ):
 
         self._closed = False
@@ -44,16 +45,16 @@ async def create_async_writer():
             create_async_writer(), self._loop
         ).result()
 
-    def _call(self, coro, *args, **kwargs):
+    def _call(self, coro):
         if self._closed:
             raise TopicWriterError("writer is closed")
 
         return asyncio.run_coroutine_threadsafe(coro, self._loop)
 
-    def _call_sync(self, coro: Coroutine, timeout, *args, **kwargs):
-        f = self._call(coro, *args, **kwargs)
+    def _call_sync(self, coro: Coroutine, timeout):
+        f = self._call(coro)
         try:
-            return f.result()
+            return f.result(timeout)
         except TimeoutError:
             f.cancel()
             raise

ydb/topic.py

@@ -7,13 +7,10 @@
 
 from ._topic_reader.topic_reader import (
     PublicReaderSettings as TopicReaderSettings,
-    Reader as TopicReader,
-    Selector as TopicSelector,
-    Events as TopicReaderEvents,
-    RetryPolicy as TopicReaderRetryPolicy,
-    StubEvent as TopicReaderStubEvent,
 )
 
+from ._topic_reader.topic_reader_sync import TopicReaderSync as TopicReader
+
 from ._topic_reader.topic_reader_asyncio import (
     PublicAsyncIOReader as TopicReaderAsyncIO,
 )
@@ -240,26 +237,27 @@ def drop_topic(self, path: str):
 
     def topic_reader(
         self,
-        topic: Union[str, TopicSelector, List[Union[str, TopicSelector]]],
         consumer: str,
-        commit_batch_time: Union[float, None] = 0.1,
-        commit_batch_count: Union[int, None] = 1000,
+        topic: str,
         buffer_size_bytes: int = 50 * 1024 * 1024,
-        sync_commit: bool = False,  # reader.commit(...) will wait commit ack from server
-        on_commit: Callable[["TopicReaderStubEvent"], None] = None,
-        on_get_partition_start_offset: Callable[
-            ["TopicReaderEvents.OnPartitionGetStartOffsetRequest"],
-            "TopicReaderEvents.OnPartitionGetStartOffsetResponse",
-        ] = None,
-        on_init_partition: Callable[["StubEvent"], None] = None,
-        on_shutdown_partition: Callable[["StubEvent"], None] = None,
-        decoder: Union[Mapping[int, Callable[[bytes], bytes]], None] = None,
-        deserializer: Union[Callable[[bytes], Any], None] = None,
-        one_attempt_connection_timeout: Union[float, None] = 1,
-        connection_timeout: Union[float, None] = None,
-        retry_policy: Union["TopicReaderRetryPolicy", None] = None,
+        # on_commit: Callable[["Events.OnCommit"], None] = None
+        # on_get_partition_start_offset: Callable[
+        #     ["Events.OnPartitionGetStartOffsetRequest"],
+        #     "Events.OnPartitionGetStartOffsetResponse",
+        # ] = None
+        # on_partition_session_start: Callable[["StubEvent"], None] = None
+        # on_partition_session_stop: Callable[["StubEvent"], None] = None
+        # on_partition_session_close: Callable[["StubEvent"], None] = None  # todo?
+        # decoder: Union[Mapping[int, Callable[[bytes], bytes]], None] = None
+        # deserializer: Union[Callable[[bytes], Any], None] = None
+        # one_attempt_connection_timeout: Union[float, None] = 1
+        # connection_timeout: Union[float, None] = None
+        # retry_policy: Union["RetryPolicy", None] = None
     ) -> TopicReader:
-        raise NotImplementedError()
+        args = locals()
+        del args["self"]
+        settings = TopicReaderSettings(**args)
+        return TopicReader(self._driver, settings)
 
     def topic_writer(
         self,