Fetch completed async results without an extra status request

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.

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:

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:

tests/unit/test_kernel_client.py

@@ -1285,20 +1285,20 @@ 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
         return 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()
 
 
 # ---------------------------------------------------------------------------