working example

Author: ptiurin

src/firebolt/async_db/connection.py

@@ -5,7 +5,8 @@
 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
@@ -78,6 +79,9 @@ class Connection(BaseConnection):
         "engine_url",
         "api_endpoint",
         "_is_closed",
+        "_transaction_id",
+        "_transaction_sequence_id",
+        "_transaction_lock",
         "client_class",
         "cursor_type",
         "id",
@@ -99,6 +103,7 @@ def __init__(
         self._cursors: List[Cursor] = []
         self._client = client
         self.id = id
+        self._transaction_lock: trio.Lock = trio.Lock()
         self.init_parameters = init_parameters or {}
         if database:
             self.init_parameters["database"] = database
@@ -192,6 +197,25 @@ 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_headers(request)
+        response = await self._client.send(request, stream=True)
+        self._handle_transaction_updates(response.headers)
+        return response
+
+    async def _execute_query(self, request: Request) -> Response:
+        if self.in_transaction():
+            async with self._transaction_lock:
+                return await self._execute_query_impl(request)
+        else:
+            return await self._execute_query_impl(request)
+
+    async def commit(self) -> None:
+        await self.cursor().execute("COMMIT")
+
+    async def rollback(self) -> None:
+        await self.cursor().execute("ROLLBACK")
+
     # Context manager support
     async def __aenter__(self) -> Connection:
         if self.closed:

src/firebolt/async_db/cursor.py

@@ -29,9 +29,7 @@
 )
 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 (
@@ -58,6 +56,10 @@
 )
 from firebolt.utils.timeout_controller import TimeoutController
 from firebolt.utils.urls import DATABASES_URL, ENGINES_URL
+from firebolt.utils.util import (
+    _parse_remove_parameters,
+    _parse_update_parameters,
+)
 
 if TYPE_CHECKING:
     from firebolt.async_db.connection import Connection
@@ -135,7 +137,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()
 

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,8 @@ 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
 
     def _remove_cursor(self, cursor: Any) -> None:
         # This way it's atomic
@@ -76,17 +91,62 @@ def _remove_cursor(self, cursor: Any) -> None:
         except ValueError:
             pass
 
+    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_headers(self, request: Request) -> None:
+        transaction_params = self.create_transaction_params()
+        for key, value in transaction_params.items():
+            request.headers[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 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,12 +25,19 @@ 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
+IMMUTABLE_PARAMETER_LIST = (
+    USE_PARAMETER_LIST + DISALLOWED_PARAMETER_LIST + TRANSACTION_PARAMETER_LIST
+)
 UPDATE_ENDPOINT_HEADER = "Firebolt-Update-Endpoint"
 UPDATE_PARAMETERS_HEADER = "Firebolt-Update-Parameters"
 RESET_SESSION_HEADER = "Firebolt-Reset-Session"

src/firebolt/common/cursor/base_cursor.py

@@ -33,22 +33,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _parse_update_parameters(parameter_header: str) -> Dict[str, str]:
-    """Parse update parameters and set them as attributes."""
-    # parse key1=value1,key2=value2 comma separated string into dict
-    param_dict = dict(item.split("=") for item in parameter_header.split(","))
-    # strip whitespace from keys and values
-    param_dict = {key.strip(): value.strip() for key, value in param_dict.items()}
-    return param_dict
-
-
-def _parse_remove_parameters(parameter_header: str) -> List[str]:
-    """Parse remove parameters header and return list of parameter names to remove."""
-    # parse key1,key2,key3 comma separated string into list
-    param_list = [item.strip() for item in parameter_header.split(",")]
-    return param_list
-
-
 def _parse_update_endpoint(
     new_engine_endpoint_header: str,
 ) -> Tuple[str, Dict[str, str]]:

src/firebolt/db/connection.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 import logging
+import threading
 from ssl import SSLContext
 from types import TracebackType
 from typing import Any, Dict, List, Optional, Type, Union
 from uuid import uuid4
 from warnings import warn
 
-from httpx import Timeout, codes
+from httpx import Request, Response, Timeout, codes
 
 from firebolt.client import DEFAULT_API_URL, Client, ClientV1, ClientV2
 from firebolt.client.auth import Auth
@@ -213,6 +214,9 @@ class Connection(BaseConnection):
         "engine_url",
         "api_endpoint",
         "_is_closed",
+        "_transaction_id",
+        "_transaction_sequence_id",
+        "_transaction_lock",
         "client_class",
         "cursor_type",
         "id",
@@ -234,6 +238,7 @@ def __init__(
         self._cursors: List[Cursor] = []
         self._client = client
         self.id = id
+        self._transaction_lock: threading.Lock = threading.Lock()
         self.init_parameters = init_parameters or {}
         if database:
             self.init_parameters["database"] = database
@@ -264,6 +269,25 @@ def close(self) -> None:
         self._client.close()
         self._is_closed = True
 
+    def _execute_query_impl(self, request: Request) -> Response:
+        self._add_transaction_headers(request)
+        response = self._client.send(request, stream=True)
+        self._handle_transaction_updates(response.headers)
+        return response
+
+    def _execute_query(self, request: Request) -> Response:
+        if self.in_transaction():
+            with self._transaction_lock:
+                return self._execute_query_impl(request)
+        else:
+            return self._execute_query_impl(request)
+
+    def commit(self) -> None:
+        self.cursor().execute("COMMIT")
+
+    def rollback(self) -> None:
+        self.cursor().execute("ROLLBACK")
+
     # Server-side async methods
 
     def get_async_query_info(self, token: str) -> List[AsyncQueryInfo]:

src/firebolt/db/cursor.py

@@ -37,9 +37,7 @@
 )
 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 (
@@ -67,7 +65,12 @@
 )
 from firebolt.utils.timeout_controller import TimeoutController
 from firebolt.utils.urls import DATABASES_URL, ENGINES_URL
-from firebolt.utils.util import Timer, raise_error_from_response
+from firebolt.utils.util import (
+    Timer,
+    _parse_remove_parameters,
+    _parse_update_parameters,
+    raise_error_from_response,
+)
 
 if TYPE_CHECKING:
     from firebolt.db.connection import Connection
@@ -163,7 +166,7 @@ def _api_request(
                 content=query,
                 timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT,
             )
-            return self._client.send(req, stream=True)
+            return self.connection._execute_query(req)
         except TimeoutException:
             raise QueryTimeoutError()
 

src/firebolt/utils/util.py

@@ -189,6 +189,22 @@ def raise_error_from_response(resp: Response) -> None:
     resp.raise_for_status()
 
 
+def _parse_update_parameters(parameter_header: str) -> Dict[str, str]:
+    """Parse update parameters and set them as attributes."""
+    # parse key1=value1,key2=value2 comma separated string into dict
+    param_dict = dict(item.split("=") for item in parameter_header.split(","))
+    # strip whitespace from keys and values
+    param_dict = {key.strip(): value.strip() for key, value in param_dict.items()}
+    return param_dict
+
+
+def _parse_remove_parameters(parameter_header: str) -> List[str]:
+    """Parse remove parameters header and return list of parameter names to remove."""
+    # parse key1,key2,key3 comma separated string into list
+    param_list = [item.strip() for item in parameter_header.split(",")]
+    return param_list
+
+
 class Timer:
     def __init__(self, message: str = ""):
         self._message = message