async topic tx with listener pattern

Author: vgvoleg

tests/topics/test_topic_transactions.py

@@ -38,14 +38,14 @@ async def callee(tx: ydb.aio.QueryTxContext):
                 assert msg.data.decode() == "123"
 
 
-@pytest.mark.skip("Not implemented yet.")
+# @pytest.mark.skip("Not implemented yet.")
 class TestTopicTransactionalWriter:
     async def test_commit(self, driver: ydb.aio.Driver, topic_path, topic_reader: ydb.TopicReaderAsyncIO):
         async with ydb.aio.QuerySessionPool(driver) as pool:
 
             async def callee(tx: ydb.aio.QueryTxContext):
                 tx_writer = driver.topic_client.tx_writer(tx, topic_path)
-                tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
+                await tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
 
             await pool.retry_tx_async(callee)
 
@@ -57,7 +57,7 @@ async def test_rollback(self, driver: ydb.aio.Driver, topic_path, topic_reader:
 
             async def callee(tx: ydb.aio.QueryTxContext):
                 tx_writer = driver.topic_client.tx_writer(tx, topic_path)
-                tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
+                await tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
 
                 await tx.rollback()
 

ydb/_grpc/grpcwrapper/ydb_topic.py

@@ -252,7 +252,7 @@ def to_proto(self) -> ydb_topic_pb2.StreamWriteMessage.WriteRequest:
             proto.codec = self.codec
 
             if self.tx_identity is not None:
-                proto.tx = self.tx_identity.to_proto()
+                proto.tx.CopyFrom(self.tx_identity.to_proto())
 
             for message in self.messages:
                 proto_mess = proto.messages.add()
@@ -314,6 +314,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")
 
@@ -326,6 +328,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"

ydb/_topic_writer/topic_writer.py

@@ -54,8 +54,12 @@ class Written:
     class Skipped:
         pass
 
+    @dataclass(eq=True)
+    class WrittenInTx:
+        pass
+
 
-PublicWriteResultTypes = Union[PublicWriteResult.Written, PublicWriteResult.Skipped]
+PublicWriteResultTypes = Union[PublicWriteResult.Written, PublicWriteResult.Skipped, PublicWriteResult.WrittenInTx]
 
 
 class WriterSettings(PublicWriterSettings):

ydb/_topic_writer/topic_writer_asyncio.py

@@ -44,6 +44,8 @@
     GrpcWrapperAsyncIO,
 )
 
+from ..query.base import TxListenerAsyncIO
+
 if typing.TYPE_CHECKING:
     from ..query.transaction import BaseQueryTxContext
 
@@ -168,12 +170,12 @@ async def wait_init(self) -> PublicWriterInitInfo:
         return await self._reconnector.wait_init()
 
 
