[async] WIF impersonation for GCP #2496

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_connection.py

@@ -414,10 +414,24 @@ async def __open_connection(self):
                             "errno": ER_INVALID_WIF_SETTINGS,
                         },
                     )
+                if (
+                    self._workload_identity_impersonation_path
+                    and self._workload_identity_provider != AttestationProvider.GCP
+                ):
+                    Error.errorhandler_wrapper(
+                        self,
+                        None,
+                        ProgrammingError,
+                        {
+                            "msg": "workload_identity_impersonation_path is currently only supported for GCP.",
+                            "errno": ER_INVALID_WIF_SETTINGS,
+                        },
+                    )
                 self.auth_class = AuthByWorkloadIdentity(
                     provider=self._workload_identity_provider,
                     token=self._token,
                     entra_resource=self._workload_identity_entra_resource,
+                    impersonation_path=self._workload_identity_impersonation_path,
                 )
             else:
                 # okta URL, e.g., https://<account>.okta.com/

src/snowflake/connector/aio/_wif_util.py

@@ -25,6 +25,10 @@
 
 logger = logging.getLogger(__name__)
 
+GCP_METADATA_SERVICE_ACCOUNT_BASE_URL = (
+    "http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default"
+)
+
 
 async def get_aws_region() -> str:
     """Get the current AWS workload's region."""
@@ -81,30 +85,108 @@ async def create_aws_attestation() -> WorkloadIdentityAttestation:
     )
 
 
-async def create_gcp_attestation(
-    session_manager: SessionManager | None = None,
-) -> WorkloadIdentityAttestation:
-    """Tries to create a workload identity attestation for GCP.
+async def get_gcp_access_token(session_manager: SessionManager) -> str:
+    """Gets a GCP access token from the metadata server.
+
+    If the application isn't running on GCP or no credentials were found, raises an error.
+    """
+    try:
+        res = await session_manager.request(
+            method="GET",
+            url=f"{GCP_METADATA_SERVICE_ACCOUNT_BASE_URL}/token",
+            headers={
+                "Metadata-Flavor": "Google",
+            },
+        )
+
+        content = await res.content.read()
+        response_text = content.decode("utf-8")
+        return json.loads(response_text)["access_token"]
+    except Exception as e:
+        raise ProgrammingError(
+            msg=f"Error fetching GCP access token: {e}. Ensure the application is running on GCP.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
+
+
+async def get_gcp_identity_token_via_impersonation(
+    impersonation_path: list[str], session_manager: SessionManager
+) -> str:
+    """Gets a GCP identity token from the metadata server.
+
+    If the application isn't running on GCP or no credentials were found, raises an error.
+    """
+    if not impersonation_path:
+        raise ProgrammingError(
+            msg="Error: impersonation_path cannot be empty.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
+
+    current_sa_token = await get_gcp_access_token(session_manager)
+    impersonation_path = [
+        f"projects/-/serviceAccounts/{client_id}" for client_id in impersonation_path
+    ]
+    try:
+        res = await session_manager.post(
+            url=f"https://iamcredentials.googleapis.com/v1/{impersonation_path[-1]}:generateIdToken",
+            headers={
+                "Authorization": f"Bearer {current_sa_token}",
+                "Content-Type": "application/json",
+            },
+            json={
+                "delegates": impersonation_path[:-1],
+                "audience": SNOWFLAKE_AUDIENCE,
+            },
+        )
+
+        content = await res.content.read()
+        response_text = content.decode("utf-8")
+        return json.loads(response_text)["token"]
+    except Exception as e:
+        raise ProgrammingError(
+            msg=f"Error fetching GCP identity token for impersonated GCP service account '{impersonation_path[-1]}': {e}. Ensure the application is running on GCP.",
+            errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+        )
+
+
+async def get_gcp_identity_token(session_manager: SessionManager) -> str:
+    """Gets a GCP identity token from the metadata server.
 
     If the application isn't running on GCP or no credentials were found, raises an error.
     """
     try:
         res = await session_manager.request(
             method="GET",
-            url=f"http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/identity?audience={SNOWFLAKE_AUDIENCE}",
+            url=f"{GCP_METADATA_SERVICE_ACCOUNT_BASE_URL}/identity?audience={SNOWFLAKE_AUDIENCE}",
             headers={
                 "Metadata-Flavor": "Google",
             },
         )
 
         content = await res.content.read()
-        jwt_str = content.decode("utf-8")
+        return content.decode("utf-8")
     except Exception as e:
         raise ProgrammingError(
-            msg=f"Error fetching GCP metadata: {e}. Ensure the application is running on GCP.",
+            msg=f"Error fetching GCP identity token: {e}. Ensure the application is running on GCP.",
             errno=ER_WIF_CREDENTIALS_NOT_FOUND,
         )
 
+
+async def create_gcp_attestation(
+    session_manager: SessionManager,
+    impersonation_path: list[str] | None = None,
+) -> WorkloadIdentityAttestation:
+    """Tries to create a workload identity attestation for GCP.
+
+    If the application isn't running on GCP or no credentials were found, raises an error.
+    """
+    if impersonation_path:
+        jwt_str = await get_gcp_identity_token_via_impersonation(
+            impersonation_path, session_manager
+        )
+    else:
+        jwt_str = await get_gcp_identity_token(session_manager)
+
     _, subject = extract_iss_and_sub_without_signature_verification(jwt_str)
     return WorkloadIdentityAttestation(
         AttestationProvider.GCP, jwt_str, {"sub": subject}
@@ -179,6 +261,7 @@ async def create_attestation(
     provider: AttestationProvider | None,
     entra_resource: str | None = None,
     token: str | None = None,
+    impersonation_path: list[str] | None = None,
     session_manager: SessionManager | None = None,
 ) -> WorkloadIdentityAttestation:
     """Entry point to create an attestation using the given provider.
