Merge branch 'main' into SNOW-2043816-Kerberos-Proxy-Auth-Python

sfc-gh-fpawlowski · sfc-gh-fpawlowski · commit cdd70bbfe040 · 2025-06-23T08:25:57.000+02:00
diff --git a/.github/workflows/build_test.yml b/.github/workflows/build_test.yml
@@ -172,7 +172,7 @@ jobs:
         run: python -m pip install tox>=4
       - name: Run tests
        # To run a single test on GHA use the below command:
-       # run: python -m tox run -e `echo py${PYTHON_VERSION/\./}-single-ci | sed 's/ /,/g'`
+#        run: python -m tox run -e `echo py${PYTHON_VERSION/\./}-single-ci | sed 's/ /,/g'`
         run: python -m tox run -e `echo py${PYTHON_VERSION/\./}-{extras,unit,integ,pandas,sso}-ci | sed 's/ /,/g'`
 
         env:
@@ -181,7 +181,7 @@ jobs:
           PYTEST_ADDOPTS: --color=yes --tb=short
           TOX_PARALLEL_NO_SPINNER: 1
          # To specify the test name (in single test mode) pass this env variable:
-         # SINGLE_TEST_NAME: test/file/path::test_name
+#          SINGLE_TEST_NAME: test/path/filename.py::test_name
         shell: bash
       - name: Combine coverages
         run: python -m tox run -e coverage --skip-missing-interpreters false
diff --git a/DESCRIPTION.md b/DESCRIPTION.md
@@ -19,6 +19,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Fix `write_pandas` special characters usage in the location name.
   - Fix usage of `use_virtual_url` when building the location for gcs storage client.
   - Bind cryptography to <=44.0.3 to avoid issues with 45.0.0.
+  - Added support for Snowflake OAuth for local applications.
 
 - v3.15.0(Apr 29,2025)
   - Bumped up min boto and botocore version to 1.24.
diff --git a/src/snowflake/connector/auth/_oauth_base.py b/src/snowflake/connector/auth/_oauth_base.py
@@ -12,7 +12,13 @@
 from typing import TYPE_CHECKING, Any
 from urllib.error import HTTPError, URLError
 
-from ..errorcode import ER_FAILED_TO_REQUEST, ER_IDP_CONNECTION_ERROR
+from ..errorcode import (
+    ER_FAILED_TO_REQUEST,
+    ER_IDP_CONNECTION_ERROR,
+    ER_NO_CLIENT_ID,
+    ER_NO_CLIENT_SECRET,
+)
+from ..errors import Error, ProgrammingError
 from ..network import OAUTH_AUTHENTICATOR
 from ..secret_detector import SecretDetector
 from ..token_cache import TokenCache, TokenKey, TokenType
@@ -185,6 +191,33 @@ def assertion_content(self) -> str:
         """Returns the token."""
         return self._access_token or ""
 
