Update WIF integration tests to verify authenticated username + prepare for impersonation (#2510)

Author: sfc-gh-pmansour

ci/test_wif.sh

@@ -12,8 +12,11 @@ run_tests_and_set_result() {
   local host="$2"
   local snowflake_host="$3"
   local rsa_key_path="$4"
+  local snowflake_user="$5"
+  local impersonation_path="$6"
+  local snowflake_user_for_impersonation="$7"
 
-  ssh -i "$rsa_key_path" -o IdentitiesOnly=yes -p 443 "$host" env BRANCH="$BRANCH" SNOWFLAKE_TEST_WIF_HOST="$snowflake_host" SNOWFLAKE_TEST_WIF_PROVIDER="$provider" SNOWFLAKE_TEST_WIF_ACCOUNT="$SNOWFLAKE_TEST_WIF_ACCOUNT" bash << EOF
+  ssh -i "$rsa_key_path" -o IdentitiesOnly=yes -p 443 "$host" env BRANCH="$BRANCH" SNOWFLAKE_TEST_WIF_HOST="$snowflake_host" SNOWFLAKE_TEST_WIF_PROVIDER="$provider" SNOWFLAKE_TEST_WIF_ACCOUNT="$SNOWFLAKE_TEST_WIF_ACCOUNT" SNOWFLAKE_TEST_WIF_USERNAME="$snowflake_user" SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH="$impersonation_path" SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION="$snowflake_user_for_impersonation" bash << EOF
       set -e
       set -o pipefail
       docker run \
@@ -22,6 +25,9 @@ run_tests_and_set_result() {
         -e SNOWFLAKE_TEST_WIF_PROVIDER \
         -e SNOWFLAKE_TEST_WIF_HOST \
         -e SNOWFLAKE_TEST_WIF_ACCOUNT \
+        -e SNOWFLAKE_TEST_WIF_USERNAME \
+        -e SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH \
+        -e SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION \
         snowflakedb/client-python-test:1 \
           bash -c "
             echo 'Running tests on branch: \$BRANCH'
@@ -71,9 +77,9 @@ setup_parameters
 # Run tests for all cloud providers
 EXIT_STATUS=0
 set +e  # Don't exit on first failure
-run_tests_and_set_result "AZURE" "$HOST_AZURE" "$SNOWFLAKE_TEST_WIF_HOST_AZURE" "$RSA_KEY_PATH_AWS_AZURE"
-run_tests_and_set_result "AWS" "$HOST_AWS" "$SNOWFLAKE_TEST_WIF_HOST_AWS" "$RSA_KEY_PATH_AWS_AZURE"
-run_tests_and_set_result "GCP" "$HOST_GCP" "$SNOWFLAKE_TEST_WIF_HOST_GCP" "$RSA_KEY_PATH_GCP"
+run_tests_and_set_result "AZURE" "$HOST_AZURE" "$SNOWFLAKE_TEST_WIF_HOST_AZURE" "$RSA_KEY_PATH_AWS_AZURE" "$SNOWFLAKE_TEST_WIF_USERNAME_AZURE" "$SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH_AZURE" "$SNOWFLAKE_TEST_WIF_USERNAME_AZURE_IMPERSONATION"
+run_tests_and_set_result "AWS" "$HOST_AWS" "$SNOWFLAKE_TEST_WIF_HOST_AWS" "$RSA_KEY_PATH_AWS_AZURE" "$SNOWFLAKE_TEST_WIF_USERNAME_AWS" "$SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH_AWS" "$SNOWFLAKE_TEST_WIF_USERNAME_AWS_IMPERSONATION"
+run_tests_and_set_result "GCP" "$HOST_GCP" "$SNOWFLAKE_TEST_WIF_HOST_GCP" "$RSA_KEY_PATH_GCP" "$SNOWFLAKE_TEST_WIF_USERNAME_GCP" "$SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH_GCP" "$SNOWFLAKE_TEST_WIF_USERNAME_GCP_IMPERSONATION"
 set -e  # Re-enable exit on error
 echo "Exit status: $EXIT_STATUS"
 exit $EXIT_STATUS

ci/wif/parameters/parameters_wif.json.gpg

@@ -22,6 +22,9 @@
 ACCOUNT = os.getenv("SNOWFLAKE_TEST_WIF_ACCOUNT")
 HOST = os.getenv("SNOWFLAKE_TEST_WIF_HOST")
 PROVIDER = os.getenv("SNOWFLAKE_TEST_WIF_PROVIDER")
+EXPECTED_USERNAME = os.getenv("SNOWFLAKE_TEST_WIF_USERNAME")
+IMPERSONATION_PATH = os.getenv("SNOWFLAKE_TEST_WIF_IMPERSONATION_PATH")
+EXPECTED_USERNAME_IMPERSONATION = os.getenv("SNOWFLAKE_TEST_WIF_USERNAME_IMPERSONATION")
 
 
 @pytest.mark.wif
@@ -33,8 +36,8 @@ def test_wif_defined_provider():
         "workload_identity_provider": PROVIDER,
     }
     assert connect_and_execute_simple_query(
-        connection_params
-    ), "Failed to connect with using WIF - automatic provider detection"
+        connection_params, EXPECTED_USERNAME
+    ), f"Failed to connect with using WIF using provider {PROVIDER}"
 
 
 @pytest.mark.wif
@@ -51,21 +54,48 @@ def test_should_authenticate_using_oidc():
     }
 
     assert connect_and_execute_simple_query(
-        connection_params
+        connection_params, expected_user=None
     ), "Failed to connect using WIF with OIDC provider"
 
 
+@pytest.mark.wif
+@pytest.mark.skip("Impersonation is still being developed")
+def test_should_authenticate_with_impersonation():
+    if not isinstance(IMPERSONATION_PATH, str) or not IMPERSONATION_PATH:
+        pytest.skip("Skipping test - IMPERSONATION_PATH is not set")
+
+    logger.debug(f"Using impersonation path: {IMPERSONATION_PATH}")
+    impersonation_path_list = IMPERSONATION_PATH.split(",")
+
+    connection_params = {
+        "host": HOST,
+        "account": ACCOUNT,
+        "authenticator": "WORKLOAD_IDENTITY",
+        "workload_identity_provider": PROVIDER,
+        "workload_identity_impersonation_path": impersonation_path_list,
+    }
+
+    assert connect_and_execute_simple_query(
+        connection_params, EXPECTED_USERNAME_IMPERSONATION
+    ), f"Failed to connect using WIF with provider {PROVIDER}"
+
+
 def is_provider_gcp() -> bool:
     return PROVIDER == "GCP"
 
 
-def connect_and_execute_simple_query(connection_params) -> bool:
+def connect_and_execute_simple_query(connection_params, expected_user=None) -> bool:
     try:
         logger.info("Trying to connect to Snowflake")
         with snowflake.connector.connect(**connection_params) as con:
-            result = con.cursor().execute("select 1;")
-            logger.debug(result.fetchall())
-            logger.info("Successfully connected to Snowflake")
+            result = con.cursor().execute("select current_user();")
+            (user,) = result.fetchone()
+            logger.debug(user)
+            if expected_user:
+                assert (
+                    expected_user == user
+                ), f"Expected user '{expected_user}', got user '{user}'"
+            logger.info(f"Successfully connected to Snowflake as {user}")
             return True
     except Exception as e:
         logger.error(e)