group messages to batches

Author: rekby

ydb/_topic_reader/topic_reader_asyncio.py

@@ -26,7 +26,6 @@
     Codec,
 )
 from .._errors import check_retriable_error
-from .. import topic
 
 
 class TopicReaderError(YdbError):
@@ -65,11 +64,11 @@ class PublicAsyncIOReader:
     _parent: typing.Any  # need for prevent close parent client by GC
 
     def __init__(
-            self,
-            driver: Driver,
-            settings: topic_reader.PublicReaderSettings,
-            *,
-            _parent=None,
+        self,
+        driver: Driver,
+        settings: topic_reader.PublicReaderSettings,
+        *,
+        _parent=None,
     ):
         self._loop = asyncio.get_running_loop()
         self._closed = False
@@ -87,7 +86,7 @@ def __del__(self):
             self._loop.create_task(self.close(flush=False), name="close reader")
 
     async def receive_batch(
-            self,
+        self,
     ) -> typing.Union[datatypes.PublicBatch, None]:
         """
         Get one messages batch from reader.
@@ -108,7 +107,7 @@ async def receive_message(self) -> typing.Optional[datatypes.PublicMessage]:
         return self._reconnector.receive_message_nowait()
 
     def commit(
-            self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]
+        self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]
     ):
         """
         Write commit message to a buffer.
