bug

Author: vgvoleg

examples/topic/topic_transactions_async_example.py

@@ -3,6 +3,8 @@
 import logging
 import ydb
 
+# logging.basicConfig(level=logging.DEBUG)
+
 
 async def connect(endpoint: str, database: str) -> ydb.aio.Driver:
     config = ydb.DriverConfig(endpoint=endpoint, database=database)
@@ -26,13 +28,19 @@ async def write_with_tx_example(driver: ydb.aio.Driver, topic: str, message_coun
 
         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.")
+            for i in range(message_count):
+                try:
+                    print("try to execute")
+                    async with await tx.execute(query=f"select {i} as res") as result_stream:
+                        async for result_set in result_stream:
+                            print("try to write msg")
+                            await tx_writer.write(ydb.TopicWriterMessage(data=str(result_set.rows[0]["res"])))
+                            print(f"Messages {result_set.rows[0]['res']} were written with tx.")
+                except BaseException as e:
+                    print("Caught unexpected exception")
+                    print(e)
+                    raise e
 
         await session_pool.retry_tx_async(callee)
 
@@ -70,7 +78,7 @@ async def main():
     args = parser.parse_args()
 
     if args.verbose:
-        logger = logging.getLogger("topicexample")
+        logger = logging.getLogger(__name__)
         logger.setLevel(logging.DEBUG)
         logger.addHandler(logging.StreamHandler())
 
@@ -79,7 +87,7 @@ async def main():
             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)
+        # await read_with_tx_example(driver, args.path, args.consumer)
 
 
 if __name__ == "__main__":

examples/topic/topic_transactions_example.py

@@ -26,12 +26,10 @@ def write_with_tx_example(driver: ydb.Driver, topic: str, message_count: int = 1
         def callee(tx: ydb.QueryTxContext):
             tx_writer: ydb.TopicTxWriter = driver.topic_client.tx_writer(tx, topic)
             for i in range(message_count):
-                result_stream = tx.execute(query=f"select {i} as res")
-                messages = [result_set.rows[0]["res"] for result_set in result_stream]
-
-                tx_writer.write([ydb.TopicWriterMessage(data=str(message)) for message in messages])
-
-                print(f"Messages {messages} were written with tx.")
+                result_stream = tx.execute(query=f"select {i} as res;")
+                for result_set in result_stream:
+                    tx_writer.write(ydb.TopicWriterMessage(data=str(result_set.rows[0]["res"])))
+                    print(f"Messages {result_set.rows[0]['res']} were written with tx.")
 
         session_pool.retry_tx_sync(callee)
 

ydb/_topic_writer/topic_writer_asyncio.py

@@ -83,6 +83,7 @@ def __del__(self):
             logger.warning("Topic writer was not closed properly. Consider using method close().")
 
     async def close(self, *, flush: bool = True):
+        print("closing writer")
         if self._closed:
             return
 
@@ -188,24 +189,24 @@ def __init__(
         tx._add_callback(TxEvent.BEFORE_COMMIT, self._on_before_commit, self._loop)
         tx._add_callback(TxEvent.BEFORE_ROLLBACK, self._on_before_rollback, self._loop)
 
-    async def write(
-        self,
-        messages: Union[Message, List[Message]],
-    ):
-        """
-        send one or number of messages to server.
-        it put message to internal buffer
+    # async def write(
+    #     self,
+    #     messages: Union[Message, List[Message]],
+    # ):
+    #     """
+    #     send one or number of messages to server.
+    #     it put message to internal buffer
 
-        For wait with timeout use asyncio.wait_for.
-        """
-        await self.write_with_ack(messages)
+    #     For wait with timeout use asyncio.wait_for.
+    #     """
+    #     await self.write_with_ack(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

ydb/aio/query/base.py

@@ -9,3 +9,4 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         #  To close stream on YDB it is necessary to scroll through it to the end
         async for _ in self:
             pass
+        print("exit from execute context iterator")