Tests

Author: hectorcast-db

Makefile

@@ -28,7 +28,7 @@ test:
 	pytest -m 'not integration and not benchmark' --cov=databricks --cov-report html tests
 
 integration:
-	pytest -n auto -m 'integration and not benchmark' --reruns 2 --dist loadgroup --cov=databricks --cov-report html tests
+	pytest -n auto  --dist loadgroup --cov=databricks --cov-report html tests/integration/test_auth.py
 
 benchmark:
 	pytest -m 'benchmark' tests

databricks/sdk/__init__.py

@@ -61,6 +61,7 @@ class Config:
     host: str = ConfigAttribute(env="DATABRICKS_HOST")
     account_id: str = ConfigAttribute(env="DATABRICKS_ACCOUNT_ID")
     token: str = ConfigAttribute(env="DATABRICKS_TOKEN", auth="pat", sensitive=True)
+    token_audience: str = ConfigAttribute(env="DATABRICKS_TOKEN_AUDIENCE", auth="databricks-wif")
     username: str = ConfigAttribute(env="DATABRICKS_USERNAME", auth="basic")
     password: str = ConfigAttribute(env="DATABRICKS_PASSWORD", auth="basic", sensitive=True)
     client_id: str = ConfigAttribute(env="DATABRICKS_CLIENT_ID", auth="oauth")

databricks/sdk/config.py

@@ -23,6 +23,7 @@
 from .azure import add_sp_management_token, add_workspace_id_header
 from .oauth import (ClientCredentials, OAuthClient, Refreshable, Token,
                     TokenCache, TokenSource)
+from .oidc_token_supplier import GitHubOIDCTokenSupplier
 
 CredentialsProvider = Callable[[], Dict[str, str]]
 
@@ -314,6 +315,44 @@ def token() -> Token:
     return OAuthCredentialsProvider(refreshed_headers, token)
 
 
+@oauth_credentials_strategy("databricks-wif", ["host", "client_id", "token_audience"])
+def databricks_wif(cfg: "Config") -> Optional[CredentialsProvider]:
+    supplier = GitHubOIDCTokenSupplier()
+    # Try to get a token. If no supplier returns a token, we cannot use this authentication mode.
+    token = supplier.get_oidc_token(cfg.token_audience)
+    if not token:
+        return None
+
+    def token_source_for(audience: str) -> TokenSource:
+        token = supplier.get_oidc_token(audience)
+        if not token:
+            # Should not happen, since we checked it above.
+            raise Exception("Cannot get OIDC token")
+        params = {
+            "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
+            "subject_token": token,
+            "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
+        }
+        return ClientCredentials(
+            client_id=cfg.client_id,
+            client_secret="",  # we have no (rotatable) secrets in OIDC flow
+            token_url=cfg.oidc_endpoints.token_endpoint,
+            endpoint_params=params,
+            scopes=["all-apis"],
+            use_params=True,
+            disable_async=not cfg.enable_experimental_async_token_refresh,
+        )
+
+    def refreshed_headers() -> Dict[str, str]:
+        token = token_source_for(cfg.token_audience).token()
+        return {"Authorization": f"{token.token_type} {token.access_token}"}
+
+    def token() -> Token:
+        return token_source_for(cfg.token_audience).token()
+
+    return OAuthCredentialsProvider(refreshed_headers, token)
+
+
 @oauth_credentials_strategy("github-oidc-azure", ["host", "azure_client_id"])
 def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]:
     if "ACTIONS_ID_TOKEN_REQUEST_TOKEN" not in os.environ:
@@ -325,16 +364,8 @@ def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]:
     if not cfg.is_azure:
         return None
 
