Read tx context from connection on cursor execute

tests/test_connections.py

@@ -103,6 +103,77 @@ def failing_rollback(*args: any, **kwargs: any) -> None:
             session_pool.release = original_release
             tx_context.rollback = original_rollback
 
+    def _create_tx_context_table(self, connection: dbapi.Connection) -> None:
+        cur = connection.cursor()
+        maybe_await(
+            cur.execute_scheme(
+                """
+                DROP TABLE IF EXISTS test_tx_context;
+                CREATE TABLE test_tx_context (
+                id Int64 NOT NULL,
+                val Int64,
+                PRIMARY KEY(id)
+                )
+                """
+            )
+        )
+        maybe_await(cur.close())
+
+    def _count_tx_context_rows(self, connection: dbapi.Connection) -> int:
+        cur = connection.cursor()
+        maybe_await(cur.execute("SELECT COUNT(*) FROM test_tx_context"))
+        row = cur.fetchone()
+        maybe_await(cur.close())
+        assert row is not None
+        return row[0]
+
+    def _test_cursor_created_before_begin_uses_transaction(
+        self,
+        connection: dbapi.Connection,
+    ) -> None:
+        self._create_tx_context_table(connection)
+        connection.set_isolation_level(dbapi.IsolationLevel.SERIALIZABLE)
+
+        cur = connection.cursor()
+        maybe_await(connection.begin())
+        maybe_await(
+            cur.execute("INSERT INTO test_tx_context(id, val) VALUES (1, 1)")
+        )
+        maybe_await(connection.rollback())
+
+        assert self._count_tx_context_rows(connection) == 0
+
+        maybe_await(cur.close())
+        maybe_await(
+            connection.cursor().execute_scheme("DROP TABLE test_tx_context")
+        )
+
+    def _test_cursor_is_reusable_after_commit(
+        self,
+        connection: dbapi.Connection,
+    ) -> None:
+        self._create_tx_context_table(connection)
+        connection.set_isolation_level(dbapi.IsolationLevel.SERIALIZABLE)
+
+        maybe_await(connection.begin())
+        cur = connection.cursor()
+        maybe_await(
+            cur.execute("INSERT INTO test_tx_context(id, val) VALUES (1, 1)")
+        )
+        maybe_await(connection.commit())
+
+        # the finished transaction must not be reused by the same cursor
+        maybe_await(
+            cur.execute("INSERT INTO test_tx_context(id, val) VALUES (2, 2)")
+        )
+
+        assert self._count_tx_context_rows(connection) == 2
+
+        maybe_await(cur.close())
+        maybe_await(
+            connection.cursor().execute_scheme("DROP TABLE test_tx_context")
+        )
+
     def _test_connection(self, connection: dbapi.Connection) -> None:
         maybe_await(connection.commit())
         maybe_await(connection.rollback())
@@ -527,6 +598,16 @@ def test_close_releases_session_after_rollback_error(
     ) -> None:
         self._test_close_releases_session_after_rollback_error(connection)
 
+    def test_cursor_created_before_begin_uses_transaction(
+        self, connection: dbapi.Connection
+    ) -> None:
+        self._test_cursor_created_before_begin_uses_transaction(connection)
+
+    def test_cursor_is_reusable_after_commit(
+        self, connection: dbapi.Connection
+    ) -> None:
+        self._test_cursor_is_reusable_after_commit(connection)
+
     def test_connection(self, connection: dbapi.Connection) -> None:
         self._test_connection(connection)
 
@@ -649,6 +730,22 @@ async def test_close_releases_session_after_rollback_error(
             self._test_close_releases_session_after_rollback_error, connection
         )
 
+    @pytest.mark.asyncio
+    async def test_cursor_created_before_begin_uses_transaction(
+        self, connection: dbapi.AsyncConnection
+    ) -> None:
+        await greenlet_spawn(
+            self._test_cursor_created_before_begin_uses_transaction, connection
+        )
+
+    @pytest.mark.asyncio
+    async def test_cursor_is_reusable_after_commit(
+        self, connection: dbapi.AsyncConnection
+    ) -> None:
+        await greenlet_spawn(
+            self._test_cursor_is_reusable_after_commit, connection
+        )
+
     @pytest.mark.asyncio
     async def test_connection(self, connection: dbapi.AsyncConnection) -> None:
         await greenlet_spawn(self._test_connection, connection)

tests/test_cursors.py

@@ -25,10 +25,14 @@ def maybe_await(obj: callable) -> any:
 
 
 class FakeSyncConnection:
+    _tx_context = None
+
     def _invalidate_session(self) -> None: ...
 
 
 class FakeAsyncConnection:
+    _tx_context = None
+
     async def _invalidate_session(self) -> None: ...
 
 

ydb_dbapi/connections.py

@@ -265,7 +265,6 @@ def cursor(self) -> Cursor:
             connection=self,
             session_pool=self._session_pool,
             tx_mode=self._tx_mode,
-            tx_context=self._tx_context,
             table_path_prefix=self.table_path_prefix,
             request_settings=self.request_settings,
             retry_settings=self.retry_settings,
@@ -467,7 +466,6 @@ def cursor(self) -> AsyncCursor:
             connection=self,
             session_pool=self._session_pool,
             tx_mode=self._tx_mode,
-            tx_context=self._tx_context,
             table_path_prefix=self.table_path_prefix,
             request_settings=self.request_settings,
             retry_settings=self.retry_settings,

ydb_dbapi/cursors.py

@@ -206,7 +206,6 @@ def __init__(
         tx_mode: ydb.BaseQueryTxMode,
         request_settings: ydb.BaseRequestSettings,
         retry_settings: ydb.RetrySettings,
-        tx_context: ydb.QueryTxContext | None = None,
         table_path_prefix: str = "",
         pyformat: bool = False,
     ) -> None:
@@ -216,11 +215,16 @@ def __init__(
         self._tx_mode = tx_mode
         self._request_settings = request_settings
         self._retry_settings = retry_settings
-        self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
         self._pyformat = pyformat
         self._stream: Iterator | None = None
 
+    @property
+    def _tx_context(self) -> ydb.QueryTxContext | None:
+        # Read the transaction from the connection on every access: a cursor
+        # must not pin the transaction that was open when it was created.
+        return self._connection._tx_context
+
     def fetchone(self) -> tuple | None:
         return self._fetchone_from_buffer()
 
@@ -389,7 +393,6 @@ def __init__(
         tx_mode: ydb.BaseQueryTxMode,
         request_settings: ydb.BaseRequestSettings,
         retry_settings: ydb.RetrySettings,
-        tx_context: ydb.aio.QueryTxContext | None = None,
         table_path_prefix: str = "",
         pyformat: bool = False,
     ) -> None:
@@ -399,11 +402,16 @@ def __init__(
         self._tx_mode = tx_mode
         self._request_settings = request_settings
         self._retry_settings = retry_settings
-        self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
         self._pyformat = pyformat
         self._stream: AsyncIterator | None = None
 
+    @property
+    def _tx_context(self) -> ydb.aio.QueryTxContext | None:
+        # Read the transaction from the connection on every access: a cursor
+        # must not pin the transaction that was open when it was created.
+        return self._connection._tx_context
+
     def fetchone(self) -> tuple | None:
         return self._fetchone_from_buffer()