Databricks Isv Python Sdk

Databricks SDK for Python (ISV)

Use this skill when implementing or testing Databricks SDK for Python (databricks-sdk) integrations for PWAF-compliant workspace management, Unity Catalog, Jobs, and REST-style API access.

PWAF Documentation Links

Requirements

• SDK: databricks-sdk 0.20.0+
• Python: 3.8+
• Install: pip install databricks-sdk

Authentication Decision Guide

Which authentication method to use?

Production / automated workloads?
  → OAuth M2M (client credentials)  ✅ RECOMMENDED

User-interactive flows?
  → U2M Custom OAuth App (PKCE)  ✅ SUPPORTED

Already have an OAuth access token?
  → U2M Token-Env  ✅ SUPPORTED

Local development/testing only?
  → PAT (Personal Access Token)  ⚠️ LIMITED

Authentication Comparison

CLIENT_ID Distinction (CRITICAL)

Never mix these. M2M uses service principal credentials for server-to-server auth. U2M custom OAuth uses a registered app for user login.

Token Lifetime Summary

User-Agent (Required per PWAF)

Call once before creating WorkspaceClient:

from databricks.sdk import useragent

useragent.with_partner("YourCompany")
useragent.with_product("YourCompany_YourProduct", "1.0.0")

Alternative (on WorkspaceClient):

from databricks.sdk import WorkspaceClient

client = WorkspaceClient(
    host="https://myworkspace.cloud.databricks.com",
    token="dapi...",
    product="YourCompany_YourProduct",
    product_version="1.0.0"
)

Environment Variables Reference

Important:

• Use APP_AUTH_TYPE (not DATABRICKS_AUTH_TYPE) for app-level auth selection because the SDK reads DATABRICKS_AUTH_TYPE internally.
• Do not mix M2M and U2M environment variables. Use env -i for clean test environments.

Config Parameters Reference

Critical: When using OAuth M2M, always pass auth_type="oauth-m2m" to Config to avoid "cannot configure default credentials" error.

Complete Examples

PAT Authentication (Testing Only)

import os
from databricks.sdk import WorkspaceClient, useragent

useragent.with_partner("YourCompany")
useragent.with_product("YourProduct", "1.0.0")

def get_client_pat():
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    if not host.startswith("https://"):
        host = f"https://{host}"
    token = os.environ.get("DATABRICKS_TOKEN", "").strip()
    if not token:
        raise ValueError("DATABRICKS_TOKEN required")
    
    return WorkspaceClient(host=host, token=token)

client = get_client_pat()
me = client.current_user.me()
print(f"Authenticated as: {me.user_name}")

OAuth M2M (Recommended for Production)

import os
from databricks.sdk import WorkspaceClient, useragent
from databricks.sdk.config import Config

useragent.with_partner("YourCompany")
useragent.with_product("YourProduct", "1.0.0")

def get_client_m2m():
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    if not host.startswith("https://"):
        host = f"https://{host}"
    
    config = Config(
        host=host,
        client_id=os.environ.get("DATABRICKS_CLIENT_ID", ""),
        client_secret=os.environ.get("DATABRICKS_CLIENT_SECRET", ""),
        auth_type="oauth-m2m",  # CRITICAL: Required for M2M
    )
    return WorkspaceClient(config=config)

client = get_client_m2m()
# SDK automatically handles token refresh
tables = list(client.tables.list("samples", "nyctaxi"))
print(f"Found {len(tables)} tables")

OAuth U2M Custom OAuth App (PKCE)

import os
import base64
import hashlib
import secrets
import string
import webbrowser
from http.server import HTTPServer, BaseHTTPRequestHandler
from threading import Thread
from urllib.parse import parse_qs, urlparse, urlencode
import requests
from databricks.sdk import WorkspaceClient

DEFAULT_REDIRECT_URI = "http://localhost:8080/callback"

