From c5ba55a3355889bc942ce786ed562ac4396a7da8 Mon Sep 17 00:00:00 2001 From: Bogdan Markov Date: Mon, 16 Dec 2024 17:05:30 +0100 Subject: [PATCH] Fix: release sessions without transaction. --- ydb_sqlalchemy/dbapi/connection.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/ydb_sqlalchemy/dbapi/connection.py b/ydb_sqlalchemy/dbapi/connection.py index dc4ac97..471bd84 100644 --- a/ydb_sqlalchemy/dbapi/connection.py +++ b/ydb_sqlalchemy/dbapi/connection.py @@ -57,6 +57,7 @@ def __init__( self.interactive_transaction: bool = False # AUTOCOMMIT self.tx_mode: ydb.AbstractTransactionModeBuilder = ydb.SerializableReadWrite() self.tx_context: Optional[ydb.TxContext] = None + self.session: Optional[ydb.Session] = None self.use_scan_query: bool = False self.request_settings: ydb.BaseRequestSettings = ydb.BaseRequestSettings() @@ -139,22 +140,29 @@ def get_ydb_request_settings(self) -> ydb.BaseRequestSettings: def begin(self): self.tx_context = None + if self.session: + self._maybe_await(self.session_pool.release, self.session) + self.session = None if self.interactive_transaction and not self.use_scan_query: - session = self._maybe_await(self.session_pool.acquire) - self.tx_context = session.transaction(self.tx_mode) + self.session = self._maybe_await(self.session_pool.acquire) + self.tx_context = self.session.transaction(self.tx_mode) self._maybe_await(self.tx_context.begin) def commit(self): if self.tx_context and self.tx_context.tx_id: self._maybe_await(self.tx_context.commit) - self._maybe_await(self.session_pool.release, self.tx_context.session) self.tx_context = None + if self.session: + self._maybe_await(self.session_pool.release, self.session) + self.session = None def rollback(self): if self.tx_context and self.tx_context.tx_id: self._maybe_await(self.tx_context.rollback) - self._maybe_await(self.session_pool.release, self.tx_context.session) self.tx_context = None + if self.session: + self._maybe_await(self.session_pool.release, self.session) + self.session = None def close(self): self.rollback()