diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index eeb496f73..56d7ba2ab 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/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. diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py index ed98984c8..15f78bf5f 100644 --- a/src/databricks/sql/backend/kernel/result_set.py +++ b/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) diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py index 700528c82..fd109752f 100644 --- a/tests/e2e/test_kernel_backend.py +++ b/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() diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index 3253c6fff..2991ee193 100644 --- a/tests/unit/test_kernel_client.py +++ b/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() + 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 + 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 + 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,28 @@ 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() + c._kernel_session = MagicMock() + 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 +1091,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 +1113,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() # ---------------------------------------------------------------------------