add executable example

Author: vgvoleg

examples/topic/topic_transactions_example.py

@@ -1,28 +1,90 @@
+import asyncio
+import argparse
+import logging
 import ydb
 
 
-def writer_example(driver: ydb.Driver, topic: str):
-    session_pool = ydb.QuerySessionPool(driver)
+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
 
-    def callee(tx: ydb.QueryTxContext):
-        tx_writer: ydb.TopicTxWriter = driver.topic_client.tx_writer(tx, topic)  # <=======
-        # дефолт - без дедупликации, без ретраев и без producer_id.
 
-        with tx.execute(query="select 1") as result_sets:
-            messages = [result_set.rows[0] for result_set in result_sets]
+async def create_topic(driver: ydb.aio.Driver, topic: str, consumer: str):
+    try:
+        await driver.topic_client.drop_topic(topic)
+    except ydb.SchemeError:
+        pass
 
-        tx_writer.write(messages)  # вне зависимости от состояния вышестоящего стрима поведение должно быть одинаковое
+    await driver.topic_client.create_topic(topic, consumers=[consumer])
 
-    session_pool.retry_tx_sync(callee)
 
+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 reader_example(driver: ydb.Driver, reader: ydb.TopicReader):
-    session_pool = ydb.QuerySessionPool(driver)
+        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]
 
-    def callee(tx: ydb.QueryTxContext):
-        batch = reader.receive_batch_with_tx(tx, max_messages=5)  # <=======
+                await tx_writer.write([ydb.TopicWriterMessage(data=str(message)) for message in messages])
 
-        with tx.execute(query="INSERT INTO max_values(val) VALUES ($val)", parameters={"$val": max(batch)}) as _:
-            pass
+                print(f"Messages {messages} were written with tx.")
 
-    session_pool.retry_tx_sync(callee)
+        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} 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())
+
+    driver = await connect(args.endpoint, args.database)
+    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())

ydb/_topic_writer/topic_writer_asyncio.py

@@ -191,9 +191,22 @@ 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
+
+        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"):

ydb/topic.py

@@ -25,6 +25,8 @@
     "TopicWriteResult",
     "TopicWriter",
     "TopicWriterAsyncIO",
+    "TopicTxWriter",
+    "TopicTxWriterAsyncIO",
     "TopicWriterInitInfo",
     "TopicWriterMessage",
     "TopicWriterSettings",
@@ -68,7 +70,7 @@
 from ydb._topic_writer.topic_writer_asyncio import TxWriterAsyncIO as TopicTxWriterAsyncIO
 from ydb._topic_writer.topic_writer_asyncio import WriterAsyncIO as TopicWriterAsyncIO
 from ._topic_writer.topic_writer_sync import WriterSync as TopicWriter
-from ._topic_writer.topic_writer_sync import TxWriterSync as TxTopicWriter
+from ._topic_writer.topic_writer_sync import TxWriterSync as TopicTxWriter
 
 from ._topic_common.common import (
     wrap_operation as _wrap_operation,
@@ -546,7 +548,7 @@ def tx_writer(
         if not settings.encoder_executor:
             settings.encoder_executor = self._executor
 
-        return TxTopicWriter(tx, self._driver, settings, _parent=self)
+        return TopicTxWriter(tx, self._driver, settings, _parent=self)
 
     def close(self):
         if self._closed: