Skip to content

Commit 22b34a2

Browse files
sfc-gh-akolodziejczyksfc-gh-akolodziejczyk
authored
SNOW-2160717 add WIF e2e tests (#2433)
1 parent d89ebee commit 22b34a2

File tree

11 files changed

+219
-1
lines changed

11 files changed

+219
-1
lines changed

‎.gitignore

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -129,3 +129,8 @@ src/snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow_arrow_iterator.cpp
129129
# Prober files
130130
prober/parameters.json
131131
prober/snowflake_prober.egg-info/
132+
133+
# SSH private key for WIF tests
134+
ci/wif/parameters/rsa_wif_aws_azure
135+
ci/wif/parameters/rsa_wif_gcp
136+
ci/wif/parameters/parameters_wif.json

‎Jenkinsfile

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,18 @@ timestamps {
7171
'''.stripMargin()
7272
}
7373
}
74+
},
75+
'Test WIF': {
76+
stage('Test WIF') {
77+
withCredentials([
78+
string(credentialsId: 'sfctest0-parameters-secret', variable: 'PARAMETERS_SECRET')
79+
]) {
80+
sh '''\
81+
|#!/bin/bash -e
82+
|$WORKSPACE/ci/test_wif.sh
83+
'''.stripMargin()
84+
}
85+
}
7486
}
7587
)
7688
}

‎ci/container/test_authentication.sh

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ set -o pipefail
66
export WORKSPACE=${WORKSPACE:-/mnt/workspace}
77
export SOURCE_ROOT=${SOURCE_ROOT:-/mnt/host}
88

9-
MVNW_EXE=$SOURCE_ROOT/mvnw
109
AUTH_PARAMETER_FILE=./.github/workflows/parameters/private/parameters_aws_auth_tests.json
1110
eval $(jq -r '.authtestparams | to_entries | map("export \(.key)=\(.value|tostring)")|.[]' $AUTH_PARAMETER_FILE)
1211

‎ci/test_wif.sh

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
#!/bin/bash -e
2+
3+
set -o pipefail
4+
5+
export THIS_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
6+
export RSA_KEY_PATH_AWS_AZURE="$THIS_DIR/wif/parameters/rsa_wif_aws_azure"
7+
export RSA_KEY_PATH_GCP="$THIS_DIR/wif/parameters/rsa_wif_gcp"
8+
export PARAMETERS_FILE_PATH="$THIS_DIR/wif/parameters/parameters_wif.json"
9+
10+
run_tests_and_set_result() {
11+
local provider="$1"
12+
local host="$2"
13+
local snowflake_host="$3"
14+
local rsa_key_path="$4"
15+
16+
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
17+
set -e
18+
set -o pipefail
19+
docker run \
20+
--rm \
21+
-e BRANCH \
22+
-e SNOWFLAKE_TEST_WIF_PROVIDER \
23+
-e SNOWFLAKE_TEST_WIF_HOST \
24+
-e SNOWFLAKE_TEST_WIF_ACCOUNT \
25+
snowflakedb/client-python-test:1 \
26+
bash -c "
27+
echo 'Running tests on branch: \$BRANCH'
28+
if [[ \"\$BRANCH\" =~ ^PR-[0-9]+\$ ]]; then
29+
curl -L https://github.com/snowflakedb/snowflake-connector-python/archive/refs/pull/\$(echo \$BRANCH | cut -d- -f2)/head.tar.gz | tar -xz
30+
mv snowflake-connector-python-* snowflake-connector-python
31+
else
32+
curl -L https://github.com/snowflakedb/snowflake-connector-python/archive/refs/heads/\$BRANCH.tar.gz | tar -xz
33+
mv snowflake-connector-python-\$BRANCH snowflake-connector-python
34+
fi
35+
cd snowflake-connector-python
36+
bash ci/wif/test_wif.sh
37+
"
38+
EOF
39+
local status=$?
40+
41+
if [[ $status -ne 0 ]]; then
42+
echo "$provider tests failed with exit status: $status"
43+
EXIT_STATUS=1
44+
else
45+
echo "$provider tests passed"
46+
fi
47+
}
48+
49+
get_branch() {
50+
local branch
51+
branch=$(git rev-parse --abbrev-ref HEAD)
52+
if [[ "$branch" == "HEAD" ]]; then
53+
branch=$(git name-rev --name-only HEAD | sed 's#^remotes/origin/##;s#^origin/##')
54+
fi
55+
echo "$branch"
56+
}
57+
58+
setup_parameters() {
59+
gpg --quiet --batch --yes --decrypt --passphrase="$PARAMETERS_SECRET" --output "$RSA_KEY_PATH_AWS_AZURE" "${RSA_KEY_PATH_AWS_AZURE}.gpg"
60+
gpg --quiet --batch --yes --decrypt --passphrase="$PARAMETERS_SECRET" --output "$RSA_KEY_PATH_GCP" "${RSA_KEY_PATH_GCP}.gpg"
61+
chmod 600 "$RSA_KEY_PATH_AWS_AZURE"
62+
chmod 600 "$RSA_KEY_PATH_GCP"
63+
gpg --quiet --batch --yes --decrypt --passphrase="$PARAMETERS_SECRET" --output "$PARAMETERS_FILE_PATH" "${PARAMETERS_FILE_PATH}.gpg"
64+
eval $(jq -r '.wif | to_entries | map("export \(.key)=\(.value|tostring)")|.[]' $PARAMETERS_FILE_PATH)
65+
}
66+
67+
BRANCH=$(get_branch)
68+
export BRANCH
69+
setup_parameters
70+
71+
# Run tests for all cloud providers
72+
EXIT_STATUS=0
73+
set +e # Don't exit on first failure
74+
run_tests_and_set_result "AZURE" "$HOST_AZURE" "$SNOWFLAKE_TEST_WIF_HOST_AZURE" "$RSA_KEY_PATH_AWS_AZURE"
75+
run_tests_and_set_result "AWS" "$HOST_AWS" "$SNOWFLAKE_TEST_WIF_HOST_AWS" "$RSA_KEY_PATH_AWS_AZURE"
76+
run_tests_and_set_result "GCP" "$HOST_GCP" "$SNOWFLAKE_TEST_WIF_HOST_GCP" "$RSA_KEY_PATH_GCP"
77+
set -e # Re-enable exit on error
78+
echo "Exit status: $EXIT_STATUS"
79+
exit $EXIT_STATUS

‎ci/wif/parameters/parameters_wif.json.gpg

294 Bytes
Binary file not shown.

‎ci/wif/parameters/rsa_wif_aws_azure.gpg

344 Bytes
Binary file not shown.

‎ci/wif/parameters/rsa_wif_gcp.gpg

356 Bytes
Binary file not shown.

‎ci/wif/test_wif.sh

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
#!/bin/bash -e
2+
3+
set -o pipefail
4+
5+
export SF_OCSP_TEST_MODE=true
6+
export SF_ENABLE_EXPERIMENTAL_AUTHENTICATION=true
7+
export RUN_WIF_TESTS=true
8+
9+
/opt/python/cp39-cp39/bin/python -m pip install --break-system-packages -e .
10+
/opt/python/cp39-cp39/bin/python -m pip install --break-system-packages pytest
11+
/opt/python/cp39-cp39/bin/python -m pytest test/wif/*

‎test/conftest.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,10 @@ def pytest_runtest_setup(item) -> None:
149149
if os.getenv("RUN_AUTH_TESTS") != "true":
150150
pytest.skip("Skipping auth test in current environment")
151151

152+
if "wif" in test_tags:
153+
if os.getenv("RUN_WIF_TESTS") != "true":
154+
pytest.skip("Skipping WIF test in current environment")
155+
152156

153157
def get_server_parameter_value(connection, parameter_name: str) -> str | None:
154158
"""Get server parameter value, returns None if parameter doesn't exist."""

‎test/wif/test_wif.py

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
import logging.config
2+
import os
3+
import subprocess
4+
5+
import pytest
6+
7+
import snowflake.connector
8+
9+
logger = logging.getLogger(__name__)
10+
logger.setLevel(logging.INFO)
11+
12+
13+
"""
14+
Running tests locally:
15+
16+
1. Push branch to repository
17+
2. Set environment variables PARAMETERS_SECRET and BRANCH
18+
3. Run ci/test_wif.sh
19+
"""
20+
21+
22+
ACCOUNT = os.getenv("SNOWFLAKE_TEST_WIF_ACCOUNT")
23+
HOST = os.getenv("SNOWFLAKE_TEST_WIF_HOST")
24+
PROVIDER = os.getenv("SNOWFLAKE_TEST_WIF_PROVIDER")
25+
26+
27+
@pytest.mark.wif
28+
def test_wif_provider_autodetection():
29+
connection_params = {
30+
"account": ACCOUNT,
31+
"authenticator": "WORKLOAD_IDENTITY",
32+
"host": HOST,
33+
}
34+
assert connect_and_execute_simple_query(
35+
connection_params
36+
), "Failed to connect with using WIF - automatic provider detection"
37+
38+
39+
@pytest.mark.wif
40+
def test_wif_defined_provider():
41+
connection_params = {
42+
"host": HOST,
43+
"account": ACCOUNT,
44+
"authenticator": "WORKLOAD_IDENTITY",
45+
"workload_identity_provider": PROVIDER,
46+
}
47+
assert connect_and_execute_simple_query(
48+
connection_params
49+
), "Failed to connect with using WIF - automatic provider detection"
50+
51+
52+
@pytest.mark.wif
53+
def test_should_authenticate_using_oidc():
54+
if not is_provider_gcp():
55+
pytest.skip("Skipping test - not running on GCP")
56+
57+
connection_params = {
58+
"host": HOST,
59+
"account": ACCOUNT,
60+
"authenticator": "WORKLOAD_IDENTITY",
61+
"workload_identity_provider": "OIDC",
62+
"token": get_gcp_access_token(),
63+
}
64+
65+
assert connect_and_execute_simple_query(
66+
connection_params
67+
), "Failed to connect using WIF with OIDC provider"
68+
69+
70+
def is_provider_gcp() -> bool:
71+
return PROVIDER == "GCP"
72+
73+
74+
def connect_and_execute_simple_query(connection_params) -> bool:
75+
try:
76+
logger.info("Trying to connect to Snowflake")
77+
with snowflake.connector.connect(**connection_params) as con:
78+
result = con.cursor().execute("select 1;")
79+
logger.debug(result.fetchall())
80+
logger.info("Successfully connected to Snowflake")
81+
return True
82+
except Exception as e:
83+
logger.error(e)
84+
return False
85+
86+
87+
def get_gcp_access_token() -> str:
88+
try:
89+
command = (
90+
'curl -H "Metadata-Flavor: Google" '
91+
'"http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/identity?audience=snowflakecomputing.com"'
92+
)
93+
94+
result = subprocess.run(
95+
["bash", "-c", command], capture_output=True, text=True, check=False
96+
)
97+
98+
if result.returncode == 0 and result.stdout and result.stdout.strip():
99+
return result.stdout.strip()
100+
else:
101+
raise RuntimeError(
102+
f"Failed to retrieve GCP access token, exit code: {result.returncode}"
103+
)
104+
105+
except Exception as e:
106+
raise RuntimeError(f"Error executing GCP metadata request: {e}")

0 commit comments

Comments
 (0)