`yggdrasil.databricks.sql.engine`¶

engine ¶

Databricks SQL engine — Spark + warehouse dual-path execution.

SQLEngine ¶

SQLEngine(
    client=None,
    catalog_name: str | None = None,
    schema_name: str | None = None,
    default_warehouse: Optional[SQLWarehouse] = None,
    spark: Optional[DatabricksSparkStatementExecutor] = None,
)

Bases: DatabricksService, StatementExecutor

Unified SQL execution and Delta-table write engine for Databricks.

Composes two inner executors:

spark — :class:SparkStatementExecutor for Spark-side execution.
the warehouse handle resolved by :meth:warehouse — :class:SQLWarehouse, also a :class:StatementExecutor.

Routing (:meth:_pick_engine): explicit override → caller-supplied session → executor's own session → environment session → fall back to warehouse API.

Singleton-cached by (client, catalog_name, schema_name, default_warehouse) so two callers asking for the same scope share the same Spark sub-executor, the same lazy default_warehouse resolution, and the same sub-service caches (catalogs / schemas / tables route through self.client, which is itself singleton- cached). The spark field doesn't participate in identity — callers that pass a custom Spark executor onto an existing engine re-bind it in place.

opened `property` ¶

opened: bool

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

closed `property` ¶

closed: bool

Inverse of :attr:opened.

wheels `property` ¶

wheels: 'Wheels'

Wheel registry service (shorthand for client.wheels).

environments `property` ¶

environments: 'Environments'

Base-environment service (shorthand for client.environments).

views `property` ¶

views: 'Tables'

Alias for :attr:tables — :class:Table covers both managed/external tables and view-shaped securables.

volumes `property` ¶

volumes: 'Volumes'

Collection-level Unity Catalog volume service (shorthand for client.volumes).

genie `property` ¶

genie: 'Genie'

Genie service (shorthand for client.genie).

ai `property` ¶

ai: 'DatabricksAI'

Databricks AI umbrella service (shorthand for client.ai).

warehouse ¶

warehouse(
    warehouse_id: str | None = None, warehouse_name: str | None = None
) -> SQLWarehouse

Resolve the warehouse used by this engine.

execute ¶

execute(
    statement: str | PreparedStatement | StatementResult,
    *,
    row_limit: int | None = None,
    catalog_name: str | None = None,
    schema_name: str | None = None,
    wait: WaitingConfigArg = True,
    raise_error: bool = True,
    engine: Optional[Literal["spark", "api"]] = None,
    warehouse_id: str | None = None,
    warehouse_name: str | None = None,
    byte_limit: int | None = None,
    spark_session: Optional["SparkSession"] = None,
    external_data: Mapping[str, "VolumePath | Any"] | None = None,
    parameters: Mapping[str, Any] | None = None,
    retry: Optional[WaitingConfigArg] = None
) -> StatementResult

Execute a SQL statement through Spark or the Databricks SQL API.

retry controls result-level retry on the warehouse path (what :meth:StatementResult.retry does after a terminal failure). Has no effect on the Spark path here.

execute_many ¶

execute_many(
    statements: (
        Iterable[str | PreparedStatement | StatementResult]
        | Mapping[str, str | PreparedStatement | StatementResult]
    ),
    *,
    wait: WaitingConfigArg = True,
    raise_error: bool = True,
    parallel: Optional[int] = None,
    engine: Optional[Literal["spark", "api"]] = None,
    warehouse_id: str | None = None,
    warehouse_name: str | None = None,
    spark_session: Optional["SparkSession"] = None,
    retry: Optional[WaitingConfigArg] = None
) -> StatementBatch

Run a collection of statements; return per-statement results in order.

Statements that already carry per-statement external_volume_paths get their {alias} substitution from the warehouse-batch coercer at submit time — the engine doesn't manage a parallel registry.

retry is broadcast onto each warehouse statement before submission (Spark statements pass through untouched). None leaves whatever the statement already says intact; False explicitly clears any existing policy.

statement_result ¶

statement_result(
    statement_id: str,
    *,
    warehouse_id: str | None = None,
    warehouse_name: str | None = None
) -> "WarehouseStatementResult"

Re-attach to an already-executed statement by its statement_id.

The Databricks Statement Execution API keeps a finished statement's result available for a window after it runs. This binds a warehouse handle (the engine default unless one is named) to statement_id and returns a readable :class:WarehouseStatementResult — call wait() / to_arrow_table() / to_polars() / … to materialise it without re-running the query.

The statement text isn't needed (the statement already ran); a bare prepared statement is attached only to carry result-read config (disposition / format).

table ¶

table(
    location: str | None = None,
    *,
    catalog_name: str | None = None,
    schema_name: str | None = None,
    table_name: str | None = None
) -> Table

Resolve a table handle.

insert_into ¶

