Merge pull request #223 fixed use transaction object when commit with flag

Author: rekby

Diff for: CHANGELOG.md

@@ -1,3 +1,5 @@
+* fixed use transaction object when commit with flag
+
 ## 2.13.2 ##
 * fix snapshot attribute in class _ResultSet
 

Diff for: tests/aio/test_tx.py

@@ -97,7 +97,7 @@ async def test_tx_snapshot_ro(driver, database):
 
 
 @pytest.mark.asyncio
-async def test_split_transactions_deny_split(driver, table_name):
+async def test_split_transactions_deny_split_explicit_commit(driver, table_name):
     async with ydb.aio.SessionPool(driver, 1) as pool:
 
         async def check_transaction(s: ydb.aio.table.Session):
@@ -117,6 +117,28 @@ async def check_transaction(s: ydb.aio.table.Session):
         await pool.retry_operation(check_transaction)
 
 
+@pytest.mark.asyncio
+async def test_split_transactions_deny_split_flag_commit(driver, table_name):
+    async with ydb.aio.SessionPool(driver, 1) as pool:
+
+        async def check_transaction(s: ydb.aio.table.Session):
+            async with s.transaction(allow_split_transactions=False) as tx:
+                await tx.execute(
+                    "INSERT INTO %s (id) VALUES (1)" % table_name, commit_tx=True
+                )
+
+                with pytest.raises(RuntimeError):
+                    await tx.execute("INSERT INTO %s (id) VALUES (2)" % table_name)
+
+                await tx.commit()
+
+            async with s.transaction() as tx:
+                rs = await tx.execute("SELECT COUNT(*) as cnt FROM %s" % table_name)
+                assert rs[0].rows[0].cnt == 1
+
+        await pool.retry_operation(check_transaction)
+
+
 @pytest.mark.asyncio
 async def test_split_transactions_allow_split(driver, table_name):
     async with ydb.aio.SessionPool(driver, 1) as pool:

Diff for: tests/table/test_tx.py

@@ -91,7 +91,7 @@ def test_tx_snapshot_ro(driver_sync, database):
     assert data[0].rows == [{"value": 2}]
 
 
-def test_split_transactions_deny_split(driver_sync, table_name):
+def test_split_transactions_deny_split_explicit_commit(driver_sync, table_name):
     with ydb.SessionPool(driver_sync, 1) as pool:
 
         def check_transaction(s: ydb.table.Session):
@@ -111,6 +111,27 @@ def check_transaction(s: ydb.table.Session):
         pool.retry_operation_sync(check_transaction)
 
 
+def test_split_transactions_deny_split_flag_commit(driver_sync, table_name):
+    with ydb.SessionPool(driver_sync, 1) as pool:
+
+        def check_transaction(s: ydb.table.Session):
+            with s.transaction(allow_split_transactions=False) as tx:
+                tx.execute(
+                    "INSERT INTO %s (id) VALUES (1)" % table_name, commit_tx=True
+                )
+
+                with pytest.raises(RuntimeError):
+                    tx.execute("INSERT INTO %s (id) VALUES (2)" % table_name)
+
+                tx.commit()
+
+            with s.transaction() as tx:
+                rs = tx.execute("SELECT COUNT(*) as cnt FROM %s" % table_name)
+                assert rs[0].rows[0].cnt == 1
+
+        pool.retry_operation_sync(check_transaction)
+
+
 def test_split_transactions_allow_split(driver_sync, table_name):
     with ydb.SessionPool(driver_sync, 1) as pool:
 

Diff for: ydb/aio/table.py

@@ -194,8 +194,6 @@ async def execute(
         self, query, parameters=None, commit_tx=False, settings=None
     ):  # pylint: disable=W0236
 
-        self._check_split()
-
         return await super().execute(query, parameters, commit_tx, settings)
 
     async def commit(self, settings=None):  # pylint: disable=W0236

Diff for: ydb/table.py

@@ -2316,6 +2316,8 @@ def execute(self, query, parameters=None, commit_tx=False, settings=None):
         """
 
         self._check_split()
+        if commit_tx:
+            self._set_finish(self._COMMIT)
 
         return self._driver(
             _tx_ctx_impl.execute_request_factory(