feat(FIR-46457): Add full transaction support (#471)

Author: ptiurin

src/firebolt/async_db/connection.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import logging
 from ssl import SSLContext
 from types import TracebackType
 from typing import Any, Dict, List, Optional, Type, Union
 from uuid import uuid4
 
-from httpx import Timeout, codes
+import trio
+from httpx import Request, Response, Timeout, codes
 
 from firebolt.async_db.cursor import Cursor, CursorV1, CursorV2
 from firebolt.client import DEFAULT_API_URL
@@ -45,6 +47,8 @@
     validate_engine_name_and_url_v1,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class Connection(BaseConnection):
     """
@@ -78,9 +82,13 @@ class Connection(BaseConnection):
         "engine_url",
         "api_endpoint",
         "_is_closed",
+        "_transaction_id",
+        "_transaction_sequence_id",
+        "_transaction_lock",
         "client_class",
         "cursor_type",
         "id",
+        "_autocommit",
     )
 
     def __init__(
@@ -92,14 +100,17 @@ def __init__(
         api_endpoint: str,
         init_parameters: Optional[Dict[str, Any]] = None,
         id: str = uuid4().hex,
+        autocommit: bool = True,
     ):
         super().__init__(cursor_type)
         self.api_endpoint = api_endpoint
         self.engine_url = engine_url
         self._cursors: List[Cursor] = []
         self._client = client
         self.id = id
+        self._transaction_lock: trio.Lock = trio.Lock()
         self.init_parameters = init_parameters or {}
+        self._autocommit = autocommit
         if database:
             self.init_parameters["database"] = database
 
@@ -192,6 +203,44 @@ async def cancel_async_query(self, token: str) -> None:
         cursor = self.cursor()
         await cursor.execute(ASYNC_QUERY_CANCEL, [async_query_info[0].query_id])
 
+    async def _execute_query_impl(self, request: Request) -> Response:
+        self._add_transaction_params(request)
+        response = await self._client.send(request, stream=True)
+        if not self.autocommit:
+            self._handle_transaction_updates(response.headers)
+        return response
+
+    async def _begin_nolock(self, request: Request) -> None:
+        """Begin a transaction without a lock. Used internally."""
+        # Create a copy of the request with "BEGIN" as the body content
+        begin_request = self._client.build_request(
+            request.method, request.url, content="BEGIN"
+        )
+        response = await self._client.send(begin_request, stream=True)
+        self._handle_transaction_updates(response.headers)
+
+    async def _execute_query(self, request: Request) -> Response:
+        if self.in_transaction or not self.autocommit:
+            async with self._transaction_lock:
+                # If autocommit is off we need to explicitly begin a transaction
+                if not self.in_transaction:
+                    await self._begin_nolock(request)
+                return await self._execute_query_impl(request)
+        else:
+            return await self._execute_query_impl(request)
+
+    async def commit(self) -> None:
+        if self.closed:
+            raise ConnectionClosedError("Unable to commit: Connection closed.")
+        # Commit is a no-op for V1
+        if self.cursor_type != CursorV1:
+            await self.cursor().execute("COMMIT")
+
+    async def rollback(self) -> None:
+        if self.closed:
+            raise ConnectionClosedError("Unable to rollback: Connection closed.")
+        await self.cursor().execute("ROLLBACK")
+
     # Context manager support
     async def __aenter__(self) -> Connection:
         if self.closed:
@@ -203,6 +252,14 @@ async def aclose(self) -> None:
         if self.closed:
             return
 
+        # Only rollback if we have a transaction and autocommit is off
+        if self.in_transaction and not self.autocommit:
+            try:
+                await self.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
         # closing lock is held, and later connection will be marked as closed
@@ -217,6 +274,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()
 
 
@@ -229,6 +290,7 @@ async def connect(
     api_endpoint: str = DEFAULT_API_URL,
     disable_cache: bool = False,
     url: Optional[str] = None,
+    autocommit: bool = True,
     additional_parameters: Dict[str, Any] = {},
 ) -> Connection:
     # auth parameter is optional in function signature
@@ -256,6 +318,7 @@ async def connect(
             user_agent_header=user_agent_header,
             database=database,
             connection_url=url,
+            autocommit=autocommit,
         )
     elif auth_version == FireboltAuthVersion.V2:
         assert account_name is not None
@@ -268,6 +331,7 @@ async def connect(
             api_endpoint=api_endpoint,
             connection_id=connection_id,
             disable_cache=disable_cache,
+            autocommit=autocommit,
         )
     elif auth_version == FireboltAuthVersion.V1:
         return await connect_v1(
@@ -293,6 +357,7 @@ async def connect_v2(
     engine_name: Optional[str] = None,
     api_endpoint: str = DEFAULT_API_URL,
     disable_cache: bool = False,
+    autocommit: bool = True,
 ) -> Connection:
     """Connect to Firebolt.
 
@@ -356,6 +421,7 @@ async def connect_v2(
             api_endpoint,
             cursor.parameters | cursor._set_parameters,
             connection_id,
+            autocommit,
         )
 
 
@@ -423,6 +489,7 @@ def connect_core(
     user_agent_header: str,
     database: Optional[str] = None,
     connection_url: Optional[str] = None,
+    autocommit: bool = True,
 ) -> Connection:
     """Connect to Firebolt Core.
 
@@ -460,6 +527,7 @@ def connect_core(
         client=client,
         cursor_type=CursorV2,
         api_endpoint=verified_url,
+        autocommit=autocommit,
     )
 
 

src/firebolt/async_db/cursor.py

@@ -8,30 +8,13 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union
 from urllib.parse import urljoin
 
-from httpx import (
-    URL,
-    USE_CLIENT_DEFAULT,
-    Headers,
-    Response,
-    TimeoutException,
-    codes,
-)
+from httpx import URL, USE_CLIENT_DEFAULT, Response, TimeoutException, codes
 
 from firebolt.client.client import AsyncClient, AsyncClientV1, AsyncClientV2
 from firebolt.common._types import ColType, ParameterType, SetParameter
-from firebolt.common.constants import (
-    JSON_OUTPUT_FORMAT,
-    REMOVE_PARAMETERS_HEADER,
-    RESET_SESSION_HEADER,
-    UPDATE_ENDPOINT_HEADER,
-    UPDATE_PARAMETERS_HEADER,
-    CursorState,
-)
+from firebolt.common.constants import JSON_OUTPUT_FORMAT, CursorState
 from firebolt.common.cursor.base_cursor import (
     BaseCursor,
-    _parse_remove_parameters,
-    _parse_update_endpoint,
-    _parse_update_parameters,
     _raise_if_internal_set_parameter,
 )
 from firebolt.common.cursor.decorators import (
@@ -135,7 +118,7 @@ async def _api_request(
                 content=query,
                 timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT,
             )
-            return await self._client.send(req, stream=True)
+            return await self.connection._execute_query(req)
         except TimeoutException:
             raise QueryTimeoutError()
 
@@ -181,25 +164,6 @@ async def _validate_set_parameter(
         # append empty result set
         await self._append_row_set_from_response(None)
 
-    async def _parse_response_headers(self, headers: Headers) -> None:
-        if headers.get(UPDATE_ENDPOINT_HEADER):
-            endpoint, params = _parse_update_endpoint(
-                headers.get(UPDATE_ENDPOINT_HEADER)
-            )
-            self._update_set_parameters(params)
-            self.engine_url = endpoint
-
-        if headers.get(RESET_SESSION_HEADER):
-            self.flush_parameters()
-
-        if headers.get(UPDATE_PARAMETERS_HEADER):
-            param_dict = _parse_update_parameters(headers.get(UPDATE_PARAMETERS_HEADER))
-            self._update_set_parameters(param_dict)
-
-        if headers.get(REMOVE_PARAMETERS_HEADER):
-            param_list = _parse_remove_parameters(headers.get(REMOVE_PARAMETERS_HEADER))
-            self._remove_set_parameters(param_list)
-
     async def _close_rowset_and_reset(self) -> None:
         """Reset cursor state."""
         if self._row_set is not None:
@@ -305,7 +269,7 @@ async def _execute_single_query(
             await resp.aread()
             self._parse_async_response(resp)
         else:
-            await self._parse_response_headers(resp.headers)
+            self._parse_response_headers(resp.headers)
             await self._append_row_set_from_response(resp)
 
         if not async_execution:

src/firebolt/common/base_connection.py

@@ -1,19 +1,32 @@
 from collections import namedtuple
 from typing import Any, Dict, List, Optional, Tuple, Type
 
+from httpx import Headers, Request
+
 from firebolt.client.auth.base import Auth
 from firebolt.common._types import ColType
+from firebolt.common.constants import (
+    REMOVE_PARAMETERS_HEADER,
+    RESET_SESSION_HEADER,
+    TRANSACTION_ID_SETTING,
+    TRANSACTION_SEQUENCE_ID_SETTING,
+    UPDATE_PARAMETERS_HEADER,
+)
 from firebolt.utils.cache import (
     ConnectionInfo,
     EngineInfo,
     SecureCacheKey,
     _firebolt_cache,
 )
-from firebolt.utils.exception import ConnectionClosedError, FireboltError
+from firebolt.utils.exception import FireboltError
 from firebolt.utils.usage_tracker import (
     get_cache_tracking_params,
     get_user_agent_header,
 )
+from firebolt.utils.util import (
+    _parse_remove_parameters,
+    _parse_update_parameters,
+)
 
 ASYNC_QUERY_STATUS_RUNNING = "RUNNING"
 ASYNC_QUERY_STATUS_SUCCESSFUL = "ENDED_SUCCESSFULLY"
@@ -68,6 +81,9 @@ def __init__(self, cursor_type: Type) -> None:
         self.cursor_type = cursor_type
         self._cursors: List[Any] = []
         self._is_closed = False
+        self._transaction_id: Optional[str] = None
+        self._transaction_sequence_id: Optional[str] = None
+        self._autocommit: bool = True
 
     def _remove_cursor(self, cursor: Any) -> None:
         # This way it's atomic
@@ -76,17 +92,70 @@ def _remove_cursor(self, cursor: Any) -> None:
         except ValueError:
             pass
 
+    @property
+    def in_transaction(self) -> bool:
+        """`True` if connection is in a transaction; `False` otherwise."""
+        return self._transaction_id is not None
+
+    def _parse_response_headers_transaction(self, headers: Headers) -> None:
+        parameters_header = headers.get(UPDATE_PARAMETERS_HEADER)
+        if not parameters_header:
+            return
+        parameters = _parse_update_parameters(parameters_header)
+        transaction_id = parameters.get(TRANSACTION_ID_SETTING)
+        if transaction_id:
+            self._transaction_id = transaction_id
+        sequence_id = parameters.get(TRANSACTION_SEQUENCE_ID_SETTING)
+        if sequence_id:
+            self._transaction_sequence_id = sequence_id
+
+    def _parse_remove_headers_transaction(self, headers: Headers) -> None:
+        parameters_header = headers.get(REMOVE_PARAMETERS_HEADER)
+        if not parameters_header:
+            return
+        parameters = _parse_remove_parameters(parameters_header)
+        for param in parameters:
+            if param == TRANSACTION_ID_SETTING:
+                self._transaction_id = None
+            elif param == TRANSACTION_SEQUENCE_ID_SETTING:
+                self._transaction_sequence_id = None
+
+    def _reset_transaction_state(self) -> None:
+        self._transaction_id = None
+        self._transaction_sequence_id = None
+
+    def create_transaction_params(self) -> Dict[str, str]:
+        params: Dict[str, str] = {}
+        if self._transaction_id:
+            params[TRANSACTION_ID_SETTING] = self._transaction_id
+        if self._transaction_sequence_id is not None:
+            params[TRANSACTION_SEQUENCE_ID_SETTING] = str(self._transaction_sequence_id)
+        return params
+
+    def _add_transaction_params(self, request: Request) -> None:
+        transaction_params = self.create_transaction_params()
+        for key, value in transaction_params.items():
+            request.url = request.url.copy_add_param(key, value)
+
+    def _handle_transaction_updates(self, headers: Headers) -> None:
+        self._parse_response_headers_transaction(headers)
+        if headers.get(RESET_SESSION_HEADER):
+            self._reset_transaction_state()
+        if headers.get(REMOVE_PARAMETERS_HEADER):
+            self._parse_remove_headers_transaction(headers)
+
+    @property
+    def autocommit(self) -> bool:
+        """
+        `True` if connection is in autocommit mode; `False` otherwise.
+        """
+        return self._autocommit
+
     @property
     def closed(self) -> bool:
         """`True` if connection is closed; `False` otherwise."""
         return self._is_closed
 
-    def commit(self) -> None:
-        """Does nothing since Firebolt doesn't have transactions."""
-
-        if self.closed:
-            raise ConnectionClosedError("Unable to commit: Connection closed.")
-
 
 def get_cached_system_engine_info(
     auth: Auth,

src/firebolt/common/constants.py

@@ -25,10 +25,15 @@ class ParameterStyle(Enum):
     FB_NUMERIC = "fb_numeric"  # $1, $2, ... as placeholders (server-side)
 
 
+TRANSACTION_ID_SETTING = "transaction_id"
+TRANSACTION_SEQUENCE_ID_SETTING = "transaction_sequence_id"
+
 # Parameters that should be set using USE instead of SET
 USE_PARAMETER_LIST = ["database", "engine"]
 # parameters that can only be set by the backend
 DISALLOWED_PARAMETER_LIST = ["output_format"]
+# Connection level transaction management
+TRANSACTION_PARAMETER_LIST = [TRANSACTION_ID_SETTING, TRANSACTION_SEQUENCE_ID_SETTING]
 # parameters that are set by the backend and should not be set by the user
 IMMUTABLE_PARAMETER_LIST = USE_PARAMETER_LIST + DISALLOWED_PARAMETER_LIST
 UPDATE_ENDPOINT_HEADER = "Firebolt-Update-Endpoint"