@@ -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+
131141def _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