Merge branch 'main' into SNOW-2176203-partial-chains

Author: sfc-gh-snoonan

DESCRIPTION.md

@@ -9,6 +9,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 # Release Notes
 - v3.18.0(TBD)
   - Added the `workload_identity_impersonation_path` parameter to support service account impersonation for Workload Identity Federation on GCP workloads only
+  - Fixed `get_results_from_sfqid` when using `DictCursor` and executing multiple statements at once
   - Added support for intermediate cetificates as roots when they are stored in the trust store
 
 - v3.17.3(September 02,2025)

ci/wif/parameters/parameters_wif.json.gpg

@@ -1358,14 +1358,18 @@ def __open_connection(self):
                     )
                 if (
                     self._workload_identity_impersonation_path
-                    and self._workload_identity_provider != AttestationProvider.GCP
+                    and self._workload_identity_provider
+                    not in (
+                        AttestationProvider.GCP,
+                        AttestationProvider.AWS,
+                    )
                 ):
                     Error.errorhandler_wrapper(
                         self,
                         None,
                         ProgrammingError,
                         {
-                            "msg": "workload_identity_impersonation_path is currently only supported for GCP.",
+                            "msg": "workload_identity_impersonation_path is currently only supported for GCP and AWS.",
                             "errno": ER_INVALID_WIF_SETTINGS,
                         },
                     )

src/snowflake/connector/connection.py

@@ -1741,8 +1741,7 @@ def wait_until_ready() -> None:
         self.connection.get_query_status_throw_if_error(
             sfqid
         )  # Trigger an exception if query failed
-        klass = self.__class__
-        self._inner_cursor = klass(self.connection)
+        self._inner_cursor = SnowflakeCursor(self.connection)
         self._sfqid = sfqid
         self._prefetch_hook = wait_until_ready
 

src/snowflake/connector/cursor.py

@@ -145,15 +145,37 @@ def get_aws_sts_hostname(region: str, partition: str) -> str:
         )
 
 
+def get_aws_session(impersonation_path: list[str] | None = None):
+    """Creates a boto3 session with the appropriate credentials.
+
+    If impersonation_path is provided, this uses the role at the end of the path. Otherwise, this uses the role attached to the current workload.
+    """
+    session = boto3.session.Session()
+
+    impersonation_path = impersonation_path or []
+    for arn in impersonation_path:
+        response = session.client("sts").assume_role(
+            RoleArn=arn, RoleSessionName="identity-federation-session"
+        )
+        creds = response["Credentials"]
+        session = boto3.session.Session(
+            aws_access_key_id=creds["AccessKeyId"],
+            aws_secret_access_key=creds["SecretAccessKey"],
+            aws_session_token=creds["SessionToken"],
+        )
+    return session
+
+
 def create_aws_attestation(
-    session_manager: SessionManager | None = None,
+    impersonation_path: list[str] | None = None,
 ) -> WorkloadIdentityAttestation:
     """Tries to create a workload identity attestation for AWS.
 
     If the application isn't running on AWS or no credentials were found, raises an error.
     """
     # TODO: SNOW-2223669 Investigate if our adapters - containing settings of http traffic - should be passed here as boto urllib3session. Those requests go to local servers, so they do not need Proxy setup or Headers customization in theory. But we may want to have all the traffic going through one class (e.g. Adapter or mixin).
-    session = boto3.session.Session()
+    session = get_aws_session(impersonation_path)
+
     aws_creds = session.get_credentials()
     if not aws_creds:
         raise ProgrammingError(
@@ -387,7 +409,7 @@ def create_attestation(
     )
 
     if provider == AttestationProvider.AWS:
-        return create_aws_attestation(session_manager)
+        return create_aws_attestation(impersonation_path)
     elif provider == AttestationProvider.AZURE:
         return create_azure_attestation(entra_resource, session_manager)
     elif provider == AttestationProvider.GCP:

src/snowflake/connector/wif_util.py

@@ -40,9 +40,8 @@ def gen_dummy_id_token(
     )
 
 
-def gen_dummy_access_token(sub="test-subject") -> str:
+def gen_dummy_access_token(sub="test-subject", key="secret") -> str:
     """Generates a dummy access token using the given subject."""
-    key = "secret"
     logger.debug(f"Generating dummy access token for subject {sub}")
     return (sub + key).encode("utf-8").hex()
 
@@ -368,6 +367,11 @@ class FakeAwsEnvironment:
     def __init__(self):
         # Defaults used for generating a token. Can be overriden in individual tests.
         self.arn = "arn:aws:sts::123456789:assumed-role/My-Role/i-34afe100cad287fab"
+        # Path of roles that can be assumed. Empty if no impersonation is allowed.
+        # Can be overriden in individual tests.
+        self.assumption_path = []
+        self.assume_role_call_count = 0
+
         self.caller_identity = {"Arn": self.arn}
         self.region = "us-east-1"
         self.credentials = Credentials(access_key="ak", secret_key="sk")
@@ -376,6 +380,25 @@ def __init__(self):
         )
         self.metadata_token = "test-token"
 
+    def assume_role(self, **kwargs):
+        if (
+            self.assumption_path
+            and kwargs["RoleArn"] == self.assumption_path[self.assume_role_call_count]
+        ):
+            arn = self.assumption_path[self.assume_role_call_count]
+            self.assume_role_call_count += 1
+            return {
+                "Credentials": {
+                    "AccessKeyId": "access_key",
+                    "SecretAccessKey": "secret_key",
+                    "SessionToken": "session_token",
+                    "Expiration": int(time()) + 60 * 60,
+                },
+                "AssumedRoleUser": {"AssumedRoleId": hash(arn), "Arn": arn},
+                "ResponseMetadata": {},
+            }
+        return {}
+
     def get_region(self):
         return self.region
 
@@ -399,6 +422,7 @@ def fetcher_fetch_metadata_token(self):
     def boto3_client(self, *args, **kwargs):
         mock_client = mock.Mock()
         mock_client.get_caller_identity.return_value = self.caller_identity
+        mock_client.assume_role = self.assume_role
         return mock_client
 
     def __enter__(self):
@@ -439,6 +463,9 @@ def __enter__(self):
                 side_effect=self.boto3_client,
             )
         )
+        self.patchers.append(
+            mock.patch("boto3.session.Session.client", side_effect=self.boto3_client)
+        )
         for patcher in self.patchers:
             patcher.__enter__()
         return self

test/csp_helpers.py

@@ -7,6 +7,7 @@
 import pytest
 
 from snowflake.connector import DatabaseError, ProgrammingError
+from snowflake.connector.cursor import DictCursor, SnowflakeCursor
 
 # Mark all tests in this file to time out after 2 minutes to prevent hanging forever
 pytestmark = [pytest.mark.timeout(120), pytest.mark.skipolddriver]
@@ -17,14 +18,15 @@
     QueryStatus = None
 
 
-def test_simple_async(conn_cnx):
+@pytest.mark.parametrize("cursor_class", [SnowflakeCursor, DictCursor])
+def test_simple_async(conn_cnx, cursor_class):
     """Simple test to that shows the most simple usage of fire and forget.
 
     This test also makes sure that wait_until_ready function's sleeping is tested and
     that some fields are copied over correctly from the original query.
     """
     with conn_cnx() as con:
-        with con.cursor() as cur:
+        with con.cursor(cursor_class) as cur:
             cur.execute_async("select count(*) from table(generator(timeLimit => 5))")
             cur.get_results_from_sfqid(cur.sfqid)
             assert len(cur.fetchall()) == 1

test/integ/test_async.py

@@ -14,6 +14,7 @@
 
 import snowflake.connector.cursor
 from snowflake.connector import ProgrammingError, errors
+from snowflake.connector.cursor import DictCursor, SnowflakeCursor
 
 try:  # pragma: no cover
     from snowflake.connector.constants import (
@@ -153,10 +154,11 @@ def test_binding_multi(conn_cnx, style: str, skip_to_last_set: bool):
             )
 
 
-def test_async_exec_multi(conn_cnx, skip_to_last_set: bool):
+@pytest.mark.parametrize("cursor_class", [SnowflakeCursor, DictCursor])
+def test_async_exec_multi(conn_cnx, cursor_class, skip_to_last_set: bool):
     """Tests whether async execution query works within a multi-statement"""
     with conn_cnx() as con:
