diff --git a/.github/workflows/kernel-e2e.yml b/.github/workflows/kernel-e2e.yml index 3bc5d7579..59da13648 100644 --- a/.github/workflows/kernel-e2e.yml +++ b/.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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 7e5b4bd37..219de41c0 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/KERNEL_REV b/KERNEL_REV index 124a7df3e..98f923813 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -d72c26492ad2fe2385b10eccd02e4ab462eec465 +cb5b11c1e6991e0b1a61f6ae67f7c7dcecfd8a23 diff --git a/README.md b/README.md index 223791562..36162e69b 100644 --- a/README.md +++ b/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 diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index c53fcde5a..0ca7ba6a8 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/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,7 +905,7 @@ 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()`` @@ -902,35 +913,39 @@ def get_execution_result( # 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. diff --git a/tests/unit/test_kernel_async_readers.py b/tests/unit/test_kernel_async_readers.py new file mode 100644 index 000000000..bb770368d --- /dev/null +++ b/tests/unit/test_kernel_async_readers.py @@ -0,0 +1,213 @@ +"""Public async cursor reads through the native kernel and a local HTTP server.""" + +import base64 +import json +import threading +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from types import SimpleNamespace +from urllib.parse import urlsplit + +import pytest + +pa = pytest.importorskip("pyarrow") +pytestmark = pytest.mark.realkernel + + +@pytest.fixture(params=[False, True], ids=["standard", "realtime"]) +def async_connection(request): + kernel = pytest.importorskip("databricks_sql_kernel") + assert getattr( + kernel, "__file__", None + ), "A fake module is shadowing the native kernel" + assert callable( + getattr(kernel.ExecutedAsyncStatement, "reader", None) + ), "These tests require the kernel source revision in KERNEL_REV" + + from databricks import sql + + sessions_path = "/api/2.0/sql/warehouses/test-wh/sessions" + base_path = f"{sessions_path}/test-session" if request.param else "/api/2.0/sql" + statement_path = f"{base_path}/statements/test-statement" + state = SimpleNamespace( + requests=[], + statement_path=statement_path, + hold_next=False, + started=threading.Event(), + release=threading.Event(), + lock=threading.Lock(), + errors=[], + ) + chunks = [] + for index, values in enumerate(([1, 2], [3, 4])): + batch = pa.record_batch([pa.array(values, type=pa.int64())], names=["n"]) + buffer = pa.BufferOutputStream() + with pa.ipc.new_stream(buffer, batch.schema) as writer: + writer.write_batch(batch) + chunk = { + "chunk_index": index, + "row_offset": index * 2, + "row_count": 2, + "attachment": base64.b64encode(buffer.getvalue()).decode(), + } + if index == 0: + chunk["next_chunk_index"] = 1 + chunks.append(chunk) + result = { + "statement_id": "test-statement", + "status": {"state": "SUCCEEDED"}, + "manifest": { + "format": "ARROW_STREAM", + "schema": { + "column_count": 1, + "columns": [ + { + "name": "n", + "type_name": "BIGINT", + "type_text": "BIGINT", + "position": 0, + } + ], + }, + "total_chunk_count": 2, + "total_row_count": 4, + }, + "result": chunks[0], + } + + class Handler(BaseHTTPRequestHandler): + def handle_request(self): + payload = self.rfile.read(int(self.headers.get("Content-Length", 0))) + body = json.loads(payload) if payload else None + path = urlsplit(self.path).path + with state.lock: + state.requests.append((self.command, path, body)) + hold = ( + self.command == "GET" and path == statement_path and state.hold_next + ) + if hold: + state.hold_next = False + if hold: + state.started.set() + if not state.release.wait(10): + state.errors.append("Timed out waiting to release the owner poll") + + status = 200 + realtime = False + if self.command == "GET" and path.startswith("/api/2.0/connector-service/"): + response = {"flags": []} + elif self.command == "POST" and path == sessions_path: + response = {"session_id": "test-session"} + realtime = request.param + elif self.command == "POST" and path == f"{base_path}/statements": + response = { + "statement_id": "test-statement", + "status": {"state": "PENDING"}, + } + elif self.command == "GET" and path == statement_path: + response = result + elif self.command == "GET" and path == f"{statement_path}/result/chunks/1": + response = chunks[1] + elif self.command == "DELETE": + response = {} + else: + status = 404 + response = { + "error_code": "NOT_FOUND", + "message": "statement unavailable", + } + data = json.dumps(response).encode() + self.send_response(status) + self.send_header("Content-Length", str(len(data))) + self.send_header("Content-Type", "application/json") + if realtime: + self.send_header("x-databricks-rt-session", "true") + self.end_headers() + self.wfile.write(data) + + do_POST = handle_request + do_GET = handle_request + do_DELETE = handle_request + + def log_message(self, *_args): + pass + + with ThreadingHTTPServer(("127.0.0.1", 0), Handler) as server: + worker = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + conn = None + try: + conn = sql.connect( + server_hostname=f"http://127.0.0.1:{server.server_port}", + http_path="/sql/1.0/warehouses/test-wh", + access_token="test-token", + use_kernel=True, + enable_telemetry=False, + ) + yield conn, state + finally: + state.release.set() + if conn is not None: + conn.close() + server.shutdown() + worker.join(timeout=5) + assert not state.errors + assert ( + sum( + method == "DELETE" and path == statement_path + for method, path, _ in state.requests + ) + == 1 + ) + + +def _assert_read_routes(state): + assert { + path + for method, path, _ in state.requests + if method == "GET" and path.startswith("/api/2.0/sql/") + } == { + state.statement_path, + f"{state.statement_path}/result/chunks/1", + } + assert not any(method == "DELETE" for method, _, _ in state.requests) + execute_bodies = [ + body + for method, path, body in state.requests + if method == "POST" and path.endswith("/statements") + ] + assert len(execute_bodies) == 1 + assert execute_bodies[0]["wait_timeout"] == "0s" + + +def test_async_cursor_refetch_preserves_session(async_connection): + from databricks.sql.backend.types import CommandState + + conn, state = async_connection + cur = conn.cursor() + cur.execute_async("SELECT id AS n FROM range(1, 5)") + for _ in range(2): + cur.get_async_execution_result() + assert cur.fetchall_arrow().to_pydict() == {"n": [1, 2, 3, 4]} + assert cur.get_query_state() == CommandState.SUCCEEDED + _assert_read_routes(state) + + +def test_async_cursor_reads_overlap_owner_poll(async_connection): + from databricks.sql.backend.types import CommandState + + conn, state = async_connection + cur = conn.cursor() + cur.execute_async("SELECT id AS n FROM range(1, 5)") + state.hold_next = True + with ThreadPoolExecutor(max_workers=1) as pool: + monitor = pool.submit(cur.get_query_state) + try: + assert state.started.wait(5), "Owner poll did not reach the server" + assert cur.get_query_state() == CommandState.SUCCEEDED + cur.get_async_execution_result() + assert cur.fetchall_arrow().to_pydict() == {"n": [1, 2, 3, 4]} + _assert_read_routes(state) + finally: + state.release.set() + assert monitor.result(timeout=5) == CommandState.SUCCEEDED diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index cd90b9398..b285ce6e1 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -1116,6 +1116,25 @@ def test_get_query_state_returns_succeeded_when_server_404s(): assert c.get_query_state(cid) == CommandState.SUCCEEDED +@pytest.mark.parametrize("read_result", [False, True], ids=["status", "result"]) +def test_independent_async_read_requires_reader_api(read_result): + c = _make_client() + c._kernel_session = MagicMock() + owner = MagicMock(spec=["status", "await_result"]) + cid = CommandId.from_sea_statement_id("tracked-async") + c._async_handles[cid.guid] = owner + c._async_result_stream_started.add(cid.guid) + + with pytest.raises(NotSupportedError, match="newer databricks-sql-kernel"): + if read_result: + c.get_execution_result(cid, cursor=MagicMock()) + else: + c.get_query_state(cid) + + c._kernel_session.attach_async_statement.assert_not_called() + assert c._async_handles[cid.guid] is owner + + def test_get_query_state_returns_closed_from_server(): """A closed-but-not-yet-GC'd async command: the server still answers GetStatementStatus with state=CLOSED (200), which flows through the @@ -1155,19 +1174,19 @@ def test_get_query_state_uses_retained_owning_handle_before_result_stream(): handle.status.assert_called_once_with() -def test_get_query_state_concurrent_poll_routes_to_attach_by_id(): +def test_get_query_state_concurrent_poll_uses_reader(): """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() + is claimed) derives a reader 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 + reader = MagicMock() + reader.status.return_value = ("Running", None) + owning_handle.reader.return_value = reader cid = CommandId.from_sea_statement_id("async-status-concurrent") c._async_handles[cid.guid] = owning_handle @@ -1178,30 +1197,30 @@ 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) + assert cid.guid in c._async_status_in_flight 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. + # The re-entrant (concurrent) poll uses a reader derived from the owner. assert reentrant_state["state"] == CommandState.RUNNING - c._kernel_session.attach_async_statement.assert_called_once_with( - "async-status-concurrent" - ) + owning_handle.reader.assert_called_once_with() + c._kernel_session.attach_async_statement.assert_not_called() 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(): +def test_get_query_state_uses_reader_after_result_stream_started(): """Once get_execution_result has claimed the owning handle for result - streaming, status polling falls back to attach-by-id.""" + streaming, status polling derives a reader.""" 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 + reader = MagicMock() + reader.status.return_value = ("Succeeded", None) + owning_handle.reader.return_value = reader cid = CommandId.from_sea_statement_id("async-status-attached") c._async_handles[cid.guid] = owning_handle c._async_result_stream_started.add(cid.guid) @@ -1209,10 +1228,9 @@ def test_get_query_state_attaches_by_id_after_result_stream_started(): 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() + owning_handle.reader.assert_called_once_with() + c._kernel_session.attach_async_statement.assert_not_called() + reader.status.assert_called_once_with() def test_get_execution_result_uses_retained_owning_handle_first(): @@ -1276,12 +1294,12 @@ def test_get_execution_result_owning_handle_failure_can_retry_owning_handle(): c._kernel_session.attach_async_statement.assert_not_called() -def test_get_execution_result_construction_failure_retains_marker_and_attaches_by_id(): +def test_get_execution_result_construction_failure_retains_marker_and_uses_reader(): """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 + through the derived reader fallback rather than re-awaiting the owning handle.""" c = _make_client() c._kernel_session = MagicMock() @@ -1302,13 +1320,13 @@ def test_get_execution_result_construction_failure_retains_marker_and_attaches_b 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 + # A retry derives a reader (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 + reader = MagicMock() + reader.await_result.return_value = retry_stream + owning_handle.reader.return_value = reader cursor = MagicMock() cursor.arraysize = 100 cursor.buffer_size_bytes = 1024 @@ -1316,10 +1334,9 @@ def test_get_execution_result_construction_failure_retains_marker_and_attaches_b 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() + owning_handle.reader.assert_called_once_with() + c._kernel_session.attach_async_statement.assert_not_called() + reader.await_result.assert_called_once_with() # The owning handle was not re-awaited on the retry. owning_handle.await_result.assert_called_once_with() @@ -1550,13 +1567,13 @@ def test_kernel_error_during_result_set_construction_is_mapped(): # --------------------------------------------------------------------------- -# get_execution_result uses the owning handle once, then attach-by-id +# get_execution_result uses the owning handle once, then derived readers # --------------------------------------------------------------------------- def test_get_execution_result_is_re_callable(): """The first result fetch uses the owning handle for telemetry; a - second fetch for the same async command re-attaches by id so + second fetch for the same async command derives a reader so Thrift-parity re-fetch still works.""" c = _make_client() c._kernel_session = MagicMock() @@ -1564,9 +1581,9 @@ def test_get_execution_result_is_re_callable(): fake_stream.arrow_schema.return_value = pa.schema([("n", pa.int64())]) 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 + reader = MagicMock() + reader.await_result.return_value = fake_stream + owning_handle.reader.return_value = reader cid = CommandId.from_sea_statement_id("async-recall-twice") c._async_handles[cid.guid] = owning_handle cursor = MagicMock() @@ -1578,10 +1595,9 @@ def test_get_execution_result_is_re_callable(): assert rs1 is not None and rs2 is not None 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() + owning_handle.reader.assert_called_once_with() + c._kernel_session.attach_async_statement.assert_not_called() + reader.await_result.assert_called_once_with() # ---------------------------------------------------------------------------