feat: Add cancel and status queries to server-side async execution

src/firebolt/async_db/cursor.py

@@ -61,6 +61,19 @@ class CursorState(Enum):
     CLOSED = 4
 
 
+class QueryStatus(Enum):
+    """Enumeration of query responses on server-side async queries."""
+
+    RUNNING = 1
+    ENDED_SUCCESSFULLY = 2
+    ENDED_UNSUCCESSFULLY = 3
+    NOT_READY = 4
+    STARTED_EXECUTION = 5
+    PARSE_ERROR = 6
+    CANCELED_EXECUTION = 7
+    EXECUTION_ERROR = 8
+
+
 class Statistics(BaseModel):
     """
     Class for query execution statistics.
@@ -273,16 +286,6 @@ def _reset(self) -> None:
         self._next_set_idx = 0
         self._query_id = ""
 
-    def _query_id_from_response_async(self, response: Response) -> str:
-        if response.headers.get("content-length", "") == "0":
-            raise OperationalError("No response to asynchronous query.")
-        query_data = response.json()
-        if "query_id" not in query_data:
-            raise OperationalError(
-                "Invalid response to asynchronous query: missing query_id."
-            )
-        return query_data["query_id"]
-
     def _row_set_from_response(
         self, response: Response
     ) -> Tuple[
@@ -328,27 +331,43 @@ def _append_row_set(
 
     async def _api_request(
         self,
-        query: str,
-        set_parameters: Optional[dict] = None,
+        query: Optional[str] = "",
+        parameters: Optional[dict] = {},
+        path: Optional[str] = "",
+        use_set_parameters: Optional[bool] = True,
     ) -> Response:
+        """
+        Query API, return Response object.
+
+        Args:
+            query (str): SQL query
+            parameters (Optional[Sequence[ParameterType]]): A sequence of substitution
+                parameters. Used to replace '?' placeholders inside a query with
+                actual values. Note: In order to "output_format" dict value, it
+                    must be an empty string. If no value not specified,
+                    JSON_OUTPUT_FORMAT will be used.
+            path (str): endpoint suffix, for example "cancel" or "status"
+        """
+        if not parameters:
+            parameters = {"output_format": JSON_OUTPUT_FORMAT}
+        elif "output_format" not in parameters:
+            parameters["output_format"] = JSON_OUTPUT_FORMAT
+        elif parameters["output_format"] == "":
+            # In the case that output_format should be missing, send in "";
+            # we remove it here so the API endpoint string won't be malformed.
+            del parameters["output_format"]
+        if use_set_parameters:
+            parameters = {**self._set_parameters, **parameters}
         return await self._client.request(
-            url="/",
+            url=f"/{path}",
             method="POST",
             params={
                 "database": self.connection.database,
-                "output_format": JSON_OUTPUT_FORMAT,
-                **self._set_parameters,
-                **(set_parameters or dict()),
+                **parameters,
             },
             content=query,
         )
 
-    async def _async_execution_api_request(self, query: str) -> Response:
-        """Do query request using SET async_execution=1."""
-        return await self._api_request(
-            query, {"async_execution": 1, "advanced_mode": 1}
-        )
-
     async def _validate_set_parameter(self, parameter: SetParameter) -> None:
         """Validate parameter by executing simple query with it."""
         if parameter.name == "async_execution":
@@ -357,7 +376,9 @@ async def _validate_set_parameter(self, parameter: SetParameter) -> None:
                 "Instead, pass it as an argument to the execute() or "
                 "executemany() function."
             )
-        resp = await self._api_request("select 1", {parameter.name: parameter.value})
+        resp = await self._api_request(
+            "select 1", {parameter.name: parameter.value, "output_format": ""}
+        )
         # Handle invalid set parameter
         if resp.status_code == codes.BAD_REQUEST:
             raise OperationalError(resp.text)
@@ -366,6 +387,35 @@ async def _validate_set_parameter(self, parameter: SetParameter) -> None:
         # set parameter passed validation
         self._set_parameters[parameter.name] = parameter.value
 
+    def _validate_ss_async_settings(
+        self,
+        parameters: Sequence[Sequence[ParameterType]],
+        queries: List[Union[SetParameter, str]],
+        skip_parsing: bool = False,
+        async_execution: Optional[bool] = False,
+    ) -> None:
+        if async_execution and self._set_parameters.get("use_standard_sql", "1") == "0":
+            raise AsyncExecutionUnavailableError(
+                "It is not possible to execute queries asynchronously if "
+                "use_standard_sql=0."
+            )
+        if parameters and skip_parsing:
+            logger.warning(
+                "Query formatting parameters are provided but skip_parsing "
+                "is specified. They will be ignored."
+            )
+
+        # Allow users to manually skip parsing for performance improvement.
+        non_set_queries = 0
+        for query in queries:
+            if type(query) is not SetParameter:
+                non_set_queries += 1
+        if non_set_queries > 1 and async_execution:
+            raise AsyncExecutionUnavailableError(
+                "It is not possible to execute multi-statement "
+                "queries asynchronously."
+            )
+
     async def _do_execute(
         self,
         raw_query: str,
@@ -374,34 +424,14 @@ async def _do_execute(
         async_execution: Optional[bool] = False,
     ) -> None:
         self._reset()
+        queries: List[Union[SetParameter, str]] = (
+            [raw_query] if skip_parsing else split_format_sql(raw_query, parameters)
+        )
         try:
-            if (
-                async_execution
-                and self._set_parameters.get("use_standard_sql", "0") == "1"
-            ):
-                raise AsyncExecutionUnavailableError(
-                    "It is not possible to execute queries asynchronously if "
-                    "use_standard_sql is in use."
-                )
-            if parameters and skip_parsing:
-                logger.warning(
-                    "Query formatting parameters are provided with skip_parsing."
-                    " They will be ignored."
-                )
-
-            # Allow users to manually skip parsing for performance improvement
-            queries: List[Union[SetParameter, str]] = (
-                [raw_query] if skip_parsing else split_format_sql(raw_query, parameters)
-            )
-            if len(queries) > 1 and async_execution:
-                raise AsyncExecutionUnavailableError(
-                    "It is not possible to execute multi-statement "
-                    "queries asynchronously."
-                )
             for query in queries:
 
                 start_time = time.time()
-                # our CREATE EXTERNAL TABLE queries currently require credentials,
+                # Our CREATE EXTERNAL TABLE queries currently require credentials,
                 # so we will skip logging those queries.
                 # https://docs.firebolt.io/sql-reference/commands/ddl-commands#create-external-table
                 if isinstance(query, SetParameter) or not re.search(
@@ -419,9 +449,24 @@ async def _do_execute(
                 if isinstance(query, SetParameter):
                     await self._validate_set_parameter(query)
                 elif async_execution:
-                    resp = await self._async_execution_api_request(query)
-                    await self._raise_if_error(resp)
-                    self._query_id = self._query_id_from_response_async(resp)
+                    self._validate_ss_async_settings(
+                        parameters,
+                        queries,
+                        skip_parsing,
+                        async_execution,
+                    )
+                    response = await self._api_request(
+                        query, {"async_execution": 1, "advanced_mode": 1}
+                    )
+                    await self._raise_if_error(response)
+                    if response.headers.get("content-length", "") == "0":
+                        raise OperationalError("No response to asynchronous query.")
+                    resp = response.json()
+                    if "query_id" not in resp or resp["query_id"] == "":
+                        raise OperationalError(
+                            "Invalid response to asynchronous query: missing query_id."
+                        )
+                    self._query_id = resp["query_id"]
                 else:
                     resp = await self._api_request(query)
                     await self._raise_if_error(resp)
@@ -431,7 +476,7 @@ async def _do_execute(
 
                 logger.info(
                     f"Query fetched {self.rowcount} rows in"
-                    f" {time.time() - start_time} seconds"
+                    f" {time.time() - start_time} seconds."
                 )
 
             self._state = CursorState.DONE
@@ -584,6 +629,44 @@ def setinputsizes(self, sizes: List[int]) -> None:
     def setoutputsize(self, size: int, column: Optional[int] = None) -> None:
         """Set a column buffer size for fetches of large columns (does nothing)."""
 
+    @check_not_closed
+    async def cancel(self, query_id: str) -> None:
+        """Cancel a server-side async query."""
+        await self._api_request(
+            parameters={"query_id": query_id, "output_format": ""},
+            path="cancel",
+            use_set_parameters=False,
+        )
+
+    @check_not_closed
+    async def get_status(self, query_id: str) -> QueryStatus:
+        """Get status of a server-side async query. Return the state of the query."""
+        try:
+            resp = await self._api_request(
+                # output_format must be empty for status to work correctly.
+                # And set parameters will cause 400 errors.
+                parameters={"query_id": query_id, "output_format": ""},
+                path="status",
+                use_set_parameters=False,
+            )
+            if resp.status_code == codes.BAD_REQUEST:
+                raise OperationalError(
+                    f"Asynchronous query {query_id} status check failed: "
+                    f"{resp.status_code}."
+                )
+            resp_json = resp.json()
+            if "status" not in resp_json:
+                raise OperationalError(
+                    f"Invalid response to asynchronous query: missing status."
+                )
+        except Exception:
+            self._state = CursorState.ERROR
+            raise
+        # Remember that query_id might be empty.
+        if resp_json["status"] == "":
+            return QueryStatus.NOT_READY
+        return QueryStatus[resp_json["status"]]
+
     # Context manager support
     @check_not_closed
     def __enter__(self) -> BaseCursor:
@@ -636,31 +719,33 @@ async def executemany(
         parameters_seq: Sequence[Sequence[ParameterType]],
         async_execution: Optional[bool] = False,
     ) -> int:
-        async with self._async_query_lock.writer:
-            return await super().executemany(query, parameters_seq, async_execution)
         """
