Fix: release sessions even in case of error.

ydb-platform · Dec 17, 2024 · 37740c2 · 37740c2
1 parent 0034d94
commit 37740c2
Showing 1 changed file with 16 additions and 12 deletions.
diff --git a/ydb_sqlalchemy/dbapi/connection.py b/ydb_sqlalchemy/dbapi/connection.py
@@ -149,20 +149,24 @@ def begin(self):
             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.tx_context = None
-        if self.session:
-            self._maybe_await(self.session_pool.release, self.session)
-            self.session = None
+        try:
+            if self.tx_context and self.tx_context.tx_id:
+                self._maybe_await(self.tx_context.commit)
+                self.tx_context = None
+        finally:
+            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.tx_context = None
-        if self.session:
-            self._maybe_await(self.session_pool.release, self.session)
-            self.session = None
+        try:
+            if self.tx_context and self.tx_context.tx_id:
+                self._maybe_await(self.tx_context.rollback)
+                self.tx_context = None
+        finally:
+            if self.session:
+                self._maybe_await(self.session_pool.release, self.session)
+                self.session = None
 
     def close(self):
         self.rollback()