diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 01274078f..93b0a98a4 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,27 @@ 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`` 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() + # Async ids whose owning-handle ``status()`` poll is currently in + # flight. A second concurrent poll of the same id (before result + # streaming is claimed) is routed to the attach-by-id fallback so + # it gets a fresh kernel handle instead of racing ``status()`` on + # the shared owning handle. Guarded by ``_async_handles_lock``; + # each entry is transient (added before the poll, discarded in a + # ``finally``). + self._async_status_in_flight: 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 @@ -406,6 +417,8 @@ 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() + self._async_status_in_flight.clear() for _, handle in tracked: # Per-handle close errors are non-fatal — PEP 249 # discourages raising from session close — so log and @@ -657,6 +670,8 @@ 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) + self._async_status_in_flight.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 @@ -686,18 +701,53 @@ 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. The owning-handle path + # is per-connection, not per-cursor: any cursor on the submitting + # connection (including a fresh cursor resuming the id) resolves + # the same owning handle until result streaming is claimed — see + # the concurrency note below for the limits that places on + # concurrent polling. Once result streaming has been claimed, or + # when this connector genuinely never held the handle (a + # cross-process / restarted-process resume), 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.") + # Concurrency note: the lock guards the _async_handles / + # _async_result_stream_started / _async_status_in_flight bookkeeping only. + # The retained owning handle it returns is a shared object, and + # handle.status() below runs OUTSIDE the lock, so it is not safe to invoke + # status() on one owning handle from two threads at once. Rather than leave + # concurrent in-process polling of a single async id "unsupported" and + # undefined, we reserve the owning handle for the first poller via + # _async_status_in_flight: a second concurrent poll of the same id (before + # result streaming is claimed) sees the id already in flight and falls + # through to the attach-by-id path, getting its own fresh kernel handle — + # preserving the pre-change behaviour where every caller re-attached by id + # and status() ran on distinct objects. The reservation is transient + # (discarded in the finally below), so serial polls still take the + # telemetry-preserving owning-handle path. + with self._async_handles_lock: + handle = ( + None + if ( + command_id.guid in self._async_result_stream_started + or command_id.guid in self._async_status_in_flight + ) + else self._async_handles.get(command_id.guid) + ) + reserved_owning_handle = handle is not None + if reserved_owning_handle: + self._async_status_in_flight.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) state, failure = handle.status() except Exception as exc: if _is_not_found(exc): @@ -721,6 +771,14 @@ def get_query_state(self, command_id: CommandId) -> CommandState: # sync-fall-through behaviour. return CommandState.SUCCEEDED raise _wrap_kernel_exception("get_query_state", exc) from exc + finally: + # Release the owning-handle reservation once this poll's + # status() has completed (or raised). Only the reserver clears + # it, so a concurrent poll that fell through to attach-by-id + # never touches another poller's reservation. + if reserved_owning_handle: + with self._async_handles_lock: + self._async_status_in_flight.discard(command_id.guid) if state == "Failed" and failure is not None: # Surface server-reported failure as a database error so # the cursor's polling loop terminates with the right @@ -743,28 +801,68 @@ 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. + # + # Concurrency: the owning handle is shared, and ``await_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 + # poll has the owning handle reserved (guid in + # ``_async_status_in_flight``), fall through to attach-by-id and + # get a fresh kernel handle, exactly as an in-flight peer poll + # does. In the normal serial flow (poll to terminal, then fetch) + # the reservation is already discarded, so the fetch still takes + # the telemetry-preserving owning-handle path. # - # ``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, or a concurrent poll + # holds it), ``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 + or command_id.guid in self._async_status_in_flight + ) + 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. + # + # Unlike the ``await_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 + # owning handle's result stream has been started (and may be + # partially consumed); re-awaiting that same handle on a retry is + # not safe. Leaving the marker set routes any retry through the + # attach-by-id 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. 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 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 6e217f26a..acae90819 100644 --- a/tests/e2e/test_kernel_backend.py +++ b/tests/e2e/test_kernel_backend.py @@ -462,14 +462,14 @@ def test_dml_rowcount_wiring_does_not_break_dml(conn): cur.execute(f"DROP TABLE IF EXISTS {tbl}") -# ── Async execution: state + result come from the server (attach-by-id) ── +# ── Async execution: owning handle first, attach-by-id for re-fetch/resume ── 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``. The first + in-process flow 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 @@ -482,10 +482,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..3eb5a9006 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -799,9 +799,110 @@ 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_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_concurrent_poll_routes_to_attach_by_id(): + """A second concurrent poll of the same async id (while the first + poll's owning-handle status() is in flight, before result streaming + is claimed) falls back to attach-by-id rather than racing status() + on the shared owning handle. The owning-handle reservation is + released once the first poll returns, so a later serial poll takes + the owning-handle path again.""" + c = _make_client() + c._kernel_session = MagicMock() + owning_handle = MagicMock() + attached_handle = MagicMock() + attached_handle.status.return_value = ("Running", None) + c._kernel_session.attach_async_statement.return_value = attached_handle + cid = CommandId.from_sea_statement_id("async-status-concurrent") + c._async_handles[cid.guid] = owning_handle + + # Simulate the first poll being mid-flight: its reservation is set. + reentrant_state = {} + + def owning_status(): + # A concurrent poll arriving while this one holds the reservation + # must not touch the owning handle. + reentrant_state["state"] = c.get_query_state(cid) + return ("Running", None) + + owning_handle.status.side_effect = owning_status + + assert c.get_query_state(cid) == CommandState.RUNNING + # The re-entrant (concurrent) poll fell back to attach-by-id. + assert reentrant_state["state"] == CommandState.RUNNING + c._kernel_session.attach_async_statement.assert_called_once_with( + "async-status-concurrent" + ) + owning_handle.status.assert_called_once_with() + # Reservation released after the first poll returns. + assert cid.guid not in c._async_status_in_flight + + +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.""" + 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 +915,76 @@ 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_construction_failure_retains_marker_and_attaches_by_id(): + """If the owning handle's ``await_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 attach-by-id fallback rather than re-awaiting 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. + owning_stream.arrow_schema.side_effect = _FakeKernelError(code="Internal") + owning_handle = MagicMock() + owning_handle.await_result.return_value = owning_stream + cid = CommandId.from_sea_statement_id("async-construct-fail") + c._async_handles[cid.guid] = owning_handle + + with pytest.raises(DatabaseError): + c.get_execution_result(cid, cursor=MagicMock()) + + # Marker stays set even though construction failed. + assert cid.guid in c._async_result_stream_started + owning_handle.await_result.assert_called_once_with() + c._kernel_session.attach_async_statement.assert_not_called() + + # A retry now attaches by id (fresh stream) instead of re-awaiting + # the partially-consumed owning handle. + retry_stream = MagicMock() + retry_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())]) + attached_handle = MagicMock() + attached_handle.await_result.return_value = retry_stream + c._kernel_session.attach_async_statement.return_value = attached_handle + cursor = MagicMock() + cursor.arraysize = 100 + cursor.buffer_size_bytes = 1024 + + rs = c.get_execution_result(cid, cursor=cursor) + + assert rs is not None + c._kernel_session.attach_async_statement.assert_called_once_with( + "async-construct-fail" + ) + attached_handle.await_result.assert_called_once_with() + # The owning handle was not re-awaited on the retry. + owning_handle.await_result.assert_called_once_with() + + 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.""" @@ -1044,24 +1211,25 @@ def test_kernel_error_during_result_set_construction_is_mapped(): # --------------------------------------------------------------------------- -# get_execution_result is re-callable via attach-by-id +# get_execution_result uses the owning handle once, then attach-by-id # --------------------------------------------------------------------------- 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 +1238,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() # ---------------------------------------------------------------------------