From 9818c5bca57e36323eb5b82b7673fc03d70537d3 Mon Sep 17 00:00:00 2001 From: Roman Tretiak Date: Mon, 16 Dec 2024 16:29:45 +0100 Subject: [PATCH] Hotfix: release session on error --- ydb_sqlalchemy/dbapi/connection.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/ydb_sqlalchemy/dbapi/connection.py b/ydb_sqlalchemy/dbapi/connection.py index dc4ac97..6375353 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() @@ -140,21 +141,25 @@ def get_ydb_request_settings(self) -> ydb.BaseRequestSettings: def begin(self): self.tx_context = 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()