snowflake-connector-python 3.16.0__cp313-cp313-win_amd64.whl

snowflake/connector/auth/by_plugin.py

@@ -0,0 +1,219 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+"""This module implements the base class for authenticator classes.
+
+Note:
+ **kwargs are added to most functions so that child classes can safely ignore extra in
+  arguments in case of a caller API change and named arguments are enforced to prevent
+  issues with argument being sent in out of order.
+"""
+
+import logging
+import time
+from abc import ABC, abstractmethod
+from enum import Enum, unique
+from os import getenv
+from typing import TYPE_CHECKING, Any, Iterator
+
+from ..errorcode import ER_FAILED_TO_CONNECT_TO_DB
+from ..errors import DatabaseError, Error, OperationalError
+from ..sqlstate import SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED
+from ..time_util import TimeoutBackoffCtx
+
+if TYPE_CHECKING:
+    from .. import SnowflakeConnection
+
+logger = logging.getLogger(__name__)
+
+"""
+Default value for max retry is 1 because
+Python requests module already tries twice
+by default. Unlike JWT where we need to refresh
+token every 10 seconds, general authenticators
+wait for 60 seconds before connection timeout
+per attempt totaling a 240 sec wait time for a non
+JWT based authenticator which is more than enough.
+This can be changed ofcourse using MAX_CNXN_RETRY_ATTEMPTS
+env variable.
+"""
+DEFAULT_MAX_CON_RETRY_ATTEMPTS = 1
+DEFAULT_AUTH_CLASS_TIMEOUT = 120
+
+
+@unique
+class AuthType(Enum):
+    DEFAULT = "SNOWFLAKE"  # default authenticator name
+    EXTERNAL_BROWSER = "EXTERNALBROWSER"
+    KEY_PAIR = "SNOWFLAKE_JWT"
+    OAUTH = "OAUTH"
+    ID_TOKEN = "ID_TOKEN"
+    USR_PWD_MFA = "USERNAME_PASSWORD_MFA"
+    OKTA = "OKTA"
+    PAT = "PROGRAMMATIC_ACCESS_TOKEN'"
+    NO_AUTH = "NO_AUTH"
+    WORKLOAD_IDENTITY = "WORKLOAD_IDENTITY"
+    PAT_WITH_EXTERNAL_SESSION = "PAT_WITH_EXTERNAL_SESSION"
+
+
+class AuthByPlugin(ABC):
+    """External Authenticator interface."""
+
+    def __init__(
+        self,
+        timeout: int | None = None,
+        backoff_generator: Iterator | None = None,
+        **kwargs,
+    ) -> None:
+        self.consent_cache_id_token = False
+
+        self._retry_ctx = TimeoutBackoffCtx(
+            timeout=timeout if timeout is not None else DEFAULT_AUTH_CLASS_TIMEOUT,
+            max_retry_attempts=kwargs.get(
+                "max_retry_attempts",
+                int(getenv("MAX_CON_RETRY_ATTEMPTS", DEFAULT_MAX_CON_RETRY_ATTEMPTS)),
+            ),
+            backoff_generator=backoff_generator,
+        )
+
+        # some authenticators may want to override socket level timeout
+        # for example, AuthByKeyPair will set this to ensure JWT tokens are refreshed in time
+        # if not None, this will override socket_timeout specified in connection
+        self._socket_timeout = None
+
+    @property
+    def timeout(self) -> int:
+        """The timeout of _retry_ctx is guaranteed not to be None during AuthByPlugin initialization"""
+        return self._retry_ctx.timeout
+
+    @timeout.setter
+    def timeout(self) -> None:
+        logger.warning(
+            "Attempting to mutate timeout of AuthByPlugin. Create a new instance with desired parameters instead."
+        )
+
+    @property
+    @abstractmethod
+    def type_(self) -> AuthType:
+        """Return the Snowflake friendly name of auth class."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def assertion_content(self) -> str:
+        """Return a safe version of the information used to authenticate with Snowflake.
+
+        This is used for logging, useful for printing temporary tokens, but make sure to
+        mask secrets.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def prepare(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        authenticator: str,
+        service_name: str | None,
+        account: str,
+        user: str,
+        password: str | None,
+        **kwargs: Any,
+    ) -> str | None:
+        """Prepare for authentication.
+
+        This function is useful for situations where we need to reach out to a 3rd-party
+        service before authenticating with Snowflake.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def update_body(self, body: dict[Any, Any]) -> None:
+        """Update the body of the authentication request."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def reset_secrets(self) -> None:
+        """Reset secret members."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def reauthenticate(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Re-perform authentication.
+
+        The difference between this and authentication is that secrets will be removed
+        from memory by the time this gets called.
+        """
+        raise NotImplementedError
+
+    def _handle_failure(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        ret: dict[Any, Any],
+        **kwargs: Any,
+    ) -> None:
+        """Handles a failure when an issue happens while connecting to Snowflake.
+
+        If the user returns from this function execution will continue. The argument
+        data can be manipulated from within this function and so recovery is possible
+        from here.
+        """
+        Error.errorhandler_wrapper(
+            conn,
+            None,
+            DatabaseError,
+            {
+                "msg": "Failed to connect to DB: {host}:{port}, {message}".format(
+                    host=conn._rest._host,
+                    port=conn._rest._port,
+                    message=ret["message"],
+                ),
+                "errno": int(ret.get("code", -1)),
+                "sqlstate": SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
+            },
+        )
+
+    def handle_timeout(
+        self,
+        *,
+        authenticator: str,
+        service_name: str | None,
+        account: str,
+        user: str,
+        password: str,
+        **kwargs: Any,
+    ) -> None:
+        """Default timeout handler.
+
+        This will trigger if the authenticator
+        hasn't implemented one. By default we retry on timeouts and use
+        jitter to deduce the time to sleep before retrying. The sleep
+        time ranges between 1 and 16 seconds.
+        """
+
+        # Some authenticators may not want to delete the parameters to this function
+        # Currently, the only authenticator where this is the case is AuthByKeyPair
+        if kwargs.pop("delete_params", True):
+            del authenticator, service_name, account, user, password
+
+        logger.debug("Default timeout handler invoked for authenticator")
+        if not self._retry_ctx.should_retry:
+            error = OperationalError(
+                msg=f"Could not connect to Snowflake backend after {self._retry_ctx.current_retry_count + 1} attempt(s)."
+                "Aborting",
+                errno=ER_FAILED_TO_CONNECT_TO_DB,
+            )
+            raise error
+        else:
+            logger.debug(
+                f"Hit connection timeout, attempt number {self._retry_ctx.current_retry_count + 1}."
+                " Will retry in a bit..."
+            )
+            time.sleep(float(self._retry_ctx.current_sleep_time))
+            self._retry_ctx.increment()

