move change to update_body() and modify unit tests

Author: sfc-gh-eqin

src/snowflake/connector/auth/_auth.py

@@ -226,13 +226,6 @@ def authenticate(
         if session_parameters:
             body["data"]["SESSION_PARAMETERS"] = session_parameters
 
-        from . import AuthByWorkloadIdentity
-
-        if isinstance(auth_instance, AuthByWorkloadIdentity):
-            body["data"]["IMPERSONATION_PATH_LENGTH"] = len(
-                auth_instance.impersonation_path or []
-            )
-
         logger.debug(
             "body['data']: %s",
             {

src/snowflake/connector/auth/workload_identity.py

@@ -78,6 +78,9 @@ def update_body(self, body: dict[typing.Any, typing.Any]) -> None:
             self.attestation
         ).value
         body["data"]["TOKEN"] = self.attestation.credential
+        body["data"].setdefault("CLIENT_ENVIRONMENT", {})[
+            "IMPERSONATION_PATH_LENGTH"
+        ] = len(self.impersonation_path or [])
 
     def prepare(
         self, *, conn: SnowflakeConnection | None, **kwargs: typing.Any

test/unit/test_auth_workload_identity.py

@@ -136,6 +136,7 @@ def test_explicit_oidc_valid_inline_token_plumbed_to_api():
         "AUTHENTICATOR": "WORKLOAD_IDENTITY",
         "PROVIDER": "OIDC",
         "TOKEN": dummy_token,
+        "CLIENT_ENVIRONMENT": {"IMPERSONATION_PATH_LENGTH": 0},
     }
 
 
@@ -191,6 +192,7 @@ def test_explicit_aws_encodes_audience_host_signature_to_api(
     data = extract_api_data(auth_class)
     assert data["AUTHENTICATOR"] == "WORKLOAD_IDENTITY"
     assert data["PROVIDER"] == "AWS"
+    assert data["CLIENT_ENVIRONMENT"]["IMPERSONATION_PATH_LENGTH"] == 0
     verify_aws_token(data["TOKEN"], fake_aws_environment.region)
 
 
@@ -308,6 +310,7 @@ def test_explicit_gcp_plumbs_token_to_api(
         "AUTHENTICATOR": "WORKLOAD_IDENTITY",
         "PROVIDER": "GCP",
         "TOKEN": fake_gce_metadata_service.token,
+        "CLIENT_ENVIRONMENT": {"IMPERSONATION_PATH_LENGTH": 0},
     }
 
 
@@ -357,6 +360,7 @@ def test_gcp_calls_correct_apis_and_populates_auth_data_for_final_sa(
         "AUTHENTICATOR": "WORKLOAD_IDENTITY",
         "PROVIDER": "GCP",
         "TOKEN": sa3_id_token,
+        "CLIENT_ENVIRONMENT": {"IMPERSONATION_PATH_LENGTH": 2},
     }
 
 
@@ -409,6 +413,7 @@ def test_explicit_azure_plumbs_token_to_api(fake_azure_metadata_service):
         "AUTHENTICATOR": "WORKLOAD_IDENTITY",
         "PROVIDER": "AZURE",
         "TOKEN": fake_azure_metadata_service.token,
+        "CLIENT_ENVIRONMENT": {"IMPERSONATION_PATH_LENGTH": 0},
     }