Topic transactions feature (#559)

Author: vgvoleg

.github/workflows/tests.yaml

@@ -18,21 +18,15 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [3.8, 3.9]
-        environment: [py-proto5, py-tls-proto5, py-proto4, py-tls-proto4, py-proto3, py-tls-proto3]
-        folder: [ydb, tests --ignore=tests/topics, tests/topics]
+        environment: [py, py-tls, py-proto4, py-tls-proto4, py-proto3, py-tls-proto3]
+        folder: [ydb, tests]
         exclude:
-          - environment: py-tls-proto5
+          - environment: py-tls
             folder: ydb
           - environment: py-tls-proto4
             folder: ydb
           - environment: py-tls-proto3
             folder: ydb
-          - environment: py-tls-proto5
-            folder: tests/topics
-          - environment: py-tls-proto4
-            folder: tests/topics
-          - environment: py-tls-proto3
-            folder: tests/topics
 
     steps:
     - uses: actions/checkout@v1

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

@@ -0,0 +1,85 @@
+import argparse
+import logging
+import ydb
+
+
+def connect(endpoint: str, database: str) -> ydb.Driver:
+    config = ydb.DriverConfig(endpoint=endpoint, database=database)
+    config.credentials = ydb.credentials_from_env_variables()
+    driver = ydb.Driver(config)
+    driver.wait(5, fail_fast=True)
+    return driver
+
+
+def create_topic(driver: ydb.Driver, topic: str, consumer: str):
+    try:
+        driver.topic_client.drop_topic(topic)
+    except ydb.SchemeError:
+        pass
+
+    driver.topic_client.create_topic(topic, consumers=[consumer])
+
+
+def write_with_tx_example(driver: ydb.Driver, topic: str, message_count: int = 10):
+    with ydb.QuerySessionPool(driver) as session_pool:
+
+        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;")
+                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.")
+
+        session_pool.retry_tx_sync(callee)
+
+
+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):
+
+                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.")
+
+                session_pool.retry_tx_sync(callee)
+
+
+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())
+
+    with connect(args.endpoint, args.database) as driver:
+        if not args.skip_drop_and_create_topic:
+            create_topic(driver, args.path, args.consumer)
+
+        write_with_tx_example(driver, args.path)
+        read_with_tx_example(driver, args.path, args.consumer)
+
+
+if __name__ == "__main__":
+    main()

tests/query/test_query_transaction.py

@@ -92,3 +92,15 @@ def test_execute_two_results(self, tx: QueryTxContext):
 
         assert res == [[1], [2]]
         assert counter == 2
+
+    def test_tx_identity_before_begin_raises(self, tx: QueryTxContext):
+        with pytest.raises(RuntimeError):
+            tx._tx_identity()
+
+    def test_tx_identity_after_begin_works(self, tx: QueryTxContext):
+        tx.begin()
+
+        identity = tx._tx_identity()
+
+        assert identity.tx_id == tx.tx_id
+        assert identity.session_id == tx.session_id

tests/topics/test_topic_reader.py

@@ -174,12 +174,13 @@ def test_read_and_commit_with_close_reader(self, driver_sync, topic_with_message
             assert message != message2
 
     def test_read_and_commit_with_ack(self, driver_sync, topic_with_messages, topic_consumer):
-        reader = driver_sync.topic_client.reader(topic_with_messages, topic_consumer)
-        message = reader.receive_message()
-        reader.commit_with_ack(message)
+        with driver_sync.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            message = reader.receive_message()
+            reader.commit_with_ack(message)
+
+        with driver_sync.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            batch = reader.receive_batch()
 
-        reader = driver_sync.topic_client.reader(topic_with_messages, topic_consumer)
-        batch = reader.receive_batch()
         assert message != batch.messages[0]
 
     def test_read_compressed_messages(self, driver_sync, topic_path, topic_consumer):
@@ -247,3 +248,6 @@ async def wait(fut):
         datas.sort()
 
         assert datas == ["10", "11"]
+
+        await reader0.close()
+        await reader1.close()