SNOW-2026797 Adding tests for PAT (#2265)

Author: sfc-gh-pcyrek

.github/workflows/parameters/private/parameters_aws_auth_tests.json.gpg

@@ -1,16 +1,14 @@
-try:
-    import os
-    import sys
+import os
+import sys
+from typing import Union
 
-    from cryptography.hazmat.backends import default_backend
-    from cryptography.hazmat.primitives import serialization
-except ImportError:
-    pass
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
 
 sys.path.append(os.path.abspath(os.path.dirname(__file__)))
 
 
-def get_oauth_token_parameters():
+def get_oauth_token_parameters() -> dict[str, str]:
     return {
         "auth_url": _get_env_variable("SNOWFLAKE_AUTH_TEST_OAUTH_URL"),
         "oauth_client_id": _get_env_variable("SNOWFLAKE_AUTH_TEST_OAUTH_CLIENT_ID"),
@@ -30,14 +28,14 @@ def _get_env_variable(name: str, required: bool = True) -> str:
     return value
 
 
-def get_okta_login_credentials():
+def get_okta_login_credentials() -> dict[str, str]:
     return {
         "login": _get_env_variable("SNOWFLAKE_AUTH_TEST_OKTA_USER"),
         "password": _get_env_variable("SNOWFLAKE_AUTH_TEST_OKTA_PASS"),
     }
 
 
-def get_soteria_okta_login_credentials():
+def get_soteria_okta_login_credentials() -> dict[str, str]:
     return {
         "login": _get_env_variable("SNOWFLAKE_AUTH_TEST_EXTERNAL_OAUTH_OKTA_CLIENT_ID"),
         "password": _get_env_variable(
@@ -56,6 +54,13 @@ def get_rsa_private_key_for_key_pair(
         return private_key
 
 
+def get_pat_setup_command_variables() -> dict[str, Union[str, bool, int]]:
+    return {
+        "snowflake_user": _get_env_variable("SNOWFLAKE_AUTH_TEST_SNOWFLAKE_USER"),
+        "role": _get_env_variable("SNOWFLAKE_AUTH_TEST_INTERNAL_OAUTH_SNOWFLAKE_ROLE"),
+    }
+
+
 class AuthConnectionParameters:
     def __init__(self):
         self.basic_config = {
@@ -69,7 +74,7 @@ def __init__(self):
             "CLIENT_STORE_TEMPORARY_CREDENTIAL": False,
         }
 
-    def get_base_connection_parameters(self):
+    def get_base_connection_parameters(self) -> dict[str, Union[str, bool, int]]:
         return self.basic_config
 
     def get_key_pair_connection_parameters(self):
@@ -79,15 +84,15 @@ def get_key_pair_connection_parameters(self):
 
         return config
 
-    def get_external_browser_connection_parameters(self):
+    def get_external_browser_connection_parameters(self) -> dict[str, str]:
         config = self.basic_config.copy()
 
         config["user"] = _get_env_variable("SNOWFLAKE_AUTH_TEST_BROWSER_USER")
         config["authenticator"] = "externalbrowser"
 
         return config
 
-    def get_store_id_token_connection_parameters(self):
+    def get_store_id_token_connection_parameters(self) -> dict[str, str]:
         config = self.get_external_browser_connection_parameters()
 
         config["CLIENT_STORE_TEMPORARY_CREDENTIAL"] = _get_env_variable(
@@ -96,7 +101,7 @@ def get_store_id_token_connection_parameters(self):
 
         return config
 
-    def get_okta_connection_parameters(self):
+    def get_okta_connection_parameters(self) -> dict[str, str]:
         config = self.basic_config.copy()
 
         config["user"] = _get_env_variable("SNOWFLAKE_AUTH_TEST_BROWSER_USER")
@@ -105,15 +110,17 @@ def get_okta_connection_parameters(self):
 
         return config
 
-    def get_oauth_connection_parameters(self, token: str):
+    def get_oauth_connection_parameters(self, token: str) -> dict[str, str]:
         config = self.basic_config.copy()
 
         config["user"] = _get_env_variable("SNOWFLAKE_AUTH_TEST_BROWSER_USER")
         config["authenticator"] = "OAUTH"
         config["token"] = token
         return config
 
-    def get_oauth_external_authorization_code_connection_parameters(self):
+    def get_oauth_external_authorization_code_connection_parameters(
+        self,
+    ) -> dict[str, Union[str, bool, int]]:
         config = self.basic_config.copy()
 
         config["authenticator"] = "OAUTH_AUTHORIZATION_CODE"
@@ -136,7 +143,9 @@ def get_oauth_external_authorization_code_connection_parameters(self):
 
         return config
 
-    def get_snowflake_authorization_code_connection_parameters(self):
+    def get_snowflake_authorization_code_connection_parameters(
+        self,
+    ) -> dict[str, Union[str, bool, int]]:
         config = self.basic_config.copy()
 
         config["authenticator"] = "OAUTH_AUTHORIZATION_CODE"
@@ -158,7 +167,9 @@ def get_snowflake_authorization_code_connection_parameters(self):
 
         return config
 
-    def get_snowflake_wildcard_external_authorization_code_connection_parameters(self):
+    def get_snowflake_wildcard_external_authorization_code_connection_parameters(
+        self,
+    ) -> dict[str, Union[str, bool, int]]:
         config = self.basic_config.copy()
 
         config["authenticator"] = "OAUTH_AUTHORIZATION_CODE"
@@ -177,7 +188,9 @@ def get_snowflake_wildcard_external_authorization_code_connection_parameters(sel
 
         return config
 
-    def get_oauth_external_client_credential_connection_parameters(self):
+    def get_oauth_external_client_credential_connection_parameters(
+        self,
+    ) -> dict[str, str]:
         config = self.basic_config.copy()
 
         config["authenticator"] = "OAUTH_CLIENT_CREDENTIALS"
@@ -195,3 +208,11 @@ def get_oauth_external_client_credential_connection_parameters(self):
         )
 
         return config
+
+    def get_pat_connection_parameters(self) -> dict[str, str]:
+        config = self.basic_config.copy()
+
+        config["authenticator"] = "PROGRAMMATIC_ACCESS_TOKEN"
+        config["user"] = _get_env_variable("SNOWFLAKE_AUTH_TEST_BROWSER_USER")
+
+        return config

test/auth/authorization_parameters.py

@@ -4,6 +4,7 @@
 import threading
 import webbrowser
 from enum import Enum
+from typing import Union
 
 import requests
 
@@ -125,3 +126,19 @@ def _provide_credentials(self, scenario: Scenario, login: str, password: str):
         except Exception as e:
             self.error_msg = e
             raise RuntimeError(e)
+
+    def connect_using_okta_connection_and_execute_custom_command(
+        self, command: str, return_token: bool = False
+    ) -> Union[bool, str]:
+        try:
+            logger.info("Setup PAT")
+            with snowflake.connector.connect(**self.configuration) as con:
+                result = con.cursor().execute(command)
+                token = result.fetchall()[0][1]
+        except Exception as e:
+            self.error_msg = e
+            logger.error(e)
+            return False
+        if return_token:
+            return token
+        return False

test/auth/authorization_test_helper.py

@@ -0,0 +1,82 @@
+from datetime import datetime
+from test.auth.authorization_parameters import (
+    AuthConnectionParameters,
+    get_pat_setup_command_variables,
+)
+from typing import Union
+
+import pytest
+from authorization_test_helper import AuthorizationTestHelper
+
+
+@pytest.mark.auth
+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 = get_pat_token(pat_command_variables)
+        connection_parameters["token"] = pat_command_variables["token"]
+        test_helper.connect_and_execute_simple_query()
+    finally:
+        remove_pat_token(pat_command_variables)
+    assert test_helper.get_error_msg() == "", "Error message should be empty"
+
+
+@pytest.mark.auth
+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 = get_pat_token(pat_command_variables)
+        connection_parameters["token"] = pat_command_variables["token"]
+        test_helper.connect_and_execute_simple_query()
+    finally:
+        remove_pat_token(pat_command_variables)
+
+    assert "Programmatic access token is invalid" in test_helper.get_error_msg()
+
+
+@pytest.mark.auth
+def test_authenticate_with_pat_invalid_token() -> None:
+    connection_parameters = AuthConnectionParameters().get_pat_connection_parameters()
+    connection_parameters["token"] = "invalidToken"
+    test_helper = AuthorizationTestHelper(connection_parameters)
+    test_helper.connect_and_execute_simple_query()
+    assert "Programmatic access token is invalid" in test_helper.get_error_msg()
+
+
+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"] = (
+        test_helper.connect_using_okta_connection_and_execute_custom_command(
+            command, True
+        )
+    )
+    return pat_command_variables
+
+
+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)
+    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")