> For the complete documentation index, see [llms.txt](https://docs.espresso.ai/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://docs.espresso.ai/databricks-optimizer/databricks-sql-onboarding-1.md).

# Databricks SQL Onboarding

Espresso AI makes real-time decisions powered by ML models to optimize your Databricks SQL workloads.

Please run the following commands in Databricks SQL. This will set up a user account we can use to access the metadata we need for our models as well as optimize your Databricks SQL account's operations.

To manage onboarding as code instead, use the [Espresso Databricks Terraform module](/databricks-optimizer/databricks-terraform-onboarding.md).

Note that we **never** access, log, or store any data from Databricks SQL. We only look at metadata.

Prerequisites: Make sure you are a Databricks workspace admin, account admin, and metastore admin.

## Steps

1. Open your Databricks workspace, create a Python notebook, and attach it to Serverless compute.
2. In the first cell, install/upgrade the SDK:

```python
%pip install databricks-sdk --upgrade
```

3. When finished, create a second cell and restart the Python kernel:

```python
%restart_python
```

4. In the next cell, copy-paste the code below. When you run this, if you have multiple workspaces, you will be prompted to grant a temporary service principal account admin credentials. This simply lets us grant permissions across workspaces, and we delete that temporary service principal immediately in the script. If you only have one workspace to manage, change the first line to `MULTI_WORKSPACE = False` before running:

```python
MULTI_WORKSPACE = True

import json
import time
from datetime import datetime

from databricks.sdk import AccountClient, WorkspaceClient
from databricks.sdk.errors import BadRequest
from databricks.sdk.service import catalog, jobs
from databricks.sdk.service.compute import ClusterSource, ListClustersFilterBy
from databricks.sdk.service.iam import (
    AccessControlRequest,
    ComplexValue,
    Patch,
    PatchOp,
    PatchSchema,
    PermissionLevel,
    WorkspacePermission,
)
from databricks.sdk.service.workspace import ImportFormat, Language
from pydantic import BaseModel, field_validator


class DatabricksOAuthToken(BaseModel):
    id: str
    oauth_secret: str
    created_at: datetime
    expires_at: datetime
    client_id: str

    @field_validator("oauth_secret", "client_id", mode="before")
    @classmethod
    def strip_whitespace(cls, v: str) -> str:
        """Strip whitespace from string fields."""
        return v.strip() if isinstance(v, str) else v


class DatabricksCredentials(BaseModel):
    oauth_token: DatabricksOAuthToken
    workspace_url: str
    service_principal_name: str = "espresso-ai-optimizer"
    service_principal_id: str
    warehouse_id: str | None = None
    warehouse_name: str | None = None
    workspace_id: str | None = None
    workspace_name: str | None = None

    def to_json(self) -> str:
        return json.dumps(self.model_dump(mode="json"))


SYNC_SP_NAME = "espresso-ai-permission-sync"

ALL_PURPOSE_ONLY = ListClustersFilterBy(cluster_sources=[ClusterSource.UI, ClusterSource.API])

MANAGED_RESOURCES = [
    ("warehouses", "warehouses", "list", "id"),
    ("clusters", "clusters", "list", "cluster_id"),
    ("jobs", "jobs", "list", "job_id"),
    ("pipelines", "pipelines", "list_pipelines", "pipeline_id"),
    ("instance-pools", "instance_pools", "list", "instance_pool_id"),
]

SYNC_NOTEBOOK = """# Databricks notebook source
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.compute import ClusterSource, ListClustersFilterBy
from databricks.sdk.service.iam import AccessControlRequest, PermissionLevel

APP_ID = "{app_id}"
ALL_PURPOSE_ONLY = ListClustersFilterBy(cluster_sources=[ClusterSource.UI, ClusterSource.API])
client = WorkspaceClient()
acl = [
    AccessControlRequest(
        service_principal_name=APP_ID, permission_level=PermissionLevel.CAN_MANAGE
    )
]
for object_type, service, list_method, id_attr in [
{resources}
]:
    lister = getattr(getattr(client, service), list_method)
    for obj in (lister(filter_by=ALL_PURPOSE_ONLY) if object_type == "clusters" else lister()):
        if getattr(obj, "creator_user_name", None) == APP_ID:
            continue
        if object_type == "jobs" and obj.settings.name == "{sync_job}":
            continue
        try:
            client.permissions.update(
                request_object_type=object_type,
                request_object_id=str(getattr(obj, id_attr)),
                access_control_list=acl,
            )
        except Exception as e:
            print(f"failed to grant on {{object_type}} {{getattr(obj, id_attr)}}: {{e}}")
"""


def _is_duplicate_error(e):
    msg = str(e).lower()
    return "already exists" in msg or "duplicate" in msg


def _patch_sp_entitlements(client, sp):
    client.service_principals.patch(
        id=sp.id,
        operations=[
            Patch(
                op=PatchOp.ADD,
                path="entitlements",
                value=[
                    {"value": "databricks-sql-access"},
                    {"value": "allow-cluster-create"},
                ],
            ),
        ],
        schemas=[PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP],
    )


def get_or_create_service_principal(client, name="espresso-ai-optimizer"):
    if sps := list(client.service_principals.list(filter=f"displayName eq '{name}'")):
        sp = sps[0]
        _patch_sp_entitlements(client, sp)
        return sp

    return client.service_principals.create(
        display_name=name,
        active=True,
        entitlements=[
            ComplexValue(value="allow-cluster-create"),
            ComplexValue(value="databricks-sql-access"),
        ],
    )


def find_service_principal(client, name="espresso-ai-optimizer", timeout_secs=60):
    deadline = time.monotonic() + timeout_secs
    last_err = None
    while time.monotonic() < deadline:
        try:
            sps = list(client.service_principals.list(filter=f"displayName eq '{name}'"))
            if sps:
                return sps[0]
        except Exception as e:
            # A freshly-assigned principal may not be able to authenticate to the
            # workspace yet; keep polling until the assignment propagates.
            last_err = e
        time.sleep(5)
    raise RuntimeError(
        f"Service principal {name!r} not visible after {timeout_secs}s. "
        "Workspace assignment may not have propagated."
        + (f" Last error: {last_err}" if last_err else "")
    )


def create_oauth_token(client, service_principal):
    token = client.service_principal_secrets_proxy.create(
        service_principal_id=service_principal.id,
        lifetime=f"{2 * 365 * 24 * 60 * 60}s",  # 2 years
    )
    return DatabricksOAuthToken(
        id=token.id,
        oauth_secret=token.secret,
        created_at=token.create_time,
        expires_at=token.expire_time,
        client_id=service_principal.application_id,
    )


def grant_service_principal_resource_permissions(client, service_principal):
    acl = [
        AccessControlRequest(
            service_principal_name=service_principal.application_id,
            permission_level=PermissionLevel.CAN_MANAGE,
        )
    ]
    for object_type, service, list_method, id_attr in MANAGED_RESOURCES:
        try:
            lister = getattr(getattr(client, service), list_method)
            objects = list(
                lister(filter_by=ALL_PURPOSE_ONLY) if object_type == "clusters" else lister()
            )
        except Exception as e:
            print(f"  {object_type}: ⚠️  could not list ({e})")
            continue

        granted = 0
        errors = []
        for obj in objects:
            object_id = getattr(obj, id_attr, None)
            if object_id is None or (
                object_type == "jobs" and obj.settings.name == SYNC_SP_NAME
            ):
                continue
            try:
                client.permissions.update(
                    request_object_type=object_type,
                    request_object_id=str(object_id),
                    access_control_list=acl,
                )
                granted += 1
            except Exception as e:
                errors.append(f"{object_type} {object_id}: {e}")
        for err in errors:
            print(f"    ⚠️  failed to grant on {err}")


def make_service_principal_workspace_admin(client, sp):
    try:
        client.groups.patch(
            id=next(client.groups.list(filter="displayName eq 'admins'")).id,
            operations=[Patch(op=PatchOp.ADD, value={"members": [{"value": sp.id}]})],
            schemas=[PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP],
        )
    except Exception as e:
        if not _is_duplicate_error(e):
            raise


def schedule_permission_sync(client, service_principal):
    sync_sp = get_or_create_service_principal(client, SYNC_SP_NAME)
    make_service_principal_workspace_admin(client, sync_sp)
    for old in list(client.jobs.list(name=SYNC_SP_NAME)):
        client.jobs.delete(job_id=old.job_id)
    secret = client.service_principal_secrets_proxy.create(
        service_principal_id=sync_sp.id, lifetime=f"{60 * 60}s"
    )
    sync_client = WorkspaceClient(
        host=client.config.host,
        client_id=sync_sp.application_id,
        client_secret=secret.secret,
    )
    find_service_principal(sync_client, SYNC_SP_NAME)

    sync_client.workspace.mkdirs(home := f"/Users/{sync_sp.application_id}")
    notebook_path = f"{home}/{SYNC_SP_NAME}"
    sync_client.workspace.upload(
        notebook_path,
        SYNC_NOTEBOOK.format(
            app_id=service_principal.application_id,
            sync_job=SYNC_SP_NAME,
            resources="\n".join(f"    {r}," for r in MANAGED_RESOURCES),
        ).encode(),
        format=ImportFormat.SOURCE,
        language=Language.PYTHON,
        overwrite=True,
    )
    sync_client.jobs.create(
        name=SYNC_SP_NAME,
        tasks=[
            jobs.Task(
                task_key="sync",
                notebook_task=jobs.NotebookTask(notebook_path=notebook_path),
            )
        ],
        schedule=jobs.CronSchedule(quartz_cron_expression="0 0 * * * ?", timezone_id="UTC"),
        performance_target=jobs.PerformanceTarget.STANDARD,
    )
    client.service_principal_secrets_proxy.delete(
        service_principal_id=sync_sp.id, secret_id=secret.id
    )
    print(f"  Hourly permission sync scheduled, owned by {SYNC_SP_NAME}")


def get_or_create_warehouse(client):
    for warehouse in client.warehouses.list():
        if warehouse.name == "ESPRESSO_AI_WAREHOUSE":
            return warehouse.id

    return client.warehouses.create_and_wait(
        name="ESPRESSO_AI_WAREHOUSE",
        cluster_size="X-Small",
        auto_stop_mins=1,
        enable_serverless_compute=True,
        min_num_clusters=1,
        max_num_clusters=1,
    ).id


def allow_service_principal_to_read_system_logs(client, service_principal):
    errors = []

    def grant(asset_name, asset_type, privilege):
        try:
            client.grants.update(
                full_name=asset_name,
                securable_type=asset_type.value,
                changes=[
                    catalog.PermissionsChange(
                        add=[privilege], principal=service_principal.application_id
                    )
                ],
            )
        except Exception as e:
            errors.append(str(e).lower())

    grant("system", catalog.SecurableType.CATALOG, catalog.Privilege.USE_CATALOG)
    for schema in client.schemas.list(catalog_name="system"):
        if schema.name in ["data_classification", "data_quality_monitoring"]:
            continue
        schema_full_name = f"system.{schema.name}"
        grant(schema_full_name, catalog.SecurableType.SCHEMA, catalog.Privilege.USE_SCHEMA)
        grant(schema_full_name, catalog.SecurableType.SCHEMA, catalog.Privilege.SELECT)

    if not errors:
        return

    if any("account admin" in e for e in errors):
        print("\n⚠️  ACCOUNT ADMIN required: Ask an account admin to grant you access.")
    if any("manage on catalog" in e or "metastore" in e for e in errors):
        accounts_url = client.config.environment.deployment_url("accounts")
        print(f"⚠️  METASTORE ADMIN required: Visit {accounts_url}/data")

    raise RuntimeError(
        "Failed to grant system table access (resolve the above and re-run):\n  - "
        + "\n  - ".join(errors)
    )


def wait_for_account_admin(accounts_url, account_id, client_id, secret):
    deadline = time.monotonic() + 300  # 5 minutes
    while time.monotonic() < deadline:
        try:
            ac = AccountClient(
                host=accounts_url,
                account_id=account_id,
                client_id=client_id,
                client_secret=secret,
            )
            workspaces = [
                (ws.workspace_id, ws.workspace_name, ws.deployment_name)
                for ws in ac.workspaces.list()
            ]
            print("Account admin access confirmed.")
            return ac, workspaces
        except Exception:
            time.sleep(5)
    raise TimeoutError(
        "Account admin not granted within 5 minutes. Grant the role via the link above and re-run."
    )


def filter_to_current_metastore(account_client, workspace_client, workspaces):
    metastore_id = workspace_client.metastores.current().metastore_id
    metastore_ws_ids = set(account_client.metastore_assignments.list(metastore_id))
    return [ws for ws in workspaces if ws[0] in metastore_ws_ids]


def check_workspaces_federated(account_client, workspaces):
    not_federated = []
    for ws_id, ws_name, _ in workspaces:
        try:
            account_client.workspace_assignment.list(workspace_id=ws_id)
        except BadRequest as e:
            if "permission assignment apis are not available" not in str(e).lower():
                raise
            not_federated.append((ws_id, ws_name))
    if not_federated:
        details = "\n  - ".join(f"{name} ({wid})" for wid, name in not_federated)
        raise PermissionError(
            "IDENTITY FEDERATION required on the following workspaces:\n"
            f"  - {details}\n\n"
            "Enable identity federation via the account console:\n"
            "  Workspaces → <workspace> → Configuration tab → "
            "'Identity federation' must show 'Enabled'.\n"
        )


def _assign_workspace_permission(account_client, ws_id, principal_id, permission):
    deadline = time.monotonic() + 60
    while True:
        try:
            account_client.workspace_assignment.update(
                workspace_id=int(ws_id),
                principal_id=int(principal_id),
                permissions=[permission],
            )
            return
        except Exception as e:
            if _is_duplicate_error(e):
                return
            if time.monotonic() >= deadline:
                raise
            time.sleep(5)


def assign_sp_to_workspaces(
    account_client, sp_id, workspaces, permission=WorkspacePermission.USER
):
    for ws_id, ws_name, _ in workspaces:
        _assign_workspace_permission(account_client, ws_id, sp_id, permission)
        print(f"  Assigned {permission.value} on {ws_name} ({ws_id})")


def setup_workspace(
    account_client,
    temp_sp,
    temp_secret,
    workspace_url,
    workspace_id,
    sp_name,
):
    _assign_workspace_permission(
        account_client, workspace_id, temp_sp.id, WorkspacePermission.ADMIN
    )
    admin_client = WorkspaceClient(
        host=workspace_url,
        client_id=temp_sp.application_id,
        client_secret=temp_secret,
    )
    sp = find_service_principal(admin_client, sp_name)
    _patch_sp_entitlements(admin_client, sp)
    wh_id = get_or_create_warehouse(admin_client)
    grant_service_principal_resource_permissions(admin_client, sp)
    schedule_permission_sync(admin_client, sp)
    return wh_id


if __name__ == "__main__":
    from importlib import import_module

    SparkSession = import_module("pyspark.sql").SparkSession
    spark = SparkSession.getActiveSession() or SparkSession.builder.getOrCreate()
    workspace_url = f"https://{spark.conf.get('spark.databricks.workspaceUrl') or ''}"
    client = WorkspaceClient()
    environment = client.config.environment
    accounts_url = environment.deployment_url("accounts")
    service_principal = get_or_create_service_principal(client)
    oauth_token = create_oauth_token(client, service_principal)
    warehouse_id = get_or_create_warehouse(client)
    grant_service_principal_resource_permissions(client, service_principal)
    schedule_permission_sync(client, service_principal)
    allow_service_principal_to_read_system_logs(client, service_principal)

    credentials_by_workspace: dict[str, DatabricksCredentials] = {}

    if not MULTI_WORKSPACE:
        workspace_id = str(client.get_workspace_id())
        workspace_name = spark.sql(
            f"SELECT workspace_name FROM system.access.workspaces_latest "
            f"WHERE workspace_id = {workspace_id}"
        ).collect()[0][0]
        credentials_by_workspace[workspace_name] = DatabricksCredentials(
            oauth_token=oauth_token,
            workspace_url=workspace_url,
            workspace_id=workspace_id,
            workspace_name=workspace_name,
            service_principal_name=service_principal.display_name,
            service_principal_id=service_principal.id,
            warehouse_id=warehouse_id,
            warehouse_name="ESPRESSO_AI_WAREHOUSE",
        )
    else:
        account_id = spark.sql(
            "SELECT account_id FROM system.billing.usage LIMIT 1"
        ).collect()[0][0]

        temp_sp = get_or_create_service_principal(client, "espresso-ai-temp")
        temp_secret = client.service_principal_secrets_proxy.create(
            service_principal_id=temp_sp.id, lifetime=f"{60 * 60}s"
        ).secret
        roles_url = (
            f"{accounts_url}/user-management/"
            f"serviceprincipals/{temp_sp.id}/roles?account_id={account_id}"
        )
        print(
            f"\n📋 To continue, grant the temporary service principal account admin access."
            f"\n   Open this link and add the 'Account admin' role:\n\n   {roles_url}\n"
        )

        account_client, workspaces = wait_for_account_admin(
            accounts_url, account_id, temp_sp.application_id, temp_secret
        )
        workspaces = filter_to_current_metastore(account_client, client, workspaces)
        check_workspaces_federated(account_client, workspaces)
        assign_sp_to_workspaces(account_client, int(service_principal.id), workspaces)

        current_ws_id = str(client.get_workspace_id())
        for ws_id, ws_name, deployment_name in workspaces:
            ws_id_str = str(ws_id)
            workspace_url_for_ws = environment.deployment_url(deployment_name)
            if ws_id_str == current_ws_id:
                ws_warehouse_id = warehouse_id
            else:
                ws_warehouse_id = setup_workspace(
                    account_client=account_client,
                    temp_sp=temp_sp,
                    temp_secret=temp_secret,
                    workspace_url=workspace_url_for_ws,
                    workspace_id=ws_id_str,
                    sp_name=service_principal.display_name,
                )
            credentials_by_workspace[ws_name] = DatabricksCredentials(
                oauth_token=oauth_token,
                workspace_url=workspace_url_for_ws,
                workspace_id=ws_id_str,
                workspace_name=ws_name,
                service_principal_name=service_principal.display_name,
                service_principal_id=service_principal.id,
                warehouse_id=ws_warehouse_id,
                warehouse_name="ESPRESSO_AI_WAREHOUSE" if ws_warehouse_id else None,
            )

        delete_url = (
            f"{accounts_url}/user-management/"
            f"serviceprincipals/{temp_sp.id}?account_id={account_id}"
        )
        print(
            f"\n🧹 When you're done, delete the temporary 'espresso-ai-temp' service"
            f"\n   principal — it still has account admin. Open this link, then"
            f"\n   ⋮ menu → Delete → Confirm delete:\n\n   {delete_url}\n"
        )

    print("\n🎉 Setup complete! Here are the Databricks credentials to send to Espresso AI:")
    print("=" * 50)
    print(
        json.dumps(
            {
                ws: cred.model_dump(mode="json")
                for ws, cred in credentials_by_workspace.items()
            },
            indent=2,
        )
    )
    print("=" * 50)

```

5. Copy the JSON credentials printed between the ===== lines and share it securely with Espresso AI.

You can securely [upload the output here](https://www.dropbox.com/request/IGcHPWby1x9tPPv8hXWr).

## What the Script Does:

* Create or reuse a service principal named "espresso-ai-optimizer" for Espresso AI.
* Grant that service principal the ability to manage SQL Warehouses.
* Grant SELECT access on Databricks system tables used for usage and cost analysis (e.g. query history, warehouse events, node timeline).
* If we're granting access to multiple workspaces, we create a temporary service principal "espresso-ai-temp" with account admin privileges to grant "espresso-ai-optimizer" the same permissions across the workspaces. We then delete "espresso-ai-temp."
* Create an hourly job that checks for any new warehouses and grants CAN\_MANAGE permissions to our service principal.
* Print a JSON blob with credentials and your workspace URL for you to share securely with Espresso AI.

## Troubleshooting

The script provides error messages for permission issues:

1. "You are not a workspace admin" → Contact one of the listed admins for access
2. "ACCOUNT ADMIN PERMISSIONS REQUIRED" → You need account admin to create service principals
3. "METASTORE ADMIN PERMISSIONS REQUIRED" →You need metastore admin to grant system table access
4. "Account admin not granted within 5 minutes. Grant the role via the link above and re-run." → You need to grant the temporary service principal account admin privileges. This SP will be deleted.

Quick permission checks:

* Verify that you started a serverless notebook, not a notebook on a specific cluster
* Verify you're a workspace admin:
  * Click on the circle with your initials
  * Click on "Settings"
  * You should see "Workspace Admin" section in addition to a "User" section.
* Verify you're an account admin:
  * You should be able to login to <https://accounts.cloud.databricks.com/> and see the console to manage your account.
* Verify you're a metastore admin:
  * Visit <https://accounts.cloud.databricks.com/data>
  * Click on the name of the metastore listed, which will open up a "Configuration" page.
  * Set yourself as the "Metastore Admin".

Please rerun the script once you have the required permissions.

## Questions?

[Book a Call](https://espresso.ai/demo)