@@ -119,7 +118,7 @@ def commit(
         self._reconnector.commit(batch)
 
     async def commit_with_ack(
-            self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]
+        self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]
     ):
         """
         write commit message to a buffer and wait ack from the server.
@@ -204,7 +203,7 @@ def receive_message_nowait(self):
         return self._stream_reader.receive_message_nowait()
 
     def commit(
-            self, batch: datatypes.ICommittable
+        self, batch: datatypes.ICommittable
     ) -> datatypes.PartitionSession.CommitAckWaiter:
         return self._stream_reader.commit(batch)
 
@@ -263,10 +262,10 @@ class ReaderStream:
     _get_token_function: Callable[[], str]
 
     def __init__(
-            self,
-            reader_reconnector_id: int,
-            settings: topic_reader.PublicReaderSettings,
-            get_token_function: Optional[Callable[[], str]] = None,
+        self,
+        reader_reconnector_id: int,
+        settings: topic_reader.PublicReaderSettings,
+        get_token_function: Optional[Callable[[], str]] = None,
     ):
         self._loop = asyncio.get_running_loop()
         self._id = ReaderStream._static_id_counter.inc_and_get()
@@ -295,9 +294,9 @@ def __init__(
 
     @staticmethod
     async def create(
-            reader_reconnector_id: int,
-            driver: SupportedDriverType,
-            settings: topic_reader.PublicReaderSettings,
+        reader_reconnector_id: int,
+        driver: SupportedDriverType,
+        settings: topic_reader.PublicReaderSettings,
     ) -> "ReaderStream":
         stream = GrpcWrapperAsyncIO(StreamReadMessage.FromServer.from_proto)
 
@@ -315,7 +314,7 @@ async def create(
         return reader
 
     async def _start(
-            self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMessage.InitRequest
+        self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMessage.InitRequest
     ):
         if self._started:
             raise TopicReaderError("Double start ReaderStream")
@@ -381,7 +380,7 @@ def receive_message_nowait(self):
         return message
 
     def commit(
-            self, batch: datatypes.ICommittable
+        self, batch: datatypes.ICommittable
     ) -> datatypes.PartitionSession.CommitAckWaiter:
         partition_session = batch._commit_get_partition_session()
 
@@ -435,19 +434,19 @@ async def _read_messages_loop(self):
                     self._on_read_response(message.server_message)
 
                 elif isinstance(
-                        message.server_message, StreamReadMessage.CommitOffsetResponse
+                    message.server_message, StreamReadMessage.CommitOffsetResponse
                 ):
                     self._on_commit_response(message.server_message)
 
                 elif isinstance(
-                        message.server_message,
-                        StreamReadMessage.StartPartitionSessionRequest,
+                    message.server_message,
+                    StreamReadMessage.StartPartitionSessionRequest,
                 ):
                     self._on_start_partition_session(message.server_message)
 
                 elif isinstance(
-                        message.server_message,
-                        StreamReadMessage.StopPartitionSessionRequest,
+                    message.server_message,
+                    StreamReadMessage.StopPartitionSessionRequest,
                 ):
                     self._on_partition_session_stop(message.server_message)
 
@@ -479,12 +478,12 @@ async def _update_token(self, token: str):
             self._update_token_event.clear()
 
     def _on_start_partition_session(
-            self, message: StreamReadMessage.StartPartitionSessionRequest
+        self, message: StreamReadMessage.StartPartitionSessionRequest
     ):
         try:
             if (
-                    message.partition_session.partition_session_id
-                    in self._partition_sessions
+                message.partition_session.partition_session_id
+                in self._partition_sessions
             ):
                 raise TopicReaderError(
                     "Double start partition session: %s"
@@ -515,7 +514,7 @@ def _on_start_partition_session(
             self._set_first_error(err)
 
     def _on_partition_session_stop(
-            self, message: StreamReadMessage.StopPartitionSessionRequest
+        self, message: StreamReadMessage.StopPartitionSessionRequest
     ):
         if message.partition_session_id not in self._partition_sessions:
             # may if receive stop partition with graceful=false after response on stop partition
@@ -563,7 +562,7 @@ def _buffer_release_bytes(self, bytes_size):
         )
 
     def _read_response_to_batches(
-            self, message: StreamReadMessage.ReadResponse
+        self, message: StreamReadMessage.ReadResponse
     ) -> typing.List[datatypes.PublicBatch]:
         batches = []
 
@@ -573,7 +572,7 @@ def _read_response_to_batches(
 
         bytes_per_batch = message.bytes_size // batch_count
         additional_bytes_to_last_batch = (
-                message.bytes_size - bytes_per_batch * batch_count
+            message.bytes_size - bytes_per_batch * batch_count
         )
 
         for partition_data in message.partition_data:

ydb/_topic_writer/topic_writer_asyncio.py

@@ -40,8 +40,6 @@
     GrpcWrapperAsyncIO,
 )
 
-from .. import topic
-
 logger = logging.getLogger(__name__)
 
 
@@ -52,10 +50,10 @@ class WriterAsyncIO:
     _parent: typing.Any  # need for prevent close parent client by GC
 
     def __init__(
-            self,
-            driver: SupportedDriverType,
-            settings: PublicWriterSettings,
-            _client=None,
+        self,
+        driver: SupportedDriverType,
+        settings: PublicWriterSettings,
+        _client=None,
     ):
         self._loop = asyncio.get_running_loop()
         self._closed = False

ydb/_topic_writer/topic_writer_sync.py

@@ -5,7 +5,6 @@
 from concurrent.futures import Future
 from typing import Union, List, Optional
 
-from .. import topic
 from .._grpc.grpcwrapper.common_utils import SupportedDriverType
 from .topic_writer import (
     PublicWriterSettings,
@@ -30,12 +29,12 @@ class WriterSync:
     _parent: typing.Any  # need for prevent close parent client by GC
 
     def __init__(
-            self,
-            driver: SupportedDriverType,
-            settings: PublicWriterSettings,
-            *,
-            eventloop: Optional[asyncio.AbstractEventLoop] = None,
-            _parent=None
+        self,
+        driver: SupportedDriverType,
+        settings: PublicWriterSettings,
+        *,
+        eventloop: Optional[asyncio.AbstractEventLoop] = None,
+        _parent=None,
     ):
 
         self._closed = False
@@ -101,17 +100,17 @@ def wait_init(self, *, timeout: TimeoutType = None) -> PublicWriterInitInfo:
         )
 
     def write(
-            self,
-            messages: Union[Message, List[Message]],
-            timeout: TimeoutType = None,
+        self,
+        messages: Union[Message, List[Message]],
+        timeout: TimeoutType = None,
     ):
         self._check_closed()
 
         self._caller.safe_call_with_result(self._async_writer.write(messages), timeout)
 
     def async_write_with_ack(
-            self,
-            messages: Union[Message, List[Message]],
+        self,
+        messages: Union[Message, List[Message]],
     ) -> Future[Union[PublicWriteResult, List[PublicWriteResult]]]:
         self._check_closed()
 
@@ -120,9 +119,9 @@ def async_write_with_ack(
         )
 
     def write_with_ack(
-            self,
-            messages: Union[Message, List[Message]],
-            timeout: Union[float, None] = None,
+        self,
+        messages: Union[Message, List[Message]],
+        timeout: Union[float, None] = None,
     ) -> Union[PublicWriteResult, List[PublicWriteResult]]:
         self._check_closed()