Intermediate changes

commit_hash:bdf980a10265fd16f4aecae0dff78846b425bd8b

Author: robot-piglet

contrib/python/ydb/py3/.dist-info/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ydb
-Version: 3.19.3
+Version: 3.20.1
 Summary: YDB Python SDK
 Home-page: http://github.com/ydb-platform/ydb-python-sdk
 Author: Yandex LLC

contrib/python/ydb/py3/ya.make

@@ -2,7 +2,7 @@
 
 PY3_LIBRARY()
 
-VERSION(3.19.3)
+VERSION(3.20.1)
 
 LICENSE(Apache-2.0)
 

contrib/python/ydb/py3/ydb/_apis.py

@@ -115,6 +115,7 @@ class TopicService(object):
     DropTopic = "DropTopic"
     StreamRead = "StreamRead"
     StreamWrite = "StreamWrite"
+    UpdateOffsetsInTransaction = "UpdateOffsetsInTransaction"
 
 
 class QueryService(object):

contrib/python/ydb/py3/ydb/_errors.py

@@ -5,6 +5,7 @@
 
 _errors_retriable_fast_backoff_types = [
     issues.Unavailable,
+    issues.ClientInternalError,
 ]
 _errors_retriable_slow_backoff_types = [
     issues.Aborted,

contrib/python/ydb/py3/ydb/_grpc/grpcwrapper/ydb_topic.py

@@ -141,6 +141,18 @@ def from_proto(msg: ydb_topic_pb2.UpdateTokenResponse) -> typing.Any:
 ########################################################################################################################
 
 
+@dataclass
+class TransactionIdentity(IToProto):
+    tx_id: str
+    session_id: str
+
+    def to_proto(self) -> ydb_topic_pb2.TransactionIdentity:
+        return ydb_topic_pb2.TransactionIdentity(
+            id=self.tx_id,
+            session=self.session_id,
+        )
+
+
 class StreamWriteMessage:
     @dataclass()
     class InitRequest(IToProto):
@@ -199,6 +211,7 @@ def from_proto(
     class WriteRequest(IToProto):
         messages: typing.List["StreamWriteMessage.WriteRequest.MessageData"]
         codec: int
+        tx_identity: Optional[TransactionIdentity]
 
         @dataclass
         class MessageData(IToProto):
@@ -237,6 +250,9 @@ def to_proto(self) -> ydb_topic_pb2.StreamWriteMessage.WriteRequest:
             proto = ydb_topic_pb2.StreamWriteMessage.WriteRequest()
             proto.codec = self.codec
 
+            if self.tx_identity is not None:
+                proto.tx.CopyFrom(self.tx_identity.to_proto())
+
             for message in self.messages:
                 proto_mess = proto.messages.add()
                 proto_mess.CopyFrom(message.to_proto())
@@ -297,6 +313,8 @@ def from_proto(cls, proto_ack: ydb_topic_pb2.StreamWriteMessage.WriteResponse.Wr
                         )
                     except ValueError:
                         message_write_status = reason
+                elif proto_ack.HasField("written_in_tx"):
+                    message_write_status = StreamWriteMessage.WriteResponse.WriteAck.StatusWrittenInTx()
                 else:
                     raise NotImplementedError("unexpected ack status")
 
@@ -309,6 +327,9 @@ def from_proto(cls, proto_ack: ydb_topic_pb2.StreamWriteMessage.WriteResponse.Wr
             class StatusWritten:
                 offset: int
 
+            class StatusWrittenInTx:
+                pass
+
             @dataclass
             class StatusSkipped:
                 reason: "StreamWriteMessage.WriteResponse.WriteAck.StatusSkipped.Reason"
@@ -1196,6 +1217,52 @@ def to_public(self) -> ydb_topic_public_types.PublicMeteringMode:
             return ydb_topic_public_types.PublicMeteringMode.UNSPECIFIED
 
 
+@dataclass
+class UpdateOffsetsInTransactionRequest(IToProto):
+    tx: TransactionIdentity
+    topics: List[UpdateOffsetsInTransactionRequest.TopicOffsets]
+    consumer: str
+
+    def to_proto(self):
+        return ydb_topic_pb2.UpdateOffsetsInTransactionRequest(
+            tx=self.tx.to_proto(),
+            consumer=self.consumer,
+            topics=list(
+                map(
+                    UpdateOffsetsInTransactionRequest.TopicOffsets.to_proto,
+                    self.topics,
+                )
+            ),
+        )
+
+    @dataclass
+    class TopicOffsets(IToProto):
+        path: str
+        partitions: List[UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets]
+
+        def to_proto(self):
+            return ydb_topic_pb2.UpdateOffsetsInTransactionRequest.TopicOffsets(
+                path=self.path,
+                partitions=list(
+                    map(
+                        UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets.to_proto,
+                        self.partitions,
+                    )
+                ),
+            )
+
+        @dataclass
+        class PartitionOffsets(IToProto):
+            partition_id: int
+            partition_offsets: List[OffsetsRange]
+
+            def to_proto(self) -> ydb_topic_pb2.UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets:
+                return ydb_topic_pb2.UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets(
+                    partition_id=self.partition_id,
+                    partition_offsets=list(map(OffsetsRange.to_proto, self.partition_offsets)),
+                )
+
+
 @dataclass
 class CreateTopicRequest(IToProto, IFromPublic):
     path: str

contrib/python/ydb/py3/ydb/_topic_reader/datatypes.py

@@ -108,6 +108,9 @@ def ack_notify(self, offset: int):
             waiter = self._ack_waiters.popleft()
             waiter._finish_ok()
 
+    def _update_last_commited_offset_if_needed(self, offset: int):
+        self.committed_offset = max(self.committed_offset, offset)
+
     def close(self):
         if self.closed:
             return
@@ -211,3 +214,9 @@ def _pop_batch(self, message_count: int) -> PublicBatch:
         self._bytes_size = self._bytes_size - new_batch._bytes_size
 
         return new_batch
+
+    def _update_partition_offsets(self, tx, exc=None):
+        if exc is not None:
+            return
+        offsets = self._commit_get_offsets_range()
+        self._partition_session._update_last_commited_offset_if_needed(offsets.end)

contrib/python/ydb/py3/ydb/_topic_reader/topic_reader_asyncio.py

@@ -5,7 +5,7 @@
 import gzip
 import typing
 from asyncio import Task
-from collections import OrderedDict
+from collections import defaultdict, OrderedDict
 from typing import Optional, Set, Dict, Union, Callable
 
 import ydb
@@ -19,17 +19,24 @@
 from .._grpc.grpcwrapper.common_utils import (
     IGrpcWrapperAsyncIO,
     SupportedDriverType,
+    to_thread,
     GrpcWrapperAsyncIO,
 )
 from .._grpc.grpcwrapper.ydb_topic import (
     StreamReadMessage,
     UpdateTokenRequest,
     UpdateTokenResponse,
+    UpdateOffsetsInTransactionRequest,
     Codec,
 )
 from .._errors import check_retriable_error
 import logging
 
+from ..query.base import TxEvent
+
+if typing.TYPE_CHECKING:
+    from ..query.transaction import BaseQueryTxContext
+
 logger = logging.getLogger(__name__)
 
 
@@ -77,7 +84,7 @@ def __init__(
     ):
         self._loop = asyncio.get_running_loop()
         self._closed = False
-        self._reconnector = ReaderReconnector(driver, settings)
+        self._reconnector = ReaderReconnector(driver, settings, self._loop)
         self._parent = _parent
 
     async def __aenter__(self):
@@ -88,8 +95,12 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
     def __del__(self):
         if not self._closed:
-            task = self._loop.create_task(self.close(flush=False))
-            topic_common.wrap_set_name_for_asyncio_task(task, task_name="close reader")
+            try:
+                logger.warning("Topic reader was not closed properly. Consider using method close().")
+                task = self._loop.create_task(self.close(flush=False))
+                topic_common.wrap_set_name_for_asyncio_task(task, task_name="close reader")
+            except BaseException:
+                logger.warning("Something went wrong during reader close in __del__")
 
     async def wait_message(self):
         """
@@ -112,6 +123,23 @@ async def receive_batch(
             max_messages=max_messages,
         )
 
+    async def receive_batch_with_tx(
+        self,
+        tx: "BaseQueryTxContext",
+        max_messages: typing.Union[int, None] = None,
+    ) -> typing.Union[datatypes.PublicBatch, None]:
+        """
+        Get one messages batch with tx from reader.
+        All messages in a batch from same partition.
+
+        use asyncio.wait_for for wait with timeout.
+        """
+        await self._reconnector.wait_message()
+        return self._reconnector.receive_batch_with_tx_nowait(
+            tx=tx,
+            max_messages=max_messages,
+        )
+
     async def receive_message(self) -> typing.Optional[datatypes.PublicMessage]:
         """
         Block until receive new message
@@ -165,18 +193,27 @@ class ReaderReconnector:
     _state_changed: asyncio.Event
     _stream_reader: Optional["ReaderStream"]
     _first_error: asyncio.Future[YdbError]
+    _tx_to_batches_map: Dict[str, typing.List[datatypes.PublicBatch]]
 
-    def __init__(self, driver: Driver, settings: topic_reader.PublicReaderSettings):
+    def __init__(
+        self,
+        driver: Driver,
+        settings: topic_reader.PublicReaderSettings,
+        loop: Optional[asyncio.AbstractEventLoop] = None,
+    ):
         self._id = self._static_reader_reconnector_counter.inc_and_get()
         self._settings = settings
         self._driver = driver
+        self._loop = loop if loop is not None else asyncio.get_running_loop()
         self._background_tasks = set()
 
         self._state_changed = asyncio.Event()
         self._stream_reader = None
         self._background_tasks.add(asyncio.create_task(self._connection_loop()))
         self._first_error = asyncio.get_running_loop().create_future()
 
+        self._tx_to_batches_map = dict()
+
     async def _connection_loop(self):
         attempt = 0
         while True:
@@ -190,6 +227,7 @@ async def _connection_loop(self):
                 if not retry_info.is_retriable:
                     self._set_first_error(err)
                     return
+
                 await asyncio.sleep(retry_info.sleep_timeout_seconds)
 
                 attempt += 1
@@ -222,9 +260,87 @@ def receive_batch_nowait(self, max_messages: Optional[int] = None):
             max_messages=max_messages,
         )
 
+    def receive_batch_with_tx_nowait(self, tx: "BaseQueryTxContext", max_messages: Optional[int] = None):
+        batch = self._stream_reader.receive_batch_nowait(
+            max_messages=max_messages,
+        )
+
+        self._init_tx(tx)
+
+        self._tx_to_batches_map[tx.tx_id].append(batch)
+
+        tx._add_callback(TxEvent.AFTER_COMMIT, batch._update_partition_offsets, self._loop)
+
+        return batch
+
     def receive_message_nowait(self):
         return self._stream_reader.receive_message_nowait()
 
+    def _init_tx(self, tx: "BaseQueryTxContext"):
+        if tx.tx_id not in self._tx_to_batches_map:  # Init tx callbacks
+            self._tx_to_batches_map[tx.tx_id] = []
+            tx._add_callback(TxEvent.BEFORE_COMMIT, self._commit_batches_with_tx, self._loop)
+            tx._add_callback(TxEvent.AFTER_COMMIT, self._handle_after_tx_commit, self._loop)
+            tx._add_callback(TxEvent.AFTER_ROLLBACK, self._handle_after_tx_rollback, self._loop)
+
+    async def _commit_batches_with_tx(self, tx: "BaseQueryTxContext"):
+        grouped_batches = defaultdict(lambda: defaultdict(list))
+        for batch in self._tx_to_batches_map[tx.tx_id]:
+            grouped_batches[batch._partition_session.topic_path][batch._partition_session.partition_id].append(batch)
+
+        request = UpdateOffsetsInTransactionRequest(tx=tx._tx_identity(), consumer=self._settings.consumer, topics=[])
+
+        for topic_path in grouped_batches:
+            topic_offsets = UpdateOffsetsInTransactionRequest.TopicOffsets(path=topic_path, partitions=[])
+            for partition_id in grouped_batches[topic_path]:
+                partition_offsets = UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets(
+                    partition_id=partition_id,
+                    partition_offsets=[
+                        batch._commit_get_offsets_range() for batch in grouped_batches[topic_path][partition_id]
+                    ],
+                )
+                topic_offsets.partitions.append(partition_offsets)
+            request.topics.append(topic_offsets)
+
+        try:
+            return await self._do_commit_batches_with_tx_call(request)
+        except BaseException:
+            exc = issues.ClientInternalError("Failed to update offsets in tx.")
+            tx._set_external_error(exc)
+            self._stream_reader._set_first_error(exc)
+        finally:
+            del self._tx_to_batches_map[tx.tx_id]
+
+    async def _do_commit_batches_with_tx_call(self, request: UpdateOffsetsInTransactionRequest):
+        args = [
+            request.to_proto(),
+            _apis.TopicService.Stub,
+            _apis.TopicService.UpdateOffsetsInTransaction,
+            topic_common.wrap_operation,
+        ]
+
+        if asyncio.iscoroutinefunction(self._driver.__call__):
+            res = await self._driver(*args)
+        else:
+            res = await to_thread(self._driver, *args, executor=None)
+
+        return res
+
+    async def _handle_after_tx_rollback(self, tx: "BaseQueryTxContext", exc: Optional[BaseException]) -> None:
+        if tx.tx_id in self._tx_to_batches_map:
+            del self._tx_to_batches_map[tx.tx_id]
+        exc = issues.ClientInternalError("Reconnect due to transaction rollback")
+        self._stream_reader._set_first_error(exc)
+
+    async def _handle_after_tx_commit(self, tx: "BaseQueryTxContext", exc: Optional[BaseException]) -> None:
+        if tx.tx_id in self._tx_to_batches_map:
+            del self._tx_to_batches_map[tx.tx_id]
+
+        if exc is not None:
+            self._stream_reader._set_first_error(
+                issues.ClientInternalError("Reconnect due to transaction commit failed")
+            )
+
     def commit(self, batch: datatypes.ICommittable) -> datatypes.PartitionSession.CommitAckWaiter:
         return self._stream_reader.commit(batch)