Connect Data in Databricks

Connect your Databricks notebook to source files or Unity Catalog tables. Complete Databricks Notebook Setup first. For S3, include the S3 integration with kumo-engine[databricks,s3].

Choose a source

Source	Connection
Files in Amazon S3	Direct S3 access
Tables in Unity Catalog	Notebook authentication
ADLS files registered as tables	ADLS through Unity Catalog

Running Python on Linux instead? Use the Linux Databricks connection.

Amazon S3 from Databricks

On AWS classic compute with an authorized instance profile already attached, use the AWS credential chain directly. The role needs permission to list the prefix and read its objects:

import kumo.engine as kumo
connector = kumo.S3Connector(
    root_dir="s3://<bucket>/<prefix>/",
    region="<aws-region>",
)
print(connector.table_names())
print(list(connector["<table-name-from-listing>"].column_dict))

Replace the bucket, prefix, region, and table name. Expect a listing and the selected table’s column names. Unity Catalog grants alone do not authorize this direct S3 read. For access failures, check the compute attachment, role, bucket policy, and any KMS key policy.

If credentials are not configured, follow IAM Instance Profile for rotating credentials on AWS classic compute, or Temporary Credentials in Databricks Secrets for a short session. Continue to Tables after the schema check succeeds.

Azure Data Lake Storage from Databricks

Connect to ADLS data through tables registered in Unity Catalog. If your files are not registered yet, have your data administrator configure an ADLS external location and register external tables.

Then follow Unity Catalog from a Databricks Notebook below with the registered catalog, schema, and table names. Kumo uses DatabricksConnector; Unity Catalog manages the underlying storage credentials.

Unity Catalog from a Databricks Notebook

Create a connector to an existing Unity Catalog schema after Databricks Notebook Setup. Choose notebook authentication below when the Databricks SDK can resolve the workspace identity from your execution environment. Alternatively, use the explicit OAuth service principal procedure. Neither choice grants access to compute or data. Kumo sends data queries through the Databricks SQL connector to the warehouse or cluster you select. An active Spark session alone does not verify this access.

For either choice, the resolved principal needs CAN USE on the SQL warehouse, plus the following Unity Catalog grants. A catalog or data owner runs these in SQL Editor. Replace <principal> with the user’s email address or the service principal’s application ID, and repeat the table grant for each source table:

GRANT USE CATALOG ON CATALOG `<catalog>` TO `<principal>`;
GRANT USE SCHEMA ON SCHEMA `<catalog>`.`<source-schema>` TO `<principal>`;
GRANT SELECT ON TABLE `<catalog>`.`<source-schema>`.`<source-table>` TO `<principal>`;

If you select an all-purpose cluster instead, grant the principal access to that compute resource. The notebook’s artifact-volume permissions are separate from these source read grants.

Use Notebook Authentication

Install the Databricks integration through the setup guide and use a notebook environment where the Databricks SDK can authenticate. The resolved principal needs access to the selected compute resource, USE CATALOG, USE SCHEMA, and SELECT on the source tables. Replace the catalog, schema, and warehouse ID with resources that principal can access:

import kumo.engine as kumo
connector = kumo.DatabricksConnector.from_notebook(
    catalog="<catalog>",
    schema_name="<source-schema>",
    warehouse_id="<sql-warehouse-id>",
)

Specify exactly one of warehouse_id or cluster_id; use cluster_id="<all-purpose-cluster-id>" instead of the warehouse argument to select a cluster. Kumo does not infer a query target from the notebook’s attached compute. from_notebook() delegates authentication to the Databricks SDK. Creating the connector alone does not authenticate or validate access. For explicit OAuth credentials, use Prepare the Service Principal. Continue to Check Source Access after creating this connector. The output grants apply to both authentication choices.

Check Source Access

Run this check after either authentication choice. List available tables, then inspect a table you intend to use:

print(connector.table_names())
source = connector["<source-table-name>"]
print([(column.name, column.dtype, column.stype) for column in source.columns])

Expect the available table names followed by the selected table’s column names and types. Use a table name within the configured schema. If listing or schema inspection fails, check the identity resolved by the Databricks SDK (or the explicit host and OAuth credentials), compute availability, Unity Catalog grants, and network access from the notebook. A successful schema read does not prove that a later full graph snapshot or output write will succeed.

Continue to Tables when the expected columns appear.

Advanced configuration

IAM Instance Profile

Use this path on Databricks AWS classic compute. An instance profile supplies rotating AWS credentials to the cluster. You need an AWS administrator to create the role and profile, a Databricks workspace administrator to register it, and permission to edit the cluster.

Create, register, and verify the instance profile

1. Create the S3 access role. On your administration machine, sign in to the AWS CLI with an administrator profile for the account running Databricks compute. Check the account with aws sts get-caller-identity. The following commands use that configured AWS CLI identity.

Choose an unused role and profile name, shown here as kumo-s3-read. Save this EC2 trust policy as ec2-trust.json:

{
  "Version": "2012-10-17",
  "Statement": [{
    "Effect": "Allow",
    "Principal": {"Service": "ec2.amazonaws.com"},
    "Action": "sts:AssumeRole"
  }]
}

Save the following as s3-read.json. Replace <bucket> with your bucket name and <prefix> with the source directory, without leading or trailing slashes. This policy permits listing and reading only that prefix:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": "s3:GetBucketLocation",
      "Resource": "arn:aws:s3:::<bucket>"
    },
    {
      "Effect": "Allow",
      "Action": "s3:ListBucket",
      "Resource": "arn:aws:s3:::<bucket>",
      "Condition": {"StringLike": {"s3:prefix": ["<prefix>", "<prefix>/*"]}}
    },
    {
      "Effect": "Allow",
      "Action": "s3:GetObject",
      "Resource": "arn:aws:s3:::<bucket>/<prefix>/*"
    }
  ]
}

Create the role, attach the policy, and place the role in an instance profile:

aws iam create-role --role-name kumo-s3-read \
  --assume-role-policy-document file://ec2-trust.json
aws iam put-role-policy --role-name kumo-s3-read \
  --policy-name KumoS3Read --policy-document file://s3-read.json
aws iam create-instance-profile --instance-profile-name kumo-s3-read
aws iam add-role-to-instance-profile \
  --instance-profile-name kumo-s3-read --role-name kumo-s3-read
aws iam get-instance-profile --instance-profile-name kumo-s3-read \
  --query 'InstanceProfile.{Arn:Arn,Roles:Roles[].Arn}'

Save the returned instance profile ARN and role ARN. For a bucket in another AWS account, its owner must also allow this role in the bucket policy. For objects encrypted with a customer-managed KMS key, grant the role kms:Decrypt on that key and allow it in the key policy. These source grants do not permit output writes.

2. Allow the workspace to pass the role. In the following policy, replace <s3-role-arn> with the role ARN from step 1. Save it as pass-role.json:

{
  "Version": "2012-10-17",
  "Statement": [{
    "Effect": "Allow",
    "Action": "iam:PassRole",
    "Resource": "<s3-role-arn>"
  }]
}

Apply it to the workspace deployment IAM role, replacing <workspace-deployment-role> with its name. This is the role Databricks uses to launch compute, separate from kumo-s3-read:

aws iam put-role-policy --role-name '<workspace-deployment-role>' \
  --policy-name PassKumoS3Role --policy-document file://pass-role.json

If your organization manages this role through infrastructure as code, apply the same permission there.

3. Register and attach the profile. As a Databricks workspace administrator, open Settings > Security > Instance profiles > Manage, select Add instance profile, and enter the profile ARN from step 1. Select the registered profile and grant access to the user or group running the notebook. See Manage instance profiles.

Edit the notebook’s classic compute, select the profile under Advanced options > Instances > Instance profile, save, and restart the compute. Allow a few minutes for IAM changes to propagate. Restrict cluster access to users authorized for the role’s data.

4. Verify the role and connect. Reattach the notebook and rerun Databricks Notebook Setup after the restart. Run this cell, replacing the bucket, prefix, and bucket region:

import boto3
import kumo.engine as kumo
print("AWS identity:", boto3.client("sts").get_caller_identity()["Arn"])
connector = kumo.S3Connector(
    root_dir="s3://<bucket>/<prefix>/",
    region="<aws-region>",
)
print(connector.table_names())

The identity should name kumo-s3-read, and the listing should contain your source files or table directories. If the identity differs, check the notebook’s compute attachment and instance profile. If access is denied, check the role, bucket, prefix, and KMS policies. Unity Catalog grants alone do not authorize this direct S3 connection.

Choose a table from the listing and verify schema access:

source = connector["<table-name-from-listing>"]
print(list(source.column_dict))

A column listing confirms that Kumo can inspect the selected source. Continue to Tables.

Temporary Credentials in Databricks Secrets

Use this path for a short notebook session with an approved temporary AWS identity. For scheduled or long-running workloads on AWS, use the instance profile path so credentials rotate automatically.

Store and use a temporary AWS session

1. Verify AWS access locally. On your administration machine, install AWS CLI v2 and Databricks CLI 0.205 or later. Configure an approved AWS profile that supplies temporary credentials. For IAM Identity Center, use aws configure sso once, then sign in:

aws sso login --profile '<aws-profile>'
aws sts get-caller-identity --profile '<aws-profile>'
aws s3api list-objects-v2 --profile '<aws-profile>' \
  --bucket '<bucket>' --prefix '<prefix>/' \
  --region '<aws-region>' --max-items 5

