Skip to content

yggdrasil.databricks.client

client

DatabricksClient

DatabricksClient(
    *,
    host: Any = ...,
    account_id: Any = ...,
    workspace_id: Any = ...,
    token: Any = ...,
    client_id: Any = ...,
    client_secret: Any = ...,
    token_audience: Any = ...,
    cluster_id: Any = ...,
    serverless_compute_id: Any = ...,
    azure_workspace_resource_id: Any = ...,
    azure_use_msi: Any = ...,
    azure_client_secret: Any = ...,
    azure_client_id: Any = ...,
    azure_tenant_id: Any = ...,
    azure_environment: Any = ...,
    google_credentials: Any = ...,
    google_service_account: Any = ...,
    profile: Any = ...,
    config_file: Any = ...,
    auth_type: Any = ...,
    http_timeout_seconds: Any = ...,
    retry_timeout_seconds: Any = ...,
    debug_truncate_bytes: Any = ...,
    debug_headers: Any = ...,
    rate_limit: Any = ...,
    max_connection_pools: Any = ...,
    max_connections_per_pool: Any = ...,
    product: Any = ...,
    product_version: Any = ...,
    skip_verify: Any = ...,
    catalog_name: Any = ...,
    schema_name: Any = ...,
    singleton_ttl: "int | None" = ...
)

Bases: Singleton, URLBased

Thin wrapper around databricks.sdk.config.Config.

URL-addressable through the :class:URLBased base: cls.scheme is :attr:Scheme.DATABRICKS (dbks), so a single dbks://[client_id[:secret]@]<host>[?profile=...&account_id=...] URL round-trips a client through :meth:from_url / :meth:to_url. The userinfo carries the credential — a bare :<token>@ for a PAT, <client_id>:<client_secret>@ for an OAuth service principal — and the query string carries every other DatabricksClient field that __init__ accepts. Sensitive fields (host, token, client_id, client_secret) are kept out of the query so they don't get duplicated when the URL is logged or persisted.

Public state is intentionally minimal
  • config: Config

Extra wrapper-only metadata is kept separately.

Cross-process / cross-host serialization is supported via :meth:__getstate__ / :meth:__setstate__. SDK clients, Configs, and lazy sub-service caches are dropped (all rebuild on demand from the dataclass fields). A best-effort session token snapshot is carried alongside so the receiving side can warm-start without re-running the auth dance (browser flow, MSI probe, gcloud, etc).

If the deserializing host is itself a Databricks runtime (driver node), the carried credentials are discarded and auth_type is forced to "runtime" — DBR's notebook-scoped auth is short-lived, identity-correct, and the right default in that environment.

explore_url property

explore_url: URL

Workspace UI root for the Catalog Explorer (/explore/data).

Mirrors :attr:Catalog.explore_url / :attr:Schema.explore_url so the whole resource hierarchy advertises a deep-link in one place.

project property writable

project: Optional[str]

The client project — an alias of :attr:product, always lowercased (the canonical project identifier). yggdrasil and ygg are the same project: yggdrasil is the name (and the import package), ygg is only its PyPI distribution alias. That distribution mapping lives in the wheel / environment layer (:func:~yggdrasil.databricks.wheels.service.distribution_for), not here — so client.project stays the human project name while the deployed wheels / base environment land under the ygg distribution folder. Defaults to "yggdrasil"; set it (or product) to name a project and it persists with the client (product is one of :data:_INIT_NAMES, so it rides config, session snapshots, and clones).

product_name property

product_name: Optional[str]

A nice, capitalized display name for the client :attr:project — the real project name (yggdrasilYggdrasil, my-appMy App), or None when unset. The project's default warehouse and cluster are named for this. The ygg PyPI alias is not applied here — it belongs to the wheel / distribution layer, not the project's identity.

wheels property

wheels: 'Wheels'

Wheel registry service — build/upload/deploy/browse wheels.

environments property

environments: 'Environments'

Base-environment service — assemble/deploy serverless + cluster images, deploy projects.

compute property

compute: 'Compute'

Default cluster helper for this client.

secrets property

secrets: 'Secrets'

Default secrets helper for this client.

tables property

tables: 'Tables'

Collection-level Unity Catalog table service for this client.

views property

views: 'Tables'

Alias for :attr:tables — Unity Catalog stores views in the same tables API, and :class:Table handles both shapes.

