Adding support for Snowflake Workload Identity Federation

dbt-snowflake/pyproject.toml

@@ -25,8 +25,7 @@ classifiers = [
 dependencies = [
     "dbt-common>=1.10,<2.0",
     "dbt-adapters>=1.16,<2.0",
-    # lower bound pin due to CVE-2025-24794
-    "snowflake-connector-python[secure-local-storage]>=3.13.1,<4.0.0",
+    "snowflake-connector-python[secure-local-storage]>=3.17.3,<4.0.0",
     "certifi<2025.4.26",
     # add dbt-core to ensure backwards compatibility of installation, this is not a functional dependency
     "dbt-core>=1.10.0rc0",

dbt-snowflake/src/dbt/adapters/snowflake/connections.py

@@ -38,6 +38,8 @@
     BindUploadError,
 )
 
+from snowflake.connector.network import WORKLOAD_IDENTITY_AUTHENTICATOR
+
 from dbt_common.exceptions import (
     DbtInternalError,
     DbtRuntimeError,
@@ -114,6 +116,8 @@ class SnowflakeCredentials(Credentials):
     # this needs to default to `None` so that we can tell if the user set it; see `__post_init__()`
     reuse_connections: Optional[bool] = None
     s3_stage_vpce_dns_name: Optional[str] = None
+    workload_identity_provider: Optional[str] = None
+    workload_identity_entra_resource: Optional[str] = None
 
     def __post_init__(self):
         if self.authenticator != "oauth" and (self.oauth_client_secret or self.oauth_client_id):
@@ -182,6 +186,8 @@ def _connection_keys(self):
             "insecure_mode",
             "reuse_connections",
             "s3_stage_vpce_dns_name",
+            "workload_identity_provider",
+            "workload_identity_entra_resource",
         )
 
     def auth_args(self):
@@ -232,6 +238,22 @@ def auth_args(self):
                 result["token"] = self.token
                 result["authenticator"] = "oauth"
 
+            elif self.authenticator.lower() == "workload_identity":
+                result["authenticator"] = WORKLOAD_IDENTITY_AUTHENTICATOR
+
+                if not self.workload_identity_provider:
+                    raise DbtConfigError(
+                        "workload_identity_provider must be set if authenticator='workload_identity'!"
+                    )
+                result["workload_identity_provider"] = self.workload_identity_provider
+
+                if self.token:
+                    result["token"] = self.token
+                if self.workload_identity_entra_resource:
+                    result["workload_identity_entra_resource"] = (
+                        self.workload_identity_entra_resource
+                    )
+
             # enable id token cache for linux
             result["client_store_temporary_credential"] = True
             # enable mfa token cache for linux

dbt-snowflake/src/dbt/include/snowflake/profile_template.yml

@@ -20,6 +20,10 @@ prompts:
       authenticator:
         hint: "'externalbrowser' or a valid Okta URL"
         default: 'externalbrowser'
+    workload_identity:
+      _fixed_authenticator: workload_identity
+      workload_identity_provider:
+        hint: Must be one of the following - [OIDC, AWS, AZURE, GCP]
   role:
     hint: 'dev role'
   warehouse:

dbt-snowflake/tests/functional/auth_tests/test_workload_identity_federation_aws.py

@@ -0,0 +1,87 @@
+"""
+Functional tests for Snowflake Workload Identity Federation (WIF) with AWS authentication.
+Prerequisites for testing WIF with AWS:
+1. **AWS IAM Configuration:**
+   Create an IAM role that can be assumed by the EC2 service. An example trust policy below:
+   ```
+   {
+     "Version": "2012-10-17",
+     "Statement": [
+       {
+         "Effect": "Allow",
+         "Principal": {
+           "Service": "ec2.amazonaws.com"  // or your specific service
+         },
+         "Action": "sts:AssumeRole"
+       }
+     ]
+   }
+   ```
+2. **EC2 Instance:**
+    Launch an EC2 instance with the IAM role attached as an instance profile.
+    Connect to the EC2 instance and
+
+
+3. **Snowflake User Configuration:**
+   Create a service user in Snowflake with WIF enabled:
+   ```sql
+   CREATE USER <username>
+     WORKLOAD_IDENTITY = (
+       TYPE = AWS
+       ARN = '<amazon_iam_role_arn>'
+     )
+     TYPE = SERVICE
+     DEFAULT_ROLE = <role>;
+   ```
+   Replace `<username>` with your desired username and `<amazon_iam_role_arn>`
+   with the ARN of your AWS IAM role.
+4. **AWS Environment:**
+   This test must run from within the configured EC2 environment.
+   Connect to the EC2 instance using SSH or similar.
+   Clone this repository, run the setup, and execute this test e.g.
+   `hatch run pytest tests/functional/auth_tests/test_workload_identity_federation_aws.py::test_snowflake_wif_basic_functionality`
+5. **Environment Variables:**
+   Set the following environment variables for testing:
+   - SNOWFLAKE_TEST_ACCOUNT: Your Snowflake account identifier
+   - SNOWFLAKE_TEST_WIF_USER: The Snowflake service user created for WIF
+   - SNOWFLAKE_TEST_DATABASE: Test database name
+   - SNOWFLAKE_TEST_WAREHOUSE: Test warehouse name
+   - SNOWFLAKE_TEST_ROLE: Snowflake Role for the user (optional)
+   - SNOWFLAKE_TEST_SCHEMA: Schema for testing (optional, defaults to schema in profile)
+Note: WIF authentication relies on being in the AWS environment, so these tests can't be run locally or in the CI/CD pipeline.
+"""
+
+import os
+from dbt.tests.util import run_dbt
+import pytest
+
+
+_MODELS__MODEL_1_SQL = """
+select 1 as id, 'wif_test' as source
+"""
+
+
+class TestSnowflakeWorkloadIdentityFederation:
+    @pytest.fixture(scope="class", autouse=True)
+    def dbt_profile_target(self):
+        return {
+            "type": "snowflake",
+            "threads": 4,
+            "account": os.getenv("SNOWFLAKE_TEST_ACCOUNT"),
+            "user": os.getenv("SNOWFLAKE_TEST_WIF_USER"),
+            "database": os.getenv("SNOWFLAKE_TEST_DATABASE"),
+            "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"),
+            "role": os.getenv("SNOWFLAKE_TEST_ROLE"),
+            "authenticator": "workload_identity",
+            "workload_identity_provider": "aws",
+        }
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "model_1.sql": _MODELS__MODEL_1_SQL,
+        }
+
+    def test_snowflake_wif_basic_functionality(self, project):
+        """Test basic dbt functionality with WIF authentication"""
+        run_dbt()

