Add support for workload identity federation (#2203)

Author: sfc-gh-pmansour

setup.cfg

@@ -44,6 +44,8 @@ python_requires = >=3.8
 packages = find_namespace:
 install_requires =
     asn1crypto>0.24.0,<2.0.0
+    boto3>=1.0
+    botocore>=1.0
     cffi>=1.9,<2.0.0
     cryptography>=3.1.0
     pyOpenSSL>=22.0.0,<26.0.0

src/snowflake/connector/auth/__init__.py

@@ -15,6 +15,7 @@
 from .pat import AuthByPAT
 from .usrpwdmfa import AuthByUsrPwdMfa
 from .webbrowser import AuthByWebBrowser
+from .workload_identity import AuthByWorkloadIdentity
 
 FIRST_PARTY_AUTHENTICATORS = frozenset(
     (
@@ -26,6 +27,7 @@
         AuthByWebBrowser,
         AuthByIdToken,
         AuthByPAT,
+        AuthByWorkloadIdentity,
         AuthNoAuth,
     )
 )
@@ -39,6 +41,7 @@
     "AuthByOkta",
     "AuthByUsrPwdMfa",
     "AuthByWebBrowser",
+    "AuthByWorkloadIdentity",
     "AuthNoAuth",
     "Auth",
     "AuthType",

src/snowflake/connector/auth/by_plugin.py

@@ -56,6 +56,7 @@ class AuthType(Enum):
     OKTA = "OKTA"
     PAT = "PROGRAMMATIC_ACCESS_TOKEN'"
     NO_AUTH = "NO_AUTH"
+    WORKLOAD_IDENTITY = "WORKLOAD_IDENTITY"
 
 
 class AuthByPlugin(ABC):

src/snowflake/connector/auth/workload_identity.py

@@ -0,0 +1,98 @@
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+import json
+import typing
+from enum import Enum, unique
+
+from ..network import WORKLOAD_IDENTITY_AUTHENTICATOR
+from ..wif_util import (
+    AttestationProvider,
+    WorkloadIdentityAttestation,
+    create_attestation,
+)
+from .by_plugin import AuthByPlugin, AuthType
+
+
+@unique
+class ApiFederatedAuthenticationType(Enum):
+    """An API-specific enum of the WIF authentication type."""
+
+    AWS = "AWS"
+    AZURE = "AZURE"
+    GCP = "GCP"
+    OIDC = "OIDC"
+
+    @staticmethod
+    def from_attestation(
+        attestation: WorkloadIdentityAttestation,
+    ) -> ApiFederatedAuthenticationType:
+        """Maps the internal / driver-specific attestation providers to API authenticator types.
+
+        The AttestationProvider is related to how the driver fetches the credential, while the API authenticator
+        type is related to how the credential is verified. In most current cases these may be the same, though
+        in the future we could have, for example, multiple AttestationProviders that all fetch an OIDC ID token.
+        """
+        if attestation.provider == AttestationProvider.AWS:
+            return ApiFederatedAuthenticationType.AWS
+        if attestation.provider == AttestationProvider.AZURE:
+            return ApiFederatedAuthenticationType.AZURE
+        if attestation.provider == AttestationProvider.GCP:
+            return ApiFederatedAuthenticationType.GCP
+        if attestation.provider == AttestationProvider.OIDC:
+            return ApiFederatedAuthenticationType.OIDC
+        return ValueError(f"Unknown attestation provider '{attestation.provider}'")
+
+
+class AuthByWorkloadIdentity(AuthByPlugin):
+    """Plugin to authenticate via workload identity."""
+
+    def __init__(
+        self,
+        *,
+        provider: AttestationProvider | None = None,
+        token: str | None = None,
+        entra_resource: str | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.provider = provider
+        self.token = token
+        self.entra_resource = entra_resource
+
+        self.attestation: WorkloadIdentityAttestation | None = None
+
+    def type_(self) -> AuthType:
+        return AuthType.WORKLOAD_IDENTITY
+
+    def reset_secrets(self) -> None:
+        self.attestation = None
+
+    def update_body(self, body: dict[typing.Any, typing.Any]) -> None:
+        body["data"]["AUTHENTICATOR"] = WORKLOAD_IDENTITY_AUTHENTICATOR
+        body["data"]["PROVIDER"] = ApiFederatedAuthenticationType.from_attestation(
+            self.attestation
+        ).value
+        body["data"]["TOKEN"] = self.attestation.credential
+
+    def prepare(self, **kwargs: typing.Any) -> None:
+        """Fetch the token."""
+        self.attestation = create_attestation(
+            self.provider, self.entra_resource, self.token
+        )
+
+    def reauthenticate(self, **kwargs: typing.Any) -> dict[str, bool]:
+        """This is only relevant for AuthByIdToken, which uses a web-browser based flow. All other auth plugins just call authenticate() again."""
+        return {"success": False}
+
+    @property
+    def assertion_content(self) -> str:
+        """Returns the CSP provider name and an identifier. Used for logging purposes."""
+        if not self.attestation:
+            return ""
+        properties = self.attestation.user_identifier_components
+        properties["_provider"] = self.attestation.provider.value
+        return json.dumps(properties, sort_keys=True, separators=(",", ":"))

src/snowflake/connector/connection.py

@@ -44,6 +44,7 @@
     AuthByPlugin,
     AuthByUsrPwdMfa,
     AuthByWebBrowser,
+    AuthByWorkloadIdentity,
     AuthNoAuth,
 )
 from .auth.idtoken import AuthByIdToken
@@ -55,6 +56,7 @@
 from .constants import (
     _CONNECTIVITY_ERR_MSG,
     _DOMAIN_NAME_MAP,
+    ENV_VAR_EXPERIMENTAL_AUTHENTICATION,
     ENV_VAR_PARTNER,
     PARAMETER_AUTOCOMMIT,
     PARAMETER_CLIENT_PREFETCH_THREADS,
@@ -87,6 +89,7 @@
     ER_FAILED_TO_CONNECT_TO_DB,
     ER_INVALID_BACKOFF_POLICY,
     ER_INVALID_VALUE,
+    ER_INVALID_WIF_SETTINGS,
     ER_NO_ACCOUNT_NAME,
     ER_NO_NUMPY,
     ER_NO_PASSWORD,
@@ -104,6 +107,7 @@
     PROGRAMMATIC_ACCESS_TOKEN,
     REQUEST_ID,
     USR_PWD_MFA_AUTHENTICATOR,
+    WORKLOAD_IDENTITY_AUTHENTICATOR,
     ReauthenticationRequest,
     SnowflakeRestful,
 )
@@ -112,6 +116,7 @@
 from .time_util import HeartBeatTimer, get_time_millis
 from .url_util import extract_top_level_domain_from_hostname
 from .util_text import construct_hostname, parse_account, split_statements
+from .wif_util import AttestationProvider
 
 DEFAULT_CLIENT_PREFETCH_THREADS = 4
 MAX_CLIENT_PREFETCH_THREADS = 10
@@ -188,12 +193,14 @@ def _get_private_bytes_from_file(
     "private_key": (None, (type(None), bytes, str, RSAPrivateKey)),
     "private_key_file": (None, (type(None), str)),
     "private_key_file_pwd": (None, (type(None), str, bytes)),
-    "token": (None, (type(None), str)),  # OAuth/JWT/PAT Token
+    "token": (None, (type(None), str)),  # OAuth/JWT/PAT/OIDC Token
     "token_file_path": (
         None,
         (type(None), str, bytes),
-    ),  # OAuth/JWT/PAT Token file path
+    ),  # OAuth/JWT/PAT/OIDC Token file path
     "authenticator": (DEFAULT_AUTHENTICATOR, (type(None), str)),
+    "workload_identity_provider": (None, (type(None), AttestationProvider)),
+    "workload_identity_entra_resource": (None, (type(None), str)),
     "mfa_callback": (None, (type(None), Callable)),
     "password_callback": (None, (type(None), Callable)),
     "auth_class": (None, (type(None), AuthByPlugin)),
@@ -1140,6 +1147,29 @@ def __open_connection(self):
                 if not self._token and self._password:
                     self._token = self._password
                 self.auth_class = AuthByPAT(self._token)
+            elif self._authenticator == WORKLOAD_IDENTITY_AUTHENTICATOR:
+                if ENV_VAR_EXPERIMENTAL_AUTHENTICATION not in os.environ:
+                    Error.errorhandler_wrapper(
+                        self,
+                        None,
+                        ProgrammingError,
+                        {
+                            "msg": f"Please set the '{ENV_VAR_EXPERIMENTAL_AUTHENTICATION}' environment variable to use the '{WORKLOAD_IDENTITY_AUTHENTICATOR}' authenticator.",
+                            "errno": ER_INVALID_WIF_SETTINGS,
+                        },
+                    )
+                # Standardize the provider enum.
+                if self._workload_identity_provider and isinstance(
+                    self._workload_identity_provider, str
+                ):
+                    self._workload_identity_provider = AttestationProvider.from_string(
+                        self._workload_identity_provider
+                    )
+                self.auth_class = AuthByWorkloadIdentity(
+                    provider=self._workload_identity_provider,
+                    token=self._token,
+                    entra_resource=self._workload_identity_entra_resource,
+                )
             else:
                 # okta URL, e.g., https://<account>.okta.com/
                 self.auth_class = AuthByOkta(
@@ -1268,6 +1298,7 @@ def __config(self, **kwargs):
                 KEY_PAIR_AUTHENTICATOR,
                 OAUTH_AUTHENTICATOR,
                 USR_PWD_MFA_AUTHENTICATOR,
+                WORKLOAD_IDENTITY_AUTHENTICATOR,
             ]:
                 self._authenticator = auth_tmp
 
@@ -1278,14 +1309,18 @@ def __config(self, **kwargs):
                 self._token = f.read()
 
         # Set of authenticators allowing empty user.
-        empty_user_allowed_authenticators = {OAUTH_AUTHENTICATOR, NO_AUTH_AUTHENTICATOR}
+        empty_user_allowed_authenticators = {
+            OAUTH_AUTHENTICATOR,
+            NO_AUTH_AUTHENTICATOR,
+            WORKLOAD_IDENTITY_AUTHENTICATOR,
+        }
 
         if not (self._master_token and self._session_token):
             if (
                 not self.user
                 and self._authenticator not in empty_user_allowed_authenticators
             ):
-                # OAuth and NoAuth Authentications does not require a username
+                # Some authenticators do not require a username
                 Error.errorhandler_wrapper(
                     self,
                     None,
@@ -1296,6 +1331,25 @@ def __config(self, **kwargs):
             if self._private_key or self._private_key_file:
                 self._authenticator = KEY_PAIR_AUTHENTICATOR
 
+            workload_identity_dependent_options = [
+                "workload_identity_provider",
+                "workload_identity_entra_resource",
+            ]
+            for dependent_option in workload_identity_dependent_options:
+                if (
+                    self.__getattribute__(f"_{dependent_option}") is not None
+                    and self._authenticator != WORKLOAD_IDENTITY_AUTHENTICATOR
+                ):
+                    Error.errorhandler_wrapper(
+                        self,
+                        None,
+                        ProgrammingError,
+                        {
+                            "msg": f"{dependent_option} was set but authenticator was not set to {WORKLOAD_IDENTITY_AUTHENTICATOR}",
+                            "errno": ER_INVALID_WIF_SETTINGS,
+                        },
+                    )
+
             if (
                 self.auth_class is None
                 and self._authenticator
@@ -1304,6 +1358,7 @@ def __config(self, **kwargs):
                     OAUTH_AUTHENTICATOR,
                     KEY_PAIR_AUTHENTICATOR,
                     PROGRAMMATIC_ACCESS_TOKEN,
+                    WORKLOAD_IDENTITY_AUTHENTICATOR,
                 )
                 and not self._password
             ):

src/snowflake/connector/constants.py

@@ -430,6 +430,7 @@ class IterUnit(Enum):
 # TODO: all env variables definitions should be here
 ENV_VAR_PARTNER = "SF_PARTNER"
 ENV_VAR_TEST_MODE = "SNOWFLAKE_TEST_MODE"
+ENV_VAR_EXPERIMENTAL_AUTHENTICATION = "SF_ENABLE_EXPERIMENTAL_AUTHENTICATION"  # Needed to enable new strong auth features during the private preview.
 
 
 _DOMAIN_NAME_MAP = {_DEFAULT_HOSTNAME_TLD: "GLOBAL", _CHINA_HOSTNAME_TLD: "CHINA"}

src/snowflake/connector/errorcode.py

@@ -31,6 +31,8 @@
 ER_JWT_RETRY_EXPIRED = 251010
 ER_CONNECTION_TIMEOUT = 251011
 ER_RETRYABLE_CODE = 251012
+ER_INVALID_WIF_SETTINGS = 251013
+ER_WIF_CREDENTIALS_NOT_FOUND = 251014
 
 # cursor
 ER_FAILED_TO_REWRITE_MULTI_ROW_INSERT = 252001

src/snowflake/connector/network.py

@@ -189,6 +189,7 @@
 USR_PWD_MFA_AUTHENTICATOR = "USERNAME_PASSWORD_MFA"
 PROGRAMMATIC_ACCESS_TOKEN = "PROGRAMMATIC_ACCESS_TOKEN"
 NO_AUTH_AUTHENTICATOR = "NO_AUTH"
+WORKLOAD_IDENTITY_AUTHENTICATOR = "WORKLOAD_IDENTITY"
 
 
 def is_retryable_http_code(code: int) -> bool: