Implements Token Federation for Python Driver #1
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 |