dbt-snowflake/tests/functional/auth_tests/test_workload_identity_federation_oidc.py

@@ -0,0 +1,106 @@
+"""
+Functional tests for Snowflake Workload Identity Federation (WIF) with OIDC authentication.
+Prerequisites for testing WIF with OIDC:
+
+1. **Create a Snowflake User with OIDC Auth**
+
+  Create a service user in Snowflake with WIF enabled:
+  ```sql
+  CREATE USER <username>
+    TYPE = SERVICE
+    WORKLOAD_IDENTITY = (
+      TYPE = OIDC,
+      ISSUER = 'https://token.actions.githubusercontent.com',
+      SUBJECT = 'repo:<REPO_OWNER>/dbt-adapters:ref:refs/heads/main',
+      OIDC_AUDIENCE_LIST = ('snowflakecomputing.com')
+    );
+  ```
+
+2. **Create a GitHub Actions that generates the OIDC token and runs the test **
+
+  ```yaml
+
+name: Run Snowflake Workload Identity Federation (WIF) Test
+on:
+  workflow_dispatch:
+  push:
+    branches: [ main ]
+
+permissions:
+  contents: read
+  id-token: write
+
+jobs:
+  run-snowflake:
+    runs-on: ubuntu-latest
+    env:
+      SNOWFLAKE_TEST_ACCOUNT: <ACCOUNT_ID>
+      SNOWFLAKE_TEST_DATABASE: <DB_NAME>
+      SNOWFLAKE_TEST_WAREHOUSE: <WH_NAME>
+      SNOWFLAKE_TEST_ROLE: <ROLE_NAME>
+      SNOWFLAKE_TEST_WIF_USER: <USERNAME>
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Get OIDC token for Snowflake
+        id: oidc
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const token = await core.getIDToken('snowflakecomputing.com');
+            core.setOutput('id_token', token);
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - uses: pypa/hatch@install
+
+      - run: hatch run setup
+        working-directory: ./dbt-snowflake
+
+      - run: hatch run python -m pytest tests/functional/auth_tests/test_workload_identity_federation_oidc.py
+        working-directory: ./dbt-snowflake
+        env:
+          ODIC_TOKEN: ${{ steps.oidc.outputs.id_token }}
+  ```
+
+"""
+
+import os
+from dbt.tests.util import run_dbt
+import pytest
+
+
+_MODELS__MODEL_1_SQL = """
+select 1 as id, 'wif_test' as source
+"""
+
+
+class TestSnowflakeWorkloadIdentityFederation:
+    @pytest.fixture(scope="class", autouse=True)
+    def dbt_profile_target(self):
+        return {
+            "type": "snowflake",
+            "threads": 4,
+            "account": os.getenv("SNOWFLAKE_TEST_ACCOUNT"),
+            "user": os.getenv("SNOWFLAKE_TEST_WIF_USER"),
+            "database": os.getenv("SNOWFLAKE_TEST_DATABASE"),
+            "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"),
+            "role": os.getenv("SNOWFLAKE_TEST_ROLE"),
+            "authenticator": "workload_identity",
+            "workload_identity_provider": "oidc",
+            "token": os.getenv("ODIC_TOKEN"),
+        }
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "model_1.sql": _MODELS__MODEL_1_SQL,
+        }
+
+    def test_snowflake_wif_basic_functionality(self, project):
+        """Test basic dbt functionality with WIF authentication"""
+        run_dbt()