-        with con.cursor() as cur:
+        with con.cursor(cursor_class) as cur:
             cur.execute_async(
                 "select 1; select 2; select count(*) from table(generator(timeLimit => 1)); select 'b';",
                 num_statements=4,
@@ -165,14 +167,29 @@ def test_async_exec_multi(conn_cnx, skip_to_last_set: bool):
             assert con.is_still_running(con.get_query_status(q_id))
         _wait_while_query_running(con, q_id, sleep_time=1)
     with conn_cnx() as con:
-        with con.cursor() as cur:
+        with con.cursor(cursor_class) as cur:
             _wait_until_query_success(con, q_id, num_checks=3, sleep_per_check=1)
             assert con.get_query_status_throw_if_error(q_id) == QueryStatus.SUCCESS
 
+            if cursor_class == SnowflakeCursor:
+                expected = [
+                    [(1,)],
+                    [(2,)],
+                    lambda x: len(x) == 1 and len(x[0]) == 1 and x[0][0] > 0,
+                    [("b",)],
+                ]
+            elif cursor_class == DictCursor:
+                expected = [
+                    [{"1": 1}],
+                    [{"2": 2}],
+                    lambda x: len(x) == 1 and len(x[0]) == 1 and x[0]["COUNT(*)"] > 0,
+                    [{"'B'": "b"}],
+                ]
+
             cur.get_results_from_sfqid(q_id)
             _check_multi_statement_results(
                 cur,
-                checks=[[(1,)], [(2,)], lambda x: x > [(0,)], [("b",)]],
+                checks=expected,
                 skip_to_last_set=skip_to_last_set,
             )
 

test/integ/test_multi_statement.py

@@ -274,6 +274,22 @@ def test_get_aws_sts_hostname_invalid_inputs(region, partition):
     assert "Invalid AWS partition" in str(excinfo.value)
 
 
+def test_aws_impersonation_calls_correct_apis_for_each_role_in_impersonation_path(
+    fake_aws_environment: FakeAwsEnvironment,
+):
+    impersonation_path = [
+        "arn:aws:iam::123456789:role/role2",
+        "arn:aws:iam::123456789:role/role3",
+    ]
+    fake_aws_environment.assumption_path = impersonation_path
+    auth_class = AuthByWorkloadIdentity(
+        provider=AttestationProvider.AWS, impersonation_path=impersonation_path
+    )
+    auth_class.prepare(conn=None)
+
+    assert fake_aws_environment.assume_role_call_count == 2
+
+
 # -- GCP Tests --
 
 

test/unit/test_auth_workload_identity.py

@@ -682,16 +682,14 @@ def test_workload_identity_provider_is_required_for_wif_authenticator(
     "provider_param",
     [
         # Strongly-typed values.
-        AttestationProvider.AWS,
         AttestationProvider.AZURE,
         AttestationProvider.OIDC,
         # String values.
-        "AWS",
         "AZURE",
         "OIDC",
     ],
 )
-def test_workload_identity_impersonation_path_unsupported_for_non_gcp_providers(
+def test_workload_identity_impersonation_path_errors_for_unsupported_providers(
     monkeypatch, provider_param
 ):
     with monkeypatch.context() as m:
@@ -709,20 +707,22 @@ def test_workload_identity_impersonation_path_unsupported_for_non_gcp_providers(
                 ],
             )
         assert (
-            "workload_identity_impersonation_path is currently only supported for GCP."
+            "workload_identity_impersonation_path is currently only supported for GCP and AWS."
             in str(excinfo.value)
         )
 
 
 @pytest.mark.parametrize(
-    "provider_param",
+    "provider_param,impersonation_path",
     [
-        AttestationProvider.GCP,
-        "GCP",
+        (AttestationProvider.GCP, ["[email protected]"]),
+        (AttestationProvider.AWS, ["arn:aws:iam::1234567890:role/role2"]),
+        ("GCP", ["[email protected]"]),
+        ("AWS", ["arn:aws:iam::1234567890:role/role2"]),
     ],
 )
-def test_workload_identity_impersonation_path_supported_for_gcp_provider(
-    monkeypatch, provider_param
+def test_workload_identity_impersonation_path_populates_auth_class_for_supported_provider(
+    monkeypatch, provider_param, impersonation_path
 ):
     with monkeypatch.context() as m:
         m.setattr(
@@ -733,14 +733,9 @@ def test_workload_identity_impersonation_path_supported_for_gcp_provider(
             account="account",
             authenticator="WORKLOAD_IDENTITY",
             workload_identity_provider=provider_param,
-            workload_identity_impersonation_path=[
-                "[email protected]"
-            ],
+            workload_identity_impersonation_path=impersonation_path,
         )
-        assert conn.auth_class.provider == AttestationProvider.GCP
-        assert conn.auth_class.impersonation_path == [
-            "[email protected]"
-        ]
+        assert conn.auth_class.impersonation_path == impersonation_path
 
 
 @pytest.mark.parametrize(