diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bae4c54e..d56207546 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Release History # Unreleased +- Avoid an extra status poll when retrieving a completed asynchronous query result with the kernel backend. - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads. - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index df52f34a5..996df5d34 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -957,6 +957,7 @@ def get_execution_result( command_id: CommandId, cursor: "Cursor", ) -> "ResultSet": + """Fetch results after the cursor has confirmed successful completion.""" # Prefer the original owning async handle for the first # in-process result stream. The kernel attaches the real # ExecuteStatementAsync telemetry to that handle; attached @@ -966,7 +967,7 @@ def get_execution_result( # derive a reader for re-fetch. This preserves the Thrift-parity # behavior where results remain re-callable until explicit close. # - # Concurrency: the owning handle is shared, and ``await_result()`` + # Concurrency: the owning handle is shared, and ``fetch_result()`` # below runs OUTSIDE the lock, so it must not run on the same # handle a concurrent ``get_query_state`` poll is already using # for ``status()``. Mirror that method's guard here — if a status @@ -1004,7 +1005,8 @@ def get_execution_result( if owner is not None else self._kernel_session.attach_async_statement(command_id.guid) ) - stream = handle.await_result() + # Cursor.get_async_execution_result has already confirmed success. + stream = handle.fetch_result() except Exception as exc: if uses_owning_handle: with self._async_handles_lock: @@ -1013,16 +1015,16 @@ def get_execution_result( # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which # can raise — map that to PEP 249 too. # - # Unlike the ``await_result()`` failure above, we deliberately do + # Unlike the ``fetch_result()`` failure above, we deliberately do # NOT discard the ``_async_result_stream_started`` marker here. - # By this point ``await_result()`` has already succeeded, so the + # By this point ``fetch_result()`` has already succeeded, so the # owning handle's result stream has been started (and may be - # partially consumed); re-awaiting that same handle on a retry is + # partially consumed); re-fetching that same handle on a retry is # not safe. Leaving the marker set routes any retry through the # reader fallback, which re-materialises a fresh stream. # The trade-off is that such a retry loses the async-statement # telemetry — an accepted, narrow gap limited to the case where - # result-set construction fails after a successful await. + # result-set construction fails after a successful fetch. try: return self._make_result_set(stream, cursor, command_id) except Exception as exc: diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py index 15f78bf5f..c15d3ddd0 100644 --- a/src/databricks/sql/backend/kernel/result_set.py +++ b/src/databricks/sql/backend/kernel/result_set.py @@ -6,7 +6,7 @@ - ``ExecutedStatement`` — returned by ``Statement.execute()``. Has a ``statement_id`` and a ``cancel()`` method. - ``ResultStream`` — returned by ``Session.metadata().list_*`` and by - ``ExecutedAsyncStatement.await_result()``. No statement id; no + ``ExecutedAsyncStatement.fetch_result()`` after completion. No statement id; no cancel. Both implement the same three methods this class actually calls: diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index f94aef0c5..74951c52b 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -1285,11 +1285,11 @@ def test_close_command_tolerant_when_handle_missing(): c.close_command(fake_command_id) # must not raise -def _attach_returns(c, *, status=None, status_error=None, await_result=None): +def _attach_returns(c, *, status=None, status_error=None, fetch_result=None): """Wire ``_kernel_session.attach_async_statement`` to return a fake handle. ``status`` is a ``(state, failure)`` tuple for ``handle.status()``; ``status_error`` makes ``attach`` itself raise (e.g. NotFound); - ``await_result`` sets ``handle.await_result()`` return value.""" + ``fetch_result`` sets ``handle.fetch_result()`` return value.""" c._kernel_session = MagicMock() if status_error is not None: c._kernel_session.attach_async_statement.side_effect = status_error @@ -1297,8 +1297,8 @@ def _attach_returns(c, *, status=None, status_error=None, await_result=None): handle = MagicMock() if status is not None: handle.status.return_value = status - if await_result is not None: - handle.await_result.return_value = await_result + if fetch_result is not None: + handle.fetch_result.return_value = fetch_result c._kernel_session.attach_async_statement.return_value = handle return handle @@ -1319,7 +1319,7 @@ def test_get_query_state_returns_succeeded_when_server_404s(): def test_independent_async_read_requires_reader_api(read_result): c = _make_client() c._kernel_session = MagicMock() - owner = MagicMock(spec=["status", "await_result"]) + owner = MagicMock(spec=["status", "fetch_result"]) cid = CommandId.from_sea_statement_id("tracked-async") c._async_handles[cid.guid] = owner c._async_result_stream_started.add(cid.guid) @@ -1440,7 +1440,7 @@ def test_get_execution_result_uses_retained_owning_handle_first(): fake_stream = MagicMock() fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())]) handle = MagicMock() - handle.await_result.return_value = fake_stream + handle.fetch_result.return_value = fake_stream cursor = MagicMock() cursor.arraysize = 100 cursor.buffer_size_bytes = 1024 @@ -1452,7 +1452,7 @@ def test_get_execution_result_uses_retained_owning_handle_first(): assert rs is not None c._kernel_session.attach_async_statement.assert_not_called() - handle.await_result.assert_called_once_with() + handle.fetch_result.assert_called_once_with() assert cid.guid in c._async_result_stream_started @@ -1462,7 +1462,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle(): c = _make_client() fake_stream = MagicMock() fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())]) - handle = _attach_returns(c, await_result=fake_stream) + handle = _attach_returns(c, fetch_result=fake_stream) cursor = MagicMock() cursor.arraysize = 100 cursor.buffer_size_bytes = 1024 @@ -1472,7 +1472,7 @@ def test_get_execution_result_attaches_by_id_when_no_retained_handle(): assert rs is not None c._kernel_session.attach_async_statement.assert_called_once_with("async-1") - handle.await_result.assert_called_once_with() + handle.fetch_result.assert_called_once_with() def test_get_execution_result_owning_handle_failure_can_retry_owning_handle(): @@ -1482,7 +1482,7 @@ def test_get_execution_result_owning_handle_failure_can_retry_owning_handle(): c = _make_client() c._kernel_session = MagicMock() handle = MagicMock() - handle.await_result.side_effect = _FakeKernelError(code="Unavailable") + handle.fetch_result.side_effect = _FakeKernelError(code="Unavailable") cid = CommandId.from_sea_statement_id("async-retry-owning") c._async_handles[cid.guid] = handle @@ -1494,20 +1494,20 @@ def test_get_execution_result_owning_handle_failure_can_retry_owning_handle(): def test_get_execution_result_construction_failure_retains_marker_and_uses_reader(): - """If the owning handle's ``await_result()`` succeeds but result-set + """If the owning handle's ``fetch_result()`` succeeds but result-set construction then raises, the ``_async_result_stream_started`` marker is deliberately left set (the owning stream may be partially - consumed, so re-awaiting it is unsafe). A subsequent call must route - through the derived reader fallback rather than re-awaiting the owning + consumed, so re-fetching it is unsafe). A subsequent call must route + through the derived reader fallback rather than re-fetching the owning handle.""" c = _make_client() c._kernel_session = MagicMock() owning_stream = MagicMock() # ``KernelResultSet.__init__`` calls ``arrow_schema()``; make that - # raise so ``_make_result_set`` fails after a successful await. + # raise so ``_make_result_set`` fails after a successful fetch. owning_stream.arrow_schema.side_effect = _FakeKernelError(code="Internal") owning_handle = MagicMock() - owning_handle.await_result.return_value = owning_stream + owning_handle.fetch_result.return_value = owning_stream cid = CommandId.from_sea_statement_id("async-construct-fail") c._async_handles[cid.guid] = owning_handle @@ -1516,15 +1516,15 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade # Marker stays set even though construction failed. assert cid.guid in c._async_result_stream_started - owning_handle.await_result.assert_called_once_with() + owning_handle.fetch_result.assert_called_once_with() c._kernel_session.attach_async_statement.assert_not_called() - # A retry derives a reader (fresh stream) instead of re-awaiting + # A retry derives a reader (fresh stream) instead of re-fetching # the partially-consumed owning handle. retry_stream = MagicMock() retry_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())]) reader = MagicMock() - reader.await_result.return_value = retry_stream + reader.fetch_result.return_value = retry_stream owning_handle.reader.return_value = reader cursor = MagicMock() cursor.arraysize = 100 @@ -1535,9 +1535,9 @@ def test_get_execution_result_construction_failure_retains_marker_and_uses_reade assert rs is not None owning_handle.reader.assert_called_once_with() c._kernel_session.attach_async_statement.assert_not_called() - reader.await_result.assert_called_once_with() - # The owning handle was not re-awaited on the retry. - owning_handle.await_result.assert_called_once_with() + reader.fetch_result.assert_called_once_with() + # The owning handle was not re-fetched on the retry. + owning_handle.fetch_result.assert_called_once_with() def test_get_execution_result_maps_not_found_to_programming_error(): @@ -1779,9 +1779,9 @@ def test_get_execution_result_is_re_callable(): fake_stream = MagicMock() fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())]) owning_handle = MagicMock() - owning_handle.await_result.return_value = fake_stream + owning_handle.fetch_result.return_value = fake_stream reader = MagicMock() - reader.await_result.return_value = fake_stream + reader.fetch_result.return_value = fake_stream owning_handle.reader.return_value = reader cid = CommandId.from_sea_statement_id("async-recall-twice") c._async_handles[cid.guid] = owning_handle @@ -1793,10 +1793,10 @@ def test_get_execution_result_is_re_callable(): rs2 = c.get_execution_result(cid, cursor=cursor) assert rs1 is not None and rs2 is not None - owning_handle.await_result.assert_called_once_with() + owning_handle.fetch_result.assert_called_once_with() owning_handle.reader.assert_called_once_with() c._kernel_session.attach_async_statement.assert_not_called() - reader.await_result.assert_called_once_with() + reader.fetch_result.assert_called_once_with() # ---------------------------------------------------------------------------