columns property

columns: 'Columns'

Collection-level Unity Catalog column service for this client.

catalogs property

catalogs: 'Catalogs'

Collection-level Unity Catalog hierarchy service for this client.

Provides dict-like access to catalogs, schemas, and tables::

client.catalogs["main"]                   # Catalog
client.catalogs["main"]["sales"]          # Schema
client.catalogs["main"]["sales"]["orders"]  # Table

external property

external: 'DatabricksExternal'

Unity Catalog external data umbrella service for this client.

Groups the securables that bind UC to outside storage — external locations and storage credentials::

client.external.locations["raw_zone"]          # ExternalLocation
client.external.credentials.create_aws("prod_s3", "arn:aws:iam::123:role/R")
client.external.credentials["prod_s3"].aws_client(region="us-east-1")

external_locations property

external_locations: 'ExternalLocations'

Unity Catalog external-locations service for this client.

Flat alias onto :attr:externalclient.external.locations::

client.external_locations["raw_zone"]          # ExternalLocation
client.external_locations.list()               # Iterator[ExternalLocation]
client.external_locations.create(name, url, credential_name)

credentials property

credentials: 'Credentials'

Unity Catalog storage-credentials service for this client.

Flat alias onto :attr:externalclient.external.credentials::

client.credentials.create_aws("prod_s3", "arn:aws:iam::123:role/R")
client.credentials["prod_s3"].aws_client(region="us-east-1")  # refreshable

schemas property

schemas: 'Schemas'

Collection-level Unity Catalog schema service for this client.

Provides dict-like access to schemas and tables::

client.schemas["main.sales"]             # Schema
client.schemas["main.sales.orders"]      # Table
client.schemas(catalog_name="main")      # Schemas scoped to "main"

volumes property

volumes: 'Volumes'

Collection-level Unity Catalog volume service for this client.

Provides dict-like access to volumes::

client.volumes["main.sales.uploads"]                # Volume
client.volumes(catalog_name="main", schema_name="sales")["uploads"]
client.volumes.list(catalog_name="main")            # Iterator[Volume]

genie property

genie: 'Genie'

Databricks AI/BI Genie service — manage spaces + ask questions by code.

dbc.genie.spaces() dbc.genie["01ef…"].ask("revenue by region last month").to_polars()

ai property

ai: 'DatabricksAI'

Databricks AI umbrella service (vector search today, serving/registry next).

Reach the concrete services through it::

client.ai.vector_search.endpoint("rag").ensure_created()
client.ai.vector_search.index("main.rag.docs").query(
    query_text="…", columns=["id", "text"],
)

is_serverless_compute property

is_serverless_compute: bool

True when this client explicitly targets serverless compute.

Only returns True when serverless_compute_id was set by the caller. A bare client with no cluster_id and no serverless_compute_id is NOT serverless — it simply has no compute target and will resolve one lazily when needed.

to_url

to_url(scheme: str | None = None) -> URL

Render this client as a dbks://... URL.

Pack everything __init__ would need to rebuild the client into the URL: the workspace host as the URL host, the credential (PAT or OAuth client_id/secret) as userinfo, and every other non-default field as query items. Sensitive fields (host, token, client_id, client_secret) are intentionally kept out of the query so they don't get duplicated alongside the userinfo.

scheme overrides :attr:scheme for callers that want a different URL scheme (e.g. "https" for the bare workspace URL); defaults to :attr:Scheme.DATABRICKS.

from_url classmethod

from_url(url: 'URL | str', **kwargs: Any) -> TC

Build a client from a dbks://... URL.

Reads:

  • the workspace host from url.host (preferred) or a host= query param;
  • credentials from url.userinfo<client_id>:<client_secret>@ for OAuth, :<token>@ (or anything-as-password) for a PAT;
  • every other init field of :class:DatabricksClient from the query string (profile, auth_type, account_id, workspace_id, http_timeout_seconds, …).

kwargs overrides anything the URL provides so callers can layer programmatic overrides on top of a parsed URL without an extra replace call.

files_session

files_session() -> 'HTTPSession'

Authenticated :class:HTTPSession bound to this workspace host.

Volume / Files-API traffic routes through yggdrasil's own HTTP transport instead of the SDK's requests-based client: the :class:HTTPSession owns a per-host keep-alive connection pool, status-aware tiered retry (429 / 5xx with backoff), and — via the :class:HTTPStream response body — transparent resume-on-disconnect for SSL UNEXPECTED_EOF / connection-reset mid-download, the failure modes the SDK's Files client handles poorly.

