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",