def pkce_verifier_and_challenge():
    """Generate PKCE code_verifier and code_challenge per RFC 7636."""
    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, scope: str = "all-apis") -> dict:
    """Run OAuth 2.0 authorization code flow with PKCE."""
    redirect_uri = redirect_uri or DEFAULT_REDIRECT_URI
    if not host.startswith("https://"):
        host = f"https://{host}"
    
    code_verifier, code_challenge = pkce_verifier_and_challenge()
    state = secrets.token_urlsafe(32)
    
    auth_url = f"{host}/oidc/v1/authorize?" + urlencode({
        "response_type": "code",
        "client_id": client_id,
        "redirect_uri": redirect_uri,
        "scope": scope,
        "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.send_header("Content-type", "text/html; charset=utf-8")
            self.end_headers()
            self.wfile.write(b"<html><body><p>Authentication successful.</p></body></html>")
        def log_message(self, *args): pass
    
    server = HTTPServer(("localhost", port), Handler)
    Thread(target=server.handle_request, daemon=True).start()
    
    webbrowser.open(auth_url)
    print(f"Waiting for redirect at {redirect_uri}...")
    server.handle_request()
    
    if not code_holder[0]:
        raise RuntimeError("No authorization code received")
    
    token_data = {
        "grant_type": "authorization_code",
        "code": code_holder[0],
        "redirect_uri": redirect_uri,
        "client_id": client_id,
        "code_verifier": code_verifier,
    }
    if client_secret:
        token_data["client_secret"] = client_secret
    
    resp = requests.post(f"{host}/oidc/v1/token", data=token_data)
    resp.raise_for_status()
    return resp.json()

def get_client_u2m():
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    client_id = os.environ.get("DATABRICKS_U2M_CLIENT_ID", "").strip()
    redirect_uri = os.environ.get("DATABRICKS_REDIRECT_URI", "") or DEFAULT_REDIRECT_URI
    client_secret = os.environ.get("DATABRICKS_U2M_CLIENT_SECRET", "").strip() or None
    
    tokens = run_u2m_flow(host, client_id, redirect_uri, client_secret)
    access_token = tokens.get("access_token")
    
    if not host.startswith("https://"):
        host = f"https://{host}"
    
    return WorkspaceClient(
        host=host,
        token=access_token,
        product="YourCompany_YourProduct",
        product_version="1.0.0"
    )

OAuth U2M Token-Env (Pre-obtained Token)

import os
from databricks.sdk import WorkspaceClient, useragent

useragent.with_partner("YourCompany")
useragent.with_product("YourProduct", "1.0.0")

def get_client_token_env():
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    if not host.startswith("https://"):
        host = f"https://{host}"
    
    # Use pre-obtained OAuth token
    access_token = os.environ.get("DATABRICKS_ACCESS_TOKEN", "").strip()
    if not access_token:
        raise ValueError("DATABRICKS_ACCESS_TOKEN required")
    
    return WorkspaceClient(host=host, token=access_token)

Error Handling Patterns

from databricks.sdk.core import DatabricksError

def classify_error(e: Exception) -> tuple[str, bool, str]:
    """Classify error and determine if retryable.
    
    Returns: (category, retryable, message)
    """
    if isinstance(e, DatabricksError):
        error_code = getattr(e, 'error_code', '')
        status_code = getattr(e, 'status_code', 0)
        
        if status_code == 401:
            return ("AUTHENTICATION_ERROR", False, "Check credentials")
        if status_code == 403:
            return ("PERMISSION_DENIED", False, "Check user/SP permissions")
        if status_code == 404:
            return ("NOT_FOUND", False, "Resource doesn't exist")
        if status_code == 429:
            return ("RATE_LIMITED", True, "Retry with exponential backoff")
        if status_code >= 500:
            return ("SERVER_ERROR", True, "Retry")
    
    msg = str(e).lower()
    if "more than one authorization method" in msg:
        return ("AUTH_CONFLICT", False, "Use env -i for clean environment")
    if "cannot configure default credentials" in msg:
        return ("CONFIG_ERROR", False, "Pass auth_type='oauth-m2m' to Config")
    if "timeout" in msg or "connection" in msg:
        return ("NETWORK_ERROR", True, "Check connectivity")
    
    return ("UNKNOWN_ERROR", False, str(e))

Retry with Exponential Backoff

import time
import random

def retry_with_backoff(operation, max_retries: int = 3):
    """Execute operation with exponential backoff on retryable errors."""
    for attempt in range(max_retries + 1):
        try:
            return operation()
        except Exception as e:
            category, retryable, _ = classify_error(e)
            if not retryable or attempt >= max_retries:
                raise
            backoff = (2 ** attempt) + random.uniform(0, 1)
            print(f"Retry {attempt + 1}/{max_retries} in {backoff:.1f}s ({category})")
            time.sleep(backoff)

Multi-Auth Pattern

import os
from databricks.sdk import WorkspaceClient, useragent
from databricks.sdk.config import Config

useragent.with_partner("YourCompany")
useragent.with_product("YourProduct", "1.0.0")

def create_client(auth_type: str = None) -> WorkspaceClient:
    """Factory for creating WorkspaceClient with different auth methods."""
    auth_type = auth_type or os.environ.get("APP_AUTH_TYPE", "oauth_m2m")
    host = os.environ.get("DATABRICKS_HOST", "").strip()
    if not host.startswith("https://"):
        host = f"https://{host}"
    
    if auth_type == "pat":
        return WorkspaceClient(
            host=host,
            token=os.environ.get("DATABRICKS_TOKEN", "")
        )
    
    if auth_type == "oauth_m2m":
        config = Config(
            host=host,
            client_id=os.environ.get("DATABRICKS_CLIENT_ID", ""),
            client_secret=os.environ.get("DATABRICKS_CLIENT_SECRET", ""),
            auth_type="oauth-m2m",
        )
        return WorkspaceClient(config=config)
    
    if auth_type == "u2m_token_env":
        return WorkspaceClient(
            host=host,
            token=os.environ.get("DATABRICKS_ACCESS_TOKEN", "")
        )
    
    raise ValueError(f"Unknown auth_type: {auth_type}")

Auth Isolation

Use env -i to ensure each authentication method is tested in isolation:

PAT Test

env -i PATH=$PATH HOME=$HOME \
  DATABRICKS_HOST=$DATABRICKS_HOST \
  DATABRICKS_TOKEN=$DATABRICKS_TOKEN \
  python your_script.py

OAuth M2M Test

env -i PATH=$PATH HOME=$HOME \
  DATABRICKS_HOST=$DATABRICKS_HOST \
  DATABRICKS_CLIENT_ID=$DATABRICKS_CLIENT_ID \
  DATABRICKS_CLIENT_SECRET=$DATABRICKS_CLIENT_SECRET \
  python your_script.py

U2M Custom OAuth App Test

env -i PATH=$PATH HOME=$HOME USER=$USER DISPLAY=$DISPLAY \
  DATABRICKS_HOST=$DATABRICKS_HOST \
  DATABRICKS_U2M_CLIENT_ID=$DATABRICKS_U2M_CLIENT_ID \
  python your_script.py

Note: U2M requires DISPLAY, HOME, USER for browser launching.

U2M Token-Env Test

env -i PATH=$PATH HOME=$HOME \
  DATABRICKS_HOST=$DATABRICKS_HOST \
  DATABRICKS_ACCESS_TOKEN=$DATABRICKS_ACCESS_TOKEN \
  python your_script.py

Redirect URI Reference

Registering Custom Redirect URIs:

1. Databricks Account Console → Settings → App connections
2. Select your OAuth app → Edit
3. Add redirect URI under "Redirect URIs"
4. Set DATABRICKS_REDIRECT_URI environment variable to match

Validation Query

Use this query to verify authentication and User-Agent:

SELECT event_time, user_agent, action_name, request_params
FROM system.access.audit
WHERE event_time > current_timestamp() - INTERVAL 1 HOUR
  AND lower(user_agent) LIKE '%yourcompany%'
ORDER BY event_time DESC
LIMIT 10;

SDK Validation

def validate_connection(client: WorkspaceClient) -> bool:
    """Validate SDK connection by calling current_user API."""
    try:
        me = client.current_user.me()
        print(f"Connected as: {me.user_name}")
        return True
    except Exception as e:
        print(f"Validation failed: {e}")
        return False

U2M Token Refresh Pattern

import time
import threading

class TokenRefreshManager:
    """Manages automatic U2M token refresh."""
    
    def __init__(self, authenticator, refresh_interval_sec: int = 45 * 60):
        self.authenticator = authenticator
        self.refresh_interval = refresh_interval_sec
        self.current_token = None
        self._stop = False
        self._thread = None
    
    def start(self):
        self._refresh()
        self._thread = threading.Thread(target=self._refresh_loop, daemon=True)
        self._thread.start()
    
    def _refresh_loop(self):
        while not self._stop:
            time.sleep(self.refresh_interval)
            if not self._stop:
                self._refresh()
    
    def _refresh(self):
        try:
            self.current_token = self.authenticator()
            print("Token refreshed successfully")
        except Exception as e:
            print(f"Token refresh failed: {e}")
    
    def get_token(self) -> str:
        return self.current_token
    
    def stop(self):
        self._stop = True

Troubleshooting

Host Normalization

def normalize_host(host: str) -> str:
    """Ensure host has https:// prefix."""
    host = host.strip()
    if not host.startswith("https://") and not host.startswith("http://"):
        host = f"https://{host}"
    return host.rstrip("/")

Implementation Learnings

1. Config import matters: Use from databricks.sdk.config import Config (not from databricks.sdk import Config)
2. auth_type is critical for M2M: Always pass auth_type="oauth-m2m" when using client_id/client_secret
3. UserAgent is global: Call useragent.with_partner() and useragent.with_product() once before creating any client
4. M2M auto-refreshes: SDK handles token lifecycle automatically for M2M
5. U2M requires manual refresh: Implement refresh timer for long-running U2M applications
6. Host normalization: SDK accepts both https://host and host, but be explicit

PKCE Flow Implementation Details

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

Use Event for callback signaling:

from threading import Thread, Event

callback_received = Event()

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 requests (ignore favicon):

class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        req_path = urlparse(self.path).path
        if req_path == callback_path:  # Only process OAuth callback
            # Extract code and signal completion
            callback_received.set()
        else:
            self.send_response(404)  # Ignore favicon, etc.
            self.end_headers()

Use output flushing for immediate feedback:

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

Set server timeout for responsiveness:

server = HTTPServer((host, port), Handler)
server.timeout = 1  # 1 second timeout per handle_request call

Wait with timeout:

CALLBACK_TIMEOUT = 120  # seconds
if not callback_received.wait(timeout=CALLBACK_TIMEOUT):
    raise RuntimeError(f"Timed out after {CALLBACK_TIMEOUT}s")

Key insight: Using server.handle_request() in a loop (instead of just once) ensures that browser prefetch requests (like favicon) don't consume the single handler and block the OAuth callback.

SDK Workarounds Summary

Key Differences from SQL Connectors