:class:HTTPSession is itself a process-wide singleton keyed by (base_url, verify, …), so repeated calls collapse onto one shared pool. skip_verify flows through to verify=False.

files_authorization

files_authorization() -> str

Fresh Authorization header value for Files-API requests.

Delegates to the SDK Config's auth flow (:meth:databricks.sdk.config.Config.authenticate) so every supported credential type — PAT, OAuth M2M, Azure SP, GCP — and the SDK's own token-refresh caching apply unchanged; only the wire transport is ours. Raises when the resolved auth produces no bearer header (e.g. a misconfigured profile).

files_headers

files_headers() -> dict[str, str]

Base header set for Files-API requests, matching the SDK transport.

The SDK's ApiClient stamps User-Agent (product + SDK version + platform + auth type) on every call and X-Databricks-Workspace-Id when the config carries a workspace id — Databricks' edge uses both to attribute and rate-limit traffic, so requests without them are classified as anonymous and throttled (429) far more aggressively. Mirror them here since Files traffic bypasses the SDK transport (:meth:files_session).

default_tags

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

Return default resource tags for Databricks assets.

On create (update=False) the tag set is enriched with environment-derived owner metadata pulled from :class:~yggdrasil.environ.UserInfo:

  • Product / ProductVersion from the client config.
  • Owner — UserInfo email when available.
  • Hostname — local hostname so per-user pools / clusters are distinguishable in shared workspaces.
  • Userwhoami key, useful when no email is reachable.

Returns:

Type Description
dict[str, str]

A dict of default tags.

user_scoped_name

user_scoped_name(
    base: str, *, separator: str = "-", max_length: int = 80
) -> str

Return base suffixed with a stable per-user slug.

Resolution order for the slug:

  1. UserInfo.email local part (alice@example.comalice).
  2. UserInfo.key (whoami).
  3. UserInfo.hostname.

Each candidate is sanitized with :meth:safe_tag_value so the result is a legal Databricks resource name. Falls back to base unchanged when no candidate is available — useful in test harnesses where the environment carries no identity. The result is truncated to max_length characters (default 80, well under the Databricks cluster / pool name cap of 100).

safe_tag_value staticmethod

safe_tag_value(value: str, *, repl: str = '-') -> str

Sanitize a tag string to match: ^[\d \w+-=:.:/@]*$ Replaces any illegal characters with repl and collapses repeats.

dbfs_path

dbfs_path(
    parts: Union[list[str], str], temporary: bool = False
) -> DatabricksPath

Create a DatabricksPath in this workspace.

.. deprecated:: 0.8.31 Use :meth:open for byte IO (client.open("/Volumes/cat/sch/vol/x", "rb")), or :meth:path when you need the (non-opened) resource itself (client.path("/Volumes/cat/sch/vol/x")).

Parameters:

Name Type Description Default
parts Union[list[str], str]

Path parts or string to parse.

required
temporary bool

Temporary path

False

Returns:

Type Description
DatabricksPath

A DatabricksPath instance.

open

open(path: Any, mode: 'Mode | str | None' = None, **kwargs: Any) -> 'IO'

Open path against this workspace and return an :class:IO cursor.

Defaults to a :class:DatabricksPath bound to this client — strings like /Volumes/cat/sch/vol/x or dbfs+volume:///cat/sch/vol/x dispatch to the right concrete subclass (DBFS / Volumes / Workspace). A pre-built :class:~yggdrasil.path.Path is opened verbatim so callers can mix in S3/HTTP/local paths without losing the workspace binding.

mode and **kwargs ride straight through to :meth:Path.open (which forwards to :meth:IO.open).

path

path(path: Any, *, temporary: bool = False, **kwargs: Any) -> 'Path'

Build a path bound to this workspace — without opening it.

The non-opening companion to :meth:open: strings like /Volumes/cat/sch/vol/x or dbfs+volume:///cat/sch/vol/x dispatch to the right concrete :class:DatabricksPath subclass (DBFS / Volumes / Workspace), while a pre-built :class:~yggdrasil.path.Path is returned verbatim so callers can mix in S3 / HTTP / local paths without losing the workspace binding.