+    @staticmethod
+    def _validate_client_credentials_present(
+        client_id: str, client_secret: str, connection: SnowflakeConnection
+    ) -> tuple[str, str]:
+        if client_id is None or client_id == "":
+            Error.errorhandler_wrapper(
+                connection,
+                None,
+                ProgrammingError,
+                {
+                    "msg": "Oauth code flow requirement 'client_id' is empty",
+                    "errno": ER_NO_CLIENT_ID,
+                },
+            )
+        if client_secret is None or client_secret == "":
+            Error.errorhandler_wrapper(
+                connection,
+                None,
+                ProgrammingError,
+                {
+                    "msg": "Oauth code flow requirement 'client_secret' is empty",
+                    "errno": ER_NO_CLIENT_SECRET,
+                },
+            )
+
+        return client_id, client_secret
+
     def reauthenticate(
         self,
         *,
diff --git a/src/snowflake/connector/auth/oauth_code.py b/src/snowflake/connector/auth/oauth_code.py
@@ -18,11 +18,13 @@
 from ..compat import parse_qs, urlparse, urlsplit
 from ..constants import OAUTH_TYPE_AUTHORIZATION_CODE
 from ..errorcode import (
+    ER_INVALID_VALUE,
     ER_OAUTH_CALLBACK_ERROR,
     ER_OAUTH_SERVER_TIMEOUT,
     ER_OAUTH_STATE_CHANGED,
     ER_UNABLE_TO_OPEN_BROWSER,
 )
+from ..errors import Error, ProgrammingError
 from ..token_cache import TokenCache
 from ._http_server import AuthHttpServer
 from ._oauth_base import AuthByOAuthBase
@@ -45,6 +47,8 @@ def _get_query_params(
 class AuthByOauthCode(AuthByOAuthBase):
     """Authenticates user by OAuth code flow."""
 
+    _LOCAL_APPLICATION_CLIENT_CREDENTIALS = "LOCAL_APPLICATION"
+
     def __init__(
         self,
         application: str,
@@ -54,13 +58,27 @@ def __init__(
         token_request_url: str,
         redirect_uri: str,
         scope: str,
+        host: str,
         pkce_enabled: bool = True,
         token_cache: TokenCache | None = None,
         refresh_token_enabled: bool = False,
         external_browser_timeout: int | None = None,
         enable_single_use_refresh_tokens: bool = False,
+        connection: SnowflakeConnection | None = None,
         **kwargs,
     ) -> None:
+        authentication_url, redirect_uri = self._validate_oauth_code_uris(
+            authentication_url, redirect_uri, connection
+        )
+        client_id, client_secret = self._validate_client_credentials_with_defaults(
+            client_id,
+            client_secret,
+            authentication_url,
+            token_request_url,
+            host,
+            connection,
+        )
+
         super().__init__(
             client_id=client_id,
             client_secret=client_secret,
@@ -385,3 +403,77 @@ def _parse_authorization_redirected_request(
                 },
             )
         return parsed.get("code", [None])[0], parsed.get("state", [None])[0]
+
+    @staticmethod
+    def _is_snowflake_as_idp(
+        authentication_url: str, token_request_url: str, host: str
+    ) -> bool:
+        return (authentication_url == "" or host in authentication_url) and (
+            token_request_url == "" or host in token_request_url
+        )
+
+    def _eligible_for_default_client_credentials(
+        self,
+        client_id: str,
+        client_secret: str,
+        authorization_url: str,
+        token_request_url: str,
+        host: str,
+    ) -> bool:
+        return (
+            (client_id == "" or client_secret is None)
+            and (client_secret == "" or client_secret is None)
+            and self.__class__._is_snowflake_as_idp(
+                authorization_url, token_request_url, host
+            )
+        )
+
+    def _validate_client_credentials_with_defaults(
+        self,
+        client_id: str,
+        client_secret: str,
+        authorization_url: str,
+        token_request_url: str,
+        host: str,
+        connection: SnowflakeConnection,
+    ) -> tuple[str, str] | None:
+        if self._eligible_for_default_client_credentials(
+            client_id, client_secret, authorization_url, token_request_url, host
+        ):
+            return (
+                self.__class__._LOCAL_APPLICATION_CLIENT_CREDENTIALS,
+                self.__class__._LOCAL_APPLICATION_CLIENT_CREDENTIALS,
+            )
+        else:
+            self._validate_client_credentials_present(
+                client_id, client_secret, connection
+            )
+            return client_id, client_secret
+
+    @staticmethod
+    def _validate_oauth_code_uris(
+        authorization_url: str, redirect_uri: str, connection: SnowflakeConnection
+    ) -> tuple[str, str]:
+        if authorization_url and not authorization_url.startswith("https://"):
+            Error.errorhandler_wrapper(
+                connection,
+                None,
+                ProgrammingError,
+                {
+                    "msg": "OAuth supports only authorization urls that use 'https' scheme",
+                    "errno": ER_INVALID_VALUE,
+                },
+            )
+        if redirect_uri and not (
+            redirect_uri.startswith("http://") or redirect_uri.startswith("https://")
+        ):
+            Error.errorhandler_wrapper(
+                connection,
+                None,
+                ProgrammingError,
+                {
+                    "msg": "OAuth supports only authorization urls that use 'http(s)' scheme",
+                    "errno": ER_INVALID_VALUE,
+                },
+            )
+        return authorization_url, redirect_uri
diff --git a/src/snowflake/connector/auth/oauth_credentials.py b/src/snowflake/connector/auth/oauth_credentials.py
@@ -29,8 +29,10 @@ def __init__(
         scope: str,
         token_cache: TokenCache | None = None,
         refresh_token_enabled: bool = False,
+        connection: SnowflakeConnection | None = None,
         **kwargs,
     ) -> None:
+        self._validate_client_credentials_present(client_id, client_secret, connection)
         super().__init__(
             client_id=client_id,
             client_secret=client_secret,
diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py
@@ -101,8 +101,6 @@
     ER_INVALID_VALUE,
     ER_INVALID_WIF_SETTINGS,
     ER_NO_ACCOUNT_NAME,
-    ER_NO_CLIENT_ID,
-    ER_NO_CLIENT_SECRET,
     ER_NO_NUMPY,
     ER_NO_PASSWORD,
     ER_NO_USER,
@@ -1278,14 +1276,14 @@ def __open_connection(self):
                     backoff_generator=self._backoff_generator,
                 )
             elif self._authenticator == OAUTH_AUTHORIZATION_CODE:
-                self._check_oauth_parameters()
                 if self._role and (self._oauth_scope == ""):
                     # if role is known then let's inject it into scope
                     self._oauth_scope = _OAUTH_DEFAULT_SCOPE.format(role=self._role)
                 self.auth_class = AuthByOauthCode(
                     application=self.application,
                     client_id=self._oauth_client_id,
                     client_secret=self._oauth_client_secret,
+                    host=self.host,
                     authentication_url=self._oauth_authorization_url.format(
                         host=self.host, port=self.port
                     ),
@@ -1305,7 +1303,6 @@ def __open_connection(self):
                     enable_single_use_refresh_tokens=self._oauth_enable_single_use_refresh_tokens,
                 )
             elif self._authenticator == OAUTH_CLIENT_CREDENTIALS:
-                self._check_oauth_parameters()
                 if self._role and (self._oauth_scope == ""):
                     # if role is known then let's inject it into scope
                     self._oauth_scope = _OAUTH_DEFAULT_SCOPE.format(role=self._role)
@@ -1323,6 +1320,7 @@ def __open_connection(self):
                         else None
                     ),
                     refresh_token_enabled=self._oauth_enable_refresh_tokens,
+                    connection=self,
                 )
             elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
                 self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (
@@ -2338,54 +2336,6 @@ def _check_experimental_authentication_flag(self) -> None:
                 },
             )
 