Use the bucket’s region and a prefix without leading or trailing slashes. Omit --prefix for the bucket root. Confirm the returned account and role are expected. Resolve AccessDenied with the data owner before continuing; a Databricks secret does not grant AWS permissions.

2. Create a private secret scope. Sign in to your workspace from the same machine, replacing <workspace-url> with its HTTPS URL and choosing a local CLI profile name:

databricks auth login --host '<workspace-url>' --profile '<databricks-profile>'
databricks current-user me --profile '<databricks-profile>'
databricks secrets create-scope '<secret-scope>' --profile '<databricks-profile>'
databricks secrets list-acls '<secret-scope>' --profile '<databricks-profile>'

The scope creator has MANAGE permission. If the scope already exists, check its access list instead of recreating it. Do not grant the workspace-wide users group access. If a different user runs the notebook, grant only that authorized user or group READ:

databricks secrets put-acl '<secret-scope>' '<user-or-group>' READ \
  --profile '<databricks-profile>'

3. Store the temporary session. Run this Python script on your administration machine, replacing the three placeholders. It obtains credentials from the AWS CLI, checks that they form a temporary session, and sends them to Databricks Secrets through standard input. It does not print credentials or write a credential file:

import json
import subprocess
from datetime import datetime, timezone
result = subprocess.run(
    ["aws", "configure", "export-credentials", "--profile", "<aws-profile>",
     "--format", "process"],
    check=True, capture_output=True, text=True,
)
credentials = json.loads(result.stdout)
required = ("AccessKeyId", "SecretAccessKey", "SessionToken", "Expiration")
if not all(credentials.get(key) for key in required):
    raise RuntimeError("Use an AWS profile with short-lived session credentials.")
expiration = datetime.fromisoformat(credentials["Expiration"].replace("Z", "+00:00"))
if expiration <= datetime.now(timezone.utc):
    raise RuntimeError("Refresh the AWS login before storing this session.")
subprocess.run(
    ["databricks", "secrets", "put-secret", "<secret-scope>", "aws-session",
     "--profile", "<databricks-profile>"],
    input=json.dumps({key: credentials[key] for key in required}),
    text=True, check=True,
)

Confirm the key exists with databricks secrets list-secrets '<secret-scope>' --profile '<databricks-profile>'. This lists metadata only.

4. Connect from the notebook. After Databricks Notebook Setup, read the stored session and create the connector:

import json
from datetime import datetime, timezone
import kumo.engine as kumo
credentials = json.loads(dbutils.secrets.get("<secret-scope>", "aws-session"))
expiration = datetime.fromisoformat(credentials["Expiration"].replace("Z", "+00:00"))
if expiration <= datetime.now(timezone.utc):
    raise RuntimeError("Refresh the AWS session secret, then rerun this cell.")
identity = kumo.S3Identity(
    access_key_id=credentials["AccessKeyId"],
    secret_access_key=credentials["SecretAccessKey"],
    session_token=credentials["SessionToken"],
)
connector = kumo.S3Connector(
    root_dir="s3://<bucket>/<prefix>/",
    region="<aws-region>",
    identity=identity,
)
print(connector.table_names())

Use s3://<bucket>/ for the bucket root. Select a table from the listing and run print(list(connector["<table-name>"].column_dict)) to check schema access. Do not print the credentials or identity. Continue to Tables when the expected columns appear.

If the session expires or AWS returns ExpiredToken, refresh your AWS login, repeat step 3, and rerun step 4. Updating the secret alone does not refresh an existing connector. For a missing secret or a permissions error, check the scope, key name, and notebook user’s scope access.

Prepare the Service Principal

Use explicit OAuth credentials when the connector should act as a service principal instead of the notebook identity. Grant that principal the compute and source permissions shown above.

Create the principal, store credentials, and connect

1. Create a Databricks service principal. A workspace administrator opens Settings > Identity and access > Service principals > Manage > Add service principal > Add new, names the principal, and adds it to the workspace. Use a Databricks-managed principal for this OAuth flow. If one already exists, add that principal to the workspace instead. See Manage service principals.

Open the principal’s Secrets tab and select Generate secret. Set a lifetime appropriate to your organization, generate the OAuth secret, and save the displayed client ID and secret securely. The secret appears only once. See Create an OAuth secret.

2. Store the credentials. On your administration machine, use Databricks CLI 0.205 or later. Sign in and create a private scope. Replace the workspace URL, local CLI profile name, and scope name:

databricks auth login --host '<workspace-url>' --profile '<databricks-profile>'
databricks secrets create-scope '<secret-scope>' --profile '<databricks-profile>'
databricks secrets list-acls '<secret-scope>' --profile '<databricks-profile>'
databricks secrets put-secret '<secret-scope>' client-id --profile '<databricks-profile>'
databricks secrets put-secret '<secret-scope>' client-secret --profile '<databricks-profile>'

