please help

Author: vgvoleg

ydb/_topic_reader/datatypes.py

@@ -163,6 +163,7 @@ class PublicBatch(ICommittable, ISessionAlive):
     _partition_session: PartitionSession
     _bytes_size: int
     _codec: Codec
+    _commited_with_tx: bool = False
 
     def _commit_get_partition_session(self) -> PartitionSession:
         return self.messages[0]._commit_get_partition_session()

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 OrderedDict, defaultdict
 from typing import Optional, Set, Dict, Union, Callable
 
 import ydb
@@ -140,7 +140,8 @@ async def receive_batch_with_tx(
         use asyncio.wait_for for wait with timeout.
         """
         await self._reconnector.wait_message()
-        return await self._reconnector.receive_batch_with_tx_nowait(
+        tx._add_listener(self)
+        return self._reconnector.receive_batch_with_tx_nowait(
             tx,
             max_messages=max_messages,
         )
@@ -177,11 +178,14 @@ async def close(self, flush: bool = True):
         self._closed = True
         await self._reconnector.close(flush)
 
-    def _on_after_commit(self, exc):
-        return super()._on_after_commit(exc)
+    async def _on_before_commit(self, tx):
+        await self._reconnector._on_before_commit(tx)
 
-    def _on_after_rollback(self, exc):
-        return super()._on_after_rollback(exc)
+    async def _on_after_commit(self, tx, exc):
+        await self._reconnector._on_after_commit(tx, exc)
+
+    async def _on_after_rollback(self, tx, exc):
+        await self._reconnector._on_after_rollback(tx, exc)
 
 
 class ReaderReconnector:
@@ -195,8 +199,10 @@ class ReaderReconnector:
     _state_changed: asyncio.Event
     _stream_reader: Optional["ReaderStream"]
     _first_error: asyncio.Future[YdbError]
-    _batches_to_commit: asyncio.Queue
-    _wait_executor: Optional[concurrent.futures.ThreadPoolExecutor]
+
+    _batches_to_commit_with_tx: asyncio.Queue
+    _tx_to_batches: typing.Dict[str, typing.List[datatypes.PublicBatch]]
+    _wait_executor: concurrent.futures.ThreadPoolExecutor
 
     def __init__(self, driver: Driver, settings: topic_reader.PublicReaderSettings):
         self._id = self._static_reader_reconnector_counter.inc_and_get()
@@ -210,7 +216,9 @@ def __init__(self, driver: Driver, settings: topic_reader.PublicReaderSettings):
         self._first_error = asyncio.get_running_loop().create_future()
         self._wait_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
 
-        self._batches_to_commit = asyncio.Queue()
+        self._batches_to_commit_with_tx = asyncio.Queue()
+        self._tx_to_batches = defaultdict(list)
+        self._background_tasks.add(asyncio.create_task(self._update_offsets_in_tx_loop()))
 
     async def _connection_loop(self):
         attempt = 0
@@ -263,19 +271,21 @@ def receive_message_nowait(self):
     def commit(self, batch: datatypes.ICommittable) -> datatypes.PartitionSession.CommitAckWaiter:
         return self._stream_reader.commit(batch)
 
-    async def _commit_with_tx(self, tx: "BaseQueryTxContext", batch: datatypes.ICommittable) -> None:
-        pass
-
-    async def receive_batch_with_tx_nowait(self, tx: "BaseQueryTxContext", max_messages: Optional[int] = None):
+    def receive_batch_with_tx_nowait(self, tx: "BaseQueryTxContext", max_messages: Optional[int] = None):
         batch = self.receive_batch_nowait(max_messages=max_messages)
-        tx._add_listener(batch)
-        await self._update_offsets_in_tx_call(self._driver, tx, batch)
+        self._tx_to_batches[tx.tx_id].append(batch)
+
+        self._batches_to_commit_with_tx.put_nowait((tx, batch))
+
+        print("batch recieved")
+
         return batch
-        # self._batches_to_commit.put_nowait((tx, batch))
 
     async def _update_offsets_in_tx_loop(self):
         while True:
-            tx, batch = self._batches_to_commit.get()
+            print("_update_offsets_in_tx_loop")
+
+            tx, batch = await self._batches_to_commit_with_tx.get()
             await self._update_offsets_in_tx_call(self._driver, tx, batch)
 
     async def _update_offsets_in_tx_call(self, driver: SupportedDriverType, tx: "BaseQueryTxContext", batch: datatypes.ICommittable) -> None:
@@ -309,9 +319,25 @@ async def _update_offsets_in_tx_call(self, driver: SupportedDriverType, tx: "Bas
             else:
                 res = await to_thread(driver, *args, executor=self._wait_executor)
 
+            batch._commited_with_tx = True
+
             return res
         except BaseException as e:
-            self._set_first_error(e)
+            self._stream_reader._set_first_error(e)
+
+    async def _ensure_all_batches_commited_with_tx(self, tx: "BaseQueryTxContext"):
+        while True:
+            print("_ensure_all_batches_commited_with_tx")
+            if tx.tx_id not in self._tx_to_batches:
+                # we should not be here
+                return True
+            batches = self._tx_to_batches.get(tx.tx_id)
+            everything_commited = True
+            for batch in batches:
+                everything_commited = everything_commited and batch._commited_with_tx
+            if everything_commited:
+                return True
+            await asyncio.sleep(0.001)
 
     async def close(self, flush: bool):
         if self._stream_reader:
@@ -336,6 +362,26 @@ def _set_first_error(self, err: issues.Error):
             # skip if already has result
             pass
 
+    async def _on_before_commit(self, tx):
+        print("on before commit")
+        await asyncio.wait_for(self._ensure_all_batches_commited_with_tx(tx), 1)
+        pass
+
+    async def _on_after_commit(self, tx, exc):
+        print(f"on after commit, exc = {exc is not None}")
+
+        if exc:
+            self._stream_reader._set_first_error(exc)
+        for batch in self._tx_to_batches[tx.tx_id]:
+            batch._partition_session.committed_offset = max(batch._partition_session.committed_offset, batch._commit_get_offsets_range().end)
+        del self._tx_to_batches[tx.tx_id]
+
+    async def _on_after_rollback(self, tx, exc):
+        print(f"on after rollback, exc = {exc is not None}")
+        print(exc)
+        exc = exc if exc is not None else issues.InternalError("tx rollback failed")
+        self._stream_reader._set_first_error(exc)
+
 
 class ReaderStream:
     _static_id_counter = AtomicCounter()

ydb/_topic_reader/topic_reader_sync.py

@@ -20,8 +20,13 @@
     TopicReaderClosedError,
 )
 
+from ..query.base import TxListener
 
-class TopicReaderSync:
+if typing.TYPE_CHECKING:
+    from ..query.transaction import BaseQueryTxContext
+
+
+class TopicReaderSync(TxListener):
     _caller: CallFromSyncToAsync
     _async_reader: PublicAsyncIOReader
     _closed: bool
@@ -155,3 +160,23 @@ def close(self, *, flush: bool = True, timeout: TimeoutType = None):
     def _check_closed(self):
         if self._closed:
             raise TopicReaderClosedError()
+
+    def _on_before_commit(self, tx: "BaseQueryTxContext"):
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_result(self._async_reader._on_before_commit(tx), 5)
+
+    def _on_after_commit(self, tx: "BaseQueryTxContext", exc):
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_result(self._async_reader._on_after_commit(tx, exc), 5)
+
+    def _on_before_rollback(self, tx: "BaseQueryTxContext"):
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_result(self._async_reader._on_before_rollback(tx), 5)
+
+    def _on_after_rollback(self, tx: "BaseQueryTxContext", exc):
+        self._check_closed()
+
+        return self._caller.unsafe_call_with_result(self._async_reader._on_after_rollback(tx, exc), 5)

ydb/_topic_writer/topic_writer_asyncio.py

@@ -187,10 +187,10 @@ def __init__(
         self._parent = _client
         self._tx._add_listener(self)
 
-    async def _on_before_commit(self):
+    async def _on_before_commit(self, **kwargs):
         await self.close()
 
-    async def _on_before_rollback(self):
+    async def _on_before_rollback(self, **kwargs):
         await self.close(flush=False)
 
 

ydb/_topic_writer/topic_writer_sync.py

@@ -160,8 +160,8 @@ async def create_async_writer():
 
         tx._add_listener(self)
 
-    def _on_before_commit(self):
+    def _on_before_commit(self, **kwargs):
         self.close()
 
-    def _on_before_rollback(self):
+    def _on_before_rollback(self, **kwargs):
         self.close(flush=False)

ydb/aio/query/transaction.py

@@ -57,7 +57,7 @@ async def begin(self, settings: Optional[BaseRequestSettings] = None) -> "QueryT
         await self._begin_call(settings)
         return self
 
-    @base.with_async_transaction_events
+    @base.with_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.
@@ -77,7 +77,7 @@ async def commit(self, settings: Optional[BaseRequestSettings] = None) -> None:
 
         await self._commit_call(settings)
 
-    @base.with_async_transaction_events
+    @base.with_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.

ydb/query/base.py

@@ -200,74 +200,76 @@ def wrap_execute_query_response(
 
 
 class TxListener:
-    def _on_before_commit(self):
+    def _on_before_commit(self, tx: "BaseQueryTxContext"):
         pass
 
-    def _on_after_commit(self, exc: typing.Optional[BaseException]):
+    def _on_after_commit(self, tx: "BaseQueryTxContext", exc: typing.Optional[BaseException]):
         pass
 
-    def _on_before_rollback(self):
+    def _on_before_rollback(self, tx: "BaseQueryTxContext"):
         pass
 
-    def _on_after_rollback(self, exc: typing.Optional[BaseException]):
+    def _on_after_rollback(self, tx: "BaseQueryTxContext", exc: typing.Optional[BaseException]):
         pass
 
 
 class TxListenerAsyncIO:
-    async def _on_before_commit(self):
+    async def _on_before_commit(self, tx: "BaseQueryTxContext"):
         pass
 
-    async def _on_after_commit(self, exc: typing.Optional[BaseException]):
+    async def _on_after_commit(self, tx: "BaseQueryTxContext", exc: typing.Optional[BaseException]):
         pass
 
-    async def _on_before_rollback(self):
+    async def _on_before_rollback(self, tx: "BaseQueryTxContext"):
         pass
 
-    async def _on_after_rollback(self, exc: typing.Optional[BaseException]):
+    async def _on_after_rollback(self, tx: "BaseQueryTxContext", exc: typing.Optional[BaseException]):
         pass
 
 
 def with_transaction_events(method):
     @functools.wraps(method)
-    def wrapper(self, *args, **kwargs):
+    def wrapper(self: ListenerHandlerMixin, *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)
+        self._notify_listeners_sync(before_event, tx=self)
 
         try:
             result = method(self, *args, **kwargs)
 
-            self._notify_listeners_sync(after_event, exc=None)
+            self._notify_listeners_sync(after_event, tx=self, exc=None)
 
             return result
         except BaseException as e:
-            self._notify_listeners_sync(after_event, exc=e)
+            self._notify_listeners_sync(after_event, tx=self, exc=e)
             raise
+        finally:
+            self._clear_listeners()
 
-    return wrapper
-
-
-def with_async_transaction_events(method):
     @functools.wraps(method)
-    async def wrapper(self, *args, **kwargs):
+    async def async_wrapper(self: ListenerHandlerMixin, *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)
+        await self._notify_listeners_async(before_event, tx=self)
 
         try:
             result = await method(self, *args, **kwargs)
 
-            await self._notify_listeners_async(after_event, exc=None)
+            await self._notify_listeners_async(after_event, tx=self, exc=None)
 
             return result
         except BaseException as e:
-            await self._notify_listeners_async(after_event, exc=e)
+            await self._notify_listeners_async(after_event, tx=self, exc=e)
             raise
+        finally:
+            self._clear_listeners()
 
+    if asyncio.iscoroutinefunction(method):
+        return async_wrapper
     return wrapper
 
 
@@ -289,24 +291,16 @@ def _clear_listeners(self):
         self.listeners.clear()
         return self
 
-    def _notify_sync_listeners(self, event_name: str, **kwargs) -> None:
+    def _notify_listeners_sync(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:
+    async def _notify_listeners_async(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)
+            await asyncio.gather(*coros)