SNOW-1825473 adding pat authentication integration (#2122)

DESCRIPTION.md

@@ -10,6 +10,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 - v3.12.5(TBD)
   - Added a feature to limit the sizes of IO-bound ThreadPoolExecutors during PUT and GET commands.
+  - Adding support for the new PAT authentication method.
 
 - v3.12.4(December 3,2024)
   - Fixed a bug where multipart uploads to Azure would be missing their MD5 hashes.

src/snowflake/connector/auth/__init__.py

@@ -11,6 +11,7 @@
 from .keypair import AuthByKeyPair
 from .oauth import AuthByOAuth
 from .okta import AuthByOkta
+from .pat import AuthByPAT
 from .usrpwdmfa import AuthByUsrPwdMfa
 from .webbrowser import AuthByWebBrowser
 
@@ -23,13 +24,15 @@
         AuthByUsrPwdMfa,
         AuthByWebBrowser,
         AuthByIdToken,
+        AuthByPAT,
     )
 )
 
 __all__ = [
     "AuthByPlugin",
     "AuthByDefault",
     "AuthByKeyPair",
+    "AuthByPAT",
     "AuthByOAuth",
     "AuthByOkta",
     "AuthByUsrPwdMfa",

src/snowflake/connector/auth/by_plugin.py

@@ -54,6 +54,7 @@ class AuthType(Enum):
     ID_TOKEN = "ID_TOKEN"
     USR_PWD_MFA = "USERNAME_PASSWORD_MFA"
     OKTA = "OKTA"
+    PAT = "PROGRAMMATIC_ACCESS_TOKEN'"
 
 
 class AuthByPlugin(ABC):

src/snowflake/connector/auth/oauth.py

@@ -22,7 +22,7 @@ def type_(self) -> AuthType:
         return AuthType.OAUTH
 
     @property
-    def assertion_content(self) -> str:
+    def assertion_content(self) -> str | None:
         """Returns the token."""
         return self._oauth_token
 

src/snowflake/connector/auth/pat.py

@@ -0,0 +1,43 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+import typing
+
+from snowflake.connector.network import PROGRAMMATIC_ACCESS_TOKEN
+
+from .by_plugin import AuthByPlugin, AuthType
+
+
+class AuthByPAT(AuthByPlugin):
+
+    def __init__(self, pat_token: str, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._pat_token: str | None = pat_token
+
+    def type_(self) -> AuthType:
+        return AuthType.PAT
+
+    def reset_secrets(self) -> None:
+        self._pat_token = None
+
+    def update_body(self, body: dict[typing.Any, typing.Any]) -> None:
+        body["data"]["AUTHENTICATOR"] = PROGRAMMATIC_ACCESS_TOKEN
+        body["data"]["TOKEN"] = self._pat_token
+
+    def prepare(
+        self,
+        **kwargs: typing.Any,
+    ) -> None:
+        """Nothing to do here, token should be obtained outside the driver."""
+        pass
+
+    def reauthenticate(self, **kwargs: typing.Any) -> dict[str, bool]:
+        return {"success": False}
+
+    @property
+    def assertion_content(self) -> str | None:
+        """Returns the token."""
+        return self._pat_token

src/snowflake/connector/connection.py

@@ -41,6 +41,7 @@
     AuthByKeyPair,
     AuthByOAuth,
     AuthByOkta,
+    AuthByPAT,
     AuthByPlugin,
     AuthByUsrPwdMfa,
     AuthByWebBrowser,
@@ -99,6 +100,7 @@
     EXTERNAL_BROWSER_AUTHENTICATOR,
     KEY_PAIR_AUTHENTICATOR,
     OAUTH_AUTHENTICATOR,
+    PROGRAMMATIC_ACCESS_TOKEN,
     REQUEST_ID,
     USR_PWD_MFA_AUTHENTICATOR,
     ReauthenticationRequest,
@@ -185,7 +187,11 @@ def _get_private_bytes_from_file(
     "private_key": (None, (type(None), bytes, RSAPrivateKey)),
     "private_key_file": (None, (type(None), str)),
     "private_key_file_pwd": (None, (type(None), str, bytes)),
-    "token": (None, (type(None), str)),  # OAuth or JWT Token
+    "token": (None, (type(None), str)),  # OAuth/JWT/PAT Token
+    "token_file_path": (
+        None,
+        (type(None), str, bytes),
+    ),  # OAuth/JWT/PAT Token file path
     "authenticator": (DEFAULT_AUTHENTICATOR, (type(None), str)),
     "mfa_callback": (None, (type(None), Callable)),
     "password_callback": (None, (type(None), Callable)),
@@ -1097,6 +1103,8 @@ def __open_connection(self):
                     timeout=self.login_timeout,
                     backoff_generator=self._backoff_generator,
                 )
+            elif self._authenticator == PROGRAMMATIC_ACCESS_TOKEN:
+                self.auth_class = AuthByPAT(self._token)
             else:
                 # okta URL, e.g., https://<account>.okta.com/
                 self.auth_class = AuthByOkta(
@@ -1245,11 +1253,12 @@ def __config(self, **kwargs):
             if (
                 self.auth_class is None
                 and self._authenticator
-                not in [
+                not in (
                     EXTERNAL_BROWSER_AUTHENTICATOR,
                     OAUTH_AUTHENTICATOR,
                     KEY_PAIR_AUTHENTICATOR,
-                ]
+                    PROGRAMMATIC_ACCESS_TOKEN,
+                )
                 and not self._password
             ):
                 Error.errorhandler_wrapper(

src/snowflake/connector/errors.py

@@ -12,7 +12,6 @@
 from logging import getLogger
 from typing import TYPE_CHECKING, Any
 
-from .compat import BASE_EXCEPTION_CLASS
 from .secret_detector import SecretDetector
 from .telemetry import TelemetryData, TelemetryField
 from .time_util import get_time_millis
@@ -28,7 +27,7 @@
 RE_FORMATTED_ERROR = re.compile(r"^(\d{6,})(?: \((\S+)\))?:")
 
 
-class Error(BASE_EXCEPTION_CLASS):
+class Error(Exception):
     """Base Snowflake exception class."""
 
     def __init__(
@@ -361,7 +360,7 @@ def errorhandler_make_exception(
         return error_class(error_value)
 
 
-class _Warning(BASE_EXCEPTION_CLASS):
+class _Warning(Exception):
     """Exception for important warnings."""
 
     pass

src/snowflake/connector/network.py

@@ -187,6 +187,7 @@
 OAUTH_AUTHENTICATOR = "OAUTH"
 ID_TOKEN_AUTHENTICATOR = "ID_TOKEN"
 USR_PWD_MFA_AUTHENTICATOR = "USERNAME_PASSWORD_MFA"
+PROGRAMMATIC_ACCESS_TOKEN = "PROGRAMMATIC_ACCESS_TOKEN"
 
 
 def is_retryable_http_code(code: int) -> bool: