interloper-databricks 0.96.0__py3-none-any.whl

interloper_databricks/__init__.py

@@ -0,0 +1,9 @@
+"""Interloper Databricks integration: connection and destination."""
+
+from interloper_databricks.connection import DatabricksConnection
+from interloper_databricks.destination import DatabricksDestination
+
+__all__ = [
+    "DatabricksConnection",
+    "DatabricksDestination",
+]

interloper_databricks/connection.py

@@ -0,0 +1,298 @@
+"""Databricks connection resource holding service principal or personal access token credentials."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Generator
+from functools import cached_property
+from typing import Any
+from urllib.parse import urlsplit
+
+import httpx2
+from databricks import sql
+from databricks.sql.client import Connection as Session
+from interloper.connection import Connection, connection
+from interloper.resource.fields import InputField, SecretField, fetch_field_provider
+from interloper.rest import JSONCursorPaginator, RESTClient
+from pydantic import PrivateAttr, field_validator, model_validator
+from pydantic_settings import SettingsConfigDict
+
+_TOKEN_ENDPOINT = "/oidc/v1/token"
+
+#: A token this close to expiry is replaced before it is handed out, so a
+#: statement never starts on a token that lapses mid-flight.
+_REFRESH_MARGIN = 300.0
+
+#: Every REST call here serves an operator waiting on a form or a check.
+_TIMEOUT = 30.0
+
+
+class _TokenAuth(httpx2.Auth):
+    """Bearer authentication reading the token from a callable on every request."""
+
+    def __init__(self, token: Callable[[], str]) -> None:
+        """Bind the auth to its token source.
+
+        Args:
+            token: Returns a valid access token, refreshing it when due.
+        """
+        self._token = token
+
+    def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
+        """Attach the current token as a bearer header.
+
+        Args:
+            request: The request to authenticate.
+
+        Yields:
+            The authenticated request.
+        """
+        request.headers["Authorization"] = f"Bearer {self._token()}"
+        yield request
+
+
+@connection(
+    key="databricks_connection",
+    name="Databricks",
+    icon="icon:databricks",
+    tags=["Cloud"],
+)
+class DatabricksConnection(Connection):
+    """Connection resource holding Databricks workspace credentials.
+
+    Authenticates either as a service principal through OAuth
+    machine-to-machine (``client_id`` and ``client_secret``, the recommended
+    path) or with a personal access token; exactly one of the two must be
+    set. The connection's own check and pickers call the workspace REST API
+    over plain HTTP; each destination opens its own SQL session through
+    :meth:`connect`.
+    """
+
+    model_config = SettingsConfigDict(env_prefix="databricks_")
+
+    host: str = InputField(
+        label="Workspace URL",
+        description="e.g. https://dbc-a1b2345c-d6e7.cloud.databricks.com",
+    )
+    client_id: str | None = InputField(
+        default=None,
+        label="Client ID",
+        description="Service principal application ID",
+        section="Service principal (recommended)",
+    )
+    client_secret: str | None = SecretField(
+        default=None,
+        label="Client secret",
+        description="Service principal OAuth secret",
+        section="Service principal (recommended)",
+    )
+    access_token: str | None = SecretField(
+        default=None,
+        label="Access token",
+        description="Personal access token, instead of a service principal",
+        section="Personal access token",
+    )
+
+    _token: tuple[str, float] | None = PrivateAttr(default=None)
+
+    @field_validator("host")
+    @classmethod
+    def normalize_host(cls, value: str) -> str:
+        """Give the workspace URL a scheme and drop any trailing slash.
+
+        Args:
+            value: The workspace URL or bare hostname.
+
+        Returns:
+            The URL as ``https://<hostname>``.
+        """
+        value = value.strip().rstrip("/")
+        return value if "://" in value else f"https://{value}"
+
+    @field_validator("client_id", "client_secret", "access_token", mode="before")
+    @classmethod
+    def blank_is_unset(cls, value: Any) -> Any:
+        """Treat an empty credential, as a form submits it, as unset.
+
+        Args:
+            value: The raw field value.
+
+        Returns:
+            ``None`` for an empty string, the value otherwise.
+        """
+        return None if value == "" else value
+
+    @model_validator(mode="after")
+    def one_credential(self) -> DatabricksConnection:
+        """Require exactly one credential: a complete service principal or an access token.
+
+        Returns:
+            The validated connection.
+
+        Raises:
+            ValueError: If the service principal is half set, or if both or
+                neither credentials are set.
+        """
+        principal = self.client_id is not None or self.client_secret is not None
+        if principal and (self.client_id is None or self.client_secret is None):
+            raise ValueError("A service principal needs both 'client_id' and 'client_secret'")
+        if principal and self.access_token is not None:
+            raise ValueError("Set either a service principal or an access token, not both")
+        if not principal and self.access_token is None:
+            raise ValueError("Set a service principal ('client_id' and 'client_secret') or an 'access_token'")
+        return self
+
+    @property
+    def hostname(self) -> str:
+        """The workspace hostname, as the SQL connector takes it.
+
+        Returns:
+            The host without its scheme.
+        """
+        return urlsplit(self.host).netloc
+
+    # -- Credentials ---------------------------------------------------------------
+
+    def token(self) -> str:
+        """Return a bearer token for the workspace.
+
+        A personal access token is returned as is. A service principal's
+        token comes from the workspace's OAuth token endpoint (client
+        credentials, ``all-apis`` scope) and is cached until close to its
+        expiry; two threads refreshing at once only cost a second exchange.
+
+        Returns:
+            The access token.
+        """
+        if self.access_token is not None:
+            return self.access_token
+        cached = self._token
+        if cached is not None and cached[1] - time.monotonic() > _REFRESH_MARGIN:
+            return cached[0]
+        assert self.client_id is not None and self.client_secret is not None
+        with RESTClient(self.host, timeout=_TIMEOUT) as oauth:
+            response = oauth.post(
+                _TOKEN_ENDPOINT,
+                data={"grant_type": "client_credentials", "scope": "all-apis"},
+                auth=httpx2.BasicAuth(self.client_id, self.client_secret),
+            )
+        response.raise_for_status()
+        body = response.json()
+        self._token = (body["access_token"], time.monotonic() + float(body.get("expires_in", 3600)))
+        return body["access_token"]
+
+    def _authorization(self) -> dict[str, str]:
+        """Build the header the SQL connector sends with each request.
+
+        Returns:
+            The ``Authorization`` header carrying a current token.
+        """
+        return {"Authorization": f"Bearer {self.token()}"}
+
+    def _credentials_provider(self) -> Callable[[], dict[str, str]]:
+        """Hand the SQL connector its header factory.
+
+        The connector calls the provider once per session and the factory it
+        returns whenever it needs headers, so the factory refreshes the
+        service principal's token as it ages.
+
+        Returns:
+            The header factory.
+        """
+        return self._authorization
+
+    def connect(self, **session: Any) -> Session:
+        """Open a new SQL session on the workspace with this connection's credentials.
+
+        A personal access token goes to the connector as is; a service
+        principal goes through a ``credentials_provider`` built on this
+        connection's own token exchange, which spares the Databricks SDK as
+        a dependency.
+
+        Args:
+            **session: Session settings passed to the connector, such as
+                ``http_path`` and ``catalog``.
+
+        Returns:
+            The new connector session.
+        """
+        if self.access_token is not None:
+            credentials: dict[str, Any] = {"access_token": self.access_token}
+        else:
+            credentials = {"credentials_provider": self._credentials_provider}
+        return sql.connect(server_hostname=self.hostname, **credentials, **session)
+
+    # -- REST API ------------------------------------------------------------------
+
+    @cached_property
+    def client(self) -> RESTClient:
+        """The workspace REST API client the check and the pickers share.
+
+        Plain HTTP rather than the SQL connector: both run in the API
+        process, where opening a warehouse session would be slow and could
+        start a stopped warehouse.
+
+        Returns:
+            The bearer-authenticated client, cached per connection instance.
+        """
+        return RESTClient(self.host, auth=_TokenAuth(self.token), timeout=_TIMEOUT)
+
+    def _list(self, path: str, key: str, params: dict[str, str] | None = None) -> list[dict[str, Any]]:
+        """Walk a paginated list endpoint and collect its items.
+
+        Args:
+            path: The endpoint path.
+            key: The response key holding each page's items.
+            params: Static query parameters for every page.
+
+        Returns:
+            Every item across the pages.
+        """
+        pages = self.client.paginate(
+            path,
+            JSONCursorPaginator(cursor_path="next_page_token", cursor_param="page_token"),
+            params=params,
+            data_selector=lambda response: response.json().get(key, []),
+        )
+        return [item for page in pages for item in page]
+
+    @fetch_field_provider
+    def warehouses(self) -> list[dict[str, str]]:
+        """List the SQL warehouses this principal can use.
+
+        Backs the destination's ``warehouse`` ``FetchField``. The stored value
+        is the warehouse's HTTP path, which is what the SQL connector takes.
+
+        Returns:
+            Warehouse options with ``path`` and ``name``, sorted case-insensitively by name.
+        """
+        warehouses = self._list("/api/2.0/sql/warehouses", "warehouses")
+        options = [{"path": w["odbc_params"]["path"], "name": w["name"]} for w in warehouses]
+        return sorted(options, key=lambda option: option["name"].lower())
+
+    @fetch_field_provider
+    def catalogs(self) -> list[dict[str, str]]:
+        """List the Unity Catalog catalogs this principal can see.
+
+        Backs the destination's ``catalog`` ``FetchField``. ``max_results=0``
+        asks for the paginated form, which Databricks recommends over the
+        unpaginated one.
+
+        Returns:
+            Catalog options with ``name``, sorted case-insensitively.
+        """
+        catalogs = self._list("/api/2.1/unity-catalog/catalogs", "catalogs", params={"max_results": "0"})
+        return sorted(({"name": c["name"]} for c in catalogs), key=lambda option: option["name"].lower())
+
+    def check(self) -> bool:
+        """Prove the credentials work by asking the workspace who the caller is.
+
+        ``/api/2.0/preview/scim/v2/Me`` needs nothing beyond a valid token,
+        so a failure isolates a bad credential from a missing grant; for a
+        service principal it also exercises the token exchange.
+
+        Returns:
+            True; a rejected credential raises out of the HTTP call.
+        """
+        self.client.get("/api/2.0/preview/scim/v2/Me").raise_for_status()
+        return True

interloper_databricks/destination.py

@@ -0,0 +1,665 @@
+"""Databricks destination implementation."""
+
+from __future__ import annotations
+
+import datetime
+import inspect
+import io
+import json
+import math
+import threading
+import uuid
+import warnings
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from decimal import Decimal
+from functools import cached_property
+from typing import Any
+
+import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
+from databricks.sql.client import Connection as Session
+from databricks.sql.client import Cursor
+from interloper.destination import IOContext, destination
+from interloper.destination.database import DatabaseDestination, PartitionFilter
+from interloper.errors import ConfigError, DataNotFoundError
+from interloper.partitioning.base import Partition
+from interloper.representation import Representation
+from interloper.resource.fields import FetchField, InputField
+from interloper.schema import FieldSpec
+from interloper.utils.data import is_empty
+from interloper.utils.json import json_default, replace_non_finite
+from pydantic import PrivateAttr, field_validator
+from typing_extensions import Self
+
+from interloper_databricks.connection import DatabricksConnection
+from interloper_databricks.types import TIMESTAMP, VARIANT, clusterable, column_type
+
+#: Delta collects statistics on a table's first 32 columns, and a clustering
+#: key must be one of them.
+_STATS_COLUMNS = 32
+
+
+class _Lock:
+    """A re-entrant lock that a deep copy replaces with a fresh one.
+
+    A destination is deep-copied with the source it is bound to, and a bare
+    ``threading.RLock`` cannot be copied.
+    """
+
+    def __init__(self) -> None:
+        """Create the underlying lock."""
+        self._lock = threading.RLock()
+
+    def __enter__(self) -> Self:
+        """Acquire the lock.
+
+        Returns:
+            The lock.
+        """
+        self._lock.acquire()
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        """Release the lock.
+
+        Args:
+            *exc: The exception details, if the block raised; unused.
+        """
+        self._lock.release()
+
+    def __deepcopy__(self, memo: dict[int, Any]) -> _Lock:
+        """Copy as a new, unheld lock.
+
+        Args:
+            memo: The deep-copy memo; unused.
+
+        Returns:
+            A fresh lock.
+        """
+        return _Lock()
+
+
+@dataclass(frozen=True)
+class _Staged:
+    """Data uploaded to the staging volume, ready to be inserted.
+
+    Attributes:
+        path: The file's ``/Volumes/...`` path.
+        query: The ``SELECT`` reading the file, its columns cast to the table's types.
+    """
+
+    path: str
+    query: str
+
+
+@destination(
+    key="databricks_destination",
+    name="Databricks",
+    icon="icon:databricks",
+    tags=["Cloud"],
+)
+class DatabricksDestination(DatabaseDestination):
+    """Databricks destination writing Delta tables in Unity Catalog through a SQL warehouse.
+
+    A dataset is a schema inside the destination's catalog. Every write
+    uploads the data as one Parquet file to the staging volume and loads it
+    with a single statement: ``INSERT INTO ... REPLACE WHERE`` for a
+    partition or a window, ``INSERT OVERWRITE`` for an unpartitioned asset.
+
+    That is why :meth:`write` and :meth:`write_partition` are overridden
+    rather than left to the base's delete-then-insert: Databricks has no
+    generally available multi-statement transaction, so a ``DELETE``
+    followed by an insert could leave a partition empty when the insert
+    fails, while ``REPLACE WHERE`` deletes and inserts in one Delta commit.
+    The rows to replace still come from the base's partition filters, and
+    :meth:`delete`, :meth:`select` and :meth:`count` remain the base's hooks.
+
+    The destination opens one SQL session, shared by every asset written
+    through it. The connector's sessions are not thread-safe, so every
+    statement runs under the destination's lock, and a write holds it from
+    the upload to the cleanup.
+    """
+
+    connection: DatabricksConnection
+
+    warehouse: str = FetchField(
+        provider="connection.warehouses",
+        label_key="name",
+        value_key="path",
+        label="SQL warehouse",
+        description="SQL warehouse running the loads and queries",
+    )
+    catalog: str = FetchField(
+        provider="connection.catalogs",
+        label_key="name",
+        value_key="name",
+        description="Unity Catalog catalog",
+        discriminator=True,
+    )
+    default_dataset: str | None = InputField(default=None, description="Default schema for assets without a dataset")
+    staging_volume: str = InputField(
+        label="Staging volume",
+        description="Unity Catalog volume the loads stage files in, as catalog.schema.volume",
+    )
+
+    _lock: _Lock = PrivateAttr(default_factory=_Lock)
+
+    @field_validator("staging_volume")
+    @classmethod
+    def three_part_volume(cls, value: str) -> str:
+        """Require a fully qualified volume name.
+
+        Args:
+            value: The volume name.
+
+        Returns:
+            The volume name, unchanged.
+
+        Raises:
+            ValueError: If the name is not ``catalog.schema.volume``.
+        """
+        parts = value.split(".")
+        if len(parts) != 3 or not all(parts):
+            raise ValueError(f"staging_volume must be 'catalog.schema.volume', got '{value}'")
+        return value
+
+    @property
+    def volume_path(self) -> str:
+        """The staging volume's path.
+
+        Returns:
+            ``/Volumes/<catalog>/<schema>/<volume>``.
+        """
+        return "/Volumes/" + "/".join(self.staging_volume.split("."))
+
+    @cached_property
+    def client(self) -> Session:
+        """The destination's own session, on its warehouse and catalog.
+
+        Returns:
+            The connector session, cached per destination instance.
+        """
+        return self.connection.connect(http_path=self.warehouse, catalog=self.catalog)
+
+    # -- Session -----------------------------------------------------------------
+
+    def _execute(
+        self,
+        sql: str,
+        *,
+        fetch: Callable[[Cursor], Any] | None = None,
+        input_stream: io.BytesIO | None = None,
+    ) -> Any:
+        """Run one statement on a fresh cursor, under the destination's lock.
+
+        Args:
+            sql: The statement.
+            fetch: Reads the result off the cursor; ``None`` when the
+                statement returns nothing of interest.
+            input_stream: The bytes a ``PUT '__input_stream__'`` uploads.
+
+        Returns:
+            What *fetch* returns, or ``None``.
+        """
+        with self._lock:
+            cursor = self.client.cursor()
+            try:
+                cursor.execute(sql, input_stream=input_stream)
+                return fetch(cursor) if fetch is not None else None
+            finally:
+                cursor.close()
+
+    # -- Destination interface -----------------------------------------------------
+
+    def write(self, context: IOContext, data: Any) -> None:
+        """Replace the partitions the context covers with the data, in one statement.
+
+        Args:
+            context: IO context carrying the target asset, the partition or window,
+                and the effective schema.
+            data: The data to write, in its native representation.
+        """
+        if is_empty(data):
+            return
+        self._warn_missing_partition_column(data, context)
+        self._replace(context, context.partitions, data)
+
+    def write_partition(self, context: IOContext, partition: Partition | None, data: Any) -> None:
+        """Replace one partition's rows with the data, in one statement.
+
+        Args:
+            context: IO context carrying the target asset and the effective schema.
+            partition: The partition being stored, or ``None`` for the whole table.
+            data: The data to store, in its native representation.
+        """
+        self._replace(context, [partition], data)
+
+    def _replace(self, context: IOContext, partitions: Sequence[Partition | None], data: Any) -> None:
+        """Load the data in place of the rows the partitions cover.
+
+        Args:
+            context: IO context carrying the target asset and the effective schema.
+            partitions: The partitions replaced; ``[None]`` for the whole table.
+            data: The data, in its native representation.
+        """
+        table, dataset = self._target(context)
+        filters = [self._filter(context, partition) for partition in partitions]
+        if None in filters:
+            self._load(table, dataset, data, context, "OVERWRITE")
+            return
+        where = _predicate([f for f in filters if f is not None])
+        self._load(table, dataset, data, context, "INTO", f" REPLACE WHERE {where}")
+
+    # -- Naming --------------------------------------------------------------------
+
+    def _resolve_dataset(self, dataset: str | None) -> str:
+        """Return the schema to use.
+
+        Args:
+            dataset: The asset's dataset, or ``None`` to fall back to the destination's default.
+
+        Returns:
+            The resolved schema name.
+
+        Raises:
+            ConfigError: If neither the asset nor the destination names a dataset.
+        """
+        schema = dataset or self.default_dataset
+        if schema is None:
+            raise ConfigError(
+                "DatabricksDestination requires a dataset. Either set 'dataset' on the asset "
+                "or provide 'default_dataset' on the destination."
+            )
+        return schema
+
+    def _table_ref(self, table: str, schema: str) -> str:
+        """Build a fully-qualified, quoted table reference.
+
+        Args:
+            table: Table name.
+            schema: The resolved schema name.
+
+        Returns:
+            ```catalog`.`schema`.`table```.
+        """
+        return f"{_quote(self.catalog)}.{_quote(schema)}.{_quote(table)}"
+
+    def _table_exists(self, table: str, schema: str) -> bool:
+        """Check whether a table exists, through the catalog's information schema.
+
+        Unity Catalog stores object names in lower case, so the lookup does too.
+
+        Args:
+            table: Table name.
+            schema: The resolved schema name.
+
+        Returns:
+            ``True`` if the table exists, ``False`` otherwise.
+        """
+        rows = self._execute(
+            f"SELECT 1 FROM {_quote(self.catalog)}.information_schema.tables "
+            f"WHERE table_schema = {_literal(schema.lower())} AND table_name = {_literal(table.lower())}",
+            fetch=lambda cursor: cursor.fetchall(),
+        )
+        return bool(rows)
+
+    # -- Loading -------------------------------------------------------------------
+
+    def _ensure_table(self, table: str, schema: str, specs: Sequence[FieldSpec], context: IOContext) -> None:
+        """Create the schema and a typed Delta table when the table does not exist.
+
+        A partitioned asset's table is clustered on its partition column,
+        which every replace and partition read filters on, when liquid
+        clustering accepts that column as a key; Databricks recommends
+        clustering over partitioning for tables of this size. Field
+        descriptions become column comments, the asset's description the
+        table comment.
+
+        Args:
+            table: Table name.
+            schema: The resolved schema name.
+            specs: The table's field specs.
+            context: IO context carrying the asset.
+        """
+        if self._table_exists(table, schema):
+            return
+        self._execute(f"CREATE SCHEMA IF NOT EXISTS {_quote(self.catalog)}.{_quote(schema)}")
+        columns = ", ".join(_column_ddl(spec) for spec in specs)
+        sql = f"CREATE TABLE IF NOT EXISTS {self._table_ref(table, schema)} ({columns}) USING DELTA"
+        partitioning = context.asset.partitioning
+        keys = [spec for spec in specs[:_STATS_COLUMNS] if partitioning and spec.name == partitioning.column]
+        if keys and clusterable(keys[0]):
+            sql += f" CLUSTER BY ({_quote(keys[0].name)})"
+        description = _asset_description(context.asset)
+        if description:
+            sql += f" COMMENT {_literal(description)}"
+        self._execute(sql)
+
+    def _stage(self, table: str, schema: str, data: Any, context: IOContext) -> _Staged:
+        """Create what the load needs and upload the data as one Parquet file.
+
+        A new table is typed from the effective schema (declared on the asset,
+        or inferred during conform), or from a schema inferred from the data
+        when the context carries none. The frame is aligned to the schema's
+        columns: an extra column is dropped with a warning, since an existing
+        table is never altered. Each file gets a name of its own, so
+        concurrent writes never read each other's files.
+
+        Args:
+            table: Target table name.
+            schema: The resolved schema name.
+            data: The data in its native representation.
+            context: IO context carrying the asset and effective schema.
+
+        Returns:
+            Where the file was staged and the query reading it.
+        """
+        specs = (context.schema or Representation.of(data).infer()).field_specs()
+        self._ensure_table(table, schema, specs, context)
+
+        frame = Representation.of(data).to("dataframe")
+        names = [spec.name for spec in specs]
+        extras = [str(c) for c in frame.columns if str(c) not in names]
+        if extras:
+            warnings.warn(
+                f"Columns {extras} are not in the schema for '{self._table_ref(table, schema)}' "
+                "and will not be written.",
+                UserWarning,
+                stacklevel=4,
+            )
+        kept = [spec for spec in specs if spec.name in frame.columns]
+
+        path = f"{self.volume_path}/interloper/{uuid.uuid4().hex}.parquet"
+        self._execute(
+            f"PUT '__input_stream__' INTO {_literal(path)} OVERWRITE",
+            input_stream=io.BytesIO(_parquet(frame, kept)),
+        )
+        projection = ", ".join(_projection(spec) for spec in kept)
+        return _Staged(path=path, query=f"SELECT {projection} FROM read_files({_literal(path)}, format => 'parquet')")
+
+    def _load(
+        self,
+        table: str,
+        dataset: str | None,
+        data: Any,
+        context: IOContext,
+        mode: str,
+        clause: str = "",
+    ) -> None:
+        """Stage the data and insert it with one statement, removing the file afterwards.
+
+        The file is removed whether or not the insert succeeds; a failed
+        removal only warns, so it never hides the insert's own outcome.
+
+        Args:
+            table: Target table name.
+            dataset: The schema, or ``None`` for the destination's default.
+            data: The data in its native representation.
+            context: IO context carrying the asset and effective schema.
+            mode: ``INTO`` to add rows, ``OVERWRITE`` to replace the whole table.
+            clause: What follows the column matching, such as a ``REPLACE WHERE``.
+        """
+        schema = self._resolve_dataset(dataset)
+        with self._lock:
+            staged = self._stage(table, schema, data, context)
+            try:
+                self._execute(f"INSERT {mode} {self._table_ref(table, schema)} BY NAME{clause} {staged.query}")
+            finally:
+                try:
+                    self._execute(f"REMOVE {_literal(staged.path)}")
+                except Exception as error:  # noqa: BLE001
+                    warnings.warn(
+                        f"Could not remove the staged file '{staged.path}': {error}", UserWarning, stacklevel=2
+                    )
+
+    # -- DatabaseDestination hooks ---------------------------------------------------
+
+    def insert(self, table: str, dataset: str | None, data: Any, context: IOContext) -> None:
+        """Append the data to the table through a staged Parquet file.
+
+        Args:
+            table: Target table name.
+            dataset: The schema, or ``None`` for the destination's default.
+            data: The data in its native representation.
+            context: IO context carrying the asset and effective schema.
+        """
+        self._load(table, dataset, data, context, "INTO")
+
+    def delete(self, table: str, dataset: str | None, where: PartitionFilter | None) -> None:
+        """Delete the rows a filter selects, or every row.
+
+        A table that does not exist has nothing to delete.
+
+        Args:
+            table: Target table name.
+            dataset: The schema, or ``None`` for the destination's default.
+            where: The rows to delete; ``None`` for the whole table.
+        """
+        schema = self._resolve_dataset(dataset)
+        if not self._table_exists(table, schema):
+            return
+        sql = f"DELETE FROM {self._table_ref(table, schema)}"
+        self._execute(sql if where is None else f"{sql} WHERE {_predicate([where])}")
+
+    def select(self, table: str, dataset: str | None, where: PartitionFilter | None) -> pd.DataFrame:
+        """Select the rows a filter selects, or every row, as a DataFrame.
+
+        The result arrives as Arrow, so column types survive the read without
+        a pass through Python records.
+
+        Args:
+            table: Target table name.
+            dataset: The schema, or ``None`` for the destination's default.
+            where: The rows to select; ``None`` for the whole table.
+
+        Returns:
+            The selected rows.
+
+        Raises:
+            DataNotFoundError: If the table does not exist yet.
+        """
+        schema = self._resolve_dataset(dataset)
+        ref = self._table_ref(table, schema)
+        if not self._table_exists(table, schema):
+            raise DataNotFoundError(f"Table '{ref}' does not exist. Has the asset been materialized?")
+        sql = f"SELECT * FROM {ref}" if where is None else f"SELECT * FROM {ref} WHERE {_predicate([where])}"
+        return self._execute(sql, fetch=lambda cursor: cursor.fetchall_arrow().to_pandas())
+
+    def count(self, table: str, dataset: str | None, column: str) -> dict[str, int]:
+        """Return row counts grouped by a column.
+
+        Args:
+            table: Target table name.
+            dataset: The schema, or ``None`` for the destination's default.
+            column: Column to group by.
+
+        Returns:
+            Mapping from the column's value (as string) to row count.
+
+        Raises:
+            DataNotFoundError: If the table does not exist yet.
+        """
+        schema = self._resolve_dataset(dataset)
+        ref = self._table_ref(table, schema)
+        if not self._table_exists(table, schema):
+            raise DataNotFoundError(f"Table '{ref}' does not exist. Has the asset been materialized?")
+        rows = self._execute(
+            f"SELECT CAST({_quote(column)} AS STRING) AS partition_value, COUNT(*) AS cnt FROM {ref} GROUP BY 1",
+            fetch=lambda cursor: cursor.fetchall(),
+        )
+        return {row[0]: row[1] for row in rows}
+
+
+# -- Utility functions ---------------------------------------------------------------
+
+
+def _quote(identifier: str) -> str:
+    """Quote an identifier with backticks.
+
+    Args:
+        identifier: A catalog, schema, table or column name.
+
+    Returns:
+        The identifier in backticks, embedded backticks doubled.
+    """
+    return "`" + identifier.replace("`", "``") + "`"
+
+
+def _literal(value: Any) -> str:
+    """Render a value as a Databricks SQL literal.
+
+    Every value a statement here carries is a literal: ``REPLACE WHERE``
+    admits literal values in its predicate, and ``PUT``, ``REMOVE`` and
+    ``read_files`` take their paths as string literals, so one renderer
+    serves them all. A naive datetime is read as UTC, as the load writes it.
+
+    Args:
+        value: A string, date, datetime, number, boolean or ``None``.
+
+    Returns:
+        The literal: a backslash-escaped string, ``DATE'...'``,
+        ``TIMESTAMP'...'``, a bare number, ``TRUE``/``FALSE`` or ``NULL``.
+    """
+    if value is None:
+        return "NULL"
+    if isinstance(value, bool):
+        return "TRUE" if value else "FALSE"
+    if isinstance(value, (int, float, Decimal)):
+        return str(value)
+    if isinstance(value, datetime.datetime):
+        stamp = value.isoformat() if value.tzinfo is not None else f"{value.isoformat()}Z"
+        return f"TIMESTAMP'{stamp}'"
+    if isinstance(value, datetime.date):
+        return f"DATE'{value.isoformat()}'"
+    return "'" + str(value).replace("\\", "\\\\").replace("'", "\\'") + "'"
+
+
+def _predicate(filters: Sequence[PartitionFilter]) -> str:
+    """Render the rows several partition filters cover as one predicate.
+
+    Time partitions contribute their half-open bounds, merged where they
+    touch, so a contiguous window becomes the single range from its first
+    start to its last end. Other partitions match their ids by equality.
+
+    Args:
+        filters: The filters, all on the same column.
+
+    Returns:
+        The predicate text.
+    """
+    column = _quote(filters[0].column)
+    ranges: list[tuple[Any, Any]] = []
+    for start, end in sorted(f.bounds for f in filters if f.bounds is not None):
+        if ranges and start <= ranges[-1][1]:
+            ranges[-1] = (ranges[-1][0], max(ranges[-1][1], end))
+        else:
+            ranges.append((start, end))
+    parts = [f"{column} >= {_literal(start)} AND {column} < {_literal(end)}" for start, end in ranges]
+    values = [f.value for f in filters if f.bounds is None]
+    if len(values) == 1:
+        parts.append(f"{column} = {_literal(values[0])}")
+    elif values:
+        parts.append(f"{column} IN ({', '.join(_literal(v) for v in values)})")
+    return parts[0] if len(parts) == 1 else " OR ".join(f"({part})" for part in parts)
+
+
+def _column_ddl(spec: FieldSpec) -> str:
+    """Render a field spec as a column definition.
+
+    Args:
+        spec: The field spec.
+
+    Returns:
+        The quoted column and its type, with its description as a comment.
+        Every column is nullable: conform already enforces the schema's
+        nullability, and a constraint here would only turn a schema change
+        into a failed load.
+    """
+    ddl = f"{_quote(spec.name)} {column_type(spec)}"
+    return f"{ddl} COMMENT {_literal(spec.description)}" if spec.description else ddl
+
+
+def _projection(spec: FieldSpec) -> str:
+    """Read one staged column as the table's type.
+
+    Args:
+        spec: The column's field spec.
+
+    Returns:
+        ``PARSE_JSON`` of the JSON text for a ``VARIANT``, a ``CAST`` to the
+        column type otherwise, aliased to the column so ``BY NAME`` matches it.
+    """
+    name = _quote(spec.name)
+    kind = column_type(spec)
+    expression = f"PARSE_JSON({name})" if kind == VARIANT else f"CAST({name} AS {kind})"
+    return f"{expression} AS {name}"
+
+
+def _parquet(frame: pd.DataFrame, specs: Sequence[FieldSpec]) -> bytes:
+    """Encode the columns a load keeps as one Parquet file.
+
+    This only chooses the file's encoding; the values are already conformed.
+    ``VARIANT`` columns travel as JSON text for ``PARSE_JSON``, and
+    ``TIMESTAMP`` columns as UTC instants in microseconds (a naive datetime
+    read as UTC, so the session's time zone never shifts them; Databricks
+    timestamps hold microseconds, and Spark does not read nanosecond Parquet
+    timestamps as timestamps). A column with no
+    values at all is written as strings, since Parquet's null type is not
+    something every reader accepts; the load casts it to its real type.
+
+    Args:
+        frame: The data as a DataFrame.
+        specs: The kept columns' field specs, in table order.
+
+    Returns:
+        The Parquet file's bytes.
+    """
+    columns: dict[str, pd.Series] = {}
+    for spec in specs:
+        series = frame[spec.name]
+        kind = column_type(spec)
+        if kind == VARIANT:
+            series = series.map(_json)
+        elif kind == TIMESTAMP:
+            series = pd.to_datetime(series, utc=True)
+        columns[spec.name] = series
+    table = pa.Table.from_pandas(pd.DataFrame(columns, index=frame.index), preserve_index=False)
+    for index, field in enumerate(table.schema):
+        if pa.types.is_null(field.type):
+            table = table.set_column(index, field.name, table.column(index).cast(pa.string()))
+    buffer = io.BytesIO()
+    pq.write_table(table, buffer, coerce_timestamps="us", allow_truncated_timestamps=True)
+    return buffer.getvalue()
+
+
+def _json(value: Any) -> str | None:
+    """Encode one nested value as JSON text.
+
+    Args:
+        value: A dict, a list, an array or a missing value.
+
+    Returns:
+        The JSON text, or ``None`` for a missing value.
+    """
+    if value is None or value is pd.NA or (isinstance(value, float) and math.isnan(value)):
+        return None
+    if hasattr(value, "tolist"):
+        value = value.tolist()
+    return json.dumps(replace_non_finite(value), default=json_default)
+
+
+def _asset_description(asset: Any) -> str | None:
+    """Return the asset's description (its class docstring), cleaned.
+
+    This mirrors how ``Component.definition()`` derives descriptions.
+
+    Args:
+        asset: The asset whose docstring is read.
+
+    Returns:
+        The cleaned docstring, or ``None`` when the asset has none.
+    """
+    doc = type(asset).__doc__
+    return inspect.cleandoc(doc) if doc else None

interloper_databricks/types.py

@@ -0,0 +1,67 @@
+"""Databricks' view of interloper's field types."""
+
+from __future__ import annotations
+
+import datetime
+from decimal import Decimal
+
+from interloper.schema import FieldSpec
+
+VARIANT = "VARIANT"
+TIMESTAMP = "TIMESTAMP"
+
+# Ordered: the first base class that matches wins, so bool (a subclass of int)
+# and datetime (a subclass of date) must come before their parents.
+_PYTHON_TO_DATABRICKS: dict[type, str] = {
+    bool: "BOOLEAN",
+    int: "BIGINT",
+    float: "DOUBLE",
+    Decimal: "DECIMAL(38,9)",
+    datetime.datetime: TIMESTAMP,
+    datetime.date: "DATE",
+    bytes: "BINARY",
+    str: "STRING",
+    dict: VARIANT,
+    list: VARIANT,
+}
+
+_CLUSTERABLE = frozenset({"BIGINT", "DOUBLE", "DECIMAL(38,9)", TIMESTAMP, "DATE", "STRING"})
+
+
+def column_type(spec: FieldSpec) -> str:
+    """Return the Databricks column type for a field spec.
+
+    A ``datetime`` is a ``TIMESTAMP``, an absolute instant like BigQuery's:
+    ``TIMESTAMP_NTZ`` is still in Public Preview and upgrades the Delta table
+    protocol. Nested and repeated fields are ``VARIANT``, Databricks' type for
+    semi-structured data, queryable with the ``:`` path operator.
+
+    Args:
+        spec: The field spec, from :meth:`Schema.field_specs` or an inferred schema.
+
+    Returns:
+        ``VARIANT`` for a nested or repeated field, the type the spec's Python
+        type maps to otherwise, and ``STRING`` for anything unmapped
+        (``typing.Any``).
+    """
+    if spec.fields is not None or spec.repeated:
+        return VARIANT
+    if isinstance(spec.type, type):
+        for base, name in _PYTHON_TO_DATABRICKS.items():
+            if issubclass(spec.type, base):
+                return name
+    return "STRING"
+
+
+def clusterable(spec: FieldSpec) -> bool:
+    """Whether a field can be a liquid clustering key.
+
+    Args:
+        spec: The field spec.
+
+    Returns:
+        True when the field's column type is one liquid clustering accepts as a
+        key (dates, timestamps, strings and numbers; not booleans, binary or
+        ``VARIANT``).
+    """
+    return column_type(spec) in _CLUSTERABLE

interloper_databricks-0.96.0.dist-info/METADATA

@@ -0,0 +1,190 @@
+Metadata-Version: 2.3
+Name: interloper-databricks
+Version: 0.96.0
+Summary: Interloper Databricks integration: connection and destination
+Author: Guillaume Onfroy
+Author-email: Guillaume Onfroy <guillaume@digitlcloud.com>
+Requires-Dist: interloper-core
+Requires-Dist: interloper-pandas
+Requires-Dist: databricks-sql-connector[pyarrow]>=4.1.2
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# interloper-databricks
+
+Databricks integration for interloper: a `DatabricksDestination` that stores
+assets as Delta tables in Unity Catalog through a SQL warehouse, and the
+`DatabricksConnection` that holds the workspace credentials.
+
+The destination is a `DatabaseDestination`: the rows a partition covers come
+from core, exactly as for BigQuery. What differs is how they are replaced (see
+[Partitions](#partitions)).
+
+## Setup
+
+### Credentials
+
+A service principal with an OAuth secret is the recommended credential; a
+personal access token works too. The connection takes exactly one of them.
+
+1. In the account console (or the workspace admin settings), create a service
+   principal and add it to the workspace.
+2. Under its **Secrets** tab, click **Generate secret**. Note the client ID
+   and the secret; the secret is shown once. The connection asks for the
+   `all-apis` scope, so leave the secret unscoped (or select all APIs): a
+   secret restricted to some scopes cannot issue that token.
+3. Give the service principal the **Can use** permission on the SQL warehouse
+   the destination loads through.
+
+The connection exchanges the client ID and secret for a one-hour access token
+at the workspace's `/oidc/v1/token` endpoint (client credentials grant), and
+replaces the token before it expires. The connection check calls
+`/api/2.0/preview/scim/v2/Me`, and the pickers list SQL warehouses and Unity
+Catalog catalogs over the REST API, all with that token.
+
+### Grants
+
+The principal needs, in Unity Catalog:
+
+- `USE CATALOG` on the destination's catalog
+- `CREATE SCHEMA` on the catalog, or `USE SCHEMA` on every schema the assets
+  write to if you create them yourself
+- `CREATE TABLE` on those schemas
+- `MODIFY` and `SELECT` on the tables (a table the principal creates is its
+  own, so it has both)
+- `USE SCHEMA`, `READ VOLUME` and `WRITE VOLUME` on the staging volume and its
+  schema
+
+A sketch, with the service principal's application ID as the grantee:
+
+```sql
+GRANT USE CATALOG, CREATE SCHEMA ON CATALOG analytics TO `<application-id>`;
+GRANT USE SCHEMA ON SCHEMA analytics.staging TO `<application-id>`;
+GRANT READ VOLUME, WRITE VOLUME ON VOLUME analytics.staging.loads TO `<application-id>`;
+```
+
+### Staging volume
+
+Every load uploads one Parquet file to a Unity Catalog volume, reads it from
+there and removes it. Create a managed volume for it once:
+
+```sql
+CREATE SCHEMA IF NOT EXISTS analytics.staging;
+CREATE VOLUME IF NOT EXISTS analytics.staging.loads COMMENT 'interloper load staging';
+```
+
+and name it on the destination as `analytics.staging.loads`. Files go under
+`/Volumes/analytics/staging/loads/interloper/`, one per write. A file whose
+removal fails (the load's outcome stands, with a warning) is left there and
+can be deleted by hand.
+
+## Usage
+
+```python
+import interloper as il
+from interloper_databricks import DatabricksConnection, DatabricksDestination
+
+destination = DatabricksDestination(
+    connection=DatabricksConnection(
+        host="https://dbc-a1b2345c-d6e7.cloud.databricks.com",
+        client_id="...",
+        client_secret="...",
+    ),
+    warehouse="/sql/1.0/warehouses/a1b234c567d8e9fa",
+    catalog="analytics",
+    staging_volume="analytics.staging.loads",
+    default_dataset="raw",
+)
+```
+
+`warehouse` is the warehouse's HTTP path (its **Connection details** tab). The
+credentials also load from the environment (`DATABRICKS_HOST`,
+`DATABRICKS_CLIENT_ID`, `DATABRICKS_CLIENT_SECRET`, or
+`DATABRICKS_ACCESS_TOKEN` for a personal access token; note that Databricks'
+own tools name the token `DATABRICKS_TOKEN`).
+
+In a deployed instance you configure this through the UI instead: add a
+Databricks connection, then a Databricks destination, picking the warehouse
+and the catalog from the lists the connection can see.
+
+## Datasets are schemas
+
+An asset's `dataset` is the schema its table lives in, inside the
+destination's `catalog`. An asset without a dataset falls back to
+`default_dataset`; with neither, the write fails with a `ConfigError`. A
+missing schema is created on the first write, and a missing table is created
+as a Delta table typed from the asset's schema (or one inferred from the
+data), with field descriptions as column comments and the asset's description
+as the table comment:
+
+| Field type | Databricks type |
+|------------|-----------------|
+| `bool` | `BOOLEAN` |
+| `int` | `BIGINT` |
+| `float` | `DOUBLE` |
+| `Decimal` | `DECIMAL(38,9)` |
+| `datetime` | `TIMESTAMP` |
+| `date` | `DATE` |
+| `bytes` | `BINARY` |
+| `str`, `Any` | `STRING` |
+| nested models, lists, dicts | `VARIANT` |
+
+A `datetime` is a `TIMESTAMP`, an absolute instant; a naive datetime is taken
+as UTC. `VARIANT` columns are queried with the path operator
+(`SELECT payload:campaign.id FROM ...`); creating a table with one enables
+Delta's `variantType` feature, which readers need Databricks Runtime 15.4 or a
+recent Delta client for.
+
+Identifiers are quoted with backticks. Unity Catalog stores schema and table
+names in lower case and keeps column names as written; queries match either
+case-insensitively.
+
+An existing table is never altered: a column the data carries but the schema
+does not is dropped with a warning.
+
+A partitioned asset's table is clustered on its partition column with liquid
+clustering (`CLUSTER BY`), which is what Databricks recommends over
+partitioning for tables under 1 TB, when that column is one liquid clustering
+accepts as a key (a date, timestamp, string or number among the first 32
+columns). Tables are never `PARTITIONED BY`.
+
+## Partitions
+
+Each write is one statement, so a partition is replaced atomically:
+
+- a time partition:
+  ``INSERT INTO ... BY NAME REPLACE WHERE `day` >= DATE'2024-01-01' AND `day` < DATE'2024-01-02' SELECT ...``
+  (half-open bounds, so a monthly partition whose rows hold daily dates is
+  replaced whole)
+- any other partition: ``REPLACE WHERE `region` = 'eu'``
+- a window: one predicate covering every partition in it, a single range from
+  its first start to its last end when the partitions are contiguous
+- an unpartitioned asset: `INSERT OVERWRITE ... BY NAME SELECT ...`
+
+The `SELECT` reads the staged file with `read_files(..., format => 'parquet')`,
+casting each column to the table's type. `REPLACE WHERE` deletes the matching
+rows and inserts the new ones in a single Delta commit, so a failed load
+leaves the partition's old rows in place.
+
+`REPLACE WHERE` is strict: every row written must match the predicate, or the
+statement fails with `DELTA_REPLACE_WHERE_MISMATCH` and writes nothing. A row
+whose partition column is null, or falls outside the partition, fails the
+write instead of landing where the next replace would not remove it.
+
+Reads come back as a DataFrame built from the result's Arrow batches.
+
+## Notes
+
+Each destination opens its own SQL session through the connection, on its own
+warehouse and catalog, so destinations sharing a connection never affect each
+other. That session is shared by every asset the destination writes, and the
+connector's sessions are not thread-safe, so the destination runs one
+statement at a time and holds a write from its upload to its cleanup.
+
+The SQL warehouse does the loading work: it reads the staged file, casts it
+and rewrites the replaced rows, so its size bounds load throughput. Since a
+destination runs one statement at a time, size the warehouse for the largest
+single write (a partition, or a whole window) rather than for concurrency;
+several destinations, or several runs, can share a warehouse that scales out.
+A stopped warehouse adds its start-up time to the first write after it
+auto-stops, which serverless warehouses keep short.

interloper_databricks-0.96.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+interloper_databricks/__init__.py,sha256=ZUOTzYoMMHFFmOjZMl0mZC-Qr5gZ1z3vy1_aD6axwP0,276
+interloper_databricks/connection.py,sha256=paI0zI7UglgfmIvvVtHdQqMy9aI-CV45Ep0F7AaeKfE,11310
+interloper_databricks/destination.py,sha256=LV8MOE6ioex4vctmka-psicqJa1TvzkHldc4cbh7I8Y,25338
+interloper_databricks/types.py,sha256=LdYpqX0FTuz5AIKLWCo9YQ6qOKYld69OZUgZ5qELmK4,2073
+interloper_databricks-0.96.0.dist-info/WHEEL,sha256=e4_1dyBeezi8ZjfxrZ3bnVOxFDa3ksqVqH0jTHkUZ3k,81
+interloper_databricks-0.96.0.dist-info/entry_points.txt,sha256=-iwSBH6L6INfkEsrZH_S6kCznO_R8E_ObMG_TD-KaUk,60
+interloper_databricks-0.96.0.dist-info/METADATA,sha256=nmjeljZw6vDqZIUS07HknFdYY0Bs67bMZiLe1isq08Q,8044
+interloper_databricks-0.96.0.dist-info/RECORD,,

interloper_databricks-0.96.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.12.19
+Root-Is-Purelib: true
+Tag: py3-none-any

interloper_databricks-0.96.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[interloper.components]
+databricks = interloper_databricks
+