From 6e445ee698ebdc89bcbf18e05b95501111db5254 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 6 Aug 2026 13:27:43 +0300 Subject: [PATCH] Read tx context from connection on cursor execute --- tests/test_connections.py | 97 +++++++++++++++++++++++++++++++++++++++ tests/test_cursors.py | 4 ++ ydb_dbapi/connections.py | 2 - ydb_dbapi/cursors.py | 16 +++++-- 4 files changed, 113 insertions(+), 6 deletions(-) diff --git a/tests/test_connections.py b/tests/test_connections.py index 75e9bec..e70664a 100644 --- a/tests/test_connections.py +++ b/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) diff --git a/tests/test_cursors.py b/tests/test_cursors.py index 177b394..f60fd72 100644 --- a/tests/test_cursors.py +++ b/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: ... diff --git a/ydb_dbapi/connections.py b/ydb_dbapi/connections.py index 672da91..827b5bb 100644 --- a/ydb_dbapi/connections.py +++ b/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, diff --git a/ydb_dbapi/cursors.py b/ydb_dbapi/cursors.py index 0c5d0a2..53aa0fe 100644 --- a/ydb_dbapi/cursors.py +++ b/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()