Merge pull request #177 test sync topic writer, fix timeout param

Author: rekby

tests/topics/test_topic_writer.py

@@ -50,3 +50,47 @@ async def test_auto_flush_on_close(self, driver: ydb.aio.Driver, topic_path):
         ) as writer:
             init_info = await writer.wait_init()
             assert init_info.last_seqno == last_seqno
+
+
+class TestTopicWriterSync:
+    def test_send_message(self, driver_sync: ydb.Driver, topic_path):
+        writer = driver_sync.topic_client.topic_writer(
+            topic_path, producer_and_message_group_id="test"
+        )
+        writer.write(ydb.TopicWriterMessage(data="123".encode()))
+        writer.close()
+
+    def test_wait_last_seqno(self, driver_sync: ydb.Driver, topic_path):
+        with driver_sync.topic_client.topic_writer(
+            topic_path,
+            producer_and_message_group_id="test",
+            auto_seqno=False,
+        ) as writer:
+            writer.write_with_ack(ydb.TopicWriterMessage(data="123".encode(), seqno=5))
+
+        with driver_sync.topic_client.topic_writer(
+            topic_path,
+            producer_and_message_group_id="test",
+            get_last_seqno=True,
+        ) as writer2:
+            init_info = writer2.wait_init()
+            assert init_info.last_seqno == 5
+
+    def test_auto_flush_on_close(self, driver_sync: ydb.Driver, topic_path):
+        with driver_sync.topic_client.topic_writer(
+            topic_path,
+            producer_and_message_group_id="test",
+            auto_seqno=False,
+        ) as writer:
+            last_seqno = 0
+            for i in range(10):
+                last_seqno = i + 1
+                writer.write(ydb.TopicWriterMessage(data=f"msg-{i}", seqno=last_seqno))
+
+        with driver_sync.topic_client.topic_writer(
+            topic_path,
+            producer_and_message_group_id="test",
+            get_last_seqno=True,
+        ) as writer:
+            init_info = writer.wait_init()
+            assert init_info.last_seqno == last_seqno

ydb/_grpc/grpcwrapper/common_utils.py

@@ -2,7 +2,9 @@
 
 import abc
 import asyncio
+import contextvars
 import datetime