Use this when you want the path resource itself — ls / stat / write_table / child navigation — and :meth:open when you want an :class:IO byte cursor. temporary and any extra **kwargs ride straight through to :meth:DatabricksPath.from_.

tmp_path

tmp_path(
    suffix: str | None = None,
    extension: str | None = None,
    max_lifetime: float | None = None,
    catalog_name: str | None = None,
    schema_name: str | None = None,
    volume_name: str | None = None,
    base_path: str | None = None,
) -> "DatabricksPath"

Shared cache base under Volumes for the current user.

Parameters:

Name Type Description Default
suffix str | None

Optional suffix

None
extension str | None

Optional extension suffix to append.

None
max_lifetime float | None

Max lifetime of temporary path

None
catalog_name str | None

Unity catalog name for volume path

None
schema_name str | None

Unity schema name for volume path

None
volume_name str | None

Unity volume name for volume path

None
base_path str | None

Base temporary path

None

Returns:

Type Description
'DatabricksPath'

A DatabricksPath pointing at the shared cache location.

lazy_property staticmethod

lazy_property(
    *, cache_attr: str, factory: Callable[[], T], use_cache: bool
) -> T

Public helper kept for backwards compatibility.

New properties on :class:DatabricksClient inline the lookup directly through self.__dict__ (one dict get, one dict set on miss) — the function-call + lambda overhead of routing every sub-service through this helper used to cost ~1.3 us per access on the hot path. Callers outside this class keep the same surface.

dataset

dataset(sql_or_table: str, *, schema: Any = None)

Return a :class:~yggdrasil.spark.tabular.Dataset from SQL or a table name.

Resolves the Spark session via :meth:spark (Databricks Connect) and builds a :class:Dataset directly — no intermediate executor hop::

dbc = DatabricksClient()
ds = dbc.dataset("SELECT * FROM main.sales.orders")
ds = dbc.dataset("main.sales.orders")

The result is a full :class:Dataset — call .map, .filter, .to_table, .toArrow, etc. on it.

parallelize

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

Distribute function over inputs via Spark executors, or create a :class:~yggdrasil.spark.tabular.Dataset directly from inputs when no function is given::

dbc = DatabricksClient()
# With function
results = dbc.parallelize(urls, fetch, schema=output_schema)
# Without function — just wrap inputs as a Dataset
ds = dbc.parallelize(rows, schema=output_schema)

spark

spark(
    *dependencies: "Any",
    registry: "Optional[Any]" = None,
    check_public: bool = False,
    cache_dir: "Optional[Union[str, os.PathLike]]" = None,
    cluster: "str | Cluster | None" = None
)

Open a Databricks Connect session with auto-resolved deps.

Returns a live :class:pyspark.sql.SparkSession (Spark Connect variant) configured against this client's workspace host and credentials. The bound :class:DatabricksClient is stashed on the session as session.ygg_client so downstream helpers (UDFs, :class:Dataset extensions, ad-hoc resource lookups) can reach the same auth without an extra DatabricksClient.current() call.

Each dependency is classified once via :func:classify_dependency:

  • Public PyPI specs ("ygg[data,databricks]==0.7.85", "numpy>=1.0", …) ride straight to the cluster via :meth:DatabricksEnv.withDependencies. ygg is always declared via :meth:default_ygg_spec — pinned to the driver's installed version with the [data, databricks] extras so the cluster sees the exact same runtime + pandas / numpy / databricks-sdk surface the driver is using. Override by passing an explicit ygg spec (e.g. client.spark("ygg==0.7.0") or client.spark("ygg") for an editable-mode rebuild).
  • Editable installs (pip install -e .) get their local working copy built into a wheel whose version carries the local hostname (0.7.85+host.<host>). The wheel lives at /Workspace/Users/<me>/.ygg/pypi/simple/<pkg>/<wheel> (overridable via registry) and is re-uploaded on every load so the registry slot tracks the developer's working code.
  • Private / non-PyPI installs get the same wheel-build + workspace-upload treatment, but lazily — the upload is skipped when the workspace slot already exists, so a team sharing one registry path only re-uploads on version bumps.

Both editable and private wheels are then handed to :meth:DatabricksEnv.withDependencies via the local:<path> prefix Databricks Connect understands; the wheel itself is read back from the workspace into a local cache so the spec is reachable from the driver process.

