fix(kernel): preserve session routing for async readers

.github/workflows/kernel-e2e.yml

@@ -332,6 +332,9 @@ jobs:
         run: |
           $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__)"
 
+      - name: Run async reader routing tests
+        run: $CONNECTOR_VENV_PY -m pytest tests/unit/test_kernel_async_readers.py -m realkernel -v
+
       - name: Run kernel e2e tests
         run: poetry run pytest tests/e2e/test_kernel_backend.py -v
 

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Release History
 
 # Unreleased
+- Preserve the originating session when re-fetching or concurrently polling async statements on the kernel backend.
 
 # 4.8.0 (2026-10-24)
 - Upgrade Databricks SQL Kernel to 1.3.0 to address security findings.

KERNEL_REV

@@ -1 +1 @@
-d72c26492ad2fe2385b10eccd02e4ab462eec465
+cb5b11c1e6991e0b1a61f6ae67f7c7dcecfd8a23

README.md

@@ -59,6 +59,11 @@ Notes:
 - Authentication supports PAT (`access_token`), OAuth M2M/U2M, and SP-wide
   workload identity federation (`identity_federation_client_id`).
 
+For Lakehouse Real-Time, submit asynchronous statements and poll or fetch results
+through the submitting connection. Re-fetching and concurrent status polling on
+that connection are supported. Submitting on one connection and polling or
+fetching through another connection is unsupported.
+
 
 ```bash
 export DATABRICKS_HOST=********.databricks.com

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

@@ -128,6 +128,16 @@ def _is_not_found(exc: BaseException) -> bool:
     )
 
 
+def _async_statement_reader(owner: Any) -> Any:
+    reader = getattr(owner, "reader", None)
+    if not callable(reader):
+        raise NotSupportedError(
+            "Independent async reads require a newer databricks-sql-kernel "
+            "wheel that exposes ExecutedAsyncStatement.reader()."
+        )
+    return reader()
+
+
 def _catalog_or_none(value: Optional[str]) -> Optional[str]:
     """Map supported all-catalog wildcards to the kernel's unset filter."""
     if value is None or value in ("%", "*"):
@@ -336,13 +346,13 @@ def __init__(
         # ``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.
+        # object, then use derived readers for re-fetch and attach-by-id
+        # for 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
+        # streaming is claimed) is routed to a reader derived from the owner 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
@@ -796,12 +806,10 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         # 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
+        # concurrent polling. Once result streaming has been claimed,
+        # derive a reader from the retained owner. If this connector never
+        # held the owner (a cross-process / restarted-process resume),
+        # attach by statement id. 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:
@@ -815,26 +823,29 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         # 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
+        # through to a reader derived from the owner. 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:
+            owner = self._async_handles.get(command_id.guid)
             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)
+                else owner
             )
             reserved_owning_handle = handle is not None
             if reserved_owning_handle:
                 self._async_status_in_flight.add(command_id.guid)
         try:
             if handle is None:
-                handle = self._kernel_session.attach_async_statement(command_id.guid)
+                handle = (
+                    _async_statement_reader(owner)
+                    if owner is not None
+                    else self._kernel_session.attach_async_statement(command_id.guid)
+                )
             state, failure = handle.status()
         except Exception as exc:
             if _is_not_found(exc):
@@ -861,7 +872,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         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
+            # it, so a concurrent poll that used a derived reader
             # never touches another poller's reservation.
             if reserved_owning_handle:
                 with self._async_handles_lock:
@@ -894,43 +905,47 @@ def get_execution_result(
         # 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
+        # derive a reader 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
+        # ``_async_status_in_flight``), derive an independent reader,
+        # 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.
         #
-        # 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
+        # If this connection does not hold the owning handle,
+        # ``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:
+            owner = self._async_handles.get(command_id.guid)
             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)
+                else owner
             )
             uses_owning_handle = handle is not None
             if uses_owning_handle:
                 self._async_result_stream_started.add(command_id.guid)
         try:
             if handle is None:
-                handle = self._kernel_session.attach_async_statement(command_id.guid)
+                handle = (
+                    _async_statement_reader(owner)
+                    if owner is not None
+                    else self._kernel_session.attach_async_statement(command_id.guid)
+                )
             stream = handle.await_result()
         except Exception as exc:
             if uses_owning_handle:
@@ -946,7 +961,7 @@ def get_execution_result(
         # 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.
+        # reader 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.