From eedb55f51a93074b2b7f97cc9ccf716e437f6c8c Mon Sep 17 00:00:00 2001 From: Pascal Schulze Date: Wed, 21 Oct 2026 12:05:41 +0200 Subject: [PATCH 1/3] feat(NoTicket): send NULL output format for server-side async queries Newer Firebolt versions accept async SELECT queries when no result serialization is requested, so async queries are now sent with output_format=NULL. Streaming combined with async execution is rejected. Co-Authored-By: Claude Opus 5.5 --- docsrc/Connecting_and_queries.rst | 3 ++- src/firebolt/common/constants.py | 1 + .../common/cursor/statement_planners.py | 16 +++++++++--- .../dbapi/async/V2/test_server_async.py | 19 +++++++++++++- .../dbapi/sync/V2/test_server_async.py | 19 +++++++++++++- .../common/cursor/test_statement_planners.py | 26 ++++++++++++------- tests/unit/db_conftest.py | 5 ++-- 7 files changed, 71 insertions(+), 18 deletions(-) diff --git a/docsrc/Connecting_and_queries.rst b/docsrc/Connecting_and_queries.rst index 2250bf3e63b..3e011307293 100644 --- a/docsrc/Connecting_and_queries.rst +++ b/docsrc/Connecting_and_queries.rst @@ -798,7 +798,8 @@ 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. + Async SELECT queries are accepted, but there is currently no way to retrieve their results - execute_async is best suited for DML queries. + Calling fetch methods after ``execute_async`` raises an exception. Monitoring the query status ---------------------------- diff --git a/src/firebolt/common/constants.py b/src/firebolt/common/constants.py index 4203e2ed5a9..fd311ec4297 100644 --- a/src/firebolt/common/constants.py +++ b/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): diff --git a/src/firebolt/common/cursor/statement_planners.py b/src/firebolt/common/cursor/statement_planners.py index 162dacf9854..e1e4a813ff7 100644 --- a/src/firebolt/common/cursor/statement_planners.py +++ b/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 diff --git a/tests/integration/dbapi/async/V2/test_server_async.py b/tests/integration/dbapi/async/V2/test_server_async.py index 5c5c807c153..eda5e2ccc18 100644 --- a/tests/integration/dbapi/async/V2/test_server_async.py +++ b/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)) diff --git a/tests/integration/dbapi/sync/V2/test_server_async.py b/tests/integration/dbapi/sync/V2/test_server_async.py index a0b7788c9bf..91bdca691ee 100644 --- a/tests/integration/dbapi/sync/V2/test_server_async.py +++ b/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)) diff --git a/tests/unit/common/cursor/test_statement_planners.py b/tests/unit/common/cursor/test_statement_planners.py index 5cf758becf0..fe0f89a9660 100644 --- a/tests/unit/common/cursor/test_statement_planners.py +++ b/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), ], ) diff --git a/tests/unit/db_conftest.py b/tests/unit/db_conftest.py index d9e9ed646cf..69d9e02645a 100644 --- a/tests/unit/db_conftest.py +++ b/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", From 6fb16f7c0fcc476107bd08c0489c869a3d28c297 Mon Sep 17 00:00:00 2001 From: Pascal Schulze Date: Tue, 3 Nov 2026 10:20:52 +0100 Subject: [PATCH 2/3] docs: name the engine version that accepts async SELECT Co-Authored-By: Claude Opus 5.5 --- docsrc/Connecting_and_queries.rst | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docsrc/Connecting_and_queries.rst b/docsrc/Connecting_and_queries.rst index 3e011307293..a8d6fff0640 100644 --- a/docsrc/Connecting_and_queries.rst +++ b/docsrc/Connecting_and_queries.rst @@ -798,8 +798,9 @@ Trying to access `async_query_token` before calling `execute_async` will raise a separately using multiple `execute_async` calls. .. note:: - Async SELECT queries are accepted, but there is currently no way to retrieve their results - execute_async is best suited for DML queries. - Calling fetch methods after ``execute_async`` raises an exception. + 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 ---------------------------- From 291cc5edf902d2fa30d8b14e9f39f0dda8f7e5a8 Mon Sep 17 00:00:00 2001 From: Pascal Schulze Date: Tue, 3 Nov 2026 12:14:00 +0100 Subject: [PATCH 3/3] fix: raise structured errors for JSON bodies sent as text/plain Async queries sent with output_format=NULL get their JSON error bodies labelled text/plain, so they were raised as a plain FireboltError instead of FireboltStructuredError. Co-Authored-By: Claude Opus 5.5 --- src/firebolt/utils/util.py | 24 ++++++++++++++------- tests/unit/utils/test_utils.py | 38 +++++++++++++++++++++++++++++++++- 2 files changed, 54 insertions(+), 8 deletions(-) diff --git a/src/firebolt/utils/util.py b/src/firebolt/utils/util.py index 83958fa3b94..94fbf64c9c4 100644 --- a/src/firebolt/utils/util.py +++ b/src/firebolt/utils/util.py @@ -1,3 +1,4 @@ +import json import logging from functools import lru_cache from os import environ @@ -5,6 +6,7 @@ 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 diff --git a/tests/unit/utils/test_utils.py b/tests/unit/utils/test_utils.py index 5271a7ae047..e3e8b56c5c7 100644 --- a/tests/unit/utils/test_utils.py +++ b/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