OAuth 2.0

Quick Start

Make an OAuth request using your configured Service Principal:

export CLIENT_ID=<client-id>
export CLIENT_SECRET=<client-secret>

curl --request POST \
--url https://<databricks-instance>/oidc/v1/token \
--user "$CLIENT_ID:$CLIENT_SECRET" \
--data 'grant_type=client_credentials&scope=all-apis'

You will get a JSON response in this format:

{
  "access_token": "<oauth-access-token>",
  "token_type": "Bearer",
  "expires_in": 3600
}

Use the access token in future queries:

export OAUTH_TOKEN=<oauth-access-token>

curl --request GET --header "Authorization: Bearer $OAUTH_TOKEN" \
'https://<workspace-URL>/api/2.0/clusters/list'

Basics and Vocab

OAuth 2.0 refers to the Open Authorization 2 standard. It is a protocol for authorizing resource access to things like data or APIs.
- This does not mean "authentication" or verification of identity, but a way to ensure the right users can get access.
It relies on Access Tokens that act like "hall passes"
- These are sometimes represented as JSON Web Tokens (JWT) which allows embedding some data in the token
- Access Tokens can be set to expire after certain time.
Scopes refer to the reason for which access is being granted. They are specified during the authorization process.
Authorization Codes can also be returned which can be exchanged for Access Tokens. This adds security by creating another barrier to entry.
Refresh Tokens are issued with Access Tokens and can be exchanged for new Access Tokens after the time expires
- This is a potential security risk, so make sure these are stored correctly on your application side!
Service Principals are identities created to manage automated pipelines
- They're not actually users, but they kind of act like them.
- It is an API-only identity and doesn't have an associated UI.

How does OAuth 2.0 work?

Here's an outline of the basic process. But first, two things have to be accomplished:

The Resource Owner sets up an Authorization Server that handles issuing tokens
The Client (or User) needs to get their credentials off the Authorization Server:
- Client ID: Unique identifier for the client
- Client Secret: Secret password that proves the client is who they say they are

Once these two steps are complete, the general process for getting access is the same:

The Client starts by creating an access request via an app/website/API call/whatever. They must supply the following:
- Client ID and Client Secret for identification
- Scope of access
- Endpoint URI for where to send the Access Token/Authorization Code
The Authorization Server authenticates the Client.
- The process involves checking if their scope is justified/they're allowed to use the resource they requested.
The Authorization Server returns an Access Token or Authorization Code to the Client as a JWT or string.
The Client uses the Access Token/Authorization Code to access Resource Server as part of their request.

OAuth and Databricks

We are making Workspace level queries, which follow very similar query structures to performing queries with PATs (Personal Access Tokens). However, due to limitations on the Databricks side, we are prevented from using PATs.

To do the next few steps, you will need to have set up a Service Principal for Databricks. This Service Principal should have its own CLIENT_ID and CLIENT_SECRET as described above. These can be stored in an environment variables file (like .env) or configured ahead of time.

Start by getting the necessary OAuth Token:

export CLIENT_ID=<client-id>
export CLIENT_SECRET=<client-secret>

curl --request POST \
--url https://<databricks-instance>/oidc/v1/token \
--user "$CLIENT_ID:$CLIENT_SECRET" \
--data 'grant_type=client_credentials&scope=all-apis'

You will get a JSON response in this format:

{
  "access_token": "<oauth-access-token>",
  "token_type": "Bearer",
  "expires_in": 3600
}

Congratulations! You now have the necessary OAuth Token!

Query Examples

API Query

Here is an example of a query that can be made to list the available clusters on the Databricks environment:

export OAUTH_TOKEN=<oauth-access-token>

curl --request GET --header "Authorization: Bearer $OAUTH_TOKEN" \
'https://<workspace-URL>/api/2.0/clusters/list'

The <oauth-access-token> is derived from what the service principal set up
The <workspace-URL> is specific to us. Talk to a developer to get the URL if you don't already have it!

In Python, the equivalent code would be as follows:

import requests

headers = {"Authorization": f"Bearer {oauth-access-token}"}
response = requests.get('https://<workspace-URL>/api/2.0/clusters/list', headers=headers)

And in JavaScript:

fetch("https://<workspace-URL>/api/2.0/clusters/list",
    headers: {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer <oauth-access-token>'
    }
)
  .then((response) => response.json())

Model Query

Here's how to make queries to model endpoints for our chat application:

The headers are the same as above:

    headers: {
       "Content-Type": "application/json",
       "Authorization": "Bearer <oauth-access-token>"
    }

The request body should look something like this:

{
    "max_tokens": 100,
    "messages": [
        {
            "content": "<query>",
            "role": "user"
        }
    ],
    "temperature": 0.1 //You can leave this one out if you want
}

Putting it all together:

curl \
  -X POST \
  -H '{"Content-Type": "application/json", "Authorization": "Bearer <oauth-access-token>"}' \
  -d '{"max_tokens": 100, "messages": [{"content": "<query>", "role": "user"}], "temperature": 0.1}' \
  <model_url>

Or in JavaScript:

fetch(<model_url>, {
  method: "POST",
  body: {
    "max_tokens": 100,
    "messages": [
        {
            "content": "<query>",
            "role": "user"
        }
    ],
    "temperature": 0.1
  },
  headers: {
    "Content-type": "application/json, 
    "Authorization": "Bearer <oauth-access-token>"
  }
})

References

https://auth0.com/intro-to-iam/what-is-oauth-2
https://learn.microsoft.com/en-us/azure/databricks/dev-tools/auth/oauth-m2m#example-token-in-workspace
https://docs.databricks.com/api/azure/workspace/servingendpoints/query