-    def _check_oauth_parameters(self) -> None:
-        if self._oauth_client_id is None:
-            Error.errorhandler_wrapper(
-                self,
-                None,
-                ProgrammingError,
-                {
-                    "msg": "Oauth code flow requirement 'client_id' is empty",
-                    "errno": ER_NO_CLIENT_ID,
-                },
-            )
-        if self._oauth_client_secret is None:
-            Error.errorhandler_wrapper(
-                self,
-                None,
-                ProgrammingError,
-                {
-                    "msg": "Oauth code flow requirement 'client_secret' is empty",
-                    "errno": ER_NO_CLIENT_SECRET,
-                },
-            )
-        if (
-            self._oauth_authorization_url
-            and not self._oauth_authorization_url.startswith("https://")
-        ):
-            Error.errorhandler_wrapper(
-                self,
-                None,
-                ProgrammingError,
-                {
-                    "msg": "OAuth supports only authorization urls that use 'https' scheme",
-                    "errno": ER_INVALID_VALUE,
-                },
-            )
-        if self._oauth_redirect_uri and not (
-            self._oauth_redirect_uri.startswith("http://")
-            or self._oauth_redirect_uri.startswith("https://")
-        ):
-            Error.errorhandler_wrapper(
-                self,
-                None,
-                ProgrammingError,
-                {
-                    "msg": "OAuth supports only authorization urls that use 'http(s)' scheme",
-                    "errno": ER_INVALID_VALUE,
-                },
-            )
-
     @staticmethod
     def _detect_application() -> None | str:
         if ENV_VAR_PARTNER in os.environ.keys():
diff --git a/test/data/wiremock/mappings/auth/oauth/authorization_code/external_idp_custom_urls_local_application.json b/test/data/wiremock/mappings/auth/oauth/authorization_code/external_idp_custom_urls_local_application.json
@@ -0,0 +1,77 @@
+{
+  "mappings": [
+    {
+      "scenarioName": "Custom urls OAuth authorization code flow local application",
+      "requiredScenarioState": "Started",
+      "newScenarioState": "Authorized",
+      "request": {
+        "urlPathPattern": "/authorization",
+        "method": "GET",
+        "queryParameters": {
+          "response_type": {
+            "equalTo": "code"
+          },
+          "scope": {
+            "equalTo": "session:role:ANALYST"
+          },
+          "code_challenge_method": {
+            "equalTo": "S256"
+          },
+          "redirect_uri": {
+            "equalTo": "http://localhost:8009/snowflake/oauth-redirect"
+          },
+          "code_challenge": {
+            "matches": ".*"
+          },
+          "state": {
+            "matches": ".*"
+          },
+          "client_id": {
+            "equalTo": "LOCAL_APPLICATION"
+          }
+        }
+      },
+      "response": {
+        "status": 302,
+        "headers": {
+          "Location": "http://localhost:8009/snowflake/oauth-redirect?code=123&state=abc123"
+        }
+      }
+    },
+    {
+      "scenarioName": "Custom urls OAuth authorization code flow local application",
+      "requiredScenarioState": "Authorized",
+      "newScenarioState": "Acquired access token",
+      "request": {
+        "urlPathPattern": "/tokenrequest.*",
+        "method": "POST",
+        "headers": {
+          "Authorization": {
+            "contains": "Basic"
+          },
+          "Content-Type": {
+            "contains": "application/x-www-form-urlencoded; charset=UTF-8"
+          }
+        },
+        "bodyPatterns": [
+          {
+            "contains": "grant_type=authorization_code&code=123&redirect_uri=http%3A%2F%2Flocalhost%3A8009%2Fsnowflake%2Foauth-redirect&code_verifier="
+          }
+        ]
+      },
+      "response": {
+        "status": 200,
+        "jsonBody": {
+          "access_token": "access-token-123",
+          "refresh_token": "123",
+          "token_type": "Bearer",
+          "username": "user",
+          "scope": "refresh_token session:role:ANALYST",
+          "expires_in": 600,
+          "refresh_token_expires_in": 86399,
+          "idpInitiated": false
+        }
+      }
+    }
+  ]
+}
diff --git a/test/unit/test_auth_oauth_auth_code.py b/test/unit/test_auth_oauth_auth_code.py
diff --git a/test/unit/test_oauth_token.py b/test/unit/test_oauth_token.py
