`yggdrasil.databricks.cluster.statement_executor`¶

statement_executor ¶

Run SQL through a Databricks cluster as a :class:StatementExecutor.

The executor is a thin wrapper: it owns a :class:Cluster (which it delegates the actual REPL command execution to) and a :class:Volume (used to stage SELECT result data — see "SELECT handling" below). All heavy lifting lives on the cluster's :class:ExecutionContext / :class:CommandExecution plumbing; this class just adapts the yggdrasil.data.executor.StatementExecutor contract on top.

SELECT handling — INSERT OVERWRITE DIRECTORY¶

A SQL REPL on a cluster ships results back as text — fine for DDL / DML but unsafe for SELECTs (the REPL truncates above ~25 MiB and serialises Python objects rather than streaming Arrow). The warehouse path solves this with EXTERNAL_LINKS-disposition Arrow chunks; the cluster path does not have that surface.

Instead, when the executor sees a query that :meth:PreparedStatement.looks_like_query accepts, it rewrites it to

INSERT OVERWRITE DIRECTORY '<volume>/<op_id>' USING parquet
<original-query>

(the documented Databricks DML grammar — see the SQL reference at https://docs.databricks.com/aws/en/sql/language-manual/sql-ref-syntax-dml-insert-overwrite-directory). The SELECT's rows are landed as Parquet under the bound volume; the result reads them back through pyarrow.dataset from :meth:ClusterStatementResult._read_arrow_batches. The staged folder is bound to the statement so the batch's clear_temporary_resources walk unlinks it on success.

Safety / reliability guardrails baked in¶

Execution-context reuse. Each cluster caps at 145 user REPL contexts before new notebooks fail to attach (https://kb.databricks.com/clusters/too-many-execution-contexts-are-open-right-now). The executor keys the context off the bound volume's full path so every statement on the same volume shares one context. Callers that need isolation can pass an explicit context_key on the prepared statement.
Unity-Catalog-only staging. Serverless restrictions push every user toward Unity Catalog volumes (no DBFS); the executor enforces this by typing the volume parameter as :class:Volume and using /Volumes/... paths exclusively.
No silent result-size truncation. Because SELECTs go to Parquet on a volume, there is no REPL stdout cap that could truncate without raising — the Parquet read either returns every row or raises if the directory is malformed.

ClusterStatementExecutor ¶

ClusterStatementExecutor(
    cluster: "Cluster",
    volume: "Volume",
    *,
    default_language: Language = Language.SQL,
    default_context_key: Optional[str] = None
)

Bases: DatabricksResource, StatementExecutor[ClusterPreparedStatement, ClusterStatementResult, ClusterStatementBatch]

Cluster-backed :class:StatementExecutor.

Wraps a :class:Cluster (the actual execution backend) and a :class:Volume (used to stage SELECT result Parquet folders). Inherits :meth:execute / :meth:execute_many / :meth:batch from the base — only :meth:_submit_statement is implemented locally, plus the SELECT-rewrite helper.

Singleton-cached on (cluster, volume) so every statement against the same staging surface lands on one executor — same REPL context, same cluster session, same per-instance cache. The cluster cap of 145 user execution contexts is one of the main reasons for this: an unintended duplicate executor is an unintended second context.

opened `property` ¶

opened: bool

True iff :meth:_acquire has run and :meth:_release hasn't.

closed `property` ¶

closed: bool

Inverse of :attr:opened.

sql `property` ¶

sql: 'SQLEngine'

Shorthand for self.service.client.sql — the active :class:SQLEngine.

submit_command ¶

submit_command(statement: ClusterPreparedStatement) -> 'CommandExecution'

Build a :class:CommandExecution from statement (no start).

Exposed so :class:ClusterStatementResult.start can mint the command lazily — the result owns the lifecycle, the executor only owns the construction.

open ¶

open() -> 'Disposable'

Acquire the resource and cascade into owned children.

Order:

Run our own :meth:_acquire (subclass body).
Flip :attr:opened to True and mark _self_opened.
For each owned child, in registration order:
If the child is already opened, just :meth:_claim it. It stays self-opened — the existing self-open is what keeps it alive after we let go.
Otherwise, call :meth:open on the child (which recursively cascades into ITS owned children), then clear the child's _self_opened flag so the child knows its open is parent-driven, then :meth:_claim it. Without that flag clear, the eventual :meth:_unclaim would refuse to close — it would see "I'm self-opened, someone explicitly opened me, leave me alone."

Both branches record the child in our per-frame scratch list so :meth:_release knows what to unclaim.

Transactional rollback: if any child's open or claim raises, we walk back through the children we already touched (in reverse), unclaim each, then call our own :meth:_release with committed=False and re-raise the original exception. From the caller's view, the open atomically either succeeded with the whole graph live, or failed with nothing changed.

Not reentrant: raises :class:RuntimeError if already opened. Nesting is expressed via with self: blocks, not via paired :meth:open calls.

commit ¶

commit()

Commit current state

rollback ¶

rollback()

Rollback current state

close ¶

close(force: bool = False) -> None

Release the resource and cascade into owned children.

Order:

If currently held open by an outside parent claim (_claim_count > 0) AND we are not in self-opened state, this is a no-op — the parents that opened us still need us live. (Handled inside :meth:_do_close.)
Walk our scratch list of acquired children in REVERSE registration order; :meth:_unclaim each. A child whose claim count hits zero and isn't otherwise self-opened closes itself.
Run :meth:_before_release, then :meth:_release — with committed reflecting the dirty bit (cleared on exception by __exit__).

Idempotent: no-op when already closed, unless force.

force=True runs teardown even when :attr:closed. Intended for error-recovery paths where subclass state might be inconsistent.

Does NOT touch :attr:depth — the with-stack counter belongs to :meth:__enter__/:meth:__exit__ exclusively. If a caller calls :meth:close inside an active with block, the outer :meth:__exit__ will harmlessly skip the now-no-op close on unwind.

mark_dirty ¶

mark_dirty() -> None

Signal pending mutations — commit on next clean :meth:close.

to_singleton ¶

to_singleton(ttl: Any = ...) -> 'Singleton'

Promote this instance into the per-class _INSTANCES cache.

Hot listing paths (iterdir / _ls / glob) build children with singleton_ttl=False so the bounded cache doesn't fill up with thousands of short-lived entries. When a caller decides one of those children is worth keeping around (handing it to a long-running worker, returning it from an API), :meth:to_singleton registers self into the cache so the next constructor call with the same key collapses to the same instance.

ttl defaults to the subclass's _SINGLETON_TTL (... = no caching, None = process lifetime, or a seconds count). When a different instance is already cached under this key, that pre-existing one wins and is returned unchanged — the cache is the source of truth.

invalidate_singleton ¶

invalidate_singleton(remove_global: bool = True) -> None

Pop self from the per-class _INSTANCES cache.

Mutating ops on a Singleton-cached object (writes, deletes, schema invalidations on a Databricks table, put_object on an :class:S3Path) want to make sure the next caller asking for the same key gets a fresh build rather than collapsing onto this stale handle — that's what remove_global=True (the default) does. The pop is :meth:identity-guarded: only an entry that still points at self is removed, so a concurrent re-construction that already raced past this thread is left alone.

remove_global=False is a no-op. The keyword exists so subclass invalidators (invalidate_singleton, _invalidate_entity_tag_cache, …) can offer the same switch without branching at the call site.

send ¶

send(statement: 'PS | PreparedStatement | str', *, start: bool = True) -> SR

Dispatch statement and return its tracking :class:StatementResult.

Mirrors :meth:Session.send. start=True (default) fires the backend submission eagerly — the result comes back in flight (or already terminal for synchronous backends). start=False returns the idled :class:StatementResult whose backend submission is deferred until :meth:StatementResult.start fires.

The returned result is always bound to this executor — every subclass _submit_statement is supposed to thread executor=self through the constructor, but that's easy to forget and downstream code (StatementResult.wait, retry, raise_for_status) needs the back-reference. Setting it here when it's missing makes the contract enforceable from one place instead of audited per backend.

execute ¶

execute(
    statement: "PS | PreparedStatement | str",
    *,
    options: Optional[ExecutionOptions] = None,
    wait: WaitingConfigArg = True,
    raise_error: bool = True,
    start: bool = True
) -> SR

Submit a single statement and optionally wait for completion.

Two ways to pass execution policy:

Per-call kwargs wait / raise_error (ergonomic, matches the previous public API).
An :class:ExecutionOptions via options= (when you want to reuse the same policy across many calls or compose from layered defaults).

The two are merged: options provides the base, kwargs override any field they explicitly set. Unknown kwargs go nowhere — they are not forwarded to the backend. Use a typed :class:PreparedStatement subclass for backend-specific configuration (parameters, byte limits, routing, etc.).

execute_many ¶

execute_many(
    statements: Iterable["PS | PreparedStatement | str"],
    *,
    options: Optional[ExecutionOptions] = None,
    wait: WaitingConfigArg = True,
    raise_error: bool = True,
    parallel: Optional[int] = None,
    **batch_kwargs: Any
) -> SB

Run several statements as a batch and return the populated batch.

Convenience wrapper around :meth:batch: enqueues every statement, submits, and (by default) waits. parallel controls the wait phase only — submission itself is sequential, since most backends either accept fast or reject fast.

**batch_kwargs are forwarded to the batch constructor (e.g. external_paths= for :class:WarehouseStatementBatch).

batch ¶

batch(
    statements: Optional[Iterable["PS | PreparedStatement | str"]] = None,
    *,
    executor: "StatementExecutor | None" = None,
    parallel: Optional[int] = None,
    **kwargs: Any
) -> SB

Construct a batch bound to this executor.

cancel_all ¶

cancel_all() -> None

Best-effort cancel every live result this executor has produced.

yggdrasil.databricks.cluster.statement_executor¶