@@ -197,7 +280,7 @@ async def create_attestation(
     elif provider == AttestationProvider.AZURE:
         return await create_azure_attestation(entra_resource, session_manager)
     elif provider == AttestationProvider.GCP:
-        return await create_gcp_attestation(session_manager)
+        return await create_gcp_attestation(session_manager, impersonation_path)
     elif provider == AttestationProvider.OIDC:
         return create_oidc_attestation(token)
     else:

src/snowflake/connector/aio/auth/_workload_identity.py

@@ -22,6 +22,7 @@ def __init__(
         provider: AttestationProvider,
         token: str | None = None,
         entra_resource: str | None = None,
+        impersonation_path: list[str] | None = None,
         **kwargs,
     ) -> None:
         """Initializes an instance with workload identity authentication."""
@@ -30,6 +31,7 @@ def __init__(
             provider=provider,
             token=token,
             entra_resource=entra_resource,
+            impersonation_path=impersonation_path,
             **kwargs,
         )
 
@@ -44,6 +46,7 @@ async def prepare(
             self.provider,
             self.entra_resource,
             self.token,
+            self.impersonation_path,
             session_manager=(
                 conn._session_manager.clone(max_retries=0) if conn else None
             ),

test/unit/aio/test_auth_workload_identity_async.py

@@ -8,6 +8,7 @@
 import os
 from base64 import b64decode
 from unittest import mock
+from unittest.mock import AsyncMock
 from urllib.parse import parse_qs, urlparse
 
 import aiohttp
@@ -18,7 +19,7 @@
 from snowflake.connector.aio.auth import AuthByWorkloadIdentity
 from snowflake.connector.errors import ProgrammingError
 
-from ...csp_helpers import gen_dummy_id_token
+from ...csp_helpers import gen_dummy_access_token, gen_dummy_id_token
 from .csp_helpers_async import FakeAwsEnvironmentAsync, FakeGceMetadataServiceAsync
 
 logger = logging.getLogger(__name__)
@@ -279,7 +280,7 @@ async def test_explicit_gcp_metadata_server_error_bubbles_up(exception):
         with pytest.raises(ProgrammingError) as excinfo:
             await auth_class.prepare(conn=None)
 
-    assert "Error fetching GCP metadata:" in str(excinfo.value)
+    assert "Error fetching GCP identity token:" in str(excinfo.value)
     assert "Ensure the application is running on GCP." in str(excinfo.value)
 
 
@@ -307,6 +308,51 @@ async def test_explicit_gcp_generates_unique_assertion_content(
     assert auth_class.assertion_content == '{"_provider":"GCP","sub":"123456"}'
 
 
+@mock.patch("snowflake.connector.aio._session_manager.SessionManager.post")
+async def test_gcp_calls_correct_apis_and_populates_auth_data_for_final_sa(
+    mock_post_request, fake_gce_metadata_service: FakeGceMetadataServiceAsync
+):
+    fake_gce_metadata_service.sub = "sa1"
+    impersonation_path = ["sa2", "sa3"]
+    sa1_access_token = gen_dummy_access_token("sa1")
+    sa3_id_token = gen_dummy_id_token("sa3")
+
+    # Mock the POST request response
+    class AsyncResponse:
+        def __init__(self, content):
+            self._content = content
+            self.content = mock.Mock()
+            self.content.read = AsyncMock(return_value=content)
+
+    mock_post_request.return_value = AsyncResponse(
+        json.dumps({"token": sa3_id_token}).encode("utf-8")
+    )
+
+    auth_class = AuthByWorkloadIdentity(
+        provider=AttestationProvider.GCP, impersonation_path=impersonation_path
+    )
+    await auth_class.prepare(conn=None)
+
+    mock_post_request.assert_called_once_with(
+        url="https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/sa3:generateIdToken",
+        headers={
+            "Authorization": f"Bearer {sa1_access_token}",
+            "Content-Type": "application/json",
+        },
+        json={
+            "delegates": ["projects/-/serviceAccounts/sa2"],
+            "audience": "snowflakecomputing.com",
+        },
+    )
+
+    assert auth_class.assertion_content == '{"_provider":"GCP","sub":"sa3"}'
+    assert await extract_api_data(auth_class) == {
+        "AUTHENTICATOR": "WORKLOAD_IDENTITY",
+        "PROVIDER": "GCP",
+        "TOKEN": sa3_id_token,
+    }
+
+
 # -- Azure Tests --
 
 

test/unit/aio/test_connection_async_unit.py

@@ -606,6 +606,7 @@ async def test_otel_error_message_async(caplog, mock_post_requests):
             "workload_identity_entra_resource",
             "api://0b2f151f-09a2-46eb-ad5a-39d5ebef917b",
         ),
+        ("workload_identity_impersonation_path", ["subject-b", "subject-c"]),
     ],
 )
 async def test_cannot_set_dependent_params_without_wlid_authenticator(
@@ -655,6 +656,79 @@ async def test_workload_identity_provider_is_required_for_wif_authenticator(
         assert expected_error_msg in str(excinfo.value)
 
 
+@pytest.mark.parametrize(
+    "provider_param",
+    [
+        # Strongly-typed values.
+        AttestationProvider.AWS,
+        AttestationProvider.AZURE,
+        AttestationProvider.OIDC,
+        # String values.
+        "AWS",
+        "AZURE",
+        "OIDC",
+    ],
+)
+async def test_workload_identity_impersonation_path_unsupported_for_non_gcp_providers(
+    monkeypatch, provider_param
+):
+    async def mock_authenticate(*_):
+        pass
+
+    with monkeypatch.context() as m:
+        m.setattr(
+            "snowflake.connector.aio._connection.SnowflakeConnection._authenticate",
+            mock_authenticate,
+        )
+
+        with pytest.raises(ProgrammingError) as excinfo:
+            await snowflake.connector.aio.connect(
+                account="account",
+                authenticator="WORKLOAD_IDENTITY",
+                workload_identity_provider=provider_param,
+                workload_identity_impersonation_path=[
+                    "[email protected]"
+                ],
+            )
+        assert (
+            "workload_identity_impersonation_path is currently only supported for GCP."
+            in str(excinfo.value)
+        )
+
+
+@pytest.mark.parametrize(
+    "provider_param",
+    [
+        AttestationProvider.GCP,
+        "GCP",
+    ],
+)
+async def test_workload_identity_impersonation_path_supported_for_gcp_provider(
+    monkeypatch, provider_param
+):
+    async def mock_authenticate(*_):
+        pass
+
+    with monkeypatch.context() as m:
+        m.setattr(
+            "snowflake.connector.aio._connection.SnowflakeConnection._authenticate",
+            mock_authenticate,
+        )
+
+        conn = await snowflake.connector.aio.connect(
+            account="account",
+            authenticator="WORKLOAD_IDENTITY",
+            workload_identity_provider=provider_param,
+            workload_identity_impersonation_path=[
+                "[email protected]"
+            ],
+        )
+        assert conn.auth_class.provider == AttestationProvider.GCP
+        assert conn.auth_class.impersonation_path == [
+            "[email protected]"
+        ]
+
+
 @pytest.mark.parametrize(
     "provider_param, parsed_provider",
     [

test/wif/test_wif_async.py

@@ -62,7 +62,6 @@ async def test_should_authenticate_using_oidc_async():
 
 @pytest.mark.wif
 @pytest.mark.aio
-@pytest.mark.skip("Impersonation is still being developed")
 async def test_should_authenticate_with_impersonation_async():
     if not isinstance(IMPERSONATION_PATH, str) or not IMPERSONATION_PATH:
         pytest.skip("Skipping test - IMPERSONATION_PATH is not set")