Databricks

Query Databricks SQL warehouses from deployed models and write forecasts back to Unity Catalog tables.

rb.sources.databricks connects deployed Rebase code to a Databricks SQL warehouse. It shares the uniform data-source surface: read, read_bitemporal, and write.

Install

uv pip install "rebase-toolkit[databricks]"

Ships databricks-sql-connector and pyarrow. The OAuth machine-to-machine auth below also needs databricks-sdk, which the extra does not include, so install it alongside. Declare both on the image you deploy:

image = rb.Image.python("3.12").uv_pip_install("rebase-toolkit[databricks]", "databricks-sdk")

With a personal access token only, databricks-sdk is not needed.

Authenticate

OAuth machine-to-machine (a Databricks service principal) is recommended — tokens are minted automatically and never stored:

# 1. Create a service principal in your Databricks account console and
#    generate an OAuth secret for it; grant it access to the SQL warehouse
#    and catalog.

# 2. Store the credentials as a Rebase secret bundle
rebase secret create acme-databricks \
    DATABRICKS_SERVER_HOSTNAME=adb-1234567890.1.azuredatabricks.net \
    DATABRICKS_HTTP_PATH=/sql/1.0/warehouses/abc123 \
    DATABRICKS_CLIENT_ID=xxxxxxxx-... \
    DATABRICKS_CLIENT_SECRET=-  < secret.txt

A personal access token (DATABRICKS_TOKEN) works too and takes precedence when set. Find server_hostname and http_path under your SQL warehouse's Connection details tab.

Settings

Setting	Environment variable	Required
`server_hostname`	`DATABRICKS_SERVER_HOSTNAME` / `DATABRICKS_HOST`	✓
`http_path`	`DATABRICKS_HTTP_PATH`	✓
`access_token`	`DATABRICKS_TOKEN` / `DATABRICKS_ACCESS_TOKEN`	one auth method
`client_id` + `client_secret`	`DATABRICKS_CLIENT_ID`, `DATABRICKS_CLIENT_SECRET`	one auth method
`catalog`, `schema`	`DATABRICKS_CATALOG`, `DATABRICKS_SCHEMA`	optional

Any setting can also be passed to the factory directly, or namespaced per connection as REBASE_SOURCE_<CONNECTION>_<FIELD> — see credential precedence.

Read

The connector binds named parameters with %(name)s and fetches results via Arrow:

import rebase as rb

src = rb.sources.databricks(connection="acme")
df = src.read(
    "SELECT ts, issued_at, load_mw FROM energy.demand WHERE site = %(site)s",
    params={"site": "site-001"},
)

For backtesting, declare knowledge time so emflow can prove there is no leakage:

from rebase.sources import BitemporalSpec

spec = BitemporalSpec(valid_time="ts", knowledge_time="issued_at")
df = src.read_bitemporal(
    "SELECT ts, issued_at, load_mw FROM energy.demand",
    spec,
)

Write forecasts back

Writes do not work yet with the pinned databricks-sql-connector (4.x). The insert uses %s markers, which the driver does not accept when use_inline_params is off, its default. Reads are unaffected. Until this is fixed in the SDK, write through a notebook or COPY INTO.

Writes run one parameterized INSERT per row, so they suit only small, forecast-sized results. mode="replace" runs TRUNCATE TABLE first, and that is not transactional with the inserts:

src.write(forecast_df, "energy.forecasts", mode="append")

For large bulk loads, stage Parquet in cloud storage and COPY INTO instead.

Reference

Data warehouses overview — shared surface, credentials, BitemporalSpec
Backtesting — what the bitemporal columns enable