snowflake/connector/auth/default.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+from typing import Any
+
+from .by_plugin import AuthByPlugin, AuthType
+
+
+class AuthByDefault(AuthByPlugin):
+    """Default username and password authenticator."""
+
+    @property
+    def type_(self) -> AuthType:
+        return AuthType.DEFAULT
+
+    @property
+    def assertion_content(self) -> str:
+        return "*********"
+
+    def __init__(self, password: str, **kwargs) -> None:
+        """Initializes an instance with a password."""
+        super().__init__(**kwargs)
+        self._password: str | None = password
+
+    def reset_secrets(self) -> None:
+        self._password = None
+
+    def prepare(self, **kwargs: Any) -> None:
+        pass
+
+    def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
+        return {"success": False}
+
+    def update_body(self, body: dict[Any, Any]) -> None:
+        """Sets the password if available."""
+        body["data"]["PASSWORD"] = self._password

snowflake/connector/auth/idtoken.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from ..network import ID_TOKEN_AUTHENTICATOR
+from .by_plugin import AuthByPlugin, AuthType
+from .webbrowser import AuthByWebBrowser
+
+if TYPE_CHECKING:
+    from ..connection import SnowflakeConnection
+
+
+class AuthByIdToken(AuthByPlugin):
+    """Internal IdToken Based Authentication.
+
+    Works by accepting an id_toke and use that to authenticate. Only be used when users are using EXTERNAL_BROWSER_AUTHENTICATOR
+    """
+
+    @property
+    def type_(self) -> AuthType:
+        return AuthType.ID_TOKEN
+
+    @property
+    def assertion_content(self) -> str:
+        return self._id_token
+
+    def __init__(
+        self,
+        id_token: str,
+        application: str,
+        protocol: str | None,
+        host: str | None,
+        port: str | None,
+        **kwargs,
+    ) -> None:
+        """Initialized an instance with an IdToken."""
+        super().__init__(**kwargs)
+        self._id_token: str | None = id_token
+        self._application = application
+        self._protocol = protocol
+        self._host = host
+        self._port = port
+
+    def reset_secrets(self) -> None:
+        self._id_token = None
+
+    def prepare(self, **kwargs: Any) -> None:
+        pass
+
+    def reauthenticate(
+        self,
+        *,
+        conn: SnowflakeConnection,
+        **kwargs: Any,
+    ) -> dict[str, bool]:
+        conn.auth_class = AuthByWebBrowser(
+            application=self._application,
+            protocol=self._protocol,
+            host=self._host,
+            port=self._port,
+            timeout=conn.login_timeout,
+            backoff_generator=conn._backoff_generator,
+        )
+        conn._authenticate(conn.auth_class)
+        conn._auth_class.reset_secrets()
+        return {"success": True}
+
+    def update_body(self, body: dict[Any, Any]) -> None:
+        """Idtoken needs the authenticator and token attributes set."""
+        body["data"]["AUTHENTICATOR"] = ID_TOKEN_AUTHENTICATOR
+        body["data"]["TOKEN"] = self._id_token

