Use owning async handle for kernel query state

Author: jay-xiao446

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

@@ -256,13 +256,12 @@ def __init__(
         # 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`` 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.
+        # survives until an explicit close. ``get_query_state`` and
+        # ``get_execution_result`` use this owning handle before result
+        # streaming starts so kernel async statement telemetry is
+        # finalized on the original ``ExecuteStatementAsync`` telemetry
+        # object, then fall 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.
@@ -689,18 +688,28 @@ def close_command(self, command_id: CommandId) -> None:
                     pass
 
     def get_query_state(self, command_id: CommandId) -> CommandState:
-        # Server is the source of truth for async command state. Re-attach
-        # to the statement by its id and read the state the server reports
-        # — no connector-side state to drift. SEA keys GetStatementStatus
-        # purely on the id, so a statement the connector no longer holds a
-        # handle for (or never held — a different process) is still
-        # queryable. CLOSED comes straight from the server: after a
+        # Server is the source of truth for async command state. Use the
+        # retained owning handle before result streaming starts so kernel
+        # async statement telemetry is finalized on the original
+        # ExecuteStatementAsync telemetry object. Once result streaming
+        # has been claimed (or when this connector never held the handle
+        # — cross-process / fresh-cursor cases), re-attach to the
+        # statement by id. SEA keys GetStatementStatus purely on the id,
+        # so a statement the connector no longer holds a handle for is
+        # still queryable. CLOSED comes straight from the server: after a
         # statement is closed (DELETE) the server still returns 200
         # state=CLOSED until the result TTL elapses.
         if self._kernel_session is None:
             raise InterfaceError("get_query_state 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)
+            )
         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)
             state, failure = handle.status()
         except Exception as exc:
             if _is_not_found(exc):

tests/unit/test_kernel_client.py

@@ -799,6 +799,45 @@ def test_get_query_state_propagates_non_not_found_error():
         c.get_query_state(cid)
 
 
+def test_get_query_state_uses_retained_owning_handle_before_result_stream():
+    """In-process status polling uses the retained submitting handle so
+    kernel async statement telemetry stays attached to the original
+    ExecuteStatementAsync telemetry object."""
+    c = _make_client()
+    c._kernel_session = MagicMock()
+    handle = MagicMock()
+    handle.status.return_value = ("Running", None)
+    cid = CommandId.from_sea_statement_id("async-status-owning")
+    c._async_handles[cid.guid] = handle
+
+    assert c.get_query_state(cid) == CommandState.RUNNING
+
+    c._kernel_session.attach_async_statement.assert_not_called()
+    handle.status.assert_called_once_with()
+
+
+def test_get_query_state_attaches_by_id_after_result_stream_started():
+    """Once get_execution_result has claimed the owning handle for result
+    streaming, status polling falls back to attach-by-id."""
+    c = _make_client()
+    c._kernel_session = MagicMock()
+    owning_handle = MagicMock()
+    attached_handle = MagicMock()
+    attached_handle.status.return_value = ("Succeeded", None)
+    c._kernel_session.attach_async_statement.return_value = attached_handle
+    cid = CommandId.from_sea_statement_id("async-status-attached")
+    c._async_handles[cid.guid] = owning_handle
+    c._async_result_stream_started.add(cid.guid)
+
+    assert c.get_query_state(cid) == CommandState.SUCCEEDED
+
+    owning_handle.status.assert_not_called()
+    c._kernel_session.attach_async_statement.assert_called_once_with(
+        "async-status-attached"
+    )
+    attached_handle.status.assert_called_once_with()
+
+
 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."""