SNOW-2454885: Add auth tests for aio

DESCRIPTION.md

@@ -8,6 +8,9 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 # Release Notes
 - v4.2.0(TBD)
+  - 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).
   - Added support for async I/O. Asynchronous version of connector is available via `snowflake.connector.aio` module.
   - Added `SnowflakeCursor.stats` property to expose granular DML statistics (rows inserted, deleted, updated, and duplicates) for operations like CTAS where `rowcount` is insufficient.
   - Added support for injecting SPCS service identifier token (`SPCS_TOKEN`) into login requests when present in SPCS containers.

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 -vvv test/auth/aio/

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

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

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,aio-parallel,sso}-ci | sed 's/ /,/g'`
         TEST_ENVLIST=fix_lint,$TEST_LIST,py${PYTHON_VERSION/\./}-coverage
         echo "[Info] Running tox for ${TEST_ENVLIST}"
 

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-parallel}-ci | sed 's/ /,/g'`
         TEST_ENVLIST=fix_lint,$TEST_LIST,py${PYTHON_VERSION/\./}-coverage
         echo "[Info] Running tox for ${TEST_ENVLIST}"
 

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
 

src/snowflake/connector/aio/auth/_oauth_code.py

@@ -67,7 +67,15 @@ async def reset_secrets(self) -> None:
         AuthByOauthCodeSync.reset_secrets(self)
 
     async def prepare(self, **kwargs: Any) -> None:
-        AuthByOauthCodeSync.prepare(self, **kwargs)
+        """Prepare OAuth authentication by running blocking operations in executor."""
+        import asyncio
+        from functools import partial
+
+        loop = asyncio.get_event_loop()
+        # Run the blocking prepare call in a thread executor to avoid blocking the event loop
+        # Use partial to bind keyword arguments since run_in_executor doesn't accept kwargs
+        prepare_func = partial(AuthByOauthCodeSync.prepare, self, **kwargs)
+        await loop.run_in_executor(None, prepare_func)
 
     async def reauthenticate(
         self, conn: SnowflakeConnection, **kwargs: Any

src/snowflake/connector/aio/auth/_webbrowser.py

@@ -206,11 +206,14 @@ async def _receive_saml_token(
                 #   an immediate successive call to socket_client.recv gets the actual data
                 while len(raw_data) == 0 and attempts < max_attempts:
                     attempts += 1
-                    read_sockets, _write_sockets, _exception_sockets = select.select(
-                        [socket_connection], [], []
+                    # Run blocking select in executor to avoid blocking the event loop
+                    read_sockets, _write_sockets, _exception_sockets = (
+                        await self._event_loop.run_in_executor(
+                            None, select.select, [socket_connection], [], [], None
+                        )
                     )
 
-                    if read_sockets[0] is not None:
+                    if read_sockets and read_sockets[0] is not None:
                         # Receive the data in small chunks and retransmit it
                         socket_client, _ = await self._event_loop.sock_accept(
                             socket_connection

test/auth/aio/authorization_test_helper.py

@@ -0,0 +1,235 @@
+import logging.config
+import os
+import subprocess
+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=30)
+            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:
+            # Start connection as an async task in the main event loop
+            # This allows the browser to be opened from the main thread context
+            connect_task = asyncio.create_task(self.connect_and_execute_simple_query())
+
+            if self.auth_test_env == "docker":
+                # Give connection time to start and open the browser
+                await asyncio.sleep(5)
+
+                # Run browser automation in a thread
+                loop = asyncio.get_running_loop()
+                await loop.run_in_executor(
+                    None,
+                    self._provide_credentials,
+                    scenario,
+                    login,
+                    password,
+                )
+
+            # Wait for connection 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:
+                logger.info("Connected to Snowflake, executing query")
+                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"
+            )
+            _ = subprocess.run(
+                [
+                    "node",
+                    provide_browser_credentials_path,
+                    scenario.value,
+                    login,
+                    password,
+                ],
+                timeout=30,
+            )
+            # 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

test/auth/aio/conftest.py

@@ -0,0 +1,24 @@
+"""
+Pytest configuration for async authentication tests.
+
+This module applies a default timeout to all tests in the test/auth/aio/ directory
+to prevent tests from hanging indefinitely when waiting for external authentication
+services (Snowflake connections, browser interactions, MFA, OAuth flows, etc.).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+# Default timeout for all auth/aio tests (in seconds)
+# These tests involve external services and browser automation,
+# so they need sufficient time to complete but should not hang indefinitely.
+DEFAULT_AUTH_TEST_TIMEOUT = 60  # seper test
+
+
+def pytest_collection_modifyitems(items) -> None:
+    """Apply default timeout to all tests in this directory."""
+    for item in items:
+        # Apply timeout if not already set
+        if not any(mark.name == "timeout" for mark in item.iter_markers()):
+            item.add_marker(pytest.mark.timeout(DEFAULT_AUTH_TEST_TIMEOUT))