dbt-snowflake/tests/unit/test_connections.py

@@ -4,6 +4,7 @@
 from unittest.mock import Mock, patch
 import multiprocessing
 from dbt.adapters.exceptions.connection import FailedToConnectError
+from dbt_common.exceptions import DbtConfigError
 import dbt.adapters.snowflake.connections as connections
 import dbt.adapters.events.logging
 
@@ -67,3 +68,37 @@ def test_snowflake_oauth_expired_token_raises_error():
 
         with pytest.raises(FailedToConnectError):
             adapter.open()
+
+
+def test_connnections_credentials_passes_through_wif_params():
+    credentials = {
+        "account": "account_id_with_underscores",
+        "database": "database",
+        "warehouse": "warehouse",
+        "schema": "schema",
+        "authenticator": "workload_identity",
+        "workload_identity_provider": "azure",
+        "workload_identity_entra_resource": "app://123",
+        "token": "test_token",
+    }
+    auth_args = connections.SnowflakeCredentials(**credentials).auth_args()
+    assert auth_args["authenticator"] == "WORKLOAD_IDENTITY"
+    assert auth_args["workload_identity_provider"] == "azure"
+    assert auth_args["workload_identity_entra_resource"] == "app://123"
+    assert auth_args["token"] == "test_token"
+
+
+def test_connnections_credentials_wif_authenticator_fails_without_provider():
+    credentials = {
+        "account": "account_id_with_underscores",
+        "database": "database",
+        "warehouse": "warehouse",
+        "schema": "schema",
+        "authenticator": "workload_identity",
+        # Missing workload_identity_provider
+    }
+    with pytest.raises(DbtConfigError) as excinfo:
+        connections.SnowflakeCredentials(**credentials).auth_args()
+    assert "workload_identity_provider must be set if authenticator='workload_identity'" in str(
+        excinfo
+    )