diff --git a/DESCRIPTION.md b/DESCRIPTION.md index 56b93c9e06..8e4ff528c5 100644 --- a/DESCRIPTION.md +++ b/DESCRIPTION.md @@ -8,7 +8,9 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne # Release Notes - v4.2.0(TBD) - - Added support for async I/O. Asynchronous version of connector is available via `snowflake.connector.aio` module. + - Added PREVIEW support for async I/O. Asynchronous version of connector is available via `snowflake.connector.aio` module. + This is preview feature and should not be used in production code. To use this feature contact your Snowflake Sales + Representative ( Snowflake Support cannot help with this feature in the current stage, while its in preview). - v4.1.1(TBD) - Relaxed pandas dependency requirements for Python below 3.12. - Changed CRL cache cleanup background task to daemon to avoid blocking main thread. diff --git a/ci/container/test_authentication.sh b/ci/container/test_authentication.sh index 18bd6e492a..ee66e600a3 100755 --- a/ci/container/test_authentication.sh +++ b/ci/container/test_authentication.sh @@ -1,6 +1,6 @@ #!/bin/bash -e -set -o pipefail +set -ox pipefail export WORKSPACE=${WORKSPACE:-/mnt/workspace} @@ -17,6 +17,9 @@ export RUN_AUTH_TESTS=true export AUTHENTICATION_TESTS_ENV="docker" export PYTHONPATH=$SOURCE_ROOT -python3 -m pip install --break-system-packages -e . +python3 -m pip install --break-system-packages -e ".[development]" -python3 -m pytest test/auth/* +python3 -m pytest test/auth/ --ignore=test/auth/aio + +python3 -m pip install --break-system-packages -e ".[development,aio,aioboto]" +python3 -m pytest test/auth/aio/ diff --git a/ci/test_darwin.sh b/ci/test_darwin.sh index a32acc2531..b40a33ea74 100755 --- a/ci/test_darwin.sh +++ b/ci/test_darwin.sh @@ -37,7 +37,7 @@ for PYTHON_VERSION in ${PYTHON_VERSIONS}; do SHORT_VERSION=$(python3 -c "print('${PYTHON_VERSION}'.replace('.', ''))") CONNECTOR_WHL=$(ls ${CONNECTOR_DIR}/dist/snowflake_connector_python*cp${SHORT_VERSION}*.whl) # pandas not tested here because of macos issue: SNOW-1660226 - TEST_ENVLIST=$(python3 -c "print('fix_lint,' + ','.join('py${SHORT_VERSION}-' + e + '-ci' for e in ['unit','integ','sso']) + ',py${SHORT_VERSION}-coverage')") + TEST_ENVLIST=$(python3 -c "print('fix_lint,' + ','.join('py${SHORT_VERSION}-' + e + '-ci' for e in ['unit','integ','sso','aio']) + ',py${SHORT_VERSION}-coverage')") echo "[Info] Running tox for ${TEST_ENVLIST}" python3.12 -m tox run -e ${TEST_ENVLIST} --installpkg ${CONNECTOR_WHL} done diff --git a/ci/test_fips.sh b/ci/test_fips.sh index 5096ddc7c4..72f69f49e0 100755 --- a/ci/test_fips.sh +++ b/ci/test_fips.sh @@ -3,6 +3,7 @@ # Test Snowflake Connector (FIPS) # Note this is the script that test_fips_docker.sh runs inside of the docker container # +set -x # Export USE_PASSWORD only on Jenkins (not on GitHub Actions) # Jenkins FIPS tests run against mocked Snowflake with password auth @@ -41,6 +42,14 @@ pip freeze cd $CONNECTOR_DIR # Run tests in parallel using pytest-xdist -pytest -n auto -vvv --cov=snowflake.connector --cov-report=xml:coverage.xml test --ignore=test/integ/aio_it --ignore=test/unit/aio --ignore=test/wif/test_wif_async.py - +pytest -n auto -vvv --cov=snowflake.connector --cov-report=xml:coverage.xml test \ + --ignore=test/integ/aio_it \ + --ignore=test/unit/aio \ + --ignore=test/auth/aio \ + --ignore=test/wif/test_wif_async.py + +pip install "${CONNECTOR_WHL}[aio,aioboto]" +# Run aio tests separately +pytest -n auto -vvv --cov=snowflake.connector --cov-append --cov-report=xml:coverage.xml -m "aio and unit" test +pytest -n auto -vvv --cov=snowflake.connector --cov-append --cov-report=xml:coverage.xml -m "aio and integ" test deactivate diff --git a/ci/test_linux.sh b/ci/test_linux.sh index 7256538ef9..a90afdb733 100755 --- a/ci/test_linux.sh +++ b/ci/test_linux.sh @@ -41,7 +41,7 @@ else echo "[Info] Testing with ${PYTHON_VERSION}" SHORT_VERSION=$(python3.10 -c "print('${PYTHON_VERSION}'.replace('.', ''))") CONNECTOR_WHL=$(ls $CONNECTOR_DIR/dist/snowflake_connector_python*cp${SHORT_VERSION}*manylinux2014*.whl | sort -r | head -n 1) - TEST_LIST=`echo py${PYTHON_VERSION/\./}-{unit-parallel,integ,pandas-parallel,sso}-ci | sed 's/ /,/g'` + TEST_LIST=`echo py${PYTHON_VERSION/\./}-{unit-parallel,integ,pandas-parallel,sso,aio}-ci | sed 's/ /,/g'` TEST_ENVLIST=fix_lint,$TEST_LIST,py${PYTHON_VERSION/\./}-coverage echo "[Info] Running tox for ${TEST_ENVLIST}" diff --git a/ci/test_rockylinux9.sh b/ci/test_rockylinux9.sh index 3c66b5face..480d005b5b 100755 --- a/ci/test_rockylinux9.sh +++ b/ci/test_rockylinux9.sh @@ -69,7 +69,7 @@ else continue fi - TEST_LIST=`echo py${PYTHON_VERSION/\./}-{extras,unit-parallel,integ-parallel,pandas-parallel,sso}-ci | sed 's/ /,/g'` + TEST_LIST=`echo py${PYTHON_VERSION/\./}-{extras,unit-parallel,integ-parallel,pandas-parallel,aio}-ci | sed 's/ /,/g'` TEST_ENVLIST=fix_lint,$TEST_LIST,py${PYTHON_VERSION/\./}-coverage echo "[Info] Running tox for ${TEST_ENVLIST}" diff --git a/ci/test_windows.bat b/ci/test_windows.bat index dad1d38e16..4400a3c368 100644 --- a/ci/test_windows.bat +++ b/ci/test_windows.bat @@ -49,7 +49,7 @@ curl https://repo1.maven.org/maven2/org/wiremock/wiremock-standalone/3.11.0/wire set JUNIT_REPORT_DIR=%workspace% set COV_REPORT_DIR=%workspace% -set TEST_ENVLIST=fix_lint,py%pv%-unit-ci,py%pv%-integ-ci,py%pv%-pandas-ci,py%pv%-sso-ci,py%pv%-coverage +set TEST_ENVLIST=fix_lint,py%pv%-unit-ci,py%pv%-integ-ci,py%pv%-pandas-ci,py%pv%-sso-ci,py%pv%-aio-ci,py%pv%-coverage tox -e %TEST_ENVLIST% --installpkg %connector_whl% if %errorlevel% neq 0 goto :error diff --git a/test/auth/aio/__init__.py b/test/auth/aio/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/test/auth/aio/authorization_test_helper.py b/test/auth/aio/authorization_test_helper.py new file mode 100644 index 0000000000..e0943d5a1d --- /dev/null +++ b/test/auth/aio/authorization_test_helper.py @@ -0,0 +1,229 @@ +import logging.config +import os +import subprocess +import threading +import webbrowser +from enum import Enum +from typing import Union + +import requests + +import snowflake.connector.aio + +try: + from src.snowflake.connector.vendored.requests.auth import HTTPBasicAuth +except ImportError: + pass + +logger = logging.getLogger(__name__) + +logger.setLevel(logging.INFO) + + +class Scenario(Enum): + SUCCESS = "success" + FAIL = "fail" + TIMEOUT = "timeout" + EXTERNAL_OAUTH_OKTA_SUCCESS = "externalOauthOktaSuccess" + INTERNAL_OAUTH_SNOWFLAKE_SUCCESS = "internalOauthSnowflakeSuccess" + + +def get_access_token_oauth(cfg): + auth_url = cfg["auth_url"] + + data = { + "username": cfg["okta_user"], + "password": cfg["okta_pass"], + "grant_type": "password", + "scope": f"session:role:{cfg['role']}", + } + + headers = {"Content-Type": "application/x-www-form-urlencoded;charset=UTF-8"} + + auth_credentials = HTTPBasicAuth(cfg["oauth_client_id"], cfg["oauth_client_secret"]) + try: + response = requests.post( + url=auth_url, data=data, headers=headers, auth=auth_credentials + ) + response.raise_for_status() + return response.json()["access_token"] + + except requests.exceptions.HTTPError as http_err: + logger.error(f"HTTP error occurred: {http_err}") + raise + + +def clean_browser_processes(): + if os.getenv("AUTHENTICATION_TESTS_ENV") == "docker": + try: + clean_browser_processes_path = "/externalbrowser/cleanBrowserProcesses.js" + process = subprocess.run(["node", clean_browser_processes_path], timeout=15) + logger.debug(f"OUTPUT: {process.stdout}, ERRORS: {process.stderr}") + except Exception as e: + raise RuntimeError(e) + + +class AuthorizationTestHelper: + def __init__(self, configuration: dict): + self.auth_test_env = os.getenv("AUTHENTICATION_TESTS_ENV") + self.configuration = configuration + self.error_msg = "" + + def update_config(self, configuration): + self.configuration = configuration + + async def connect_and_provide_credentials( + self, scenario: Scenario, login: str, password: str + ): + import asyncio + + try: + # Use asyncio task for connection instead of thread + connect_task = asyncio.create_task(self.connect_and_execute_simple_query()) + + if self.auth_test_env == "docker": + # Browser credentials still needs to run in thread since it's sync + browser = threading.Thread( + target=self._provide_credentials, args=(scenario, login, password) + ) + browser.start() + # Wait for browser thread to complete + await asyncio.get_event_loop().run_in_executor(None, browser.join) + + # Wait for connection task to complete + await connect_task + + except Exception as e: + self.error_msg = e + logger.error(e) + + def get_error_msg(self) -> str: + return str(self.error_msg) + + async def connect_and_execute_simple_query(self): + try: + logger.info("Trying to connect to Snowflake") + async with snowflake.connector.aio.SnowflakeConnection( + **self.configuration + ) as con: + result = await con.cursor().execute("select 1;") + logger.debug(await result.fetchall()) + logger.info("Successfully connected to Snowflake") + return True + except Exception as e: + self.error_msg = e + logger.error(e) + return False + + async def connect_and_execute_set_session_state(self, key: str, value: str): + try: + logger.info("Trying to connect to Snowflake") + async with snowflake.connector.aio.SnowflakeConnection( + **self.configuration + ) as con: + result = await con.cursor().execute(f"SET {key} = '{value}'") + logger.debug(await result.fetchall()) + logger.info("Successfully SET session variable") + return True + except Exception as e: + self.error_msg = e + logger.error(e) + return False + + async def connect_and_execute_check_session_state(self, key: str): + try: + logger.info("Trying to connect to Snowflake") + async with snowflake.connector.aio.SnowflakeConnection( + **self.configuration + ) as con: + result = await con.cursor().execute(f"SELECT 1, ${key}") + value = (await result.fetchone())[1] + logger.debug(value) + logger.info("Successfully READ session variable") + return value + except Exception as e: + self.error_msg = e + logger.error(e) + return False + + def _provide_credentials(self, scenario: Scenario, login: str, password: str): + try: + webbrowser.register("xdg-open", None, webbrowser.GenericBrowser("xdg-open")) + provide_browser_credentials_path = ( + "/externalbrowser/provideBrowserCredentials.js" + ) + process = subprocess.run( + [ + "node", + provide_browser_credentials_path, + scenario.value, + login, + password, + ], + timeout=15, + ) + logger.debug(f"OUTPUT: {process.stdout}, ERRORS: {process.stderr}") + except Exception as e: + self.error_msg = e + raise RuntimeError(e) + + def get_totp(self, seed: str = "") -> []: + if self.auth_test_env == "docker": + try: + provide_totp_generator_path = "/externalbrowser/totpGenerator.js" + process = subprocess.run( + ["node", provide_totp_generator_path, seed], + timeout=40, + capture_output=True, + text=True, + ) + logger.debug(f"OUTPUT: {process.stdout}, ERRORS: {process.stderr}") + return process.stdout.strip().split() + except Exception as e: + self.error_msg = e + raise RuntimeError(e) + else: + logger.info("TOTP generation is not supported in this environment") + return "" + + async def connect_using_okta_connection_and_execute_custom_command( + self, command: str, return_token: bool = False + ) -> Union[bool, str]: + try: + logger.info("Setup PAT") + async with snowflake.connector.aio.SnowflakeConnection( + **self.configuration + ) as con: + result = await con.cursor().execute(command) + token = (await result.fetchall())[0][1] + except Exception as e: + self.error_msg = e + logger.error(e) + return False + if return_token: + return token + return False + + async def connect_and_execute_simple_query_with_mfa_token(self, totp_codes): + # Try each TOTP code until one works + for i, totp_code in enumerate(totp_codes): + logging.info(f"Trying TOTP code {i + 1}/{len(totp_codes)}") + + self.configuration["passcode"] = totp_code + self.error_msg = "" + + connection_success = await self.connect_and_execute_simple_query() + + if connection_success: + logging.info(f"Successfully connected with TOTP code {i + 1}") + return True + else: + last_error = str(self.error_msg) + logging.warning(f"TOTP code {i + 1} failed: {last_error}") + if "TOTP Invalid" in last_error: + logging.info("TOTP/MFA error detected.") + continue + else: + logging.error(f"Non-TOTP error detected: {last_error}") + break + return False diff --git a/test/auth/aio/test_external_browser.py b/test/auth/aio/test_external_browser.py new file mode 100644 index 0000000000..cba181a745 --- /dev/null +++ b/test/auth/aio/test_external_browser.py @@ -0,0 +1,102 @@ +import logging +from test.auth.aio.authorization_test_helper import ( + AuthorizationTestHelper, + Scenario, + clean_browser_processes, +) +from test.auth.authorization_parameters import ( + AuthConnectionParameters, + get_okta_login_credentials, +) + +import pytest + + +@pytest.fixture(autouse=True) +def setup_and_teardown(): + logging.info("Cleanup before test") + clean_browser_processes() + + yield + + logging.info("Teardown: Performing specific actions after the test") + clean_browser_processes() + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_external_browser_successful(): + connection_parameters = ( + AuthConnectionParameters().get_external_browser_connection_parameters() + ) + test_helper = AuthorizationTestHelper(connection_parameters) + browser_login, browser_password = get_okta_login_credentials().values() + await test_helper.connect_and_provide_credentials( + Scenario.SUCCESS, browser_login, browser_password + ) + + # Clear any error from browser automation before final assertion + test_helper.error_msg = "" + + # Verify connection succeeded by attempting a simple query + assert ( + await test_helper.connect_and_execute_simple_query() is True + ), "Connection should be established" + assert test_helper.error_msg == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_external_browser_mismatched_user(): + connection_parameters = ( + AuthConnectionParameters().get_external_browser_connection_parameters() + ) + connection_parameters["user"] = "differentUsername" + browser_login, browser_password = get_okta_login_credentials().values() + + test_helper = AuthorizationTestHelper(connection_parameters) + await test_helper.connect_and_provide_credentials( + Scenario.SUCCESS, browser_login, browser_password + ) + assert ( + "The user you were trying to authenticate as differs from the user" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.skip(reason="SNOW-2007651 Adding custom browser timeout") +@pytest.mark.asyncio +async def test_external_browser_wrong_credentials(): + connection_parameters = ( + AuthConnectionParameters().get_external_browser_connection_parameters() + ) + browser_login, browser_password = "invalidUser", "invalidPassword" + connection_parameters["external_browser_timeout"] = 10 + test_helper = AuthorizationTestHelper(connection_parameters) + await test_helper.connect_and_provide_credentials( + Scenario.FAIL, browser_login, browser_password + ) + + assert ( + "Unable to receive the OAuth message within a given timeout" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.skip(reason="SNOW-2007651 Adding custom browser timeout") +@pytest.mark.asyncio +async def test_external_browser_timeout(): + connection_parameters = ( + AuthConnectionParameters().get_external_browser_connection_parameters() + ) + test_helper = AuthorizationTestHelper(connection_parameters) + connection_parameters["external_browser_timeout"] = 1 + assert ( + not await test_helper.connect_and_execute_simple_query() + ), "Connection should not be established" + assert ( + "Unable to receive the OAuth message within a given timeout" + in test_helper.get_error_msg() + ) diff --git a/test/auth/aio/test_external_session_with_PAT.py b/test/auth/aio/test_external_session_with_PAT.py new file mode 100644 index 0000000000..2d4e915c8b --- /dev/null +++ b/test/auth/aio/test_external_session_with_PAT.py @@ -0,0 +1,83 @@ +import uuid +from test.auth.aio.authorization_test_helper import AuthorizationTestHelper +from test.auth.aio.test_pat import get_pat_token, remove_pat_token +from test.auth.authorization_parameters import ( + AuthConnectionParameters, + get_pat_setup_command_variables, +) + +import pytest + +EXTERNAL_SESSION_ID = str(uuid.uuid4()) +SESSION_VAR_KEY = "PAT_WITH_EXTERNAL_SESSION_TEST_KEY" +SESSION_VAR_VALUE = "PAT_WITH_EXTERNAL_SESSION_TEST_VALUE" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_pat_with_external_session_authN_success() -> None: + pat_command_variables = get_pat_setup_command_variables() + connection_parameters = AuthConnectionParameters().get_pat_connection_parameters() + try: + pat_command_variables = await get_pat_token(pat_command_variables) + connection_parameters["token"] = pat_command_variables["token"] + connection_parameters["external_session_id"] = EXTERNAL_SESSION_ID + connection_parameters["authenticator"] = "PAT_WITH_EXTERNAL_SESSION" + test_helper = AuthorizationTestHelper(connection_parameters) + + # Verify the SET operation succeeded + set_result = await test_helper.connect_and_execute_set_session_state( + SESSION_VAR_KEY, SESSION_VAR_VALUE + ) + assert ( + set_result is True + ), f"Failed to set session variable: {test_helper.get_error_msg()}" + + # Clear error message before the next operation + test_helper.error_msg = "" + + # Verify the GET operation succeeded + ret = await test_helper.connect_and_execute_check_session_state(SESSION_VAR_KEY) + assert ( + ret == SESSION_VAR_VALUE + ), f"Failed to get session variable. Got {ret}, error: {test_helper.get_error_msg()}" + finally: + await remove_pat_token(pat_command_variables) + assert test_helper.get_error_msg() == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_pat_with_external_session_authN_fail() -> None: + pat_command_variables = get_pat_setup_command_variables() + try: + pat_command_variables = await get_pat_token(pat_command_variables) + connection_parameters = ( + AuthConnectionParameters().get_pat_connection_parameters() + ) + connection_parameters["token"] = pat_command_variables["token"] + connection_parameters["external_session_id"] = EXTERNAL_SESSION_ID + connection_parameters["authenticator"] = "PAT_WITH_EXTERNAL_SESSION" + test_helper = AuthorizationTestHelper(connection_parameters) + + # Verify the SET operation succeeded + set_result = await test_helper.connect_and_execute_set_session_state( + SESSION_VAR_KEY, SESSION_VAR_VALUE + ) + assert ( + set_result is True + ), f"Failed to set session variable: {test_helper.get_error_msg()}" + + connection_parameters["external_session_id"] = str( + uuid.uuid4() + ) # Use different external session + test_helper = AuthorizationTestHelper(connection_parameters) + ret = await test_helper.connect_and_execute_check_session_state(SESSION_VAR_KEY) + assert ret != SESSION_VAR_VALUE + finally: + await remove_pat_token(pat_command_variables) + print(test_helper.get_error_msg()) + assert ( + f"Session variable '${SESSION_VAR_KEY}' does not exist" + in test_helper.get_error_msg() + ) diff --git a/test/auth/aio/test_key_pair.py b/test/auth/aio/test_key_pair.py new file mode 100644 index 0000000000..aee7a02376 --- /dev/null +++ b/test/auth/aio/test_key_pair.py @@ -0,0 +1,41 @@ +from test.auth.aio.authorization_test_helper import AuthorizationTestHelper +from test.auth.authorization_parameters import ( + AuthConnectionParameters, + get_rsa_private_key_for_key_pair, +) + +import pytest + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_key_pair_successful(): + connection_parameters = ( + AuthConnectionParameters().get_key_pair_connection_parameters() + ) + connection_parameters["private_key"] = get_rsa_private_key_for_key_pair( + "SNOWFLAKE_AUTH_TEST_PRIVATE_KEY_PATH" + ) + + test_helper = AuthorizationTestHelper(connection_parameters) + assert ( + await test_helper.connect_and_execute_simple_query() + ), "Failed to connect with Snowflake" + assert test_helper.error_msg == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_key_pair_invalid_key(): + connection_parameters = ( + AuthConnectionParameters().get_key_pair_connection_parameters() + ) + connection_parameters["private_key"] = get_rsa_private_key_for_key_pair( + "SNOWFLAKE_AUTH_TEST_INVALID_PRIVATE_KEY_PATH" + ) + + test_helper = AuthorizationTestHelper(connection_parameters) + assert ( + not await test_helper.connect_and_execute_simple_query() + ), "Connection to Snowflake should not be established" + assert "JWT token is invalid" in test_helper.get_error_msg() diff --git a/test/auth/aio/test_mfa.py b/test/auth/aio/test_mfa.py new file mode 100644 index 0000000000..8703748928 --- /dev/null +++ b/test/auth/aio/test_mfa.py @@ -0,0 +1,41 @@ +import logging +from test.auth.aio.authorization_test_helper import AuthorizationTestHelper +from test.auth.authorization_parameters import AuthConnectionParameters + +import pytest + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_mfa_successful(): + connection_parameters = AuthConnectionParameters().get_mfa_connection_parameters() + connection_parameters["client_request_mfa_token"] = True + test_helper = AuthorizationTestHelper(connection_parameters) + totp_codes = test_helper.get_totp() + logging.info(f"Got {len(totp_codes)} TOTP codes to try") + + connection_success = ( + await test_helper.connect_and_execute_simple_query_with_mfa_token(totp_codes) + ) + + assert ( + connection_success + ), f"Failed to connect with any of the {len(totp_codes)} TOTP codes. Last error: {test_helper.error_msg}" + assert ( + test_helper.error_msg == "" + ), f"Final error message should be empty but got: {test_helper.error_msg}" + + logging.info("Testing MFA token caching with second connection...") + + connection_parameters["passcode"] = None + cache_test_helper = AuthorizationTestHelper(connection_parameters) + cache_connection_success = ( + await cache_test_helper.connect_and_execute_simple_query() + ) + + assert ( + cache_connection_success + ), f"Failed to connect with cached MFA token. Error: {cache_test_helper.error_msg}" + assert ( + cache_test_helper.error_msg == "" + ), f"Cache test error message should be empty but got: {cache_test_helper.error_msg}" diff --git a/test/auth/aio/test_oauth.py b/test/auth/aio/test_oauth.py new file mode 100644 index 0000000000..b90d1b435c --- /dev/null +++ b/test/auth/aio/test_oauth.py @@ -0,0 +1,62 @@ +from test.auth.aio.authorization_test_helper import ( + AuthorizationTestHelper, + get_access_token_oauth, +) +from test.auth.authorization_parameters import ( + AuthConnectionParameters, + get_oauth_token_parameters, +) + +import pytest + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_oauth_successful(): + token = get_oauth_token() + connection_parameters = AuthConnectionParameters().get_oauth_connection_parameters( + token + ) + test_helper = AuthorizationTestHelper(connection_parameters) + assert ( + await test_helper.connect_and_execute_simple_query() + ), "Failed to connect with OAuth token" + assert test_helper.error_msg == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_oauth_mismatched_user(): + token = get_oauth_token() + connection_parameters = AuthConnectionParameters().get_oauth_connection_parameters( + token + ) + connection_parameters["user"] = "differentUsername" + test_helper = AuthorizationTestHelper(connection_parameters) + assert ( + await test_helper.connect_and_execute_simple_query() is False + ), "Connection should not be established" + assert ( + "The user you were trying to authenticate as differs from the user" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_oauth_invalid_token(): + token = "invalidToken" + connection_parameters = AuthConnectionParameters().get_oauth_connection_parameters( + token + ) + test_helper = AuthorizationTestHelper(connection_parameters) + assert ( + await test_helper.connect_and_execute_simple_query() is False + ), "Connection should not be established" + assert "Invalid OAuth access token" in test_helper.get_error_msg() + + +def get_oauth_token(): + oauth_config = get_oauth_token_parameters() + token = get_access_token_oauth(oauth_config) + return token diff --git a/test/auth/aio/test_okta.py b/test/auth/aio/test_okta.py new file mode 100644 index 0000000000..6b8f548f86 --- /dev/null +++ b/test/auth/aio/test_okta.py @@ -0,0 +1,62 @@ +from test.auth.aio.authorization_test_helper import AuthorizationTestHelper +from test.auth.authorization_parameters import AuthConnectionParameters + +import pytest + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_successful(): + connection_parameters = AuthConnectionParameters().get_okta_connection_parameters() + test_helper = AuthorizationTestHelper(connection_parameters) + + assert ( + await test_helper.connect_and_execute_simple_query() + ), "Failed to connect with Snowflake" + assert test_helper.error_msg == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_with_wrong_okta_username(): + connection_parameters = AuthConnectionParameters().get_okta_connection_parameters() + connection_parameters["user"] = "differentUsername" + + test_helper = AuthorizationTestHelper(connection_parameters) + assert ( + not await test_helper.connect_and_execute_simple_query() + ), "Connection to Snowflake should not be established" + assert "Failed to get authentication by OKTA" in test_helper.get_error_msg() + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_wrong_url(): + connection_parameters = AuthConnectionParameters().get_okta_connection_parameters() + + connection_parameters["authenticator"] = "https://invalid.okta.com/" + test_helper = AuthorizationTestHelper(connection_parameters) + assert ( + not await test_helper.connect_and_execute_simple_query() + ), "Connection to Snowflake should not be established" + assert ( + "The specified authenticator is not accepted by your Snowflake account configuration" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.skip(reason="SNOW-1852279 implement error handling for invalid URL") +@pytest.mark.asyncio +async def test_okta_wrong_url_2(): + connection_parameters = AuthConnectionParameters().get_okta_connection_parameters() + + connection_parameters["authenticator"] = "https://invalid.abc.com/" + test_helper = AuthorizationTestHelper(connection_parameters) + assert ( + not await test_helper.connect_and_execute_simple_query() + ), "Connection to Snowflake should not be established" + assert ( + "The specified authenticator is not accepted by your Snowflake account configuration" + in test_helper.get_error_msg() + ) diff --git a/test/auth/aio/test_okta_authorization_code.py b/test/auth/aio/test_okta_authorization_code.py new file mode 100644 index 0000000000..0b78840142 --- /dev/null +++ b/test/auth/aio/test_okta_authorization_code.py @@ -0,0 +1,110 @@ +import logging +from test.auth.aio.authorization_test_helper import ( + AuthorizationTestHelper, + Scenario, + clean_browser_processes, +) +from test.auth.authorization_parameters import ( + AuthConnectionParameters, + get_okta_login_credentials, +) + +import pytest + + +@pytest.fixture(autouse=True) +def setup_and_teardown(): + logging.info("Cleanup before test") + clean_browser_processes() + + yield + + logging.info("Teardown: Performing specific actions after the test") + clean_browser_processes() + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_authorization_code_successful(): + connection_parameters = ( + AuthConnectionParameters().get_oauth_external_authorization_code_connection_parameters() + ) + test_helper = AuthorizationTestHelper(connection_parameters) + browser_login, browser_password = get_okta_login_credentials().values() + await test_helper.connect_and_provide_credentials( + Scenario.SUCCESS, browser_login, browser_password + ) + + # Clear any error from browser automation before final assertion + test_helper.error_msg = "" + + # Verify connection succeeded by attempting a simple query + assert ( + await test_helper.connect_and_execute_simple_query() is True + ), "Connection should be established" + assert test_helper.error_msg == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_authorization_code_mismatched_user(): + connection_parameters = ( + AuthConnectionParameters().get_oauth_external_authorization_code_connection_parameters() + ) + connection_parameters["user"] = "differentUsername" + browser_login, browser_password = get_okta_login_credentials().values() + + test_helper = AuthorizationTestHelper(connection_parameters) + await test_helper.connect_and_provide_credentials( + Scenario.SUCCESS, browser_login, browser_password + ) + + assert ( + "The user you were trying to authenticate as differs from the user" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_authorization_code_timeout(): + connection_parameters = ( + AuthConnectionParameters().get_oauth_external_authorization_code_connection_parameters() + ) + test_helper = AuthorizationTestHelper(connection_parameters) + connection_parameters["external_browser_timeout"] = 1 + + assert ( + await test_helper.connect_and_execute_simple_query() is False + ), "Connection should not be established" + assert ( + "Unable to receive the OAuth message within a given timeout" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_authorization_code_with_token_cache(): + connection_parameters = ( + AuthConnectionParameters().get_oauth_external_authorization_code_connection_parameters() + ) + connection_parameters["client_store_temporary_credential"] = True + connection_parameters["external_browser_timeout"] = 10 + + test_helper = AuthorizationTestHelper(connection_parameters) + browser_login, browser_password = get_okta_login_credentials().values() + + await test_helper.connect_and_provide_credentials( + Scenario.SUCCESS, browser_login, browser_password + ) + + clean_browser_processes() + + # Clear any error from first connection before testing cache + test_helper.error_msg = "" + + assert ( + await test_helper.connect_and_execute_simple_query() is True + ), "Connection should be established" + assert test_helper.error_msg == "", "Error message should be empty" diff --git a/test/auth/aio/test_okta_client_credentials.py b/test/auth/aio/test_okta_client_credentials.py new file mode 100644 index 0000000000..4fd00febcb --- /dev/null +++ b/test/auth/aio/test_okta_client_credentials.py @@ -0,0 +1,63 @@ +import logging +from test.auth.aio.authorization_test_helper import ( + AuthorizationTestHelper, + clean_browser_processes, +) +from test.auth.authorization_parameters import AuthConnectionParameters + +import pytest + + +@pytest.fixture(autouse=True) +def setup_and_teardown(): + logging.info("Cleanup before test") + clean_browser_processes() + + yield + + logging.info("Teardown: Performing specific actions after the test") + clean_browser_processes() + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_client_credentials_successful(): + connection_parameters = ( + AuthConnectionParameters().get_oauth_external_client_credential_connection_parameters() + ) + test_helper = AuthorizationTestHelper(connection_parameters) + + await test_helper.connect_and_execute_simple_query() + + assert test_helper.error_msg == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_client_credentials_mismatched_user(): + connection_parameters = ( + AuthConnectionParameters().get_oauth_external_client_credential_connection_parameters() + ) + connection_parameters["user"] = "differentUsername" + test_helper = AuthorizationTestHelper(connection_parameters) + + await test_helper.connect_and_execute_simple_query() + + assert ( + "The user you were trying to authenticate as differs from the user" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_okta_client_credentials_unauthorized(): + connection_parameters = ( + AuthConnectionParameters().get_oauth_external_client_credential_connection_parameters() + ) + connection_parameters["oauth_client_id"] = "invalidClientID" + test_helper = AuthorizationTestHelper(connection_parameters) + + await test_helper.connect_and_execute_simple_query() + + assert "Invalid HTTP request from web browser" in test_helper.get_error_msg() diff --git a/test/auth/aio/test_pat.py b/test/auth/aio/test_pat.py new file mode 100644 index 0000000000..5b5bd41654 --- /dev/null +++ b/test/auth/aio/test_pat.py @@ -0,0 +1,85 @@ +from datetime import datetime +from test.auth.aio.authorization_test_helper import AuthorizationTestHelper +from test.auth.authorization_parameters import ( + AuthConnectionParameters, + get_pat_setup_command_variables, +) +from typing import Union + +import pytest + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_authenticate_with_pat_successful() -> None: + pat_command_variables = get_pat_setup_command_variables() + connection_parameters = AuthConnectionParameters().get_pat_connection_parameters() + test_helper = AuthorizationTestHelper(connection_parameters) + try: + pat_command_variables = await get_pat_token(pat_command_variables) + connection_parameters["token"] = pat_command_variables["token"] + await test_helper.connect_and_execute_simple_query() + finally: + await remove_pat_token(pat_command_variables) + assert test_helper.get_error_msg() == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_authenticate_with_pat_mismatched_user() -> None: + pat_command_variables = get_pat_setup_command_variables() + connection_parameters = AuthConnectionParameters().get_pat_connection_parameters() + connection_parameters["user"] = "differentUsername" + test_helper = AuthorizationTestHelper(connection_parameters) + try: + pat_command_variables = await get_pat_token(pat_command_variables) + connection_parameters["token"] = pat_command_variables["token"] + await test_helper.connect_and_execute_simple_query() + finally: + await remove_pat_token(pat_command_variables) + + assert "Programmatic access token is invalid" in test_helper.get_error_msg() + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_authenticate_with_pat_invalid_token() -> None: + connection_parameters = AuthConnectionParameters().get_pat_connection_parameters() + connection_parameters["token"] = "invalidToken" + test_helper = AuthorizationTestHelper(connection_parameters) + await test_helper.connect_and_execute_simple_query() + assert "Programmatic access token is invalid" in test_helper.get_error_msg() + + +async def get_pat_token(pat_command_variables) -> dict[str, Union[str, bool]]: + okta_connection_parameters = ( + AuthConnectionParameters().get_okta_connection_parameters() + ) + + pat_name = "PAT_PYTHON_" + generate_random_suffix() + pat_command_variables["pat_name"] = pat_name + command = ( + f"alter user {pat_command_variables['snowflake_user']} add programmatic access token {pat_name} " + f"ROLE_RESTRICTION = '{pat_command_variables['role']}' DAYS_TO_EXPIRY=1;" + ) + test_helper = AuthorizationTestHelper(okta_connection_parameters) + pat_command_variables["token"] = ( + await test_helper.connect_using_okta_connection_and_execute_custom_command( + command, True + ) + ) + return pat_command_variables + + +async def remove_pat_token(pat_command_variables: dict[str, Union[str, bool]]) -> None: + okta_connection_parameters = ( + AuthConnectionParameters().get_okta_connection_parameters() + ) + + command = f"alter user {pat_command_variables['snowflake_user']} remove programmatic access token {pat_command_variables['pat_name']};" + test_helper = AuthorizationTestHelper(okta_connection_parameters) + await test_helper.connect_using_okta_connection_and_execute_custom_command(command) + + +def generate_random_suffix() -> str: + return datetime.now().strftime("%Y%m%d%H%M%S%f") diff --git a/test/auth/aio/test_snowflake_authorization_code.py b/test/auth/aio/test_snowflake_authorization_code.py new file mode 100644 index 0000000000..af3eac4ddb --- /dev/null +++ b/test/auth/aio/test_snowflake_authorization_code.py @@ -0,0 +1,111 @@ +import logging +from test.auth.aio.authorization_test_helper import ( + AuthorizationTestHelper, + Scenario, + clean_browser_processes, +) +from test.auth.authorization_parameters import ( + AuthConnectionParameters, + get_soteria_okta_login_credentials, +) + +import pytest + + +@pytest.fixture(autouse=True) +def setup_and_teardown(): + logging.info("Cleanup before test") + clean_browser_processes() + + yield + + logging.info("Teardown: Performing specific actions after the test") + clean_browser_processes() + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_snowflake_authorization_code_successful(): + connection_parameters = ( + AuthConnectionParameters().get_snowflake_authorization_code_connection_parameters() + ) + test_helper = AuthorizationTestHelper(connection_parameters) + browser_login, browser_password = get_soteria_okta_login_credentials().values() + + await test_helper.connect_and_provide_credentials( + Scenario.INTERNAL_OAUTH_SNOWFLAKE_SUCCESS, browser_login, browser_password + ) + + # Clear any error from browser automation before final assertion + test_helper.error_msg = "" + + # Verify connection succeeded by attempting a simple query + assert ( + await test_helper.connect_and_execute_simple_query() is True + ), "Connection should be established" + assert test_helper.error_msg == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_snowflake_authorization_code_mismatched_user(): + connection_parameters = ( + AuthConnectionParameters().get_snowflake_authorization_code_connection_parameters() + ) + connection_parameters["user"] = "differentUsername" + browser_login, browser_password = get_soteria_okta_login_credentials().values() + test_helper = AuthorizationTestHelper(connection_parameters) + + await test_helper.connect_and_provide_credentials( + Scenario.INTERNAL_OAUTH_SNOWFLAKE_SUCCESS, browser_login, browser_password + ) + + assert ( + "The user you were trying to authenticate as differs from the user" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_snowflake_authorization_code_timeout(): + connection_parameters = ( + AuthConnectionParameters().get_snowflake_authorization_code_connection_parameters() + ) + test_helper = AuthorizationTestHelper(connection_parameters) + connection_parameters["external_browser_timeout"] = 1 + + assert ( + await test_helper.connect_and_execute_simple_query() is False + ), "Connection should not be established" + assert ( + "Unable to receive the OAuth message within a given timeout" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_snowflake_authorization_code_without_token_cache(): + connection_parameters = ( + AuthConnectionParameters().get_snowflake_authorization_code_connection_parameters() + ) + connection_parameters["client_store_temporary_credential"] = False + connection_parameters["external_browser_timeout"] = 15 + test_helper = AuthorizationTestHelper(connection_parameters) + browser_login, browser_password = get_soteria_okta_login_credentials().values() + + await test_helper.connect_and_provide_credentials( + Scenario.INTERNAL_OAUTH_SNOWFLAKE_SUCCESS, browser_login, browser_password + ) + + clean_browser_processes() + + assert ( + await test_helper.connect_and_execute_simple_query() is False + ), "Connection should be established" + + assert ( + "Unable to receive the OAuth message within a given timeout" + in test_helper.get_error_msg() + ), "Error message should contain timeout" diff --git a/test/auth/aio/test_snowflake_authorization_code_wildcards.py b/test/auth/aio/test_snowflake_authorization_code_wildcards.py new file mode 100644 index 0000000000..718706fac6 --- /dev/null +++ b/test/auth/aio/test_snowflake_authorization_code_wildcards.py @@ -0,0 +1,142 @@ +import logging +from test.auth.aio.authorization_test_helper import ( + AuthorizationTestHelper, + Scenario, + clean_browser_processes, +) +from test.auth.authorization_parameters import ( + AuthConnectionParameters, + get_soteria_okta_login_credentials, +) + +import pytest + + +@pytest.fixture(autouse=True) +def setup_and_teardown(): + logging.info("Cleanup before test") + clean_browser_processes() + + yield + + logging.info("Teardown: Performing specific actions after the test") + clean_browser_processes() + + +@pytest.mark.skip( + "temporarily disabled, update redirect uri for the security integration will break other drivers tests" +) +@pytest.mark.auth +@pytest.mark.asyncio +async def test_snowflake_authorization_code_wildcards_successful(): + connection_parameters = ( + AuthConnectionParameters().get_snowflake_wildcard_external_authorization_code_connection_parameters() + ) + test_helper = AuthorizationTestHelper(connection_parameters) + browser_login, browser_password = get_soteria_okta_login_credentials().values() + + await test_helper.connect_and_provide_credentials( + Scenario.INTERNAL_OAUTH_SNOWFLAKE_SUCCESS, browser_login, browser_password + ) + + # Clear any error from browser automation before final assertion + test_helper.error_msg = "" + + # Verify connection succeeded by attempting a simple query + assert ( + await test_helper.connect_and_execute_simple_query() is True + ), "Connection should be established" + assert test_helper.error_msg == "", "Error message should be empty" + + +@pytest.mark.skip( + "temporarily disabled, update redirect uri for the security integration will break other drivers tests" +) +@pytest.mark.auth +@pytest.mark.asyncio +async def test_snowflake_authorization_code_wildcards_mismatched_user(): + connection_parameters = ( + AuthConnectionParameters().get_snowflake_wildcard_external_authorization_code_connection_parameters() + ) + connection_parameters["user"] = "differentUsername" + browser_login, browser_password = get_soteria_okta_login_credentials().values() + test_helper = AuthorizationTestHelper(connection_parameters) + + await test_helper.connect_and_provide_credentials( + Scenario.INTERNAL_OAUTH_SNOWFLAKE_SUCCESS, browser_login, browser_password + ) + + assert ( + "The user you were trying to authenticate as differs from the user" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_snowflake_authorization_code_wildcards_timeout(): + connection_parameters = ( + AuthConnectionParameters().get_snowflake_wildcard_external_authorization_code_connection_parameters() + ) + test_helper = AuthorizationTestHelper(connection_parameters) + connection_parameters["external_browser_timeout"] = 1 + + assert ( + await test_helper.connect_and_execute_simple_query() is False + ), "Connection should not be established" + assert ( + "Unable to receive the OAuth message within a given timeout" + in test_helper.get_error_msg() + ) + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_snowflake_authorization_code_wildcards_with_token_cache(): + connection_parameters = ( + AuthConnectionParameters().get_snowflake_wildcard_external_authorization_code_connection_parameters() + ) + connection_parameters["external_browser_timeout"] = 15 + connection_parameters["client_store_temporary_credential"] = True + test_helper = AuthorizationTestHelper(connection_parameters) + browser_login, browser_password = get_soteria_okta_login_credentials().values() + + await test_helper.connect_and_provide_credentials( + Scenario.INTERNAL_OAUTH_SNOWFLAKE_SUCCESS, browser_login, browser_password + ) + + clean_browser_processes() + + # Clear any error from first connection before testing cache + test_helper.error_msg = "" + + assert ( + await test_helper.connect_and_execute_simple_query() is True + ), "Connection should be established" + assert test_helper.get_error_msg() == "", "Error message should be empty" + + +@pytest.mark.auth +@pytest.mark.asyncio +async def test_snowflake_authorization_code_wildcards_without_token_cache(): + connection_parameters = ( + AuthConnectionParameters().get_snowflake_wildcard_external_authorization_code_connection_parameters() + ) + connection_parameters["client_store_temporary_credential"] = False + connection_parameters["external_browser_timeout"] = 15 + test_helper = AuthorizationTestHelper(connection_parameters) + browser_login, browser_password = get_soteria_okta_login_credentials().values() + + await test_helper.connect_and_provide_credentials( + Scenario.INTERNAL_OAUTH_SNOWFLAKE_SUCCESS, browser_login, browser_password + ) + + clean_browser_processes() + + assert ( + await test_helper.connect_and_execute_simple_query() is False + ), "Connection should be established" + assert ( + "Unable to receive the OAuth message within a given timeout" + in test_helper.get_error_msg() + ), "Error message should contain timeout" diff --git a/tox.ini b/tox.ini index 9ca4e77d96..9bd617ea0a 100644 --- a/tox.ini +++ b/tox.ini @@ -45,7 +45,7 @@ setenv = SNOWFLAKE_PYTEST_COV_LOCATION = {env:JUNIT_REPORT_DIR:{toxworkdir}}/junit.{envname}-{env:cloud_provider:dev}.xml SNOWFLAKE_PYTEST_COV_CMD = --cov snowflake.connector --junitxml {env:SNOWFLAKE_PYTEST_COV_LOCATION} --cov-report= SNOWFLAKE_PYTEST_CMD = pytest {env:SNOWFLAKE_PYTEST_OPTS:} {env:SNOWFLAKE_PYTEST_COV_CMD} - SNOWFLAKE_PYTEST_CMD_IGNORE_AIO = {env:SNOWFLAKE_PYTEST_CMD} --ignore=test/integ/aio_it --ignore=test/unit/aio --ignore=test/wif/test_wif_async.py + SNOWFLAKE_PYTEST_CMD_IGNORE_AIO = {env:SNOWFLAKE_PYTEST_CMD} --ignore=test/auth/aio --ignore=test/integ/aio_it --ignore=test/unit/aio --ignore=test/wif/test_wif_async.py SNOWFLAKE_TEST_MODE = true passenv = AWS_ACCESS_KEY_ID @@ -72,6 +72,7 @@ commands = lambda: {env:SNOWFLAKE_PYTEST_CMD_IGNORE_AIO} -m "{env:SNOWFLAKE_TEST_TYPE} and lambda" {posargs:} test extras: python -m test.extras.run {posargs:} single: {env:SNOWFLAKE_PYTEST_CMD} -s "{env:SINGLE_TEST_NAME}" {posargs:} + aio: {env:SNOWFLAKE_PYTEST_CMD_IGNORE_AIO} -m "{env:SNOWFLAKE_TEST_TYPE} and aio" {posargs:} test [testenv:olddriver] basepython = python3.9 @@ -118,6 +119,23 @@ extras= aioboto pandas secure-local-storage +commands = + {env:SNOWFLAKE_PYTEST_CMD} -n auto -m "aio and unit" -vvv {posargs:} test + {env:SNOWFLAKE_PYTEST_CMD} -n auto -m "aio and integ" -vvv {posargs:} test + {env:SNOWFLAKE_PYTEST_CMD} -n auto -m "aio and auth" -vvv {posargs:} test + +[testenv:py{39,310,311,312,313}-aio] +description = Run aio tests for {basepython} +extras= + development + aio + aioboto + pandas + secure-local-storage +setenv = + {[testenv]setenv} + SNOWFLAKE_TEST_TYPE = unit +passenv = {[testenv]passenv} commands = {env:SNOWFLAKE_PYTEST_CMD} -n auto -m "aio and unit" -vvv {posargs:} test {env:SNOWFLAKE_PYTEST_CMD} -n auto -m "aio and integ" -vvv {posargs:} test