token-federation-test.yml

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:
        group: databricks-protected-runner-group
        labels: linux-ubuntu-latest
    
    steps:
      - name: Debug OIDC Claims
        uses: github/actions-oidc-debugger@main
        with:
          audience: '${{ github.server_url }}/${{ github.repository_owner }}'
          
      - 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_FOR_TF")
              http_path = os.environ.get("DATABRICKS_HTTP_PATH_FOR_TF")
              identity_federation_client_id = os.environ.get("IDENTITY_FEDERATION_CLIENT_ID_FOR_TF")
              
              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_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 test_github_token_federation.py