databricks-isv-python-sqlalchemy

<!-- skill-version: 1.0.0 --> # 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](https://databrickslabs.github.io/partner-architecture/) - [Access & Authentication](https://databrickslabs.github.io/partner-architecture/isv-partners/lakehouse-patterns/access-auth/) - [OAuth M2M](https://databrickslabs.github.io/partner-architecture/isv-partners/lakehouse-patterns/access-auth/oauth-m2m) - [OAuth U2M](https://databrickslabs.github.io/partner-architecture/isv-partners/lakehouse-patterns/access-auth/oauth-u2m) - [Telemetry Attribution](https://databrickslabs.github.io/partner-architecture/isv-partners/lakehouse-patterns/access-auth/telemetry) ## Packages ```bash 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 | Auth Type | Use Case | Token Lifetime | Refresh | PWAF Status | |-----------|----------|---------------|---------|-------------| | PAT | Testing/development | Configurable (90d default) | Manual regeneration | Limited | | OAuth M2M | Production, automated | 1 hour | Re-authenticate via SDK | Recommended | | OAuth U2M PKCE | User-interactive apps | 1 hour | Use refresh_token | Recommended | | OAuth U2M Token-Env | Headless/CI | 1 hour | External refresh | Supported | ## 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 | Token Type | Lifetime | Refresh Strategy | |------------|----------|-----------------| | PAT | Configurable (default 90 days) | Manual regeneration via workspace | | OAuth M2M access_token | 1 hour | Re-call `Config.authenticate()` | | OAuth U2M access_token | 1 hour | Use refresh_token or re-authenticate | | OAuth U2M refresh_token | Up to 90 days | Request `offline_access` scope | ## User-Agent (Required for PWAF) Format: `<isv-name>_<product-name>/<product-version>` ```python 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):** ```bash 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:** ```bash DATABRICKS_TOKEN # Personal access token (starts with dapi...) ``` **OAuth M2M Authentication:** ```bash DATABRICKS_CLIENT_ID # Service Principal UUID DATABRICKS_CLIENT_SECRET # OAuth secret for the service principal ``` **OAuth U2M Custom OAuth App:** ```bash 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:** ```bash DATABRICKS_ACCESS_TOKEN # Pre-obtained OAuth access token ``` **Multi-Auth Selector:** ```bash 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://`): ```python 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 ```python 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) ```python 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) ```python 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(