Fix kernel async statement telemetry handle

Author: jay-xiao446

src/databricks/sql/backend/kernel/client.py

@@ -25,7 +25,7 @@
 import logging
 import threading
 import uuid
-from typing import Any, Dict, List, Optional, TYPE_CHECKING, Union
+from typing import Any, Dict, List, Optional, Set, TYPE_CHECKING, Union
 
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.kernel._errors import (
@@ -251,16 +251,20 @@ def __init__(
         # concurrent cursors on the same connection don't race on submit /
         # close / close-session.
         #
-        # This is a KEEP-ALIVE registry, not a state/result lookup: the
+        # This is primarily a KEEP-ALIVE registry: the
         # submitting ``ExecutedAsyncStatement``'s ``Drop`` fires a
         # fire-and-forget ``close_statement``, which would kill the
         # still-running async query the moment the handle is dropped. We
         # retain it (and its parent ``Statement``) here so the live query
-        # survives until an explicit close. ``get_query_state`` /
-        # ``get_execution_result`` do NOT consult this map — they
-        # re-attach to the statement by id (the server is the source of
-        # truth for async state), so they work even cross-process.
+        # survives until an explicit close. ``get_query_state`` still
+        # re-attaches to the statement by id (the server is the source
+        # of truth for async state). ``get_execution_result`` uses this
+        # owning handle for the first in-process result stream so kernel
+        # async statement telemetry is finalized on the original
+        # ``ExecuteStatementAsync`` telemetry object, then falls back to
+        # attach-by-id for re-fetch / cross-process cases.
         self._async_handles: Dict[str, Any] = {}
+        self._async_result_stream_started: Set[str] = set()
         # Parent ``Statement`` objects kept alive alongside async handles.
         # On the kernel, ``Statement.close()`` flips the validity flag on
         # the produced executed handle (see kernel
@@ -403,6 +407,7 @@ def close_session(self, session_id: SessionId) -> None:
             tracked_stmts = list(self._async_statements.items())
             self._async_handles.clear()
             self._async_statements.clear()
+            self._async_result_stream_started.clear()
         for _, handle in tracked:
             # Per-handle close errors are non-fatal — PEP 249
             # discourages raising from session close — so log and
@@ -654,6 +659,7 @@ def close_command(self, command_id: CommandId) -> None:
         with self._async_handles_lock:
             handle = self._async_handles.pop(command_id.guid, None)
             stmt = self._async_statements.pop(command_id.guid, None)
+            self._async_result_stream_started.discard(command_id.guid)
         # Closing the handle below fires the server-side CloseStatement.
         # A subsequent ``get_query_state`` re-attaches by id and reads
         # ``CLOSED`` straight from the server — no connector-side
@@ -740,25 +746,39 @@ def get_execution_result(
         command_id: CommandId,
         cursor: "Cursor",
     ) -> "ResultSet":
-        # Re-attach to the statement by id and await its result. SEA keys
-        # GetStatementResult on the id, so this works whether or not the
-        # connector still holds the submitting handle — and it's
-        # inherently re-callable (each call attaches a fresh handle and
-        # re-materialises the result stream), matching the Thrift backend
-        # where the operation handle stays re-fetchable until an explicit
-        # close. No connector-side handle lookup, so no
-        # ``unknown command_id`` failure on a second call.
+        # Prefer the original owning async handle for the first
+        # in-process result stream. The kernel attaches the real
+        # ExecuteStatementAsync telemetry to that handle; attached
+        # handles intentionally use no-op telemetry, so always
+        # re-attaching loses the SEA async statement row when the result
+        # is drained. After the owning result stream has been started,
+        # attach by id for re-fetch. This preserves the Thrift-parity
+        # behavior where results remain re-callable until explicit close.
         #
-        # ``attach_async_statement`` issues a GetStatementStatus to seed
-        # the handle; a 404 (unknown / aged-out id) surfaces as a
-        # NotFound KernelError mapped to ``ProgrammingError`` below via
-        # ``_wrap_kernel_exception``.
+        # If this process does not hold the owning handle (fresh cursor,
+        # restarted process, already re-fetched), ``attach_async_statement``
+        # issues a GetStatementStatus to seed the handle; a 404 (unknown
+        # / aged-out id) surfaces as a NotFound KernelError mapped to
+        # ``ProgrammingError`` below via ``_wrap_kernel_exception``.
         if self._kernel_session is None:
             raise InterfaceError("get_execution_result requires an open session.")
+        with self._async_handles_lock:
+            handle = (
+                None
+                if command_id.guid in self._async_result_stream_started
+                else self._async_handles.get(command_id.guid)
+            )
+            uses_owning_handle = handle is not None
+            if uses_owning_handle:
+                self._async_result_stream_started.add(command_id.guid)
         try:
-            handle = self._kernel_session.attach_async_statement(command_id.guid)
+            if handle is None:
+                handle = self._kernel_session.attach_async_statement(command_id.guid)
             stream = handle.await_result()
         except Exception as exc:
+            if uses_owning_handle:
+                with self._async_handles_lock:
+                    self._async_result_stream_started.discard(command_id.guid)
             raise _wrap_kernel_exception("get_execution_result", exc) from exc
         # ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
         # can raise — map that to PEP 249 too.

src/databricks/sql/backend/kernel/result_set.py

@@ -252,13 +252,10 @@ def close(self) -> None:
             # connection close path stays clean.
             logger.warning("Error closing kernel handle: %s", exc)
         # Honor the base ``ResultSet`` contract: notify the backend.
-        # ``backend.close_command`` also drops the ``_async_handles``
-        # entry and records the guid in ``_closed_commands`` — no
-        # separate pop needed here. Sync-execute and metadata paths
-        # never registered in ``_async_handles`` to begin with, and
-        # ``get_execution_result`` pops the async path before the
-        # result set is even constructed (see the M1 fix), so this
-        # call is the single bookkeeping seam.
+        # For async results, ``backend.close_command`` drops the
+        # retained owning handle and parent Statement. Sync-execute and
+        # metadata paths never registered in ``_async_handles`` to begin
+        # with, so this call is tolerant bookkeeping for them.
         backend = cast("KernelDatabricksClient", self.backend)
         try:
             backend.close_command(self.command_id)

tests/e2e/test_kernel_backend.py

@@ -422,9 +422,10 @@ def test_dml_rowcount_wiring_does_not_break_dml(conn):
 
 def test_async_execute_polls_and_fetches_result(conn):
     """The full async CUJ: ``execute_async`` → poll
-    ``get_query_state`` → ``get_async_execution_result``. State and
-    result are read from the server by re-attaching to the statement
-    id (no connector-side state)."""
+    ``get_query_state`` → ``get_async_execution_result``. State comes
+    from the server by re-attaching to the statement id; first
+    in-process result fetch uses the retained owning handle so kernel
+    async telemetry is finalized."""
     with conn.cursor() as cur:
         cur.execute_async("SELECT 7 AS n")
         cur.get_async_execution_result()  # polls to terminal, fetches
@@ -437,10 +438,9 @@ def test_async_execute_polls_and_fetches_result(conn):
 
 
 def test_async_get_execution_result_is_re_callable(conn):
-    """``get_async_execution_result`` re-attaches by id on each call,
-    so fetching the same async command twice both succeed — the
-    connector never relied on a one-shot retained handle (Thrift-parity
-    re-fetch)."""
+    """Fetching the same async command twice succeeds: the first
+    in-process result fetch can use the owning handle, and later
+    re-fetches attach by id (Thrift-parity re-fetch)."""
     with conn.cursor() as cur:
         cur.execute_async("SELECT 11 AS n")
         cur.get_async_execution_result()

tests/unit/test_kernel_client.py

@@ -799,9 +799,33 @@ def test_get_query_state_propagates_non_not_found_error():
         c.get_query_state(cid)
 
 
-def test_get_execution_result_attaches_by_id():
-    """``get_execution_result`` re-attaches to the statement by id and
-    awaits its result — no connector-side handle lookup."""
+def test_get_execution_result_uses_retained_owning_handle_first():
+    """The first in-process result fetch uses the retained submitting
+    handle so the kernel finalizes the original async statement telemetry."""
+    c = _make_client()
+    fake_stream = MagicMock()
+    fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
+    handle = MagicMock()
+    handle.await_result.return_value = fake_stream
+    cursor = MagicMock()
+    cursor.arraysize = 100
+    cursor.buffer_size_bytes = 1024
+    cursor.row_limit = 5
+    cid = CommandId.from_sea_statement_id("async-1")
+    c._async_handles[cid.guid] = handle
+
+    rs = c.get_execution_result(cid, cursor=cursor)
+
+    assert rs is not None
+    assert rs._row_limit == 5
+    c._kernel_session.attach_async_statement.assert_not_called()
+    handle.await_result.assert_called_once_with()
+    assert cid.guid in c._async_result_stream_started
+
+
+def test_get_execution_result_attaches_by_id_when_no_retained_handle():
+    """Fallback by statement id keeps cross-process / fresh-cursor
+    result retrieval working when this connector lacks the owning handle."""
     c = _make_client()
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
@@ -814,10 +838,27 @@ def test_get_execution_result_attaches_by_id():
     rs = c.get_execution_result(cid, cursor=cursor)
 
     assert rs is not None
-    c._kernel_session.attach_async_statement.assert_called_with("async-1")
+    c._kernel_session.attach_async_statement.assert_called_once_with("async-1")
     handle.await_result.assert_called_once_with()
 
 
+def test_get_execution_result_owning_handle_failure_can_retry_owning_handle():
+    """If the owning handle's await fails before producing a result
+    stream, clear the claimed marker so a retry can still use the
+    telemetry-bearing owning handle."""
+    c = _make_client()
+    handle = MagicMock()
+    handle.await_result.side_effect = _FakeKernelError(code="Unavailable")
+    cid = CommandId.from_sea_statement_id("async-retry-owning")
+    c._async_handles[cid.guid] = handle
+
+    with pytest.raises(OperationalError):
+        c.get_execution_result(cid, cursor=MagicMock())
+
+    assert cid.guid not in c._async_result_stream_started
+    c._kernel_session.attach_async_statement.assert_not_called()
+
+
 def test_get_execution_result_maps_not_found_to_programming_error():
     """An unknown / aged-out id surfaces the kernel's NotFound as a
     mapped PEP 249 exception rather than a raw error."""
@@ -1049,19 +1090,20 @@ def test_kernel_error_during_result_set_construction_is_mapped():
 
 
 def test_get_execution_result_is_re_callable():
-    """``get_execution_result`` re-attaches by id on every call, so a
-    second fetch for the same async command succeeds (Thrift-parity
-    re-fetch). Each call attaches a fresh handle and awaits its result;
-    neither raises, and the connector never depended on a retained
-    handle. The kernel's ``await_result()`` is idempotent server-side."""
+    """The first result fetch uses the owning handle for telemetry; a
+    second fetch for the same async command re-attaches by id so
+    Thrift-parity re-fetch still works."""
     c = _make_client()
     c._kernel_session = MagicMock()
     fake_stream = MagicMock()
     fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())])
-    handle = MagicMock()
-    handle.await_result.return_value = fake_stream
-    c._kernel_session.attach_async_statement.return_value = handle
+    owning_handle = MagicMock()
+    owning_handle.await_result.return_value = fake_stream
+    attached_handle = MagicMock()
+    attached_handle.await_result.return_value = fake_stream
+    c._kernel_session.attach_async_statement.return_value = attached_handle
     cid = CommandId.from_sea_statement_id("async-recall-twice")
+    c._async_handles[cid.guid] = owning_handle
     cursor = MagicMock()
     cursor.arraysize = 100
     cursor.buffer_size_bytes = 1024
@@ -1070,10 +1112,11 @@ 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
-    # Two calls -> two attaches -> two await_results. No reliance on a
-    # connector-tracked handle.
-    assert c._kernel_session.attach_async_statement.call_count == 2
-    assert handle.await_result.call_count == 2
+    owning_handle.await_result.assert_called_once_with()
+    c._kernel_session.attach_async_statement.assert_called_once_with(
+        "async-recall-twice"
+    )
+    attached_handle.await_result.assert_called_once_with()
 
 
 # ---------------------------------------------------------------------------