insert_into(
    data: Union[
        Table,
        RecordBatch,
        RecordBatchReader,
        dict,
        list,
        str,
        PreparedStatement,
        StatementResult,
        "pandas.DataFrame",
        "polars.DataFrame",
        "pyspark.sql.DataFrame",
    ],
    *,
    mode: Mode | str | None = None,
    schema_mode: Mode | str | None = None,
    location: str | None = None,
    catalog_name: str | None = None,
    schema_name: str | None = None,
    table_name: str | None = None,
    cast_options: Optional[CastOptions] = None,
    overwrite_schema: bool | None = None,
    match_by: Optional[list[str]] = None,
    update_column_names: Optional[list[str]] = None,
    wait: WaitingConfigArg = True,
    raise_error: bool = True,
    zorder_by: Optional[list[str]] = None,
    optimize_after_merge: bool = False,
    vacuum_hours: int | None = None,
    spark_session: Optional["pyspark.sql.SparkSession"] = None,
    spark_options: Optional[Dict[str, Any]] = None,
    table: Optional[Table] = None,
    predicate: Predicate | None = None,
    retry: Optional[WaitingConfigArg] = None
) -> "StatementBatch | None"

Resolve the target :class:Table and call :meth:Table.insert_into.

arrow_insert_into ¶

arrow_insert_into(
    data,
    location: str | None = None,
    catalog_name: str | None = None,
    schema_name: str | None = None,
    table_name: str | None = None,
    mode: Mode | str | None = None,
    schema_mode: Mode | str | None = None,
    cast_options: Optional[CastOptions] = None,
    overwrite_schema: bool | None = None,
    match_by: Optional[list[str]] = None,
    update_column_names: Optional[list[str]] = None,
    wait: WaitingConfigArg = True,
    raise_error: bool = True,
    zorder_by: Optional[list[str]] = None,
    optimize_after_merge: bool = False,
    vacuum_hours: int | None = None,
    table: Optional[Table] = None,
    predicate: Predicate | None = None,
    retry: Optional[WaitingConfigArg] = None,
) -> "StatementBatch | None"

Resolve target and forward to :meth:Table.arrow_insert.

spark_insert_into ¶

spark_insert_into(
    data: Any,
    *,
    mode: Mode | str | None = None,
    schema_mode: Mode | str | None = None,
    location: str | None = None,
    catalog_name: str | None = None,
    schema_name: str | None = None,
    table_name: str | None = None,
    cast_options: Optional[CastOptions] = None,
    overwrite_schema: bool | None = None,
    match_by: Optional[list[str]] = None,
    update_column_names: Optional[list[str]] = None,
    wait: WaitingConfigArg = True,
    raise_error: bool = True,
    zorder_by: Optional[list[str]] = None,
    optimize_after_merge: bool = False,
    vacuum_hours: int | None = None,
    spark_options: Optional[Dict[str, Any]] = None,
    table: Optional[Table] = None,
    predicate: Predicate | None = None,
    spark_session: Optional["pyspark.sql.SparkSession"] = None,
    retry: Optional[WaitingConfigArg] = None
) -> "StatementBatch | None"

Resolve target and forward to :meth:Table.spark_insert.

drop_table ¶

drop_table(
    location: str | None = None,
    *,
    catalog_name: str | None = None,
    schema_name: str | None = None,
    table_name: str | None = None,
    wait: WaitingConfigArg = True,
    raise_error: bool = True
) -> None

Drop a table if it exists.

create_table ¶

create_table(
    definition: Union[Field, Schema, Any],
    *,
    full_name: str | None = None,
    catalog_name: str | None = None,
    schema_name: str | None = None,
    table_name: str | None = None,
    **kwargs
) -> Table

Create a table if it does not already exist.

dataset ¶

dataset(sql_or_table: str, *, schema: Any = None) -> 'SparkDataset'

Return a :class:Dataset from a SQL query or table name.

Auto-detects SQL (SELECT, WITH, …) vs table name. Session resolved through Databricks Connect.

parallelize ¶

parallelize(
    inputs: "Iterable",
    function: "Callable | None" = None,
    *,
    schema: Any = None,
    byte_size: int = 128 * 1024 * 1024
) -> "SparkDataset"

Distribute function over inputs via Spark executors, or create a Dataset directly from inputs when no function is given.

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

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.

prepare ¶

prepare(statement: 'PS | PreparedStatement | str') -> PS

Coerce statement into this executor's prepared-statement type.

Mirrors :meth:Session.prepare_request_before_send: takes whatever the caller passed (raw string, cross-backend :class:PreparedStatement, already-typed instance) and returns the concrete :attr:_PREPARED_CLASS every downstream hook expects. Subclasses that need to inject per-statement defaults (warehouse routing, catalog binding, SELECT-rewrite for cluster execution) override this — same shape as Session's hook.

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.

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.

default_tags ¶

default_tags(update: bool = True) -> dict[str, str]

Return default resource tags for Databricks assets.

Returns:

Type	Description
`dict[str, str]`	A dict of default tags.

yggdrasil.databricks.sql.engine¶