Add context management and improve tests

Author: ptiurin

src/firebolt/async_db/connection.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from ssl import SSLContext
 from types import TracebackType
 from typing import Any, Dict, List, Optional, Type, Union
@@ -46,6 +47,8 @@
     validate_engine_name_and_url_v1,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class Connection(BaseConnection):
     """
@@ -243,8 +246,13 @@ async def aclose(self) -> None:
         if self.closed:
             return
 
-        if self.in_transaction:
-            await self.cursor().execute("ROLLBACK")
+        # Only rollback if we have a transaction and autocommit is off
+        if self.in_transaction and not self.autocommit:
+            try:
+                await self.cursor().execute("ROLLBACK")
+            except Exception:
+                # If rollback fails during close, continue closing
+                logger.warning("Rollback failed during close")
 
         # self._cursors is going to be changed during closing cursors
         # after this point no cursors would be added to _cursors, only removed since
@@ -260,6 +268,10 @@ async def aclose(self) -> None:
     async def __aexit__(
         self, exc_type: type, exc_val: Exception, exc_tb: TracebackType
     ) -> None:
+        # If exiting normally (no exception) and we have a transaction with
+        # autocommit=False, commit the transaction before closing
+        if exc_type is None and not self.autocommit and self.in_transaction:
+            await self.commit()
         await self.aclose()
 
 

src/firebolt/db/connection.py

@@ -269,8 +269,13 @@ def close(self) -> None:
         if self.closed:
             return
 
-        if self.in_transaction:
-            self.cursor().execute("ROLLBACK")
+        # Only rollback if we have a transaction and autocommit is off
+        if self.in_transaction and not self.autocommit:
+            try:
+                self.cursor().execute("ROLLBACK")
+            except Exception:
+                # If rollback fails during close, continue closing
+                logger.warning("Rollback failed during close")
 
         cursors = self._cursors[:]
         for c in cursors:
@@ -402,6 +407,10 @@ def __enter__(self) -> Connection:
     def __exit__(
         self, exc_type: type, exc_val: Exception, exc_tb: TracebackType
     ) -> None:
+        # If exiting normally (no exception) and we have a transaction with
+        # autocommit=False, commit the transaction before closing
+        if exc_type is None and not self.autocommit and self.in_transaction:
+            self.commit()
         self.close()
 
     def __del__(self) -> None:

tests/integration/dbapi/async/V2/conftest.py

@@ -11,54 +11,41 @@
 from tests.integration.conftest import Secret
 
 
-@fixture(params=["remote", "core"])
+@fixture
 async def connection(
-    engine_name: str,
-    database_name: str,
-    auth: Auth,
-    core_auth: Auth,
-    account_name: str,
-    api_endpoint: str,
-    core_url: str,
-    request: Any,
+    connection_factory: Callable[..., Connection],
 ) -> Connection:
-    if request.param == "core":
-        kwargs = {
-            "database": "firebolt",
-            "auth": core_auth,
-            "url": core_url,
-        }
-    else:
-        kwargs = {
-            "engine_name": engine_name,
-            "database": database_name,
-            "auth": auth,
-            "account_name": account_name,
-            "api_endpoint": api_endpoint,
-        }
-    async with await connect(
-        **kwargs,
-    ) as connection:
+    async with await connection_factory() as connection:
         yield connection
 
 
-@fixture
+@fixture(params=["remote", "core"])
 async def connection_factory(
     engine_name: str,
     database_name: str,
     auth: Auth,
+    core_auth: Auth,
     account_name: str,
     api_endpoint: str,
+    core_url: str,
+    request: Any,
 ) -> Callable[..., Connection]:
     async def factory(**kwargs: Any) -> Connection:
-        return await connect(
-            engine_name=engine_name,
-            database=database_name,
-            auth=auth,
-            account_name=account_name,
-            api_endpoint=api_endpoint,
-            **kwargs,
-        )
+        if request.param == "core":
+            base_kwargs = {
+                "database": "firebolt",
+                "auth": core_auth,
+                "url": core_url,
+            }
+        else:
+            base_kwargs = {
+                "engine_name": engine_name,
+                "database": database_name,
+                "auth": auth,
+                "account_name": account_name,
+                "api_endpoint": api_endpoint,
+            }
+        return await connect(**base_kwargs, **kwargs)
 
     return factory