Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion docsrc/Connecting_and_queries.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Minor: This should say from which Firebolt version async SELECT is accepted. Older engines still reject it. With this PR as is, they would reject all async queries from this SDK version (see the comment on statement_planners.py). Also, the " - " reads like a broken dash. Two sentences would be cleaner.

Calling fetch methods after ``execute_async`` raises an exception.

Monitoring the query status
----------------------------
Expand Down
1 change: 1 addition & 0 deletions src/firebolt/common/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
16 changes: 12 additions & 4 deletions src/firebolt/common/cursor/statement_planners.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Comment on lines +140 to +141

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Major: This makes every execute_async depend on the engine accepting NULL, not just async SELECT. Engines from before the async SELECT change reject the Null output format without advanced_mode ("Output format 'NULL' is not supported, choose one of: ..."). So after the SDK upgrade, async INSERT / COPY / VACUUM break against any older engine. The "Before merging" note covers staging, but not customers. Firebolt Core users pin image tags. The SDK release is not tied to the engine rollout. I see three options. We keep JSON_Compact as the default and send NULL only for SELECT (when the statement is parsed). Or we fall back to JSON_Compact on that specific format error. Or we hold the release until every supported engine and Core image has the engine change, and document the minimum version. I'd prefer one of the first two, so the SDK stays compatible with older engines.

if streaming:
return JSON_LINES_OUTPUT_FORMAT
return JSON_OUTPUT_FORMAT
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
19 changes: 18 additions & 1 deletion tests/integration/dbapi/async/V2/test_server_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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))
Expand Down
19 changes: 18 additions & 1 deletion tests/integration/dbapi/sync/V2/test_server_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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))
Expand Down
26 changes: 17 additions & 9 deletions tests/unit/common/cursor/test_statement_planners.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
],
)
Expand Down Expand Up @@ -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),
],
)
Expand Down
5 changes: 3 additions & 2 deletions tests/unit/db_conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
},
Expand Down Expand Up @@ -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",
Expand Down
Loading