Skip to content

Commit 8974e7b

Browse files
committed
feat(kernel): consume the folded directResults execute()
The kernel's `Statement.execute()` is now the directResults call (databricks/databricks-sql-kernel#136): it returns a SINGLE `ExecutedAsyncStatement` -- seeded with the inline result when the query finished within the server inline wait (fast path, zero extra round-trips), or a poll/cancel handle when still running -- instead of always blocking to terminal. `KernelDatabricksClient.execute_command` now drives ONE uniform path (no `hasattr(await_result)` arm to feature-detect): register the handle (so `cursor.cancel()` / `close()` reach it), then `await_result()` to drive it to a ready result set -- preserving execute()'s blocking contract. `cursor.execute()` behaviour is unchanged for callers (still blocks to a ready result set); this just adapts to the kernel's new single-handle return. Mid-run cancel still works via the `_sync_cancellers` canceller registered before execute. Tested e2e against live warehouses (auto-closing and non-auto-closing): SELECT 1, range(N), and CREATE+count all succeed on both. The kernel maps an auto-closed `CLOSED` statement to `Succeeded`, so the connector sees a uniform success regardless of warehouse auto-close behaviour. Co-authored-by: Isaac Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>
1 parent 85f8ba3 commit 8974e7b

1 file changed

Lines changed: 19 additions & 24 deletions

File tree

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

Lines changed: 19 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -472,20 +472,26 @@ def execute_command(
472472
except Exception:
473473
# Canceller is best-effort; never block execute on it.
474474
pass
475-
executed = stmt.execute()
476-
# Execute succeeded: the kernel now owns the statement
477-
# lifecycle. It auto-closes the server statement when the
478-
# result stream is fully drained (``ExecutedStatement::
479-
# next_batch`` end-of-stream), with the executed handle's
480-
# ``Drop`` as the backstop for partial/abandoned reads.
481-
# So we must NOT close ``stmt`` here: a premature
482-
# ``CloseStatement`` at execute-return broke lazy
483-
# CloudFetch chunk-link fetches (``get_result_chunks``
484-
# against the live statement) for large paginated-link
485-
# results. Closing here is left ONLY for the error path
486-
# below, where no executed handle / result set was
487-
# produced to reap it.
475+
# ``execute()`` is the kernel's directResults call (Thrift /
476+
# JDBC / use_sea model), returning a single
477+
# ``ExecutedAsyncStatement`` — no arm to feature-detect. On the
478+
# fast path the handle is seeded with the inline result, so
479+
# ``await_result()`` returns it with zero extra round-trips; on
480+
# a slow query it is a poll/cancel handle. Register it so
481+
# ``cursor.cancel()`` (cancel_command) and ``close_command`` can
482+
# reach it, then drive it to terminal via ``await_result`` —
483+
# preserving ``execute()``'s blocking contract. The
484+
# ``_sync_cancellers`` canceller registered above also covers a
485+
# cancel issued during this drive.
486+
async_exec = stmt.execute()
487+
command_id = CommandId.from_sea_statement_id(async_exec.statement_id)
488+
cursor.active_command_id = command_id
489+
with self._async_handles_lock:
490+
self._async_handles[command_id.guid] = async_exec
491+
self._async_statements[command_id.guid] = stmt
488492
close_stmt = False
493+
stream = async_exec.await_result()
494+
return self._make_result_set(stream, cursor, command_id)
489495
except Exception as exc:
490496
raise _wrap_kernel_exception("execute_command", exc) from exc
491497
finally:
@@ -501,17 +507,6 @@ def execute_command(
501507
except Exception:
502508
pass
503509

504-
command_id = CommandId.from_sea_statement_id(executed.statement_id)
505-
cursor.active_command_id = command_id
506-
# ``KernelResultSet.__init__`` calls ``arrow_schema()`` which
507-
# can itself raise ``KernelError`` (or, in principle, a PyO3
508-
# native exception) — wrap the construction so callers see a
509-
# mapped PEP 249 exception.
510-
try:
511-
return self._make_result_set(executed, cursor, command_id)
512-
except Exception as exc:
513-
raise _wrap_kernel_exception("execute_command", exc) from exc
514-
515510
def cancel_command(self, command_id: CommandId) -> None:
516511
with self._async_handles_lock:
517512
handle = self._async_handles.get(command_id.guid)

0 commit comments

Comments
 (0)