-    # See https://docs.github.com/en/actions/deployment/security-hardening-your-deployments/configuring-openid-connect-in-cloud-providers
-    headers = {"Authorization": f"Bearer {os.environ['ACTIONS_ID_TOKEN_REQUEST_TOKEN']}"}
-    endpoint = f"{os.environ['ACTIONS_ID_TOKEN_REQUEST_URL']}&audience=api://AzureADTokenExchange"
-    response = requests.get(endpoint, headers=headers)
-    if not response.ok:
-        return None
-
-    # get the ID Token with aud=api://AzureADTokenExchange sub=repo:org/repo:environment:name
-    response_json = response.json()
-    if "value" not in response_json:
+    token = GitHubOIDCTokenSupplier().get_oidc_token("api://AzureADTokenExchange")
+    if not token:
         return None
 
     logger.info(
@@ -344,7 +375,7 @@ def github_oidc_azure(cfg: "Config") -> Optional[CredentialsProvider]:
     params = {
         "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
         "resource": cfg.effective_azure_login_app_id,
-        "client_assertion": response_json["value"],
+        "client_assertion": token,
     }
     aad_endpoint = cfg.arm_environment.active_directory_endpoint
     if not cfg.azure_tenant_id:
@@ -927,6 +958,7 @@ def __init__(self) -> None:
             basic_auth,
             metadata_service,
             oauth_service_principal,
+            databricks_wif,
             azure_service_principal,
             github_oidc_azure,
             azure_cli,

databricks/sdk/credentials_provider.py

@@ -0,0 +1,33 @@
+import os
+from time import sleep
+from typing import Optional
+
+import requests
+
+
+class GitHubOIDCTokenSupplier:
+    """
+    Supplies OIDC tokens from GitHub Actions.
+    """
+
+    def get_oidc_token(self, audience: str) -> Optional[str]:
+        if "ACTIONS_ID_TOKEN_REQUEST_TOKEN" not in os.environ or "ACTIONS_ID_TOKEN_REQUEST_URL" not in os.environ:
+            # not in GitHub actions
+            return None
+        # See https://docs.github.com/en/actions/deployment/security-hardening-your-deployments/configuring-openid-connect-in-cloud-providers
+        headers = {"Authorization": f"Bearer {os.environ['ACTIONS_ID_TOKEN_REQUEST_TOKEN']}"}
+        endpoint = f"{os.environ['ACTIONS_ID_TOKEN_REQUEST_URL']}&audience={audience}"
+        response = requests.get(endpoint, headers=headers)
+        if not response.ok:
+            return None
+
+        # get the ID Token with aud=api://AzureADTokenExchange sub=repo:org/repo:environment:name
+        response_json = response.json()
+        if "value" not in response_json:
+            return None
+
+        # GitHub issued time is not allways in sync, and can give tokens which are not yet valid.
+        # TODO: Remove this after Databricks API is updated to handle such cases.
+        sleep(2)
+
+        return response_json["value"]

databricks/sdk/oidc_token_supplier.py

@@ -12,6 +12,8 @@
 
 import pytest
 
+from databricks.sdk import AccountClient, WorkspaceClient
+from databricks.sdk.service import iam, oauth2
 from databricks.sdk.service.compute import (ClusterSpec, DataSecurityMode,
                                             Library, ResultType, SparkVersion)
 from databricks.sdk.service.jobs import NotebookTask, Task, ViewType
@@ -198,3 +200,72 @@ def _task_outputs(w, run):
                             output += data["data"]
         task_outputs[task_run.task_key] = output
     return task_outputs
+
+
+def test_wif_account(ucacct, env_or_skip, random):
+
+    sp = ucacct.service_principals.create(
+        active=True,
+        display_name="py-sdk-test-" + random(),
+        roles=[iam.ComplexValue(value="account_admin")],
+    )
+
+    ucacct.service_principal_federation_policy.create(
+        policy=oauth2.FederationPolicy(
+            oidc_policy=oauth2.OidcFederationPolicy(
+                issuer="https://token.actions.githubusercontent.com",
+                audiences=["https://github.com/databricks-eng"],
+                subject="repo:databricks-eng/eng-dev-ecosystem:environment:integration-tests",
+            )
+        ),
+        service_principal_id=sp.id,
+    )
+
+    ac = AccountClient(
+        host=ucacct.config.host,
+        account_id=ucacct.config.account_id,
+        client_id=sp.application_id,
+        auth_type="databricks-wif",
+        token_audience="https://github.com/databricks-eng",
+    )
+
+    groups = ac.groups.list()
+
+    next(groups)
+
+
+def test_wif_workspace(ucacct, env_or_skip, random):
+
+    workspace_id = env_or_skip("TEST_WORKSPACE_ID")
+    workspace_url = env_or_skip("TEST_WORKSPACE_URL")
+
+    sp = ucacct.service_principals.create(
+        active=True,
+        display_name="py-sdk-test-" + random(),
+    )
+
+    ucacct.service_principal_federation_policy.create(
+        policy=oauth2.FederationPolicy(
+            oidc_policy=oauth2.OidcFederationPolicy(
+                issuer="https://token.actions.githubusercontent.com",
+                audiences=["https://github.com/databricks-eng"],
+                subject="repo:databricks-eng/eng-dev-ecosystem:environment:integration-tests",
+            )
+        ),
+        service_principal_id=sp.id,
+    )
+
+    ucacct.workspace_assignment.update(
+        workspace_id=workspace_id,
+        principal_id=sp.id,
+        permissions=[iam.WorkspacePermission.ADMIN],
+    )
+
+    ws = WorkspaceClient(
+        host=workspace_url,
+        client_id=sp.application_id,
+        auth_type="databricks-wif",
+        token_audience="https://github.com/databricks-eng",
+    )
+
+    ws.current_user.me()