Fix review comments

vgvoleg · vgvoleg · commit 170827270f0b · 2025-03-25T14:49:54.000+03:00
diff --git a/examples/topic/topic_transactions_async_example.py b/examples/topic/topic_transactions_async_example.py
@@ -0,0 +1,86 @@
+import asyncio
+import argparse
+import logging
+import ydb
+
+
+async def connect(endpoint: str, database: str) -> ydb.aio.Driver:
+    config = ydb.DriverConfig(endpoint=endpoint, database=database)
+    config.credentials = ydb.credentials_from_env_variables()
+    driver = ydb.aio.Driver(config)
+    await driver.wait(5, fail_fast=True)
+    return driver
+
+
+async def create_topic(driver: ydb.aio.Driver, topic: str, consumer: str):
+    try:
+        await driver.topic_client.drop_topic(topic)
+    except ydb.SchemeError:
+        pass
+
+    await driver.topic_client.create_topic(topic, consumers=[consumer])
+
+
+async def write_with_tx_example(driver: ydb.aio.Driver, topic: str, message_count: int = 10):
+    async with ydb.aio.QuerySessionPool(driver) as session_pool:
+
+        async def callee(tx: ydb.aio.QueryTxContext):
+            tx_writer: ydb.TopicTxWriterAsyncIO = driver.topic_client.tx_writer(tx, topic)
+            for i in range(message_count):
+                result_stream = await tx.execute(query=f"select {i} as res")
+                messages = [result_set.rows[0]["res"] async for result_set in result_stream]
+
+                await tx_writer.write([ydb.TopicWriterMessage(data=str(message)) for message in messages])
+
+                print(f"Messages {messages} were written with tx.")
+
+        await session_pool.retry_tx_async(callee)
+
+
+async def read_with_tx_example(driver: ydb.aio.Driver, topic: str, consumer: str, message_count: int = 10):
+    async with driver.topic_client.reader(topic, consumer) as reader:
+        async with ydb.aio.QuerySessionPool(driver) as session_pool:
+            for _ in range(message_count):
+
+                async def callee(tx: ydb.aio.QueryTxContext):
+                    batch = await reader.receive_batch_with_tx(tx, max_messages=1)
+                    print(f"Messages [{batch.messages[0].data.decode()}] were read with tx.")
+
+                await session_pool.retry_tx_async(callee)
+
+
+async def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        description="""YDB topic basic example.\n""",
+    )
+    parser.add_argument("-d", "--database", default="/local", help="Name of the database to use")
+    parser.add_argument("-e", "--endpoint", default="grpc://localhost:2136", help="Endpoint url to use")
+    parser.add_argument("-p", "--path", default="test-topic", help="Topic name")
+    parser.add_argument("-c", "--consumer", default="consumer", help="Consumer name")
+    parser.add_argument("-v", "--verbose", default=False, action="store_true")
+    parser.add_argument(
+        "-s",
+        "--skip-drop-and-create-topic",
+        default=False,
+        action="store_true",
+        help="Use existed topic, skip remove it and re-create",
+    )
+
+    args = parser.parse_args()
+
+    if args.verbose:
+        logger = logging.getLogger("topicexample")
+        logger.setLevel(logging.DEBUG)
+        logger.addHandler(logging.StreamHandler())
+
+    async with await connect(args.endpoint, args.database) as driver:
+        if not args.skip_drop_and_create_topic:
+            await create_topic(driver, args.path, args.consumer)
+
+        await write_with_tx_example(driver, args.path)
+        await read_with_tx_example(driver, args.path, args.consumer)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
diff --git a/examples/topic/topic_transactions_example.py b/examples/topic/topic_transactions_example.py
@@ -1,59 +1,54 @@
-import asyncio
 import argparse
 import logging
 import ydb
 
 
-async def connect(endpoint: str, database: str) -> ydb.aio.Driver:
+def connect(endpoint: str, database: str) -> ydb.Driver:
     config = ydb.DriverConfig(endpoint=endpoint, database=database)
     config.credentials = ydb.credentials_from_env_variables()
-    driver = ydb.aio.Driver(config)
-    await driver.wait(5, fail_fast=True)
+    driver = ydb.Driver(config)
+    driver.wait(5, fail_fast=True)
     return driver
 
 
-async def create_topic(driver: ydb.aio.Driver, topic: str, consumer: str):
+def create_topic(driver: ydb.Driver, topic: str, consumer: str):
     try:
-        await driver.topic_client.drop_topic(topic)
+        driver.topic_client.drop_topic(topic)
     except ydb.SchemeError:
         pass
 
-    await driver.topic_client.create_topic(topic, consumers=[consumer])
+    driver.topic_client.create_topic(topic, consumers=[consumer])
 
 
-async def write_with_tx_example(driver: ydb.aio.Driver, topic: str, message_count: int = 10):
-    async with ydb.aio.QuerySessionPool(driver) as session_pool:
+def write_with_tx_example(driver: ydb.Driver, topic: str, message_count: int = 10):
+    with ydb.QuerySessionPool(driver) as session_pool:
 
