From 98293c9ebfa17347f6a253be6a8d2215b07fcfef Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 6 Aug 2026 13:24:51 +0300 Subject: [PATCH 1/2] Release session in close() even if rollback fails --- tests/test_connections.py | 44 +++++++++++++++++++++++++++++++++++++++ ydb_dbapi/connections.py | 32 +++++++++++++++++----------- 2 files changed, 64 insertions(+), 12 deletions(-) diff --git a/tests/test_connections.py b/tests/test_connections.py index d87f9ed..b75307f 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -66,6 +66,37 @@ def _test_commit_rollback_after_begin( maybe_await(connection.rollback()) + def _test_close_releases_session_after_rollback_error( + self, + connection: dbapi.Connection, + ) -> None: + connection.set_isolation_level(dbapi.IsolationLevel.SERIALIZABLE) + maybe_await(connection.begin()) + + acquired_session = connection._session + assert acquired_session is not None + + released = [] + original_release = connection._session_pool.release + + def release(session: ydb.QuerySession) -> any: + released.append(session) + return original_release(session) + + connection._session_pool.release = release + + def failing_rollback(*args: any, **kwargs: any) -> None: + raise ydb.issues.BadSession("session is invalidated") + + connection._tx_context.rollback = failing_rollback + + with pytest.raises(dbapi.Error): + maybe_await(connection.close()) + + assert released == [acquired_session] + assert connection._session is None + assert connection._tx_context is None + def _test_connection(self, connection: dbapi.Connection) -> None: maybe_await(connection.commit()) maybe_await(connection.rollback()) @@ -466,6 +497,11 @@ def test_commit_rollback_after_begin( connection, isolation_level ) + def test_close_releases_session_after_rollback_error( + self, connection: dbapi.Connection + ) -> None: + self._test_close_releases_session_after_rollback_error(connection) + def test_connection(self, connection: dbapi.Connection) -> None: self._test_connection(connection) @@ -580,6 +616,14 @@ async def test_commit_rollback_after_begin( isolation_level ) + @pytest.mark.asyncio + async def test_close_releases_session_after_rollback_error( + self, connection: dbapi.AsyncConnection + ) -> None: + await greenlet_spawn( + self._test_close_releases_session_after_rollback_error, connection + ) + @pytest.mark.asyncio async def test_connection(self, connection: dbapi.AsyncConnection) -> None: await greenlet_spawn(self._test_connection, connection) diff --git a/ydb_dbapi/connections.py b/ydb_dbapi/connections.py index cc7a117..285fd1b 100644 --- a/ydb_dbapi/connections.py +++ b/ydb_dbapi/connections.py @@ -291,14 +291,18 @@ def rollback(self) -> None: @handle_ydb_errors def close(self) -> None: - self.rollback() + try: + self.rollback() + finally: + self._tx_context = None - if self._session: - self._session_pool.release(self._session) + if self._session: + self._session_pool.release(self._session) + self._session = None - if not self._shared_session_pool: - self._session_pool.stop() - self._driver.stop() + if not self._shared_session_pool: + self._session_pool.stop() + self._driver.stop() @handle_ydb_errors def describe(self, table_path: str) -> ydb.TableSchemeEntry: @@ -489,14 +493,18 @@ async def rollback(self) -> None: @handle_ydb_errors async def close(self) -> None: - await self.rollback() + try: + await self.rollback() + finally: + self._tx_context = None - if self._session: - await self._session_pool.release(self._session) + if self._session: + await self._session_pool.release(self._session) + self._session = None - if not self._shared_session_pool: - await self._session_pool.stop() - await self._driver.stop() + if not self._shared_session_pool: + await self._session_pool.stop() + await self._driver.stop() @handle_ydb_errors async def describe(self, table_path: str) -> ydb.TableSchemeEntry: From e4b193539bbffd4463ecdac02cd36c83b475c64b Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Thu, 6 Aug 2026 15:31:50 +0300 Subject: [PATCH 2/2] Restore patched methods in close() regression test --- tests/test_connections.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/tests/test_connections.py b/tests/test_connections.py index b75307f..029a064 100644 --- a/tests/test_connections.py +++ b/tests/test_connections.py @@ -77,25 +77,31 @@ def _test_close_releases_session_after_rollback_error( assert acquired_session is not None released = [] - original_release = connection._session_pool.release + session_pool = connection._session_pool + tx_context = connection._tx_context + original_release = session_pool.release + original_rollback = tx_context.rollback def release(session: ydb.QuerySession) -> any: released.append(session) return original_release(session) - connection._session_pool.release = release - def failing_rollback(*args: any, **kwargs: any) -> None: raise ydb.issues.BadSession("session is invalidated") - connection._tx_context.rollback = failing_rollback + session_pool.release = release + tx_context.rollback = failing_rollback - with pytest.raises(dbapi.Error): - maybe_await(connection.close()) + try: + with pytest.raises(dbapi.Error): + maybe_await(connection.close()) - assert released == [acquired_session] - assert connection._session is None - assert connection._tx_context is None + assert released == [acquired_session] + assert connection._session is None + assert connection._tx_context is None + finally: + session_pool.release = original_release + tx_context.rollback = original_rollback def _test_connection(self, connection: dbapi.Connection) -> None: maybe_await(connection.commit())