Fix review comments

Author: vgvoleg

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):
+                    async with await tx.execute(query=f"select {i} as res;") as result_stream:
+                        async for result_set in result_stream:
+                            message = str(result_set.rows[0]["res"])
+                            await tx_writer.write(ydb.TopicWriterMessage(message))
+                            print(f"Message {result_set.rows[0]['res']} was 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"Message {batch.messages[0].data.decode()} was 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())

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}")
-            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])
+        def callee(tx: ydb.QueryTxContext):
+            tx_writer: ydb.TopicTxWriter = driver.topic_client.tx_writer(tx, topic)
 
-                print(f"Messages {messages} were written with tx.")
+            for i in range(message_count):
+                result_stream = tx.execute(query=f"select {i} as res;")
+                for result_set in result_stream:
+                    message = str(result_set.rows[0]["res"])
+                    tx_writer.write(ydb.TopicWriterMessage(message))
+                    print(f"Message {message} was 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"Message {batch.messages[0].data.decode()} was 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()

tests/topics/test_topic_transactions.py

@@ -357,6 +357,19 @@ async def callee(tx: ydb.aio.QueryTxContext):
             with pytest.raises(asyncio.TimeoutError):
                 await wait_for(topic_reader.receive_message(), 0.1)
 
+    async def test_writes_do_not_conflict_with_executes(
+        self, driver: ydb.aio.Driver, topic_path
+    ):
+        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)
+                for _ in range(3):
+                    async with await tx.execute("select 1"):
+                        await tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
+
+            await pool.retry_tx_async(callee, retry_settings=DEFAULT_RETRY_SETTINGS)
+
 
 class TestTopicTransactionalWriterSync:
     def test_commit(self, driver_sync: ydb.Driver, topic_path, topic_reader_sync: ydb.TopicReader):
@@ -445,3 +458,16 @@ def callee(tx: ydb.QueryTxContext):
 
             with pytest.raises(TimeoutError):
                 topic_reader_sync.receive_message(timeout=DEFAULT_TIMEOUT)
+
+    def test_writes_do_not_conflict_with_executes(
+        self, driver_sync: ydb.Driver, topic_path
+    ):
+        with ydb.QuerySessionPool(driver_sync) as pool:
+
+            def callee(tx: ydb.QueryTxContext):
+                tx_writer = driver_sync.topic_client.tx_writer(tx, topic_path)
+                for _ in range(3):
+                    with tx.execute("select 1"):
+                        tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
+
+            pool.retry_tx_sync(callee, retry_settings=DEFAULT_RETRY_SETTINGS)

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))

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
@@ -186,6 +185,10 @@ def __init__(
         self._parent = _client
         self._is_implicit = _is_implicit
 
+        # For some reason, creating partition could conflict with other session operations.
+        # Could be removed later.
+        self._first_write = True
+
         tx._add_callback(TxEvent.BEFORE_COMMIT, self._on_before_commit, self._loop)
         tx._add_callback(TxEvent.BEFORE_ROLLBACK, self._on_before_rollback, self._loop)
 
@@ -199,14 +202,18 @@ async def write(
 
         For wait with timeout use asyncio.wait_for.
         """
-        await self.write_with_ack(messages)
+        if self._first_write:
+            self._first_write = False
+            return await super().write_with_ack(messages)
+        return await super().write(messages)
+
 
     async def _on_before_commit(self, tx: "BaseQueryTxContext"):
         if self._is_implicit:
             return
-        await self.flush()
         await self.close()
 
+
     async def _on_before_rollback(self, tx: "BaseQueryTxContext"):
         if self._is_implicit:
             return