-        async def callee(tx: ydb.aio.QueryTxContext):
-            print(f"TX ID: {tx.tx_id}")
-            print(f"TX STATE: {tx._tx_state._state.value}")
-            tx_writer: ydb.TopicTxWriterAsyncIO = driver.topic_client.tx_writer(tx, topic)
-            print(f"TX ID: {tx.tx_id}")
-            print(f"TX STATE: {tx._tx_state._state.value}")
+        def callee(tx: ydb.QueryTxContext):
+            tx_writer: ydb.TopicTxWriter = driver.topic_client.tx_writer(tx, topic)
             for i in range(message_count):
-                result_stream = await tx.execute(query=f"select {i} as res")
-                messages = [result_set.rows[0]["res"] async for result_set in result_stream]
+                result_stream = tx.execute(query=f"select {i} as res")
+                messages = [result_set.rows[0]["res"] for result_set in result_stream]
 
-                await tx_writer.write([ydb.TopicWriterMessage(data=str(message)) for message in messages])
+                tx_writer.write([ydb.TopicWriterMessage(data=str(message)) for message in messages])
 
                 print(f"Messages {messages} were written with tx.")
 
-        await session_pool.retry_tx_async(callee)
+        session_pool.retry_tx_sync(callee)
 
 
-async def read_with_tx_example(driver: ydb.aio.Driver, topic: str, consumer: str, message_count: int = 10):
-    async with driver.topic_client.reader(topic, consumer) as reader:
-        async with ydb.aio.QuerySessionPool(driver) as session_pool:
+def read_with_tx_example(driver: ydb.Driver, topic: str, consumer: str, message_count: int = 10):
+    with driver.topic_client.reader(topic, consumer) as reader:
+        with ydb.QuerySessionPool(driver) as session_pool:
             for _ in range(message_count):
 
-                async def callee(tx: ydb.aio.QueryTxContext):
-                    batch = await reader.receive_batch_with_tx(tx, max_messages=1)
-                    print(f"Messages {batch.messages[0].data} were read with tx.")
+                def callee(tx: ydb.QueryTxContext):
+                    batch = reader.receive_batch_with_tx(tx, max_messages=1)
+                    print(f"Messages [{batch.messages[0].data.decode()}] were read with tx.")
 
-                await session_pool.retry_tx_async(callee)
+                session_pool.retry_tx_sync(callee)
 
 
-async def main():
+def main():
     parser = argparse.ArgumentParser(
         formatter_class=argparse.RawDescriptionHelpFormatter,
         description="""YDB topic basic example.\n""",
@@ -78,13 +73,13 @@ async def main():
         logger.setLevel(logging.DEBUG)
         logger.addHandler(logging.StreamHandler())
 
-    driver = await connect(args.endpoint, args.database)
-    if not args.skip_drop_and_create_topic:
-        await create_topic(driver, args.path, args.consumer)
+    with connect(args.endpoint, args.database) as driver:
+        if not args.skip_drop_and_create_topic:
+            create_topic(driver, args.path, args.consumer)
 
-    await write_with_tx_example(driver, args.path)
-    await read_with_tx_example(driver, args.path, args.consumer)
+        write_with_tx_example(driver, args.path)
+        read_with_tx_example(driver, args.path, args.consumer)
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    main()
diff --git a/ydb/_topic_reader/topic_reader_asyncio.py b/ydb/_topic_reader/topic_reader_asyncio.py
@@ -84,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):
@@ -190,18 +190,24 @@ class ReaderReconnector:
     _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 = defaultdict(list)
+        self._tx_to_batches_map = dict()
 
     async def _connection_loop(self):
         attempt = 0
@@ -254,22 +260,23 @@ def receive_batch_with_tx_nowait(self, tx: "BaseQueryTxContext", max_messages: O
             max_messages=max_messages,
         )
 
-        self._init_tx_if_needed(tx)
+        self._init_tx(tx)
 
         self._tx_to_batches_map[tx.tx_id].append(batch)
 
-        tx._add_callback(TxEvent.AFTER_COMMIT, batch._update_partition_offsets, None)  # probably should be current loop
+        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_if_needed(self, tx: "BaseQueryTxContext"):
+    def _init_tx(self, tx: "BaseQueryTxContext"):
         if tx.tx_id not in self._tx_to_batches_map:  # Init tx callbacks
-            tx._add_callback(TxEvent.BEFORE_COMMIT, self._commit_batches_with_tx, None)
-            tx._add_callback(TxEvent.AFTER_COMMIT, self._handle_after_tx_commit, None)
-            tx._add_callback(TxEvent.AFTER_ROLLBACK, self._handle_after_tx_rollback, None)
+            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))
diff --git a/ydb/_topic_writer/topic_writer_asyncio.py b/ydb/_topic_writer/topic_writer_asyncio.py
@@ -1,7 +1,6 @@
 import asyncio
 import concurrent.futures
 import datetime
-import functools
 import gzip
 import typing
 from collections import deque