snowflake/connector/auth/keypair.py

@@ -0,0 +1,222 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import base64
+import hashlib
+import os
+from datetime import datetime, timedelta, timezone
+from logging import getLogger
+from typing import Any
+
+import jwt
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
+from cryptography.hazmat.primitives.serialization import (
+    Encoding,
+    PublicFormat,
+    load_der_private_key,
+)
+
+from ..errorcode import ER_CONNECTION_TIMEOUT, ER_INVALID_PRIVATE_KEY
+from ..errors import OperationalError, ProgrammingError
+from ..network import KEY_PAIR_AUTHENTICATOR
+from .by_plugin import AuthByPlugin, AuthType
+
+logger = getLogger(__name__)
+
+
+class AuthByKeyPair(AuthByPlugin):
+    """Key pair based authentication."""
+
+    ALGORITHM = "RS256"
+    ISSUER = "iss"
+    SUBJECT = "sub"
+    EXPIRE_TIME = "exp"
+    ISSUE_TIME = "iat"
+    LIFETIME = 60
+    DEFAULT_JWT_RETRY_ATTEMPTS = 10
+    DEFAULT_JWT_CNXN_WAIT_TIME = 10
+
+    def __init__(
+        self,
+        private_key: bytes | str | RSAPrivateKey,
+        lifetime_in_seconds: int = LIFETIME,
+        **kwargs,
+    ) -> None:
+        """Inits AuthByKeyPair class with private key.
+
+        Args:
+            private_key: a byte array of der formats of private key, or an
+                object that implements the `RSAPrivateKey` interface.
+            lifetime_in_seconds: number of seconds the JWT token will be valid
+        """
+        super().__init__(
+            max_retry_attempts=int(
+                os.getenv(
+                    "JWT_CNXN_RETRY_ATTEMPTS", AuthByKeyPair.DEFAULT_JWT_RETRY_ATTEMPTS
+                )
+            ),
+            **kwargs,
+        )
+
+        # set internal socket timeout override
+        self._socket_timeout = int(
+            timedelta(
+                seconds=int(
+                    os.getenv(
+                        "JWT_CNXN_WAIT_TIME",
+                        AuthByKeyPair.DEFAULT_JWT_CNXN_WAIT_TIME,
+                    )
+                )
+            ).total_seconds()
+        )
+
+        self._private_key: bytes | str | RSAPrivateKey | None = private_key
+        self._jwt_token = ""
+        self._jwt_token_exp = 0
+        self._lifetime = timedelta(
+            seconds=int(os.getenv("JWT_LIFETIME_IN_SECONDS", lifetime_in_seconds))
+        )
+
+    def reset_secrets(self) -> None:
+        self._private_key = None
+
+    @property
+    def type_(self) -> AuthType:
+        return AuthType.KEY_PAIR
+
+    def prepare(
+        self,
+        *,
+        account: str,
+        user: str,
+        **kwargs: Any,
+    ) -> str:
+        if ".global" in account:
+            account = account.partition("-")[0]
+        else:
+            account = account.partition(".")[0]
+        account = account.upper()
+        user = user.upper()
+
+        now = datetime.now(timezone.utc).replace(tzinfo=None)
+
+        if isinstance(self._private_key, str):
+            try:
+                self._private_key = base64.b64decode(self._private_key)
+            except Exception as e:
+                raise ProgrammingError(
+                    msg=f"Failed to decode private key: {e}\nPlease provide a valid "
+                    "unencrypted rsa private key in base64-encoded DER format as a "
+                    "str object",
+                    errno=ER_INVALID_PRIVATE_KEY,
+                )
+
+        if isinstance(self._private_key, bytes):
+            try:
+                private_key = load_der_private_key(
+                    data=self._private_key,
+                    password=None,
+                    backend=default_backend(),
+                )
+            except Exception as e:
+                raise ProgrammingError(
+                    msg=f"Failed to load private key: {e}\nPlease provide a valid "
+                    "unencrypted rsa private key in DER format as bytes object",
+                    errno=ER_INVALID_PRIVATE_KEY,
+                )
+
+            if not isinstance(private_key, RSAPrivateKey):
+                raise ProgrammingError(
+                    msg=f"Private key type ({private_key.__class__.__name__}) not supported."
+                    "\nPlease provide a valid rsa private key in DER format as bytes "
+                    "object",
+                    errno=ER_INVALID_PRIVATE_KEY,
+                )
+        elif isinstance(self._private_key, RSAPrivateKey):
+            private_key = self._private_key
+        else:
+            raise TypeError(
+                f"Expected bytes or RSAPrivateKey, got {type(self._private_key)}"
+            )
+
+        public_key_fp = self.calculate_public_key_fingerprint(private_key)
+
+        self._jwt_token_exp = now + self._lifetime
+        payload = {
+            self.ISSUER: f"{account}.{user}.{public_key_fp}",
+            self.SUBJECT: f"{account}.{user}",
+            self.ISSUE_TIME: now,
+            self.EXPIRE_TIME: self._jwt_token_exp,
+        }
+
+        _jwt_token = jwt.encode(payload, private_key, algorithm=self.ALGORITHM)
+
+        # jwt.encode() returns bytes in pyjwt 1.x and a string
+        # in pyjwt 2.x
+        if isinstance(_jwt_token, bytes):
+            self._jwt_token = _jwt_token.decode("utf-8")
+        else:
+            self._jwt_token = _jwt_token
+
+        return self._jwt_token
+
+    def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
+        return {"success": False}
+
+    @staticmethod
+    def calculate_public_key_fingerprint(private_key):
+        # get public key bytes
+        public_key_der = private_key.public_key().public_bytes(
+            Encoding.DER, PublicFormat.SubjectPublicKeyInfo
+        )
+
+        # take sha256 on raw bytes and then do base64 encode
+        sha256hash = hashlib.sha256()
+        sha256hash.update(public_key_der)
+
+        public_key_fp = "SHA256:" + base64.b64encode(sha256hash.digest()).decode(
+            "utf-8"
+        )
+        logger.debug("Public key fingerprint is %s", public_key_fp)
+
+        return public_key_fp
+
+    def update_body(self, body: dict[Any, Any]) -> None:
+        body["data"]["AUTHENTICATOR"] = KEY_PAIR_AUTHENTICATOR
+        body["data"]["TOKEN"] = self._jwt_token
+
+    def assertion_content(self) -> str:
+        return self._jwt_token
+
+    def should_retry(self, count: int) -> bool:
+        return count < self._jwt_retry_attempts
+
+    def handle_timeout(
+        self,
+        *,
+        authenticator: str,
+        service_name: str | None,
+        account: str,
+        user: str,
+        password: str | None,
+        **kwargs: Any,
+    ) -> None:
+        logger.debug("Invoking base timeout handler")
+        super().handle_timeout(
+            authenticator=authenticator,
+            service_name=service_name,
+            account=account,
+            user=user,
+            password=password,
+            delete_params=False,
+        )
+
+        logger.debug("Base timeout handler passed, preparing new token before retrying")
+        self.prepare(account=account, user=user)
+
+    @staticmethod
+    def can_handle_exception(op: OperationalError) -> bool:
+        if op.errno is ER_CONNECTION_TIMEOUT:
+            return True
+        return False

