autocommit

Author: ptiurin

src/firebolt/async_db/connection.py

@@ -85,6 +85,7 @@ class Connection(BaseConnection):
         "client_class",
         "cursor_type",
         "id",
+        "_autocommit",
     )
 
     def __init__(
@@ -96,6 +97,7 @@ 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
@@ -105,6 +107,7 @@ def __init__(
         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
 
@@ -198,14 +201,26 @@ async def cancel_async_query(self, token: str) -> None:
         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)
+        self._add_transaction_params(request)
         response = await self._client.send(request, stream=True)
         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():
+        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)
@@ -227,6 +242,9 @@ async def aclose(self) -> None:
         if self.closed:
             return
 
+        if self.in_transaction:
+            await self.cursor().execute("ROLLBACK")
+
         # 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
@@ -253,6 +271,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
@@ -280,6 +299,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
@@ -292,6 +312,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(
@@ -317,6 +338,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.
 
@@ -380,6 +402,7 @@ async def connect_v2(
             api_endpoint,
             cursor.parameters | cursor._set_parameters,
             connection_id,
+            autocommit,
         )
 
 
@@ -447,6 +470,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.
 
@@ -484,6 +508,7 @@ def connect_core(
         client=client,
         cursor_type=CursorV2,
         api_endpoint=verified_url,
+        autocommit=autocommit,
     )
 
 

src/firebolt/common/base_connection.py

@@ -83,6 +83,7 @@ def __init__(self, cursor_type: Type) -> None:
         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
@@ -91,6 +92,7 @@ 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
@@ -130,10 +132,10 @@ def create_transaction_params(self) -> Dict[str, str]:
             params[TRANSACTION_SEQUENCE_ID_SETTING] = str(self._transaction_sequence_id)
         return params
 
-    def _add_transaction_headers(self, request: Request) -> None:
+    def _add_transaction_params(self, request: Request) -> None:
         transaction_params = self.create_transaction_params()
         for key, value in transaction_params.items():
-            request.headers[key] = value
+            request.url = request.url.copy_add_param(key, value)
 
     def _handle_transaction_updates(self, headers: Headers) -> None:
         self._parse_response_headers_transaction(headers)
@@ -142,6 +144,13 @@ def _handle_transaction_updates(self, headers: Headers) -> None:
         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."""

src/firebolt/common/constants.py

@@ -35,9 +35,7 @@ class ParameterStyle(Enum):
 # 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 + TRANSACTION_PARAMETER_LIST
-)
+IMMUTABLE_PARAMETER_LIST = USE_PARAMETER_LIST + DISALLOWED_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

@@ -15,6 +15,7 @@
     IMMUTABLE_PARAMETER_LIST,
     JSON_LINES_OUTPUT_FORMAT,
     JSON_OUTPUT_FORMAT,
+    TRANSACTION_PARAMETER_LIST,
     USE_PARAMETER_LIST,
     CursorState,
 )
@@ -207,7 +208,7 @@ def _update_set_parameters(self, parameters: Dict[str, Any]) -> None:
         user_parameters = {
             key: value
             for key, value in parameters.items()
-            if key not in IMMUTABLE_PARAMETER_LIST
+            if key not in IMMUTABLE_PARAMETER_LIST + TRANSACTION_PARAMETER_LIST
         }
 
         self.parameters.update(immutable_parameters)

src/firebolt/db/connection.py

@@ -59,6 +59,7 @@ 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
@@ -87,6 +88,7 @@ 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
@@ -99,6 +101,7 @@ def connect(
             api_endpoint=api_endpoint,
             connection_id=connection_id,
             disable_cache=disable_cache,
+            autocommit=autocommit,
         )
     elif auth_version == FireboltAuthVersion.V1:
         return connect_v1(
@@ -124,6 +127,7 @@ 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.
 
@@ -186,6 +190,7 @@ def connect_v2(
             api_endpoint,
             cursor.parameters | cursor._set_parameters,
             connection_id,
+            autocommit,
         )
 
 
@@ -220,6 +225,7 @@ class Connection(BaseConnection):
         "client_class",
         "cursor_type",
         "id",
+        "_autocommit",
     )
 
     def __init__(
@@ -231,6 +237,7 @@ def __init__(
         api_endpoint: str = DEFAULT_API_URL,
         init_parameters: Optional[Dict[str, Any]] = None,
         id: str = uuid4().hex,
+        autocommit: bool = True,
     ):
         super().__init__(cursor_type)
         self.api_endpoint = api_endpoint
@@ -240,6 +247,7 @@ def __init__(
         self.id = id
         self._transaction_lock: threading.Lock = threading.Lock()
         self.init_parameters = init_parameters or {}
+        self._autocommit = autocommit
         if database:
             self.init_parameters["database"] = database
 
@@ -261,6 +269,9 @@ def close(self) -> None:
         if self.closed:
             return
 
+        if self.in_transaction:
+            self.cursor().execute("ROLLBACK")
+
         cursors = self._cursors[:]
         for c in cursors:
             # Here c can already be closed by another thread,
@@ -270,14 +281,25 @@ def close(self) -> None:
         self._is_closed = True
 
     def _execute_query_impl(self, request: Request) -> Response:
-        self._add_transaction_headers(request)
+        self._add_transaction_params(request)
         response = self._client.send(request, stream=True)
         self._handle_transaction_updates(response.headers)
         return response
 
+    def _begin_nolock(self, request: Request) -> None:
+        """Begin a transaction without a lock. Used internally."""
+        begin_request = self._client.build_request(
+            request.method, request.url, content="BEGIN"
+        )
+        response = self._client.send(begin_request, stream=True)
+        self._handle_transaction_updates(response.headers)
+
     def _execute_query(self, request: Request) -> Response:
-        if self.in_transaction():
+        if self.in_transaction or not self.autocommit:
             with self._transaction_lock:
+                # If autocommit is off we need to explicitly begin a transaction
+                if not self.in_transaction:
+                    self._begin_nolock(request)
                 return self._execute_query_impl(request)
         else:
             return self._execute_query_impl(request)
@@ -451,6 +473,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.
 
@@ -489,6 +512,7 @@ def connect_core(
         client=client,
         cursor_type=CursorV2,
         api_endpoint=verified_url,
+        autocommit=autocommit,
     )
 
 

src/firebolt/db/cursor.py

@@ -191,6 +191,7 @@ def _validate_set_parameter(
         self._append_row_set_from_response(None)
 
     def _parse_response_headers(self, headers: Headers) -> None:
+        # TODO: merge with async in base_cursor
         if headers.get(UPDATE_ENDPOINT_HEADER):
             endpoint, params = _parse_update_endpoint(
                 headers.get(UPDATE_ENDPOINT_HEADER)