-            Prepare and execute a database query against all parameter
-            sequences provided.
+        Prepare and execute a database query against all parameter
+        sequences provided.
         """
+        async with self._async_query_lock.writer:
+            return await super().executemany(query, parameters_seq, async_execution)
 
     @wraps(BaseCursor.fetchone)
     async def fetchone(self) -> Optional[List[ColType]]:
         async with self._async_query_lock.reader:
+            """Fetch the next row of a query result set."""
             return super().fetchone()
-        """Fetch the next row of a query result set."""
 
     @wraps(BaseCursor.fetchmany)
     async def fetchmany(self, size: Optional[int] = None) -> List[List[ColType]]:
         async with self._async_query_lock.reader:
+            """
+            Fetch the next set of rows of a query result;
+            size is cursor.arraysize by default.
+            """
             return super().fetchmany(size)
-        """Fetch the next set of rows of a query result;
-          size is cursor.arraysize by default."""
 
     @wraps(BaseCursor.fetchall)
     async def fetchall(self) -> List[List[ColType]]:
         async with self._async_query_lock.reader:
+            """Fetch all remaining rows of a query result."""
             return super().fetchall()
-        """Fetch all remaining rows of a query result."""
 
     @wraps(BaseCursor.nextset)
     async def nextset(self) -> None:

src/firebolt/db/cursor.py

@@ -10,6 +10,7 @@
 from firebolt.async_db.cursor import BaseCursor as AsyncBaseCursor
 from firebolt.async_db.cursor import (
     ParameterType,
+    QueryStatus,
     check_not_closed,
     check_query_executed,
 )
@@ -101,3 +102,13 @@ def __iter__(self) -> Generator[List[ColType], None, None]:
             if row is None:
                 return
             yield row
+
+    @wraps(AsyncBaseCursor.get_status)
+    def get_status(self, query_id: str) -> QueryStatus:
+        with self._query_lock.gen_rlock():
+            return async_to_sync(super().get_status, self._async_job_thread)(query_id)
+
+    @wraps(AsyncBaseCursor.cancel)
+    def cancel(self, query_id: str) -> None:
+        with self._query_lock.gen_rlock():
+            return async_to_sync(super().cancel, self._async_job_thread)(query_id)