snowflake/connector/auth/no_auth.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+from typing import Any
+
+from .by_plugin import AuthByPlugin, AuthType
+
+
+class AuthNoAuth(AuthByPlugin):
+    """No-auth Authentication.
+
+    It is a dummy auth that requires no extra connection establishment.
+    """
+
+    @property
+    def type_(self) -> AuthType:
+        return AuthType.NO_AUTH
+
+    @property
+    def assertion_content(self) -> str | None:
+        return None
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def reset_secrets(self) -> None:
+        pass
+
+    def prepare(
+        self,
+        **kwargs: Any,
+    ) -> None:
+        pass
+
+    def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
+        return {"success": True}
+
+    def update_body(self, body: dict[Any, Any]) -> None:
+        pass

snowflake/connector/auth/oauth.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+from typing import Any
+
+from ..network import OAUTH_AUTHENTICATOR
+from .by_plugin import AuthByPlugin, AuthType
+
+
+class AuthByOAuth(AuthByPlugin):
+    """OAuth Based Authentication.
+
+    Works by accepting an OAuth token and using that to authenticate.
+    """
+
+    @property
+    def type_(self) -> AuthType:
+        return AuthType.OAUTH
+
+    @property
+    def assertion_content(self) -> str | None:
+        """Returns the token."""
+        return self._oauth_token
+
+    def __init__(self, oauth_token: str, **kwargs) -> None:
+        """Initializes an instance with an OAuth Token."""
+        super().__init__(**kwargs)
+        self._oauth_token: str | None = oauth_token
+
+    def reset_secrets(self) -> None:
+        self._oauth_token = None
+
+    def prepare(
+        self,
+        **kwargs: Any,
+    ) -> None:
+        """Nothing to do here, token should be obtained outside the driver."""
+        pass
+
+    def reauthenticate(self, **kwargs: Any) -> dict[str, bool]:
+        return {"success": False}
+
+    def update_body(self, body: dict[Any, Any]) -> None:
+        """Update some information required by OAuth.
+
+        OAuth needs the authenticator and token attributes set, as well as loginname, which is set already in auth.py.
+        """
+        body["data"]["AUTHENTICATOR"] = OAUTH_AUTHENTICATOR
+        body["data"]["TOKEN"] = self._oauth_token