fixes in ydb.aio library

Author: gridnevvvit

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 2.7.0 ##
+
+* fixes in TxContext in ydb.aio module.
+
 ## 2.6.0 ##
 
 * support `with_native_interval_in_result_sets` flag in the table client.

setup.py

@@ -6,7 +6,7 @@
 
 setuptools.setup(
     name="ydb",
-    version="2.6.0",
+    version="2.7.0",
     description="YDB Python SDK",
     author="Yandex LLC",
     author_email="[email protected]",

tests/aio/test_tx.py

@@ -0,0 +1,38 @@
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_tx_commit(driver, database):
+    session = await driver.table_client.session().create()
+    prepared = await session.prepare(
+        "DECLARE $param as Int32;\n SELECT $param as value",
+    )
+
+    tx = session.transaction()
+    await tx.execute(prepared, {"$param": 2})
+    await tx.commit()
+    await tx.commit()
+
+
+@pytest.mark.asyncio
+async def test_tx_rollback(driver, database):
+    session = await driver.table_client.session().create()
+    prepared = await session.prepare(
+        "DECLARE $param as Int32;\n SELECT $param as value",
+    )
+
+    tx = session.transaction()
+    await tx.execute(prepared, {"$param": 2})
+    await tx.rollback()
+    await tx.rollback()
+
+
+@pytest.mark.asyncio
+async def test_tx_begin(driver, database):
+    session = await driver.table_client.session().create()
+    await session.create()
+
+    tx = session.transaction()
+    await tx.begin()
+    await tx.begin()
+    await tx.rollback()

ydb/aio/table.py

@@ -54,7 +54,10 @@ async def keep_alive(self, settings=None):  # pylint: disable=W0236
         return await super().keep_alive(settings)
 
     async def create(self, settings=None):  # pylint: disable=W0236
-        return await super().create(settings)
+        res = super().create(settings)
+        if asyncio.iscoroutine(res):
+            res = await res
+        return res
 
     async def delete(self, settings=None):  # pylint: disable=W0236
         return await super().delete(settings)
@@ -184,13 +187,22 @@ async def execute(
         return await super().execute(query, parameters, commit_tx, settings)
 
     async def commit(self, settings=None):  # pylint: disable=W0236
-        return await super().commit(settings)
+        res = super().commit(settings)
+        if asyncio.iscoroutine(res):
+            res = await res
+        return res
 
     async def rollback(self, settings=None):  # pylint: disable=W0236
-        return await super().rollback(settings)
+        res = super().rollback(settings)
+        if asyncio.iscoroutine(res):
+            res = await res
+        return res
 
     async def begin(self, settings=None):  # pylint: disable=W0236
-        return await super().begin(settings)
+        res = super().begin(settings)
+        if asyncio.iscoroutine(res):
+            res = await res
+        return res
 
 
 async def retry_operation(

ydb/ydb_version.py

@@ -1 +1 @@
-VERSION = "2.6.0"
+VERSION = "2.7.0"