Skip to main content

On-behalf-of-user authentication

This recipe demonstrates how to use Databricks Apps on-behalf-of-user authentication to run a SQL query using the user's credentials instead of the app's service principal.

Code snippet​

app.py
import reflex as rx
import pandas as pd
from databricks import sql
from databricks.sdk.core import Config

cfg = Config()

_connection = None

def get_user_token(headers):
    return getattr(headers, "x_forwarded_access_token", "")

def connect_with_obo(http_path: str, user_token: str):
    global _connection
    if _connection:
        return _connection
    _connection = sql.connect(
        server_hostname=cfg.host,
        http_path=http_path,
        access_token=user_token
    )
    return _connection

def execute_query(table_name: str, conn):
    with conn.cursor() as cursor:
        query = f"SELECT * FROM {table_name} LIMIT 10"
        cursor.execute(query)
        return cursor.fetchall_arrow().to_pandas()

class OboState(rx.State):
    warehouse_paths: dict[str, str] = {}
    selected_warehouse: str = ""
    table_name: str = "samples.nyctaxi.trips"
    data: list[list] = []
    columns: list[dict] = []
    is_loading: bool = False
    error_message: str = ""

    @rx.event(background=True)
    async def run_query(self):
        async with self:
            self.is_loading = True
            self.error_message = ""
            headers = self.router.headers
            user_token = get_user_token(headers)

        if not user_token:
            async with self:
                self.error_message = "No OBO token found."
                self.is_loading = False
            return

        try:
            http_path = self.warehouse_paths.get(self.selected_warehouse)
            conn = connect_with_obo(http_path, user_token)
            df = execute_query(self.table_name, conn)

            # Process dataframe for display
            d = df.values.tolist()
            c = [{"title": col, "id": col, "type": "str"} for col in df.columns]

            async with self:
                self.data = d
                self.columns = c
        except Exception as e:
            async with self:
                self.error_message = str(e)
        finally:
            async with self:
                self.is_loading = False
info

This sample caches the SQL connection in a module-level _connection variable so it can be reused within the same app process. The app will only work when deployed to Databricks Apps with on-behalf-of-user authentication enabled.

warning

You need to enable on-behalf-of-user authentication for your application for this sample to work. When running this code locally, the X-Forwarded-Access-Token will not be present and the sample will not work as intended.

Resources​

Permissions​

For the on-behalf-of-user authentication model, permissions work as follows:

  • User's permissions: When using OBO authentication, the query runs with the end user's permissions

    • User needs SELECT permissions on the tables being queried
    • User needs CAN USE on the SQL warehouse

  • App service principal: When falling back to service principal authentication

    • Needs CAN USE on the SQL warehouse
    • Needs SELECT on the Unity Catalog tables for fallback access

See Databricks Apps authorization model for more information.

Dependencies​

requirements.txt
databricks-sdk
databricks-sql-connector
pandas
reflex