Enter the client ID and OAuth secret at the respective interactive prompts. Do not include the secret in command arguments or notebook source. For an existing scope, skip creation and check its access list. If a different user runs the notebook, grant that authorized user READ with databricks secrets put-acl '<secret-scope>' '<notebook-user>' READ --profile '<databricks-profile>'. Keep the scope restricted to users allowed to act as this principal.

3. Select the SQL warehouse and grant source access. Open SQL Warehouses, select a warehouse, and grant the service principal CAN USE in its permissions. Copy the warehouse ID from Connection details: it is the final segment of the HTTP path after /sql/1.0/warehouses/. The warehouse handles data access; attach the notebook itself to the GPU compute from setup.

Have the catalog or data owner apply the source grants above to this service principal’s application ID. The notebook execution principal still needs its own artifact-volume access from Databricks Notebook Setup. See Unity Catalog privileges.

Create the connector

Replace the workspace host (without a URL path), secret scope, warehouse ID, catalog, and source schema. Use the client-id and client-secret keys created above. Read the secret in the same Python kernel that creates the connector. Do not print the identity or secret values:

import kumo.engine as kumo
identity = kumo.DatabricksIdentity(
    host="https://<workspace-host>",
    client_id=dbutils.secrets.get("<secret-scope>", "client-id"),
    client_secret=dbutils.secrets.get("<secret-scope>", "client-secret"),
)
connector = kumo.DatabricksConnector(
    identity=identity,
    warehouse_id="<sql-warehouse-id>",
    catalog="<catalog>",
    schema_name="<schema>",
)

This connector streams source tables through the SQL warehouse. You do not need a Kumo account or Kumo API key for this local Python workflow.

Prepare Prediction Outputs

You can write predictions to Unity Catalog whether your source is Unity Catalog or S3. Source read grants do not permit prediction writes. Use a separate output schema and grant its staging volume and table permissions to the Databricks output identity. For S3 output, use the S3 output procedure with a dedicated prefix and write permissions.

Configure output permissions and connector

For prediction writes, use a separate output schema. Have its owner create the managed staging volume and grant the output connector’s principal access. That principal also needs CAN USE on the SQL warehouse selected for output writes. If your source is S3, its AWS identity supplies no Databricks permissions; grant the notebook’s Databricks identity access separately. Replace <principal> with the notebook’s resolved principal for notebook authentication, or the service principal’s application ID for explicit OAuth. Replace the catalog and output schema before running this in SQL Editor:

CREATE SCHEMA IF NOT EXISTS `<catalog>`.`<output-schema>`;
CREATE VOLUME IF NOT EXISTS `<catalog>`.`<output-schema>`.`_kumo_staging`;
GRANT USE CATALOG ON CATALOG `<catalog>` TO `<principal>`;
GRANT USE SCHEMA ON SCHEMA `<catalog>`.`<output-schema>` TO `<principal>`;
GRANT CREATE TABLE ON SCHEMA `<catalog>`.`<output-schema>` TO `<principal>`;
GRANT READ VOLUME, WRITE VOLUME
  ON VOLUME `<catalog>`.`<output-schema>`.`_kumo_staging` TO `<principal>`;

For a Unity Catalog source, create an output connector that reuses the source Databricks connector’s identity, catalog, and query target. This works with either notebook authentication or explicit OAuth:

output_connector = kumo.DatabricksConnector(
    identity=connector.identity,
    warehouse_id=connector.warehouse_id,
    cluster_id=connector.cluster_id,
    catalog=connector.catalog,
    schema_name="<output-schema>",
)

For an S3 source, keep the S3 connector for source reads and create a Databricks output connector using notebook authentication. Replace the output catalog, schema, and warehouse ID with the resources granted above:

output_connector = kumo.DatabricksConnector.from_notebook(
    catalog="<catalog>",
    schema_name="<output-schema>",
    warehouse_id="<sql-warehouse-id>",
)

This uses the notebook’s Databricks SDK identity, independently of the S3 credentials. If output writes should use a service principal instead, follow Prepare the Service Principal and construct the Databricks connector with that identity and the output catalog and schema.

Pass output_connector to PredictionOutputConfig with a new destination name in the form <output-schema>__kumo__<table>. See Making Predictions. Writing to an existing table also requires the appropriate table permissions. The _kumo_staging volume is temporary storage for SQL table writes, separate from the retained artifact volume configured during setup.

The connector is a Python object in the current process. After a Python or cluster restart, rerun the installation check, scratch configuration, and connector cells. Keep secrets in the secret scope, not in notebook source or retained artifacts.

Continue to the Databricks notebook example to define tables, build a graph, train, and predict. See Making Predictions for output tables and Publish and Retrieve Models for model retention.