Databricks Isv Python Sqlalchemy

Python SQLAlchemy + Databricks (ISV)

Use this skill when implementing SQLAlchemy with the Databricks dialect (databricks-sqlalchemy) for partner integrations. SQLAlchemy provides ORM capabilities and SQL abstraction on top of the databricks-sql-connector driver.

PWAF Documentation Links

• Partner Architecture Overview
• Access & Authentication
• OAuth M2M
• OAuth U2M
• Telemetry Attribution

Packages

pip install databricks-sqlalchemy databricks-sql-connector sqlalchemy databricks-sdk requests

• databricks-sqlalchemy: SQLAlchemy dialect (registers databricks:// URL scheme)
• databricks-sql-connector: Underlying SQL driver
• sqlalchemy 2.x: ORM and SQL toolkit
• databricks-sdk: OAuth M2M token acquisition
• requests: PKCE token exchange (U2M custom OAuth app)

Authentication Decision Guide

ISV Integration Authentication Selection:

1. Is this for automated/unattended workloads (CI, jobs, services)?
   → OAuth M2M (Client Credentials) – RECOMMENDED
   
2. Is this for user-interactive flows (web app, desktop app)?
   → OAuth U2M with Custom OAuth App (PKCE) – RECOMMENDED
   
3. Is this for quick testing/development only?
   → PAT (Personal Access Token) – LIMITED USE
   
4. Is the OAuth token obtained externally (mobile app, web frontend)?
   → OAuth U2M Token-Env (pass pre-obtained token)

Authentication Comparison

CLIENT_ID Distinction (CRITICAL)

┌─────────────────────────────────────────────────────────────────────────────┐
│ DATABRICKS_CLIENT_ID vs DATABRICKS_U2M_CLIENT_ID                            │
├─────────────────────────────────────────────────────────────────────────────┤
│ DATABRICKS_CLIENT_ID:                                                       │
│   - Service Principal UUID for OAuth M2M (machine-to-machine)              │
│   - Created in: Workspace → Settings → Identity and access → Service       │
│     principals                                                              │
│   - Used with: DATABRICKS_CLIENT_SECRET                                    │
│                                                                             │
│ DATABRICKS_U2M_CLIENT_ID:                                                   │
│   - Custom OAuth Application client ID for OAuth U2M (user-to-machine)     │
│   - Created in: Account Console → Settings → App connections               │
│   - Used with: DATABRICKS_U2M_CLIENT_SECRET (if confidential app)          │
│                                                                             │
│ NEVER mix these – they serve completely different authentication flows.    │
└─────────────────────────────────────────────────────────────────────────────┘

Token Lifetime Summary

User-Agent (Required for PWAF)

Format: <isv-name>_<product-name>/<product-version>

USER_AGENT = "YourCompany_YourProduct/1.0.0"
engine = create_engine(url, connect_args={"user_agent_entry": USER_AGENT})

Important: Use user_agent_entry (not _user_agent_entry which is deprecated).

Connection URL Format

databricks://token:<access_token>@<server_hostname>?http_path=<path>&catalog=<catalog>&schema=<schema>

• server_hostname: Workspace host without https:// (e.g., myworkspace.cloud.databricks.com)
• http_path: SQL Warehouse HTTP path (e.g., /sql/1.0/warehouses/abc123)
• catalog / schema: Optional default catalog and schema

Environment Variables Reference

Common (all auth types):

DATABRICKS_HOST          # Workspace URL (e.g., https://myworkspace.cloud.databricks.com)
DATABRICKS_HTTP_PATH     # SQL Warehouse HTTP path (e.g., /sql/1.0/warehouses/<id>)
DATABRICKS_CATALOG       # Optional: Default catalog (default: main)
DATABRICKS_SCHEMA        # Optional: Default schema (default: default)

PAT Authentication:

DATABRICKS_TOKEN         # Personal access token (starts with dapi...)

OAuth M2M Authentication:

DATABRICKS_CLIENT_ID     # Service Principal UUID
DATABRICKS_CLIENT_SECRET # OAuth secret for the service principal

OAuth U2M Custom OAuth App:

DATABRICKS_U2M_CLIENT_ID      # Custom OAuth app client ID
DATABRICKS_REDIRECT_URI       # Optional: Redirect URI (default: http://localhost:8080/callback)
DATABRICKS_U2M_CLIENT_SECRET  # Optional: Required if OAuth app is confidential

OAuth U2M Token-Env:

DATABRICKS_ACCESS_TOKEN  # Pre-obtained OAuth access token

Multi-Auth Selector:

APP_AUTH_TYPE            # pat | oauth_m2m | oauth_u2m

Note: Use APP_AUTH_TYPE (not DATABRICKS_AUTH_TYPE) to avoid conflict with SDK's internal auth_type parameter.

Host Normalization Helper

SQLAlchemy connection URL requires bare hostname (no https://):

def normalize_host(host: str) -> str:
    """Normalize host to bare hostname for SQLAlchemy URL."""
    return host.replace("https://", "").replace("http://", "").split("/")[0]

# Usage
server_hostname = normalize_host(os.environ.get("DATABRICKS_HOST", ""))
url = f"databricks://token:{token}@{server_hostname}?http_path={http_path}"

Complete Examples

Auth Type 1: PAT Authentication

import os
from sqlalchemy import create_engine, text

USER_AGENT = "YourCompany_YourProduct/1.0.0"

def get_engine():
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    http_path = os.environ.get("DATABRICKS_HTTP_PATH", "").strip()
    token = os.environ.get("DATABRICKS_TOKEN", "").strip()
    if not host or not http_path or not token:
        raise ValueError("Set DATABRICKS_HOST, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN")
    
    server_hostname = host.replace("https://", "").split("/")[0]
    catalog = os.environ.get("DATABRICKS_CATALOG", "main")
    schema = os.environ.get("DATABRICKS_SCHEMA", "default")
    
    url = f"databricks://token:{token}@{server_hostname}?http_path={http_path}&catalog={catalog}&schema={schema}"
    return create_engine(url, connect_args={"user_agent_entry": USER_AGENT})

with get_engine().connect() as conn:
    result = conn.execute(text("SELECT current_user()"))
    print(result.fetchall())

Auth Type 2: OAuth M2M (Client Credentials)

import os
from databricks.sdk.config import Config
from sqlalchemy import create_engine, text

USER_AGENT = "YourCompany_YourProduct/1.0.0"

def get_access_token():
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    host_url = host if host.startswith("https://") else f"https://{host}"
    
    config = Config(
        host=host_url,
        client_id=os.environ.get("DATABRICKS_CLIENT_ID", ""),
        client_secret=os.environ.get("DATABRICKS_CLIENT_SECRET", ""),
        auth_type="oauth-m2m",  # Required: prevents default credential resolution
    )
    # authenticate() returns headers dict, NOT an object with .token
    headers = config.authenticate()
    return headers.get("Authorization", "").replace("Bearer ", "").strip()

def get_engine():
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    http_path = os.environ.get("DATABRICKS_HTTP_PATH", "").strip()
    server_hostname = host.replace("https://", "").split("/")[0]
    
    access_token = get_access_token()
    url = f"databricks://token:{access_token}@{server_hostname}?http_path={http_path}&catalog=main&schema=default"
    return create_engine(url, connect_args={"user_agent_entry": USER_AGENT})

with get_engine().connect() as conn:
    result = conn.execute(text("SELECT current_user()"))
    print(result.fetchall())

Auth Type 3: OAuth U2M Custom OAuth App (PKCE)

import os
import base64
import hashlib
import secrets
import string
import urllib.parse
from http.server import HTTPServer, BaseHTTPRequestHandler
from threading import Thread
from urllib.parse import parse_qs, urlparse
import requests
import webbrowser
from sqlalchemy import create_engine, text

USER_AGENT = "YourCompany_YourProduct/1.0.0"
DEFAULT_REDIRECT_URI = "http://localhost:8080/callback"

def _pkce_verifier_and_challenge():
    allowed = string.ascii_letters + string.digits + "-._~"
    code_verifier = "".join(secrets.choice(allowed) for _ in range(64))
    digest = hashlib.sha256(code_verifier.encode()).digest()
    code_challenge = base64.urlsafe_b64encode(digest).decode().rstrip("=")
    return code_verifier, code_challenge

def run_u2m_flow(host: str, client_id: str, redirect_uri: str = None, client_secret: str = None) -> dict:
    redirect_uri = redirect_uri or DEFAULT_REDIRECT_URI
    host = host.rstrip("/")
    if not host.startswith("https://"):
        host = "https://" + host
    
    code_verifier, code_challenge = _pkce_verifier_and_challenge()
    state = secrets.token_urlsafe(32)
    
    auth_url = f"{host}/oidc/v1/authorize?" + urllib.parse.urlencode({
        "response_type": "code",
        "client_id": client_id,
        "redirect_uri": redirect_uri,
        "scope": "all-apis",
        "code_challenge": code_challenge,
        "code_challenge_method": "S256",
        "state": state,
    })
    
    code_holder = []
    parsed = urlparse(redirect_uri)
    port = parsed.port or 8080
    
    class Handler(BaseHTTPRequestHandler):
        def do_GET(self):
            q = parse_qs(urlparse(self.path).query)
            code_holder.append(q.get("code", [None])[0])
            self.send_response(200)
            self.end_headers()
            self.wfile.write(b"Authentication successful. Close this window.")
        def log_message(self, *args): pass
    
    server = HTTPServer(("localhost", port), Handler)
    thread = Thread(target=server.handle_request, daemon=True)
    thread.start()
    webbrowser.open(auth_url)
    thread.join(timeout=120)
    
    if not code_holder or not code_holder[0]:
        raise RuntimeError("No authorization code received")
    
    resp = requests.post(
        f"{host}/oidc/v1/token",
        data={
            "grant_type": "authorization_code",
            "code": code_holder[0],
            "redirect_uri": redirect_uri,
            "client_id": client_id,
            "code_verifier": code_verifier,
            **({"client_secret": client_secret} if client_secret else {}),
        },
    )
    resp.raise_for_status()
    return resp.json()

# Usage
tokens = run_u2m_flow(
    host=os.environ["DATABRICKS_HOST"],
    client_id=os.environ["DATABRICKS_U2M_CLIENT_ID"],
)
access_token = tokens["access_token"]
server = os.environ["DATABRICKS_HOST"].replace("https://", "").split("/")[0]
http_path = os.environ["DATABRICKS_HTTP_PATH"]
url = f"databricks://token:{access_token}@{server}?http_path={http_path}"
engine = create_engine(url, connect_args={"user_agent_entry": USER_AGENT})

Auth Type 4: OAuth U2M Token-Env (Pre-obtained Token)

import os
from sqlalchemy import create_engine, text

USER_AGENT = "YourCompany_YourProduct/1.0.0"

def get_engine():
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    http_path = os.environ.get("DATABRICKS_HTTP_PATH", "").strip()
    access_token = os.environ.get("DATABRICKS_ACCESS_TOKEN", "").strip()
    if not host or not http_path or not access_token:
        raise ValueError("Set DATABRICKS_HOST, DATABRICKS_HTTP_PATH, and DATABRICKS_ACCESS_TOKEN")
    
    server_hostname = host.replace("https://", "").split("/")[0]
    url = f"databricks://token:{access_token}@{server_hostname}?http_path={http_path}"
    return create_engine(url, connect_args={"user_agent_entry": USER_AGENT})

with get_engine().connect() as conn:
    result = conn.execute(text("SELECT current_user()"))
    print(result.fetchall())

Error Handling Patterns

Error Classification

class ErrorCategory:
    AUTH = "authentication"
    PERMISSION = "permission"
    NOT_FOUND = "not_found"
    SYNTAX = "syntax"
    RATE_LIMIT = "rate_limit"
    WAREHOUSE = "warehouse"
    NETWORK = "network"
    TIMEOUT = "timeout"
    UNKNOWN = "unknown"

def classify_error(error: Exception) -> tuple[str, str, bool]:
    """Classify error into (category, remediation, retryable)."""
    error_str = str(error).lower()
    
    # Authentication errors
    if "401" in error_str or "unauthorized" in error_str:
        return (ErrorCategory.AUTH,
                "Check credentials: verify token is valid and not expired.",
                False)
    
    # Permission errors
    if "403" in error_str or "permission_denied" in error_str:
        return (ErrorCategory.PERMISSION,
                "Grant CAN_USE on warehouse, USE_CATALOG, USE_SCHEMA, SELECT on table.",
                False)
    
    # Not found errors
    if "404" in error_str or "not found" in error_str:
        return (ErrorCategory.NOT_FOUND,
                "Verify DATABRICKS_HOST and DATABRICKS_HTTP_PATH are correct.",
                False)
    
    # SQL syntax errors
    if "analysisexception" in error_str or "syntax error" in error_str:
        return (ErrorCategory.SYNTAX,
                "Review query for typos and invalid SQL.",
                False)
    
    # Rate limiting
    if "429" in error_str or "rate limit" in error_str:
        return (ErrorCategory.RATE_LIMIT,
                "Implement exponential backoff retry.",
                True)
    
    # Warehouse errors
    if "503" in error_str or "temporarily_unavailable" in error_str:
        return (ErrorCategory.WAREHOUSE,
                "Warehouse starting. Wait and retry.",
                True)
    
    # Network errors
    if "connection" in error_str and "refused" in error_str:
        return (ErrorCategory.NETWORK,
                "Check network connectivity.",
                True)
    
    # Timeout errors
    if "timeout" in error_str:
        return (ErrorCategory.TIMEOUT,
                "Increase timeout or optimize query.",
                True)
    
    return (ErrorCategory.UNKNOWN, "Check error message.", False)

Retry with Exponential Backoff

import random
import time
from dataclasses import dataclass

@dataclass
class RetryConfig:
    max_retries: int = 3
    initial_backoff_seconds: float = 2.0
    max_backoff_seconds: float = 60.0
    backoff_multiplier: float = 2.0
    jitter_factor: float = 0.25

def is_retryable(error: Exception) -> bool:
    error_str = str(error).lower()
    retryable = ["429", "503", "rate limit", "temporarily_unavailable", "connection", "timeout"]
    non_retryable = ["401", "403", "404", "analysisexception", "syntax error"]
    for p in non_retryable:
        if p in error_str:
            return False
    for p in retryable:
        if p in error_str:
            return True
    return False

def retry_with_backoff(operation, config: RetryConfig = None):
    config = config or RetryConfig()
    for attempt in range(config.max_retries + 1):
        try:
            return operation()
        except Exception as e:
            if not is_retryable(e) or attempt >= config.max_retries:
                raise
            backoff = min(
                config.initial_backoff_seconds * (config.backoff_multiplier ** attempt),
                config.max_backoff_seconds
            )
            backoff += backoff * config.jitter_factor * random.random()
            time.sleep(backoff)

Multi-Auth Pattern

import os
from databricks.sdk.config import Config
from sqlalchemy import create_engine, text

USER_AGENT = "YourCompany_YourProduct/1.0.0"

def get_engine():
    # Use APP_AUTH_TYPE to avoid SDK internal conflict
    auth_type = os.environ.get("APP_AUTH_TYPE", "oauth_m2m").strip().lower()
    
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    http_path = os.environ.get("DATABRICKS_HTTP_PATH", "").strip()
    server_hostname = host.replace("https://", "").split("/")[0]
    
    if auth_type == "pat":
        token = os.environ.get("DATABRICKS_TOKEN", "")
    elif auth_type == "oauth_m2m":
        config = Config(
            host=host if host.startswith("https://") else f"https://{host}",
            client_id=os.environ.get("DATABRICKS_CLIENT_ID", ""),
            client_secret=os.environ.get("DATABRICKS_CLIENT_SECRET", ""),
            auth_type="oauth-m2m",
        )
        headers = config.authenticate()
        token = headers.get("Authorization", "").replace("Bearer ", "").strip()
    elif auth_type == "oauth_u2m":
        # Run PKCE flow (see Auth Type 3 example)
        token = run_u2m_flow(
            host=host,
            client_id=os.environ.get("DATABRICKS_U2M_CLIENT_ID", ""),
        )["access_token"]
    else:
        raise ValueError(f"Unknown auth type: {auth_type}")
    
    url = f"databricks://token:{token}@{server_hostname}?http_path={http_path}"
    return create_engine(url, connect_args={"user_agent_entry": USER_AGENT})

Auth Isolation Testing

Run each auth type in an isolated environment to prevent credential conflicts:

# PAT only
env -i PATH="$PATH" HOME="$HOME" \
  DATABRICKS_HOST="https://myworkspace.cloud.databricks.com" \
  DATABRICKS_HTTP_PATH="/sql/1.0/warehouses/abc123" \
  DATABRICKS_TOKEN="dapi..." \
  python your_script.py

# OAuth M2M only
env -i PATH="$PATH" HOME="$HOME" \
  DATABRICKS_HOST="https://myworkspace.cloud.databricks.com" \
  DATABRICKS_HTTP_PATH="/sql/1.0/warehouses/abc123" \
  DATABRICKS_CLIENT_ID="<sp-uuid>" \
  DATABRICKS_CLIENT_SECRET="<secret>" \
  python your_script.py

# OAuth U2M (needs DISPLAY for browser)
env -i PATH="$PATH" HOME="$HOME" DISPLAY="$DISPLAY" USER="$USER" \
  DATABRICKS_HOST="https://myworkspace.cloud.databricks.com" \
  DATABRICKS_HTTP_PATH="/sql/1.0/warehouses/abc123" \
  DATABRICKS_U2M_CLIENT_ID="<client-id>" \
  python your_script.py

Redirect URI Reference

Must be registered in Account Console → Settings → App connections → your OAuth app → Redirect URIs.

Validation Query

def validate_connection(engine):
    with engine.connect() as conn:
        result = conn.execute(text("""
            SELECT 
                current_user() AS user,
                current_catalog() AS catalog,
                current_schema() AS schema,
                current_timestamp() AS timestamp
        """))
        return dict(result.fetchone()._mapping)

Troubleshooting

Implementation Learnings

1. SDK Config Returns Headers Dict: Config.authenticate() returns {"Authorization": "Bearer <token>"}, not an object with .token. Extract with .get("Authorization", "").replace("Bearer ", "").

2. Explicit auth_type for M2M: Always set auth_type="oauth-m2m" in Config to prevent SDK from using default credential resolution (which may pick up PAT or other credentials).

3. Use APP_AUTH_TYPE: For multi-auth patterns, use APP_AUTH_TYPE instead of DATABRICKS_AUTH_TYPE to avoid conflict with SDK's internal auth_type parameter.

4. CLIENT_ID Variables: Use DATABRICKS_CLIENT_ID for service principals (M2M) and DATABRICKS_U2M_CLIENT_ID for custom OAuth apps (U2M). Never mix these.

5. Host Normalization: SQLAlchemy URL requires bare hostname. Always strip https:// from host before building URL.

6. Auth Isolation: SDK picks up credentials from environment. Run tests with env -i to ensure clean environment.

7. user_agent_entry Parameter: Use user_agent_entry (not _user_agent_entry) in connect_args for PWAF compliance.

PKCE Flow Implementation Details

When implementing custom OAuth U2M PKCE flows, follow these patterns for reliable browser callback handling:

Use threading.Event for callback signaling:

from threading import Thread, Event

callback_received = Event()
CALLBACK_TIMEOUT = 120  # seconds

def serve_until_callback():
    while not callback_received.is_set():
        server.handle_request()

thread = Thread(target=serve_until_callback, daemon=True)
thread.start()

Handle multiple HTTP requests (browsers send favicon, prefetch, etc.):

class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        req_path = urlparse(self.path).path
        if req_path == "/callback":  # Only process OAuth callback
            code = parse_qs(urlparse(self.path).query).get("code", [None])[0]
            if code:
                callback_received.set()
                self.send_success_response()
        else:
            self.send_response(404)
            self.end_headers()

Use output flushing for immediate feedback:

print("Waiting for browser authentication...", flush=True)

Wait with timeout:

if not callback_received.wait(timeout=CALLBACK_TIMEOUT):
    raise RuntimeError(f"Authentication timed out after {CALLBACK_TIMEOUT}s")

Key insight: Using server.handle_request() only once can be consumed by browser prefetch requests (like favicon), causing the actual OAuth callback to be missed. The Event-based loop pattern ensures all requests are handled until the actual callback arrives.

Key Differences from Python SQL Connector