+import functools
 import typing
 from typing import (
     Optional,
@@ -118,7 +120,7 @@ def __aiter__(self):
 
     async def __anext__(self):
         try:
-            res = await asyncio.to_thread(self._sync_iterator.__next__)
+            res = await to_thread(self._sync_iterator.__next__)
             return res
         except StopAsyncIteration:
             raise StopIteration()
@@ -180,7 +182,7 @@ async def _start_asyncio_driver(self, driver: ydb.aio.Driver, stub, method):
 
     async def _start_sync_driver(self, driver: ydb.Driver, stub, method):
         requests_iterator = AsyncQueueToSyncIteratorAsyncIO(self.from_client_grpc)
-        stream_call = await asyncio.to_thread(
+        stream_call = await to_thread(
             driver,
             requests_iterator,
             stub,
@@ -257,6 +259,25 @@ def callback_from_asyncio(
         return loop.run_in_executor(None, callback)
 
 
+async def to_thread(func, /, *args, **kwargs):
+    """Asynchronously run function *func* in a separate thread.
+
+    Any *args and **kwargs supplied for this function are directly passed
+    to *func*. Also, the current :class:`contextvars.Context` is propagated,
+    allowing context variables from the main thread to be accessed in the
+    separate thread.
+
+    Return a coroutine that can be awaited to get the eventual result of *func*.
+
+    copy to_thread from 3.10
+    """
+
+    loop = asyncio.get_running_loop()
+    ctx = contextvars.copy_context()
+    func_call = functools.partial(ctx.run, func, *args, **kwargs)
+    return await loop.run_in_executor(None, func_call)
+
+
 def proto_duration_from_timedelta(t: Optional[datetime.timedelta]) -> ProtoDuration:
     if t is None:
         return None

ydb/_topic_common/common.py

@@ -4,6 +4,7 @@
 from .._grpc.grpcwrapper.common_utils import IFromProtoWithProtoType
 
 TokenGetterFuncType = typing.Optional[typing.Callable[[], str]]
+TimeoutType = typing.Union[int, float]
 
 
 def wrap_operation(rpc_state, response_pb, driver=None):

ydb/_topic_writer/topic_writer.py

@@ -1,4 +1,3 @@
-import concurrent.futures
 import datetime
 import enum
 from dataclasses import dataclass
@@ -12,126 +11,7 @@
 from .._grpc.grpcwrapper.common_utils import IToProto
 
 
-class Writer:
-    @property
-    def last_seqno(self) -> int:
-        raise NotImplementedError()
-
-    def __init__(self, db: ydb.Driver):
-        pass
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
-
-    def close(self):
-        pass
-
-    MessageType = typing.Union["PublicMessage", "PublicMessage.SimpleMessageSourceType"]
-
-    def write(
-        self,
-        message: Union[MessageType, List[MessageType]],
-        *args: Optional[MessageType],
-        timeout: [float, None] = None,
-    ):
-        """
-        send one or number of messages to server.
-        it fast put message to internal buffer, without wait message result
-        return None
-
-        message will send independent of wait/no wait result
-
-        timeout - time for waiting for put message into internal queue.
-            if 0 or negative - non block calls
-            if None or not set - infinite wait
-            It will raise TimeoutError() exception if it can't put message to internal queue by limits during timeout.
-        """
-        raise NotImplementedError()
-
-    def async_write_with_ack(
-        self,
-        message: Union[MessageType, List[MessageType]],
-        *args: Optional[MessageType],
-        timeout: [float, None] = None,
-    ) -> concurrent.futures.Future:
-        """
-        send one or number of messages to server.
-        return feature, which can be waited for check send result: ack/duplicate/error
-
-        Usually it is fast method, but can wait if internal buffer is full.
-
-        timeout - time for waiting for put message into internal queue.
-            The method can be blocked up to timeout seconds before return future.
-
-            if 0 or negative - non block calls
-            if None or not set - infinite wait
-            It will raise TimeoutError() exception if it can't put message to internal queue by limits during timeout.
-        """
-        raise NotImplementedError()
-
-    def write_with_ack(
-        self,
-        message: Union[MessageType, List[MessageType]],
-        *args: Optional[MessageType],
-        buffer_timeout: [float, None] = None,
-    ) -> Union["MessageWriteStatus", List["MessageWriteStatus"]]:
-        """
-        IT IS SLOWLY WAY. IT IS BAD CHOISE IN MOST CASES.
-        It is recommended to use write with optionally flush or async_write_with_ack and receive acks by wait future.
-
-        send one or number of messages to server.
-        blocked until receive server ack for the message/messages.
-
-        message will send independent of wait/no wait result
-
-        buffer_timeout - time for send message to server and receive ack.
-            if 0 or negative - non block calls
-            if None or not set - infinite wait
-            It will raise TimeoutError() exception if it isn't receive ack in timeout
-        """
-        raise NotImplementedError()
-
-    def async_flush(self):
-        """
-        Force send all messages from internal buffer and wait acks from server for all
-        messages.
-
-        flush starts of flush process, and return Future for wait result.
-        messages will be flushed independent of future waiting.
-        """
-        raise NotImplementedError()
-
-    def flush(self, timeout: Optional[float] = None) -> concurrent.futures.Future:
-        """
-        Force send all messages from internal buffer and wait acks from server for all
-        messages.
-
-        timeout - time for waiting for send all messages and receive server ack.
-            if 0 or negative - non block calls
-            if None or not set - infinite wait
-            It will raise TimeoutError() exception if it isn't receive ack in timeout
-        """
-        raise NotImplementedError()
-
-    def async_wait_init(self) -> concurrent.futures.Future:
-        """
-        Return feature, which done when underling connection established
-        """
-        raise NotImplementedError()
-
-    def wait_init(self, timeout: Optional[float] = None):
-        """
-        Wait until underling connection established
-
-        timeout - time for waiting for send all messages and receive server ack.
-            if 0 or negative - non block calls
-            if None or not set - infinite wait
-            It will raise TimeoutError() exception if it isn't receive ack in timeout
-        """
-        raise NotImplementedError()
+MessageType = typing.Union["PublicMessage", "PublicMessage.SimpleMessageSourceType"]
 
 
 @dataclass

ydb/_topic_writer/topic_writer_asyncio.py

@@ -7,14 +7,14 @@
 from .topic_writer import (
     PublicWriterSettings,
     WriterSettings,
-    Writer,
     PublicMessage,
     PublicWriterInitInfo,
     InternalMessage,
     TopicWriterStopped,
     TopicWriterError,
     messages_to_proto_requests,
     PublicWriteResultTypes,
+    MessageType,
 )
 from .. import (
     _apis,
@@ -75,8 +75,8 @@ async def close(self, *, flush: bool = True):
 
     async def write_with_ack(
         self,
-        messages: Union[Writer.MessageType, List[Writer.MessageType]],
-        *args: Optional[Writer.MessageType],
+        messages: Union[MessageType, List[MessageType]],
+        *args: Optional[MessageType],
     ) -> Union[PublicWriteResultTypes, List[PublicWriteResultTypes]]:
         """
         IT IS SLOWLY WAY. IT IS BAD CHOISE IN MOST CASES.
@@ -97,8 +97,8 @@ async def write_with_ack(
 
     async def write_with_ack_future(
         self,
-        messages: Union[Writer.MessageType, List[Writer.MessageType]],
-        *args: Optional[Writer.MessageType],
+        messages: Union[MessageType, List[MessageType]],
+        *args: Optional[MessageType],
     ) -> Union[asyncio.Future, List[asyncio.Future]]:
         """
         send one or number of messages to server.
@@ -120,8 +120,8 @@ async def write_with_ack_future(
 
     async def write(
         self,
-        messages: Union[Writer.MessageType, List[Writer.MessageType]],
-        *args: Optional[Writer.MessageType],
+        messages: Union[MessageType, List[MessageType]],
+        *args: Optional[MessageType],
     ):
         """
         send one or number of messages to server.

ydb/_topic_writer/topic_writer_sync.py

@@ -11,11 +11,12 @@
     TopicWriterError,
     PublicWriterInitInfo,
     PublicMessage,
-    Writer,
     PublicWriteResult,
+    MessageType,
 )
 
 from .topic_writer_asyncio import WriterAsyncIO
+from .._topic_common.common import TimeoutType
 
 _shared_event_loop_lock = threading.Lock()
 _shared_event_loop = None  # type: Optional[asyncio.AbstractEventLoop]
@@ -78,6 +79,12 @@ async def create_async_writer():
             create_async_writer(), self._loop
         ).result()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
     def _call(self, coro, *args, **kwargs):
         if self._closed:
             raise TopicWriterError("writer is closed")
@@ -87,7 +94,7 @@ def _call(self, coro, *args, **kwargs):
     def _call_sync(self, coro: Coroutine, timeout, *args, **kwargs):
         f = self._call(coro, *args, **kwargs)
         try:
-            return f.result()
+            return f.result(timeout=timeout)
         except TimeoutError:
             f.cancel()
             raise
@@ -111,7 +118,7 @@ def flush(self, timeout=None):
     def async_wait_init(self) -> Future[PublicWriterInitInfo]:
         return self._call(self._async_writer.wait_init())
 
-    def wait_init(self, timeout) -> PublicWriterInitInfo:
+    def wait_init(self, timeout: Optional[TimeoutType] = None) -> PublicWriterInitInfo:
         return self._call_sync(self._async_writer.wait_init(), timeout)
 
     def write(
@@ -124,15 +131,15 @@ def write(
 
     def async_write_with_ack(
         self,
-        messages: Union[Writer.MessageType, List[Writer.MessageType]],
-        *args: Optional[Writer.MessageType],
+        messages: Union[MessageType, List[MessageType]],
+        *args: Optional[MessageType],
     ) -> Future[Union[PublicWriteResult, List[PublicWriteResult]]]:
         return self._call(self._async_writer.write_with_ack(messages, *args))
 
     def write_with_ack(
         self,
-        messages: Union[Writer.MessageType, List[Writer.MessageType]],
-        *args: Optional[Writer.MessageType],
+        messages: Union[MessageType, List[MessageType]],
+        *args: Optional[MessageType],
         timeout: Union[float, None] = None,
     ) -> Union[PublicWriteResult, List[PublicWriteResult]]:
         return self._call_sync(

ydb/topic.py

@@ -19,12 +19,13 @@
 )
 
 from ._topic_writer.topic_writer import (  # noqa: F401
-    Writer as TopicWriter,
     PublicWriterSettings as TopicWriterSettings,
     PublicMessage as TopicWriterMessage,
     RetryPolicy as TopicWriterRetryPolicy,
 )
 
+from ._topic_writer.topic_writer_sync import WriterSync as TopicWriter
+
 from ._topic_common.common import (
     wrap_operation as _wrap_operation,
     create_result_wrapper as _create_result_wrapper,
@@ -278,7 +279,10 @@ def topic_writer(
         get_last_seqno: bool = False,
         retry_policy: Union["TopicWriterRetryPolicy", None] = None,
     ) -> TopicWriter:
-        raise NotImplementedError()
+        args = locals()
+        del args["self"]
+        settings = TopicWriterSettings(**args)
+        return TopicWriter(self._driver, settings)
 
 
 class TopicClientSettings: