initial commit

Author: madhav-db

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

@@ -0,0 +1,166 @@
+name: Token Federation Test
+
+# This workflow tests token federation functionality with GitHub Actions OIDC tokens
+# in the databricks-sql-python connector to ensure CI/CD functionality
+
+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
+  
+  # Automatically run on PR that changes token federation files
+  pull_request:
+    branches:
+      - main
+  
+  # Run on push to main that affects token federation
+  push:
+    paths:
+      - 'src/databricks/sql/auth/token_federation.py'
+      - 'src/databricks/sql/auth/auth.py'
+      - 'examples/token_federation_*.py'
+    branches:
+      - main
+
+permissions:
+  # Required for GitHub OIDC token
+  id-token: write
+  contents: read
+
+jobs:
+  test-token-federation:
+    runs-on: 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'
+      
+      - 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')
+            core.setSecret(token)
+            core.setOutput('token', token)
+      
+      - name: Create test script
+        run: |
+          cat > test_github_token_federation.py << 'EOF'
+          #!/usr/bin/env python3
+          
+          """
+          Test script for Databricks SQL token federation with GitHub Actions OIDC tokens.
+          
+          This script demonstrates how to use the Databricks SQL connector with token federation
+          using a GitHub Actions OIDC token. It connects to a Databricks SQL warehouse,
+          runs a simple query, and shows the connected user.
+          """
+          
+          import os
+          import sys
+          import json
+          import base64
+          from databricks import sql
+          
+          def decode_jwt(token):
+              """Decode and return the claims from a JWT token."""
+              try:
+                  parts = token.split(".")
+                  if len(parts) != 3:
+                      raise ValueError("Invalid JWT format")
+                  
+                  payload = parts[1]
+                  padding = '=' * (4 - len(payload) % 4)
+                  payload += padding
+                  
+                  decoded = base64.b64decode(payload)
+                  return json.loads(decoded)
+              except Exception as e:
+                  print(f"Failed to decode token: {str(e)}")
+                  return None
+          
+          def main():
+              # Get GitHub OIDC token
+              github_token = os.environ.get("OIDC_TOKEN")
+              if not github_token:
+                  print("GitHub OIDC token not available")
+                  sys.exit(1)
+              
+              # Get Databricks connection parameters
+              host = os.environ.get("DATABRICKS_HOST")
+              http_path = os.environ.get("DATABRICKS_HTTP_PATH")
+              identity_federation_client_id = os.environ.get("IDENTITY_FEDERATION_CLIENT_ID")
+              
+              if not host or not http_path:
+                  print("Missing Databricks connection parameters")
+                  sys.exit(1)
+              
+              claims = decode_jwt(github_token)
+              if claims:
+                  print(f"Token issuer: {claims.get('iss', 'unknown')}")
+                  print(f"Token subject: {claims.get('sub', 'unknown')}")
+                  print(f"Token audience: {claims.get('aud', 'unknown')}")
+              
+              try:
+                  # Connect to Databricks using token federation
+                  print(f"Connecting to Databricks at {host}{http_path}")
+                  with sql.connect(
+                      server_hostname=host,
+                      http_path=http_path,
+                      access_token=github_token,
+                      auth_type="token-federation",
+                      identity_federation_client_id=identity_federation_client_id
+                  ) as connection:
+                      print("Connection established successfully")
+                      
+                      # Execute a simple query
+                      cursor = connection.cursor()
+                      cursor.execute("SELECT 1 + 1 as result")
+                      result = cursor.fetchall()
+                      print(f"Query result: {result[0][0]}")
+                      
+                      # Show current user
+                      cursor.execute("SELECT current_user() as user")
+                      result = cursor.fetchall()
+                      print(f"Connected as user: {result[0][0]}")
+                      
+                      print("Token federation test successful!")
+                      return True
+              except Exception as e:
+                  print(f"Error connecting to Databricks: {str(e)}")
+                  sys.exit(1)
+          
+          if __name__ == "__main__":
+              main()
+          EOF
+          chmod +x test_github_token_federation.py
+      
+      - name: Test token federation with GitHub OIDC token
+        env:
+          DATABRICKS_HOST: ${{ github.event_name == 'workflow_dispatch' && inputs.databricks_host || secrets.DATABRICKS_HOST }}
+          DATABRICKS_HTTP_PATH: ${{ github.event_name == 'workflow_dispatch' && inputs.databricks_http_path || secrets.DATABRICKS_HTTP_PATH }}
+          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 test_github_token_federation.py