Skip to content

Commit 6989189

Browse files
fix(kernel): preserve session routing for async readers (#1001)
* fix(kernel): preserve session routing for async readers Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com> * fix(kernel): update async reader kernel revision Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com> * fix(kernel): pin merged async reader revision Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com> --------- Signed-off-by: Aakash Saravanan <aakash.saravanan@databricks.com>
1 parent e7a02cd commit 6989189

7 files changed

Lines changed: 315 additions & 62 deletions

File tree

‎.github/workflows/kernel-e2e.yml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,9 @@ jobs:
332332
run: |
333333
$CONNECTOR_VENV_PY -c "import databricks_sql_kernel as k; assert k.__file__, 'kernel module has no __file__ — wheel install failed'; print('kernel ok:', k.__file__)"
334334
335+
- name: Run async reader routing tests
336+
run: $CONNECTOR_VENV_PY -m pytest tests/unit/test_kernel_async_readers.py -m realkernel -v
337+
335338
- name: Run kernel e2e tests
336339
run: poetry run pytest tests/e2e/test_kernel_backend.py -v
337340

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Release History
22

33
# Unreleased
4+
- Preserve the originating session when re-fetching or concurrently polling async statements on the kernel backend.
45

56
# 4.8.0 (2026-10-24)
67
- Upgrade Databricks SQL Kernel to 1.3.0 to address security findings.

‎KERNEL_REV‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
d72c26492ad2fe2385b10eccd02e4ab462eec465
1+
cb5b11c1e6991e0b1a61f6ae67f7c7dcecfd8a23

‎README.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,11 @@ Notes:
5959
- Authentication supports PAT (`access_token`), OAuth M2M/U2M, and SP-wide
6060
workload identity federation (`identity_federation_client_id`).
6161

62+
For Lakehouse Real-Time, submit asynchronous statements and poll or fetch results
63+
through the submitting connection. Re-fetching and concurrent status polling on
64+
that connection are supported. Submitting on one connection and polling or
65+
fetching through another connection is unsupported.
66+
6267

6368
```bash
6469
export DATABRICKS_HOST=********.databricks.com

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

Lines changed: 39 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,16 @@ def _is_not_found(exc: BaseException) -> bool:
128128
)
129129

130130

131+
def _async_statement_reader(owner: Any) -> Any:
132+
reader = getattr(owner, "reader", None)
133+
if not callable(reader):
134+
raise NotSupportedError(
135+
"Independent async reads require a newer databricks-sql-kernel "
136+
"wheel that exposes ExecutedAsyncStatement.reader()."
137+
)
138+
return reader()
139+
140+
131141
def _catalog_or_none(value: Optional[str]) -> Optional[str]:
132142
"""Map supported all-catalog wildcards to the kernel's unset filter."""
133143
if value is None or value in ("%", "*"):
@@ -336,13 +346,13 @@ def __init__(
336346
# ``get_execution_result`` use this owning handle before result
337347
# streaming starts so kernel async statement telemetry is
338348
# finalized on the original ``ExecuteStatementAsync`` telemetry
339-
# object, then fall back to attach-by-id for re-fetch /
340-
# cross-process cases.
349+
# object, then use derived readers for re-fetch and attach-by-id
350+
# for cross-process cases.
341351
self._async_handles: Dict[str, Any] = {}
342352
self._async_result_stream_started: Set[str] = set()
343353
# Async ids whose owning-handle ``status()`` poll is currently in
344354
# flight. A second concurrent poll of the same id (before result
345-
# streaming is claimed) is routed to the attach-by-id fallback so
355+
# streaming is claimed) is routed to a reader derived from the owner so
346356
# it gets a fresh kernel handle instead of racing ``status()`` on
347357
# the shared owning handle. Guarded by ``_async_handles_lock``;
348358
# each entry is transient (added before the poll, discarded in a
@@ -796,12 +806,10 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
796806
# connection (including a fresh cursor resuming the id) resolves
797807
# the same owning handle until result streaming is claimed — see
798808
# the concurrency note below for the limits that places on
799-
# concurrent polling. Once result streaming has been claimed, or
800-
# when this connector genuinely never held the handle (a
801-
# cross-process / restarted-process resume), re-attach to the
802-
# statement by id. SEA keys GetStatementStatus purely on the id,
803-
# so a statement the connector no longer holds a handle for is
804-
# still queryable. CLOSED comes straight from the server: after a
809+
# concurrent polling. Once result streaming has been claimed,
810+
# derive a reader from the retained owner. If this connector never
811+
# held the owner (a cross-process / restarted-process resume),
812+
# attach by statement id. CLOSED comes straight from the server: after a
805813
# statement is closed (DELETE) the server still returns 200
806814
# state=CLOSED until the result TTL elapses.
807815
if self._kernel_session is None:
@@ -815,26 +823,29 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
815823
# undefined, we reserve the owning handle for the first poller via
816824
# _async_status_in_flight: a second concurrent poll of the same id (before
817825
# result streaming is claimed) sees the id already in flight and falls
818-
# through to the attach-by-id path, getting its own fresh kernel handle —
819-
# preserving the pre-change behaviour where every caller re-attached by id
820-
# and status() ran on distinct objects. The reservation is transient
826+
# through to a reader derived from the owner. The reservation is transient
821827
# (discarded in the finally below), so serial polls still take the
822828
# telemetry-preserving owning-handle path.
823829
with self._async_handles_lock:
830+
owner = self._async_handles.get(command_id.guid)
824831
handle = (
825832
None
826833
if (
827834
command_id.guid in self._async_result_stream_started
828835
or command_id.guid in self._async_status_in_flight
829836
)
830-
else self._async_handles.get(command_id.guid)
837+
else owner
831838
)
832839
reserved_owning_handle = handle is not None
833840
if reserved_owning_handle:
834841
self._async_status_in_flight.add(command_id.guid)
835842
try:
836843
if handle is None:
837-
handle = self._kernel_session.attach_async_statement(command_id.guid)
844+
handle = (
845+
_async_statement_reader(owner)
846+
if owner is not None
847+
else self._kernel_session.attach_async_statement(command_id.guid)
848+
)
838849
state, failure = handle.status()
839850
except Exception as exc:
840851
if _is_not_found(exc):
@@ -861,7 +872,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
861872
finally:
862873
# Release the owning-handle reservation once this poll's
863874
# status() has completed (or raised). Only the reserver clears
864-
# it, so a concurrent poll that fell through to attach-by-id
875+
# it, so a concurrent poll that used a derived reader
865876
# never touches another poller's reservation.
866877
if reserved_owning_handle:
867878
with self._async_handles_lock:
@@ -894,43 +905,47 @@ def get_execution_result(
894905
# handles intentionally use no-op telemetry, so always
895906
# re-attaching loses the SEA async statement row when the result
896907
# is drained. After the owning result stream has been started,
897-
# attach by id for re-fetch. This preserves the Thrift-parity
908+
# derive a reader for re-fetch. This preserves the Thrift-parity
898909
# behavior where results remain re-callable until explicit close.
899910
#
900911
# Concurrency: the owning handle is shared, and ``await_result()``
901912
# below runs OUTSIDE the lock, so it must not run on the same
902913
# handle a concurrent ``get_query_state`` poll is already using
903914
# for ``status()``. Mirror that method's guard here — if a status
904915
# poll has the owning handle reserved (guid in
905-
# ``_async_status_in_flight``), fall through to attach-by-id and
906-
# get a fresh kernel handle, exactly as an in-flight peer poll
916+
# ``_async_status_in_flight``), derive an independent reader,
917+
# exactly as an in-flight peer poll
907918
# does. In the normal serial flow (poll to terminal, then fetch)
908919
# the reservation is already discarded, so the fetch still takes
909920
# the telemetry-preserving owning-handle path.
910921
#
911-
# If this process does not hold the owning handle (fresh cursor,
912-
# restarted process, already re-fetched, or a concurrent poll
913-
# holds it), ``attach_async_statement`` issues a
922+
# If this connection does not hold the owning handle,
923+
# ``attach_async_statement`` issues a
914924
# GetStatementStatus to seed the handle; a 404 (unknown / aged-out
915925
# id) surfaces as a NotFound KernelError mapped to
916926
# ``ProgrammingError`` below via ``_wrap_kernel_exception``.
917927
if self._kernel_session is None:
918928
raise InterfaceError("get_execution_result requires an open session.")
919929
with self._async_handles_lock:
930+
owner = self._async_handles.get(command_id.guid)
920931
handle = (
921932
None
922933
if (
923934
command_id.guid in self._async_result_stream_started
924935
or command_id.guid in self._async_status_in_flight
925936
)
926-
else self._async_handles.get(command_id.guid)
937+
else owner
927938
)
928939
uses_owning_handle = handle is not None
929940
if uses_owning_handle:
930941
self._async_result_stream_started.add(command_id.guid)
931942
try:
932943
if handle is None:
933-
handle = self._kernel_session.attach_async_statement(command_id.guid)
944+
handle = (
945+
_async_statement_reader(owner)
946+
if owner is not None
947+
else self._kernel_session.attach_async_statement(command_id.guid)
948+
)
934949
stream = handle.await_result()
935950
except Exception as exc:
936951
if uses_owning_handle:
@@ -946,7 +961,7 @@ def get_execution_result(
946961
# owning handle's result stream has been started (and may be
947962
# partially consumed); re-awaiting that same handle on a retry is
948963
# not safe. Leaving the marker set routes any retry through the
949-
# attach-by-id fallback, which re-materialises a fresh stream.
964+
# reader fallback, which re-materialises a fresh stream.
950965
# The trade-off is that such a retry loses the async-statement
951966
# telemetry — an accepted, narrow gap limited to the case where
952967
# result-set construction fails after a successful await.

0 commit comments

Comments
 (0)