federation follow up [draft]

.github/workflows/token-federation-test.yml

@@ -0,0 +1,78 @@
+name: Token Federation Test
+
+# Tests token federation functionality with GitHub Actions OIDC tokens
+on:
+  # Manual trigger with required inputs
+  workflow_dispatch:
+    inputs:
+      databricks_host:
+        description: 'Databricks host URL (e.g., example.cloud.databricks.com)'
+        required: true
+      databricks_http_path:
+        description: 'Databricks HTTP path (e.g., /sql/1.0/warehouses/abc123)'
+        required: true
+      identity_federation_client_id:
+        description: 'Identity federation client ID'
+        required: true
+
+  # Run on PRs that might affect token federation
+  pull_request:
+    branches: [main]
+    paths:
+      - 'src/databricks/sql/auth/**'
+      - 'examples/token_federation_*.py'
+      - 'tests/token_federation/**'
+      - '.github/workflows/token-federation-test.yml'
+
+  # Run on push to main that affects token federation
+  push:
+    branches: [main]
+    paths:
+      - 'src/databricks/sql/auth/**'
+      - 'examples/token_federation_*.py'
+      - 'tests/token_federation/**'
+      - '.github/workflows/token-federation-test.yml'
+
+permissions:
+  id-token: write  # Required for GitHub OIDC token
+  contents: read
+
+jobs:
+  test-token-federation:
+    name: Test Token Federation
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python 3.9
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.9'
+          cache: 'pip'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+          pip install pyarrow
+
+      - name: Get GitHub OIDC token
+        id: get-id-token
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const token = await core.getIDToken('https://github.com/databricks')
+            core.setSecret(token)
+            core.setOutput('token', token)
+
+      - name: Test token federation with GitHub OIDC token
+        env:
+          DATABRICKS_HOST_FOR_TF: ${{ github.event_name == 'workflow_dispatch' && inputs.databricks_host || secrets.DATABRICKS_HOST_FOR_TF }}
+          DATABRICKS_HTTP_PATH_FOR_TF: ${{ github.event_name == 'workflow_dispatch' && inputs.databricks_http_path || secrets.DATABRICKS_HTTP_PATH_FOR_TF }}
+          IDENTITY_FEDERATION_CLIENT_ID: ${{ github.event_name == 'workflow_dispatch' && inputs.identity_federation_client_id || secrets.IDENTITY_FEDERATION_CLIENT_ID }}
+          OIDC_TOKEN: ${{ steps.get-id-token.outputs.token }}
+        run: python tests/token_federation/github_oidc_test.py

examples/token_federation_examples.py

@@ -0,0 +1,109 @@
+"""
+Databricks SQL Token Federation Examples
+
+This script token federation flows:
+1. U2M + Account-wide federation
+2. U2M + Workflow-level federation  
+3. M2M + Account-wide federation
+4. M2M + Workflow-level federation
+5. Access Token + Workflow-level federation
+6. Access Token + Account-wide federation
+
+Token Federation Documentation:
+------------------------------
+For detailed setup instructions, refer to the official Databricks documentation:
+
+- General Token Federation Overview:
+  https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation.html
+
+- Token Exchange Process:
+  https://docs.databricks.com/aws/en/dev-tools/auth/oauth-federation-howto.html
+
+- Azure OAuth Token Federation:
+  https://learn.microsoft.com/en-us/azure/databricks/dev-tools/auth/oauth-federation
+
+Environment variables required:
+- DATABRICKS_HOST: Databricks workspace hostname
+- DATABRICKS_HTTP_PATH: HTTP path for the SQL warehouse
+- AZURE_TENANT_ID: Azure tenant ID
+- AZURE_CLIENT_ID: Azure client ID for service principal
+- AZURE_CLIENT_SECRET: Azure client secret
+- DATABRICKS_SERVICE_PRINCIPAL_ID: Databricks service principal ID for workflow federation
+"""
+
+import os
+from databricks import sql
+
+def run_query(connection, description):
+    cursor = connection.cursor()
+    cursor.execute("SELECT 1+1 AS result")
+    result = cursor.fetchall()
+    print(f"Query result: {result[0][0]}")
+
+    cursor.close()
+
+def demonstrate_m2m_federation(env_vars, use_workflow_federation=False):
+    """Demonstrate M2M (service principal) token federation"""
+
+    connection_params = {
+        "server_hostname": env_vars["DATABRICKS_HOST"],
+        "http_path": env_vars["DATABRICKS_HTTP_PATH"],
+        "auth_type": "client-credentials",
+        "oauth_client_id": env_vars["AZURE_CLIENT_ID"],
+        "client_secret": env_vars["AZURE_CLIENT_SECRET"],
+        "tenant_id": env_vars["AZURE_TENANT_ID"],
+        "use_token_federation": True
+    }
+
+    if use_workflow_federation and env_vars["DATABRICKS_SERVICE_PRINCIPAL_ID"]:
+        connection_params["identity_federation_client_id"] = env_vars["DATABRICKS_SERVICE_PRINCIPAL_ID"]
+        description = "M2M + Workflow-level Federation"
+    else:
+        description = "M2M + Account-wide Federation"
+
+    with sql.connect(**connection_params) as connection:
+        run_query(connection, description)
+
+
+def demonstrate_u2m_federation(env_vars, use_workflow_federation=False):
+    """Demonstrate U2M (interactive) token federation"""
+
+    connection_params = {
+        "server_hostname": env_vars["DATABRICKS_HOST"],
+        "http_path": env_vars["DATABRICKS_HTTP_PATH"],
+        "auth_type": "databricks-oauth",  # Will open browser for interactive auth
+        "use_token_federation": True
+    }
+
+    if use_workflow_federation and env_vars["DATABRICKS_SERVICE_PRINCIPAL_ID"]:
+        connection_params["identity_federation_client_id"] = env_vars["DATABRICKS_SERVICE_PRINCIPAL_ID"]
+        description = "U2M + Workflow-level Federation (Interactive)"
+    else:
+        description = "U2M + Account-wide Federation (Interactive)"
+
+    # This will open a browser for interactive auth
+    with sql.connect(**connection_params) as connection:
+        run_query(connection, description)
+
+def demonstrate_access_token_federation(env_vars):
+    """Demonstrate access token token federation"""
+
+    access_token = os.environ.get("ACCESS_TOKEN") # This is to demonstrate a token obtained from an identity provider
+
+    connection_params = {
+        "server_hostname": env_vars["DATABRICKS_HOST"],
+        "http_path": env_vars["DATABRICKS_HTTP_PATH"],
+        "access_token": access_token,
+        "use_token_federation": True
+    }
+
+    # Add workflow federation if available
+    if env_vars["DATABRICKS_SERVICE_PRINCIPAL_ID"]:
+        connection_params["identity_federation_client_id"] = env_vars["DATABRICKS_SERVICE_PRINCIPAL_ID"]
+        description = "Access Token + Workflow-level Federation"
+    else:
+        description = "Access Token + Account-wide Federation"
+
+    with sql.connect(**connection_params) as connection:
+        run_query(connection, description)
+