Serverless compute (the default — no cluster_id) wires deps through DatabricksEnv + withEnvironment; classic compute falls back to :meth:SparkSession.addArtifacts with pyfile=True since classic clusters don't honour the declarative environment.

Arguments:

  • dependencies — variadic. Each entry is anything :func:classify_dependency accepts (PyPI spec string, bare module name, :class:os.PathLike, or any object with __module__). client.spark("polars", "my_internal", Path("/some/pkg")) is the headline shape; ygg[data,databricks] is appended automatically unless the caller already provides their own ygg spec.
  • registry — a :class:WorkspacePyPIRegistry (or any shape its constructor accepts) to use as the shared wheel cache. Defaults to a registry rooted at /Workspace/Users/<me>/.ygg/pypi/simple so a single-user setup needs no configuration.
  • check_public — when True, an HTTPS probe to pypi.org decides whether an installed dist counts as public. Off by default so an offline registry stays fast; turn on when shipping mixed public / private deps.
  • cache_dir — local scratch dir used by the classic compute fallback (and for wheel materialization when no explicit registry is passed).

When a :class:pyspark.sql.SparkSession is already active in the process (notebook driver, an outer client.spark() call, a Databricks Job task), that session is returned as-is — dependency classification and wheel publishing are skipped, since the active session's environment is already fixed. The client is still stashed on it as session.ygg_client so downstream helpers find the same auth.

for_scheme classmethod

for_scheme(scheme: Any) -> 'type[URLBased]'

Return the :class:URLBased subclass registered for scheme.

Lazy: if no subclass is registered yet, this routes through :meth:Scheme.path_class which imports the backend module on demand (firing :meth:__init_subclass__ as a side effect).

Raises :class:ValueError for an unknown scheme and :class:ImportError when the backend's optional dependencies aren't installed.

dispatch classmethod

dispatch(url: Any, **kwargs: Any) -> 'URLBased'

Build the right :class:URLBased subclass from url.

Looks up the subclass via :meth:for_scheme, then delegates to that subclass's :meth:from_url. Used as the cross-cutting entry point when the caller has a URL but doesn't know (or care) which concrete class owns its scheme.

URL.from_(url).scheme drives the lookup; an empty scheme falls back to the file:// handler so bare paths work.

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.

DatabricksService

DatabricksService(client: Optional[DatabricksClient] = None)

Bases: ABC

Base class for every Databricks service wrapper.

Subclasses are plain classes (no @dataclass); they inherit the client-only constructor by default and override :meth:__init__ explicitly when they need to accept additional configuration.

wheels property

wheels: 'Wheels'

Wheel registry service (shorthand for client.wheels).

environments property

environments: 'Environments'

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

tables property

tables: 'Tables'

Collection-level Unity Catalog table service (shorthand for client.tables).

views property

views: 'Tables'

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

catalogs property

catalogs: 'Catalogs'

Collection-level Unity Catalog hierarchy service (shorthand for client.catalogs).

schemas property

schemas: 'Schemas'

Collection-level Unity Catalog schema service (shorthand for client.schemas).

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).

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.

DatabricksResource

DatabricksResource(service=None, *args, **kwargs)

Bases: ExploreUrlRepr, ABC

explore_url property

explore_url: Optional[URL]

Workspace UI deep-link for this resource, or None.

Concrete resources (:class:Catalog, :class:Schema, :class:Volume, :class:Table, :class:SQLWarehouse, :class:Job, :class:VolumePath, …) override this to return the /explore/data/... / /sql/warehouses/... / /jobs/... URL that opens the resource in the workspace UI. The inherited :class:ExploreUrlRepr keys off the override — anything that returns a non-None URL gets a ClassName(<url>) repr (and a clickable _repr_html_) for free without restating it on every subclass.

sql property

sql: 'SQLEngine'

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

current_catalog

current_catalog() -> Optional[str]

The process-/async-local default catalog name, or None.

current_schema

current_schema() -> Optional[str]

The process-/async-local default schema name, or None.

invalidate_env_defaults

invalidate_env_defaults() -> None

Drop the env-default snapshot so the next constructor re-reads.

Call after rotating DATABRICKS_* / ARM_* / GOOGLE_* env vars so a subsequent :class:DatabricksClient build picks them up. The in-process singleton cache is not cleared — entries already keyed off the previous snapshot stay live until they're replaced or evicted.