initial changes

Author: sfc-gh-eqin

src/snowflake/connector/auth/workload_identity.py

@@ -55,12 +55,14 @@ def __init__(
         provider: AttestationProvider | None = None,
         token: str | None = None,
         entra_resource: str | None = None,
+        impersonations: list[str] | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
         self.provider = provider
         self.token = token
         self.entra_resource = entra_resource
+        self.impersonations = impersonations
 
         self.attestation: WorkloadIdentityAttestation | None = None
 
@@ -85,6 +87,7 @@ def prepare(
             self.provider,
             self.entra_resource,
             self.token,
+            self.impersonations,
             session_manager=(
                 conn._session_manager.clone(max_retries=0) if conn else None
             ),

src/snowflake/connector/connection.py

@@ -214,6 +214,7 @@ def _get_private_bytes_from_file(
     "authenticator": (DEFAULT_AUTHENTICATOR, (type(None), str)),
     "workload_identity_provider": (None, (type(None), AttestationProvider)),
     "workload_identity_entra_resource": (None, (type(None), str)),
+    "workload_identity_impersonations": (None, (type(None), list)),
     "mfa_callback": (None, (type(None), Callable)),
     "password_callback": (None, (type(None), Callable)),
     "auth_class": (None, (type(None), AuthByPlugin)),
@@ -1359,6 +1360,7 @@ def __open_connection(self):
                     provider=self._workload_identity_provider,
                     token=self._token,
                     entra_resource=self._workload_identity_entra_resource,
+                    impersonations=self._workload_identity_impersonations,
                 )
             else:
                 # okta URL, e.g., https://<account>.okta.com/
@@ -1531,6 +1533,7 @@ def __config(self, **kwargs):
             workload_identity_dependent_options = [
                 "workload_identity_provider",
                 "workload_identity_entra_resource",
+                "workload_identity_impersonations",
             ]
             for dependent_option in workload_identity_dependent_options:
                 if (

src/snowflake/connector/wif_util.py

@@ -12,6 +12,8 @@
 from botocore.auth import SigV4Auth
 from botocore.awsrequest import AWSRequest
 from botocore.utils import InstanceMetadataRegionFetcher
+from google.auth import impersonated_credentials
+from google.auth.transport.requests import Request
 
 from .errorcode import ER_INVALID_WIF_SETTINGS, ER_WIF_CREDENTIALS_NOT_FOUND
 from .errors import ProgrammingError
@@ -185,6 +187,7 @@ def create_aws_attestation(
 
 
 def create_gcp_attestation(
+    impersonations: list[str] | None = None,
     session_manager: SessionManager | None = None,
 ) -> WorkloadIdentityAttestation:
     """Tries to create a workload identity attestation for GCP.
@@ -207,6 +210,24 @@ def create_gcp_attestation(
         )
 
     jwt_str = res.content.decode("utf-8")
+    if impersonations:
+        try:
+            for impersonation in impersonations:
+                jwt_str = impersonated_credentials.Credentials(
+                    source_credentials=jwt_str,
+                    target_principal=impersonation,
+                    target_audience=SNOWFLAKE_AUDIENCE,
+                )
+
+            # Refresh the last impersonated credential to get the final token
+            jwt_str.refresh(Request())
+            jwt_str = jwt_str.token
+        except Exception as e:
+            raise ProgrammingError(
+                msg=f"Error impersonating GCP service account: {e}. Ensure the service account has the 'Service Account Token Creator' role.",
+                errno=ER_WIF_CREDENTIALS_NOT_FOUND,
+            )
+
     _, subject = extract_iss_and_sub_without_signature_verification(jwt_str)
     return WorkloadIdentityAttestation(
         AttestationProvider.GCP, jwt_str, {"sub": subject}
@@ -295,6 +316,7 @@ def create_attestation(
     provider: AttestationProvider,
     entra_resource: str | None = None,
     token: str | None = None,
+    impersonations: list[str] | None = None,
     session_manager: SessionManager | None = None,
 ) -> WorkloadIdentityAttestation:
     """Entry point to create an attestation using the given provider.
@@ -313,7 +335,7 @@ def create_attestation(
     elif provider == AttestationProvider.AZURE:
         return create_azure_attestation(entra_resource, session_manager)
     elif provider == AttestationProvider.GCP:
-        return create_gcp_attestation(session_manager)
+        return create_gcp_attestation(impersonations, session_manager)
     elif provider == AttestationProvider.OIDC:
         return create_oidc_attestation(token)
     else: