Add PAT to async authentication (#2122)

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_connection.py

@@ -61,6 +61,7 @@
     EXTERNAL_BROWSER_AUTHENTICATOR,
     KEY_PAIR_AUTHENTICATOR,
     OAUTH_AUTHENTICATOR,
+    PROGRAMMATIC_ACCESS_TOKEN,
     REQUEST_ID,
     USR_PWD_MFA_AUTHENTICATOR,
     ReauthenticationRequest,
@@ -82,6 +83,7 @@
     AuthByKeyPair,
     AuthByOAuth,
     AuthByOkta,
+    AuthByPAT,
     AuthByPlugin,
     AuthByUsrPwdMfa,
     AuthByWebBrowser,
@@ -298,6 +300,8 @@ async 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)
             elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
                 self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (
                     self._client_request_mfa_token if IS_LINUX else True

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

@@ -12,6 +12,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
 
@@ -24,13 +25,15 @@
         AuthByUsrPwdMfa,
         AuthByWebBrowser,
         AuthByIdToken,
+        AuthByPAT,
     )
 )
 
 __all__ = [
     "AuthByPlugin",
     "AuthByDefault",
     "AuthByKeyPair",
+    "AuthByPAT",
     "AuthByOAuth",
     "AuthByOkta",
     "AuthByUsrPwdMfa",

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

@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+from typing import Any
+
+from ...auth.pat import AuthByPAT as AuthByPATSync
+from ._by_plugin import AuthByPlugin as AuthByPluginAsync
+
+
+class AuthByPAT(AuthByPluginAsync, AuthByPATSync):
+    def __init__(self, pat_token: str, **kwargs) -> None:
+        """Initializes an instance with a PAT Token."""
+        AuthByPATSync.__init__(self, pat_token, **kwargs)
+
+    async def reset_secrets(self) -> None:
+        AuthByPATSync.reset_secrets(self)
+
+    async def prepare(self, **kwargs: Any) -> None:
+        AuthByPATSync.prepare(self, **kwargs)
+
+    async def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
+        return AuthByPATSync.reauthenticate(self, **kwargs)
+
+    async def update_body(self, body: dict[Any, Any]) -> None:
+        AuthByPATSync.update_body(self, body)

test/unit/aio/test_auth_pat_async.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+from snowflake.connector.aio.auth import AuthByPAT
+from snowflake.connector.auth.by_plugin import AuthType
+from snowflake.connector.network import PROGRAMMATIC_ACCESS_TOKEN
+
+
+async def test_auth_pat():
+    """Simple test if AuthByPAT class."""
+    token = "patToken"
+    auth = AuthByPAT(token)
+    assert auth.type_ == AuthType.PAT
+    assert auth.assertion_content == token
+    body = {"data": {}}
+    await auth.update_body(body)
+    assert body["data"]["TOKEN"] == token, body
+    assert body["data"]["AUTHENTICATOR"] == PROGRAMMATIC_ACCESS_TOKEN, body
+
+    await auth.reset_secrets()
+    assert auth.assertion_content is None
+
+
+async def test_auth_pat_reauthenticate():
+    """Test PAT reauthenticate."""
+    token = "patToken"
+    auth = AuthByPAT(token)
+    result = await auth.reauthenticate()
+    assert result == {"success": False}
+
+
+async def test_pat_authenticator_creates_auth_by_pat(monkeypatch):
+    """Test that using PROGRAMMATIC_ACCESS_TOKEN authenticator creates AuthByPAT instance."""
+    import snowflake.connector.aio
+    from snowflake.connector.aio._network import SnowflakeRestful
+
+    # Mock the network request - this prevents actual network calls and connection errors
+    async def mock_post_request(request, url, headers, json_body, **kwargs):
+        return {
+            "success": True,
+            "message": None,
+            "data": {
+                "token": "TOKEN",
+                "masterToken": "MASTER_TOKEN",
+                "idToken": None,
+                "parameters": [{"name": "SERVICE_NAME", "value": "FAKE_SERVICE_NAME"}],
+            },
+        }
+
+    # Apply the mock using monkeypatch
+    monkeypatch.setattr(SnowflakeRestful, "_post_request", mock_post_request)
+
+    # Create connection with PAT authenticator
+    conn = snowflake.connector.aio.SnowflakeConnection(
+        user="user",
+        account="account",
+        database="TESTDB",
+        warehouse="TESTWH",
+        authenticator=PROGRAMMATIC_ACCESS_TOKEN,
+        token="test_pat_token",
+    )
+
+    await conn.connect()
+
+    # Verify that the auth_class is an instance of AuthByPAT
+    assert isinstance(conn.auth_class, AuthByPAT)
+
+    await conn.close()