Skip to main content
Last updated on

Tutorial: Scope an AI agent's access with RBAC

Use role-based access control (RBAC) to limit an AI agent's access to the data and resources it needs for a specific task. Supply the agent with an OAuth token scoped to a role. Requests made with that token use the role's permissions instead of the broader permissions of the authenticating user or Databricks service principal.

Scenario​

An AI support agent needs ticket summaries and status, but not customer contact details containing personally identifiable information (PII). Its authenticating Databricks service principal has broader access and can read both datasets.

In this example, the support.tickets schema contains ticket summaries and status without customer PII. Customer contact details are stored separately in support.customers.contact_details.

You create role-support-agent, grant it read access to support.tickets, and grant the Databricks service principal the Assume permission on it. With a token scoped to that role, the agent can read support tickets but can't access customer contact details.

The role can inherit data privileges from parent groups, and catalog or schema grants can apply to all tables beneath them. For this example, verify that those grants do not give the role access to the customer contact table. See Privilege inheritance. Changing shared grants affects other users and roles, so use a test workspace instead of removing other users' access.

This tutorial is for custom, self-hosted, or third-party agents that accept an OAuth access token that you supply. For agents deployed on Databricks Apps, use the built-in app or user authorization. See Authentication for agents.

Prerequisites​

  • A Unity Catalog-enabled workspace with RBAC enabled. See Enable RBAC.
  • Account admin or workspace admin permissions to create the role, grant the Assume permission, and assign it to the workspace. See Who can manage groups?.
  • Ownership of, or MANAGE on, the support catalog.
  • Workspace admin permissions or Can manage on the SQL warehouse used in the example.
  • Existing tables support.tickets.open_cases and support.customers.contact_details. Use fictitious test data and substitute your own names throughout if needed.
  • A Databricks service principal assigned to the workspace, with Databricks SQL access, Can use on the warehouse, and permission to query both tables. It needs an OAuth secret that permits the sql scope for machine-to-machine (M2M) authentication. See Step 1: Create an OAuth secret.
  • curl in a trusted environment where you can request and test tokens.
  • An audit reader with access to system.access.audit to verify the agent's activity under the role. See Verify audit attribution.

Step 1: Create the role​

For this tutorial, create an empty account group to serve as the role.

  1. As an account admin, log in to the account console.
  2. In the sidebar, click User management.
  3. On the Groups tab, click Add group.
  4. Enter role-support-agent. Do not add members.
  5. Click Confirm.
  6. Assign the group to the workspace with Databricks SQL access. See Assign a group to a workspace.

The Databricks SQL access entitlement grants access to Databricks SQL features, including dashboards, queries, and SQL warehouses. The example uses SQL queries. If your agent also needs core workspace features such as notebooks, jobs, models, or pipelines, grant Workspace access. See Access entitlements. You configure data privileges and warehouse permissions in Step 2.

In this example, you grant the Assume permission separately instead of adding the Databricks service principal as a member. Membership also adds the group's permissions to the Databricks service principal's own identity.

Step 2: Grant the role access to support tickets​

As the catalog owner or a user with MANAGE on support, run:

SQL
GRANT USE CATALOG ON CATALOG support TO `role-support-agent`;
GRANT USE SCHEMA, SELECT ON SCHEMA support.tickets TO `role-support-agent`;

As a workspace admin or a user with Can manage on the warehouse, grant the role access to the warehouse:

  1. In the workspace sidebar, click SQL Warehouses.
  2. Click the Kebab menu icon. kebab menu next to the warehouse, then click Permissions.
  3. Add role-support-agent with Can use.

The schema-level SELECT grant applies to all current and future tables and views in support.tickets. It does not grant access to support.customers or other schemas. To restrict the role to one table, grant SELECT on that table instead. See Manage privileges in Unity Catalog.

Step 3: Grant permission to assume the role​

Grant the Databricks service principal the Assume permission using either method:

  1. As an account admin, open the account console and click User management.
  2. On the Groups tab, click role-support-agent.
  3. Click the Permissions tab, then click Grant access.
  4. Search for and select the Databricks service principal, select Assume, and click Save.

Step 4: Request a role-scoped token​

In your trusted token-request environment, set these environment variables:

  • DATABRICKS_HOST: the workspace URL, including https:// and without a trailing slash.
  • CLIENT_ID and CLIENT_SECRET: the Databricks service principal's application ID and OAuth secret.
  • ROLE_GROUP_ID: the numeric ID of role-support-agent. An account admin can find it at the end of the group's URL in the account console. Account admins, workspace admins, and group managers can also retrieve the group's id using the Account Groups API.
  • DATABRICKS_SQL_WAREHOUSE_ID: the ID of the warehouse granted in Step 2. It follows /sql/1.0/warehouses/ in the warehouse's HTTP path. See Get connection details for a Databricks compute resource.

Request a workspace-level token with the assume_group parameter and the sql OAuth scope. The scope limits the token to SQL APIs. The role's data privileges determine which tables it can query. See Scoped OAuth secrets.

To request the token, run:

Bash
curl --request POST \
  --url "$DATABRICKS_HOST/oidc/v1/token" \
  --user "$CLIENT_ID:$CLIENT_SECRET" \
  --data-urlencode 'grant_type=client_credentials' \
  --data-urlencode 'scope=sql' \
  --data-urlencode "assume_group=$ROLE_GROUP_ID"

Set DATABRICKS_TOKEN to the access_token value from the response. Use the expires_in value to determine when to request a replacement. Account-level token requests do not support assume_group. See Assume a role.

Step 5: Verify the token's access​

Run these checks in the trusted environment before supplying the token to the agent.

First, get a baseline with a token for the Databricks service principal's own identity. Repeat the request in Step 4 without the assume_group parameter, and temporarily set DATABRICKS_TOKEN to that token. Run both queries in this step and verify that they succeed. Then set DATABRICKS_TOKEN back to the role-scoped token and repeat the checks.

Read support tickets​

Run this request with the role-scoped token:

Bash
curl --request POST \
  --url "$DATABRICKS_HOST/api/2.0/sql/statements" \
  --header "Authorization: Bearer $DATABRICKS_TOKEN" \
  --header 'Content-Type: application/json' \
  --data "{
    \"warehouse_id\": \"$DATABRICKS_SQL_WAREHOUSE_ID\",
    \"statement\": \"SELECT * FROM support.tickets.open_cases LIMIT 10\",
    \"wait_timeout\": \"50s\"
  }"

The query succeeds because the role can read support.tickets. Check that the response's status.state is SUCCEEDED. If it is PENDING or RUNNING, use the returned statement_id to retrieve the result with the same token:

Bash
curl --request GET \
  --url "$DATABRICKS_HOST/api/2.0/sql/statements/<statement-id>" \
  --header "Authorization: Bearer $DATABRICKS_TOKEN"

See Statement Execution API: Run SQL on warehouses for statement status and result handling.

Check that access to customer PII is denied​

Using the same role-scoped token and warehouse, run:

Bash
curl --request POST \
  --url "$DATABRICKS_HOST/api/2.0/sql/statements" \
  --header "Authorization: Bearer $DATABRICKS_TOKEN" \
  --header 'Content-Type: application/json' \
  --data "{
    \"warehouse_id\": \"$DATABRICKS_SQL_WAREHOUSE_ID\",
    \"statement\": \"SELECT 1 FROM support.customers.contact_details LIMIT 1\",
    \"wait_timeout\": \"50s\"
  }"

If the statement is still running, retrieve the result as described in the previous section. The query fails with a permission error in status.error, although it succeeds with the Databricks service principal's own token. An HTTP success response alone does not mean the query succeeded. A missing table, an expired token, or a warehouse error doesn't establish that the role's data privileges are configured correctly.

Check audit attribution​

Use an audit reader with the access described in Verify audit attribution. Inspect the test query events in system.access.audit: identity_metadata.run_as identifies the authorizing role, run_as_display_name contains role-support-agent, and run_by identifies the Databricks service principal that initiated the request. Audit data is updated throughout the day, so a recent event might not appear immediately. See Identity metadata reference.

Step 6: Configure the agent​

Configure the agent's Databricks client with the workspace URL and the verified role-scoped token. Supply only that token to the agent, not the Databricks service principal's OAuth secret or its unscoped token. The role-scoped token limits requests made with that token. It does not remove the Databricks service principal's permissions in other sessions.

For an agent that uses the Databricks SDK for Python, set DATABRICKS_HOST and DATABRICKS_TOKEN in the agent's environment, then initialize the client:

Python
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

The SDK reads the token from the environment using unified authentication. Supplying an access token doesn't give the SDK the credentials it needs to request a new token. For a long-running agent, arrange for a trusted component to request and supply a replacement role-scoped token before the token expires.

Optional: Use a developer's identity​

For local testing with user-to-machine (U2M) authentication, grant a participating developer the Assume permission on role-support-agent as in Step 3. The developer installs the Databricks CLI as described in Install or update the Databricks CLI, then authenticates using their own sign-in and a separate profile:

Bash
databricks auth login --host <workspace-url> --profile support-agent

The developer selects role-support-agent under Authorize as.. in the browser and clicks Continue, then retrieves the token for that profile:

Bash
databricks auth token --profile support-agent

The developer uses the returned access_token as DATABRICKS_TOKEN for the role-scoped query checks in Step 5. Audit events for these requests identify the developer as run_by. The token returned by the Databricks CLI expires. The developer retrieves a fresh token when needed. See API (OAuth access tokens).

Clean up​

For a temporary test, remove the agent's token configuration and revoke the grant of the Assume permission created for the test when you finish. Before deleting the role, check whether it owns assets or is used by other workloads.

Next steps​

On this page