feat(NoTicket): send NULL output format for server-side async queries

docsrc/Connecting_and_queries.rst

@@ -798,7 +798,9 @@ Trying to access `async_query_token` before calling `execute_async` will raise a
     separately using multiple `execute_async` calls.
 
 .. note::
-    Fetching data via SELECT is not supported and will raise an exception. execute_async is best suited for DML queries.
+    Firebolt accepts async SELECT queries from engine version 4.32.13 on. Older engines reject them.
+    There is currently no way to retrieve the results of an async SELECT, so ``execute_async`` is best
+    suited for DML queries. Calling fetch methods after ``execute_async`` raises an exception.
 
 Monitoring the query status
 ----------------------------

src/firebolt/common/constants.py

@@ -10,6 +10,7 @@
 ENGINE_STATUS_RUNNING_LIST = ["RUNNING", "Running", "ENGINE_STATE_RUNNING"]
 JSON_OUTPUT_FORMAT = "JSON_Compact"
 JSON_LINES_OUTPUT_FORMAT = "JSONLines_Compact"
+NULL_OUTPUT_FORMAT = "NULL"
 
 
 class CursorState(Enum):

src/firebolt/common/cursor/statement_planners.py

@@ -11,6 +11,7 @@
 from firebolt.common.constants import (
     JSON_LINES_OUTPUT_FORMAT,
     JSON_OUTPUT_FORMAT,
+    NULL_OUTPUT_FORMAT,
 )
 from firebolt.utils.exception import (
     ConfigurationError,
@@ -129,8 +130,15 @@ def _create_bulk_plan_impl(
         """Create parameter-style specific bulk execution plan."""
 
     @staticmethod
-    def _get_output_format(streaming: bool) -> str:
+    def _get_output_format(streaming: bool, async_execution: bool = False) -> str:
         """Get output format for the query."""
+        assert not (
+            streaming and async_execution
+        ), "Streaming and async execution are not supported together"
+        # Async queries return no rows, and the server only accepts async
+        # SELECTs when no result serialization is requested.
+        if async_execution:
+            return NULL_OUTPUT_FORMAT
         if streaming:
             return JSON_LINES_OUTPUT_FORMAT
         return JSON_OUTPUT_FORMAT
@@ -203,7 +211,7 @@ def _build_fb_numeric_query_params(
         ]
 
         query_params: Dict[str, Any] = {
-            "output_format": self._get_output_format(streaming),
+            "output_format": self._get_output_format(streaming, async_execution),
         }
         if query_parameters:
             query_params["query_parameters"] = json.dumps(query_parameters)
@@ -263,7 +271,7 @@ def _create_standard_execution_plan(
 
         # Build basic query parameters for qmark style
         query_params: Dict[str, Any] = {
-            "output_format": self._get_output_format(streaming),
+            "output_format": self._get_output_format(streaming, async_execution),
         }
         if async_execution:
             query_params["async"] = True
@@ -288,7 +296,7 @@ def _create_bulk_plan_impl(
 
         # Build query parameters for bulk insert
         query_params: Dict[str, Any] = {
-            "output_format": self._get_output_format(False),
+            "output_format": self._get_output_format(False, async_execution),
         }
         if async_execution:
             query_params["async"] = True

src/firebolt/utils/util.py

@@ -1,10 +1,12 @@
+import json
 import logging
 from functools import lru_cache
 from os import environ
 from time import time
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
+    Any,
     Callable,
     Dict,
     List,
@@ -161,6 +163,12 @@ def validate_engine_name_and_url_v1(
         )
 
 
+def _structured_error_from_body(body: Any) -> Optional[FireboltStructuredError]:
+    if isinstance(body, dict) and body.get("errors"):
+        return FireboltStructuredError(body)
+    return None
+
+
 def raise_error_from_response(resp: Response) -> None:
     """
     Raise a correct error from the response.
@@ -172,15 +180,17 @@ def raise_error_from_response(resp: Response) -> None:
         resp (Response): HTTP response
     """
     to_raise = None
-    # If error is Text - raise as is
     if "text/plain" in resp.headers.get("Content-Type", ""):
-        raise FireboltError(resp.text)
-    # If error is Json - parse it and raise
+        # Queries sent with output_format=NULL get JSON error bodies labelled
+        # as plain text, so check for a structured error before raising as is
+        try:
+            to_raise = _structured_error_from_body(json.loads(resp.text))
+        except ValueError:
+            pass
+        raise to_raise or FireboltError(resp.text)
     try:
-        decoded = resp.json()
-        if "errors" in decoded and len(decoded["errors"]) > 0:
-            # Raise later to avoid catching it in the except block
-            to_raise = FireboltStructuredError(decoded)
+        # Raise later to avoid catching it in the except block
+        to_raise = _structured_error_from_body(resp.json())
 
     except Exception:
         # If we can't parse the body, print out the error body

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

@@ -5,7 +5,11 @@
 from pytest import mark, raises
 
 from firebolt.async_db import Connection
-from firebolt.utils.exception import FireboltError, FireboltStructuredError
+from firebolt.utils.exception import (
+    FireboltError,
+    FireboltStructuredError,
+    MethodNotAllowedInAsyncError,
+)
 
 LONG_SELECT = "SELECT checksum(*) FROM GENERATE_SERIES(1, 2500000000)"  # approx 3 sec
 
@@ -53,6 +57,19 @@ async def test_insert_async(connection: Connection) -> None:
         await cursor.execute(f"DROP TABLE {table_name}")
 
 
+async def test_select_async(connection: Connection) -> None:
+    cursor = connection.cursor()
+    await cursor.execute_async("SELECT 1")
+    token = cursor.async_query_token
+    assert token is not None, "Async token was not returned"
+    # sleep for 2 sec to make sure the async query is completed
+    time.sleep(2)
+    assert await connection.is_async_query_running(token) == False
+    assert await connection.is_async_query_successful(token) == True
+    with raises(MethodNotAllowedInAsyncError):
+        await cursor.fetchall()
+
+
 async def test_insert_async_running(connection: Connection) -> None:
     cursor = connection.cursor()
     rnd_suffix = str(randint(0, 1000))

tests/integration/dbapi/sync/V2/test_server_async.py

@@ -5,7 +5,11 @@
 from pytest import mark, raises
 
 from firebolt.db import Connection
-from firebolt.utils.exception import FireboltError, FireboltStructuredError
+from firebolt.utils.exception import (
+    FireboltError,
+    FireboltStructuredError,
+    MethodNotAllowedInAsyncError,
+)
 
 LONG_SELECT = "SELECT checksum(*) FROM GENERATE_SERIES(1, 2500000000)"  # approx 3 sec
 
@@ -51,6 +55,19 @@ def test_insert_async(connection: Connection) -> None:
         cursor.execute(f"DROP TABLE {table_name}")
 
 
+def test_select_async(connection: Connection) -> None:
+    cursor = connection.cursor()
+    cursor.execute_async("SELECT 1")
+    token = cursor.async_query_token
+    assert token is not None, "Async token was not returned"
+    # sleep for 2 sec to make sure the async query is completed
+    time.sleep(2)
+    assert connection.is_async_query_running(token) == False
+    assert connection.is_async_query_successful(token) == True
+    with raises(MethodNotAllowedInAsyncError):
+        cursor.fetchall()
+
+
 def test_insert_async_running(connection: Connection) -> None:
     cursor = connection.cursor()
     rnd_suffix = str(randint(0, 1000))

tests/unit/common/cursor/test_statement_planners.py

@@ -9,6 +9,7 @@
 from firebolt.common.constants import (
     JSON_LINES_OUTPUT_FORMAT,
     JSON_OUTPUT_FORMAT,
+    NULL_OUTPUT_FORMAT,
 )
 from firebolt.common.cursor.statement_planners import (
     BaseStatementPlanner,
@@ -63,15 +64,24 @@ def test_execution_plan_defaults():
 
 # BaseStatementPlanner tests
 @pytest.mark.parametrize(
-    "streaming,expected_format",
+    "streaming,async_execution,expected_format",
     [
-        (True, JSON_LINES_OUTPUT_FORMAT),
-        (False, JSON_OUTPUT_FORMAT),
+        (True, False, JSON_LINES_OUTPUT_FORMAT),
+        (False, False, JSON_OUTPUT_FORMAT),
+        (False, True, NULL_OUTPUT_FORMAT),
     ],
 )
-def test_get_output_format(streaming, expected_format):
+def test_get_output_format(streaming, async_execution, expected_format):
     """Test output format selection."""
-    assert BaseStatementPlanner._get_output_format(streaming) == expected_format
+    assert (
+        BaseStatementPlanner._get_output_format(streaming, async_execution)
+        == expected_format
+    )
+
+
+def test_get_output_format_rejects_streaming_async():
+    with pytest.raises(AssertionError):
+        BaseStatementPlanner._get_output_format(True, True)
 
 
 # FbNumericStatementPlanner tests
@@ -126,8 +136,7 @@ def test_fb_numeric_execution_plan_no_parameters(fb_numeric_planner):
     "streaming,async_execution,expected_format,expected_async",
     [
         (True, False, JSON_LINES_OUTPUT_FORMAT, None),
-        (False, True, JSON_OUTPUT_FORMAT, True),
-        (True, True, JSON_LINES_OUTPUT_FORMAT, True),
+        (False, True, NULL_OUTPUT_FORMAT, True),
         (False, False, JSON_OUTPUT_FORMAT, None),
     ],
 )
@@ -213,8 +222,7 @@ def test_qmark_execution_plan_skip_parsing(qmark_planner):
     "streaming,async_execution,expected_format,expected_async",
     [
         (True, False, JSON_LINES_OUTPUT_FORMAT, None),
-        (False, True, JSON_OUTPUT_FORMAT, True),
-        (True, True, JSON_LINES_OUTPUT_FORMAT, True),
+        (False, True, NULL_OUTPUT_FORMAT, True),
         (False, False, JSON_OUTPUT_FORMAT, None),
     ],
 )

tests/unit/db_conftest.py

@@ -16,6 +16,7 @@
 from firebolt.common.constants import (
     JSON_LINES_OUTPUT_FORMAT,
     JSON_OUTPUT_FORMAT,
+    NULL_OUTPUT_FORMAT,
     REMOVE_PARAMETERS_HEADER,
     RESET_SESSION_HEADER,
     TRANSACTION_ID_SETTING,
@@ -319,7 +320,7 @@ def async_query_url(engine_url: str, db_name: str) -> URL:
     return URL(
         f"https://{engine_url}/",
         params={
-            "output_format": JSON_OUTPUT_FORMAT,
+            "output_format": NULL_OUTPUT_FORMAT,
             "database": db_name,
             "async": "true",
         },
@@ -852,7 +853,7 @@ def async_query_callback(async_token: str) -> Callable:
     def do_query(request: Request, **kwargs) -> Response:
         assert request.read() != b""
         assert request.method == "POST"
-        assert f"output_format={JSON_OUTPUT_FORMAT}" in str(request.url)
+        assert f"output_format={NULL_OUTPUT_FORMAT}" in str(request.url)
         assert "async=true" in str(request.url)
         query_response = {
             "message": "the query was accepted for async processing",

tests/unit/utils/test_utils.py

@@ -1,7 +1,12 @@
 import pytest
 from httpx import Response, codes
 
-from firebolt.utils.util import get_internal_error_code, parse_url_and_params
+from firebolt.utils.exception import FireboltError, FireboltStructuredError
+from firebolt.utils.util import (
+    get_internal_error_code,
+    parse_url_and_params,
+    raise_error_from_response,
+)
 
 
 @pytest.mark.parametrize(
@@ -65,3 +70,34 @@ def test_parse_url_and_params_multiple_values(url):
     with pytest.raises(ValueError) as excinfo:
         parse_url_and_params(url)
     assert "Multiple values found for key 'param1'" in str(excinfo.value)
+
+
+STRUCTURED_ERROR_BODY = (
+    '{"errors": [{"description": "Line 1, Column 1: syntax error"}],'
+    ' "query": {"query_id": "id"}}'
+)
+
+
+@pytest.mark.parametrize("content_type", ["text/plain", "application/json"])
+def test_raise_error_from_response_structured(content_type):
+    resp = Response(
+        codes.BAD_REQUEST,
+        text=STRUCTURED_ERROR_BODY,
+        headers={"Content-Type": content_type},
+    )
+    with pytest.raises(FireboltStructuredError) as exc_info:
+        raise_error_from_response(resp)
+    assert "syntax error" in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+    "text", ["Plaintext error message", '{"message": "no errors key"}', "[1]", ""]
+)
+def test_raise_error_from_response_plaintext(text):
+    resp = Response(
+        codes.BAD_REQUEST, text=text, headers={"Content-Type": "text/plain"}
+    )
+    with pytest.raises(FireboltError) as exc_info:
+        raise_error_from_response(resp)
+    assert not isinstance(exc_info.value, FireboltStructuredError)
+    assert str(exc_info.value) == text