May 7, 2025 · May 7, 2025 · May 7, 2025 · May 7, 2025 · May 7, 2025 · May 7, 2025
diff --git a/.github/workflows/token-federation-test.yml b/.github/workflows/token-federation-test.yml
 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
diff --git a/examples/token_federation_examples.py b/examples/token_federation_examples.py
 """
 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)
Original file line number	Diff line number	Diff line change
		@@ -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
Original file line number	Diff line number	Diff line change
		@@ -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)