-class TxWriterAsyncIO(WriterAsyncIO):
-    _tx: object
+class TxWriterAsyncIO(WriterAsyncIO, TxListenerAsyncIO):
+    _tx: "BaseQueryTxContext"
 
     def __init__(
         self,
-        tx,
+        tx: "BaseQueryTxContext",
         driver: SupportedDriverType,
         settings: PublicWriterSettings,
         _client=None,
@@ -183,6 +185,13 @@ def __init__(
         self._closed = False
         self._reconnector = WriterAsyncIOReconnector(driver=driver, settings=WriterSettings(settings), tx=self._tx)
         self._parent = _client
+        self._tx._add_listener(self)
+
+    async def _on_before_commit(self):
+        await self.close()
+
+    async def _on_before_rollback(self):
+        await self.close()
 
 
 class WriterAsyncIOReconnector:
@@ -560,6 +569,8 @@ def _handle_receive_ack(self, ack):
             result = PublicWriteResult.Skipped()
         elif isinstance(status, write_ack_msg.StatusWritten):
             result = PublicWriteResult.Written(offset=status.offset)
+        elif isinstance(status, write_ack_msg.StatusWrittenInTx):
+            result = PublicWriteResult.WrittenInTx()
         else:
             raise TopicWriterError("internal error - receive unexpected ack message.")
         message_future.set_result(result)
@@ -575,6 +586,7 @@ async def _send_loop(self, writer: "WriterAsyncIOStream"):
 
             while True:
                 m = await self._new_messages.get()  # type: InternalMessage
+                print("NEW MESSAGE")
                 if m.seq_no > last_seq_no:
                     writer.write([m])
         except asyncio.CancelledError:
@@ -606,6 +618,7 @@ async def flush(self):
 
         # wait last message
         await asyncio.wait(self._messages_future)
+        print("ALL MESSAGES WERE SENT TO SERVER")
 
 
 class WriterAsyncIOStream:

ydb/aio/query/pool.py

@@ -158,6 +158,8 @@ async def retry_tx_async(
         async def wrapped_callee():
             async with self.checkout() as session:
                 async with session.transaction(tx_mode=tx_mode) as tx:
+                    if tx_mode.name in ["serializable_read_write", "snapshot_read_only"]:
+                        await tx.begin()
                     result = await callee(tx, *args, **kwargs)
                     await tx.commit()
                 return result

ydb/aio/query/transaction.py

@@ -57,6 +57,7 @@ async def begin(self, settings: Optional[BaseRequestSettings] = None) -> "QueryT
         await self._begin_call(settings)
         return self
 
+    @base.with_async_transaction_events
     async def commit(self, settings: Optional[BaseRequestSettings] = None) -> None:
         """Calls commit on a transaction if it is open otherwise is no-op. If transaction execution
         failed then this method raises PreconditionFailed.
@@ -65,7 +66,7 @@ async def commit(self, settings: Optional[BaseRequestSettings] = None) -> None:
 
         :return: A committed transaction or exception if commit is failed
         """
-        if self._tx_state._already_in(QueryTxStateEnum.COMMITTED):
+        if self._tx_state._should_skip(QueryTxStateEnum.COMMITTED):
             return
 
         if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:
@@ -76,6 +77,7 @@ async def commit(self, settings: Optional[BaseRequestSettings] = None) -> None:
 
         await self._commit_call(settings)
 
+    @base.with_async_transaction_events
     async def rollback(self, settings: Optional[BaseRequestSettings] = None) -> None:
         """Calls rollback on a transaction if it is open otherwise is no-op. If transaction execution
         failed then this method raises PreconditionFailed.
@@ -84,7 +86,7 @@ async def rollback(self, settings: Optional[BaseRequestSettings] = None) -> None
 
         :return: A committed transaction or exception if commit is failed
         """
-        if self._tx_state._already_in(QueryTxStateEnum.ROLLBACKED):
+        if self._tx_state._should_skip(QueryTxStateEnum.ROLLBACKED):
             return
 
         if self._tx_state._state == QueryTxStateEnum.NOT_INITIALIZED:

ydb/query/base.py

@@ -1,4 +1,5 @@
 import abc
+import asyncio
 import enum
 import functools
 
@@ -196,3 +197,116 @@ def wrap_execute_query_response(
         return convert.ResultSet.from_message(response_pb.result_set, settings)
 
     return None
+
+
+class TxListener:
+    def _on_before_commit(self):
+        pass
+
+    def _on_after_commit(self, exc: typing.Optional[BaseException]):
+        pass
+
+    def _on_before_rollback(self):
+        pass
+
+    def _on_after_rollback(self, exc: typing.Optional[BaseException]):
+        pass
+
+
+class TxListenerAsyncIO:
+    async def _on_before_commit(self):
+        pass
+
+    async def _on_after_commit(self, exc: typing.Optional[BaseException]):
+        pass
+
+    async def _on_before_rollback(self):
+        pass
+
+    async def _on_after_rollback(self, exc: typing.Optional[BaseException]):
+        pass
+
+
+def with_transaction_events(method):
+    @functools.wraps(method)
+    def wrapper(self, *args, **kwargs):
+        method_name = method.__name__
+        before_event = f"_on_before_{method_name}"
+        after_event = f"_on_after_{method_name}"
+
+        self._notify_listeners_sync(before_event)
+
+        try:
+            result = method(self, *args, **kwargs)
+
+            self._notify_listeners_sync(after_event, exc=None)
+
+            return result
+        except BaseException as e:
+            self._notify_listeners_sync(after_event, exc=e)
+            raise
+
+    return wrapper
+
+
+def with_async_transaction_events(method):
+    @functools.wraps(method)
+    async def wrapper(self, *args, **kwargs):
+        method_name = method.__name__
+        before_event = f"_on_before_{method_name}"
+        after_event = f"_on_after_{method_name}"
+
+        await self._notify_listeners_async(before_event)
+
+        try:
+            result = await method(self, *args, **kwargs)
+
+            await self._notify_listeners_async(after_event, exc=None)
+
+            return result
+        except BaseException as e:
+            await self._notify_listeners_async(after_event, exc=e)
+            raise
+
+    return wrapper
+
+
+class ListenerHandlerMixin:
+    def _init_listener_handler(self):
+        self.listeners = []
+
+    def _add_listener(self, listener):
+        if listener not in self.listeners:
+            self.listeners.append(listener)
+        return self
+
+    def _remove_listener(self, listener):
+        if listener in self.listeners:
+            self.listeners.remove(listener)
+        return self
+
+    def _clear_listeners(self):
+        self.listeners.clear()
+        return self
+
+    def _notify_sync_listeners(self, event_name: str, **kwargs) -> None:
+        for listener in self.listeners:
+            if isinstance(listener, TxListener) and hasattr(listener, event_name):
+                getattr(listener, event_name)(**kwargs)
+
+    async def _notify_async_listeners(self, event_name: str, **kwargs) -> None:
+        coros = []
+        for listener in self.listeners:
+            if isinstance(listener, TxListenerAsyncIO) and hasattr(listener, event_name):
+                coros.append(getattr(listener, event_name)(**kwargs))
+
+        if coros:
+            await asyncio.gather(*coros)
+
+    def _notify_listeners_sync(self, event_name: str, **kwargs) -> None:
+        self._notify_sync_listeners(event_name, **kwargs)
+
+    async def _notify_listeners_async(self, event_name: str, **kwargs) -> None:
+        # self._notify_sync_listeners(event_name, **kwargs)
+
+        await self._notify_async_listeners(event_name, **kwargs)

ydb/query/pool.py

@@ -167,6 +167,8 @@ def retry_tx_sync(
         def wrapped_callee():
             with self.checkout(timeout=retry_settings.max_session_acquire_timeout) as session:
                 with session.transaction(tx_mode=tx_mode) as tx:
+                    if tx_mode.name in ["serializable_read_write", "snapshot_read_only"]:
+                        tx.begin()
                     result = callee(tx, *args, **kwargs)
                     tx.commit()
                 return result