metaflow 2.11.15__py2.py3-none-any.whl → 2.12.0__py2.py3-none-any.whl

metaflow/plugins/azure/azure_secret_manager_secrets_provider.py

@@ -0,0 +1,240 @@
+from metaflow.plugins.secrets import SecretsProvider
+import re
+import base64
+import codecs
+from urllib.parse import urlparse
+from metaflow.exception import MetaflowException
+import sys
+from metaflow.metaflow_config import AZURE_KEY_VAULT_PREFIX
+from metaflow.plugins.azure.azure_credential import (
+    create_cacheable_azure_credential,
+)
+
+
+class MetaflowAzureKeyVaultBadVault(MetaflowException):
+    """Raised when the secretid is fully qualified but does not have the right key vault domain"""
+
+
+class MetaflowAzureKeyVaultBadSecretType(MetaflowException):
+    """Raised when the secret type is anything except secrets"""
+
+
+class MetaflowAzureKeyVaultBadSecretPath(MetaflowException):
+    """Raised when the secret path does not match to expected length"""
+
+
+class MetaflowAzureKeyVaultBadSecretName(MetaflowException):
+    """Raised when the secret name does not match expected pattern"""
+
+
+class MetaflowAzureKeyVaultBadSecretVersion(MetaflowException):
+    """Raised when the secret version does not match expected pattern"""
+
+
+class MetaflowAzureKeyVaultBadSecret(MetaflowException):
+    """Raised when the secret does not match supported patterns in Metaflow"""
+
+
+class AzureKeyVaultSecretsProvider(SecretsProvider):
+    TYPE = "az-key-vault"
+    key_vault_domains = [
+        ".vault.azure.net",
+        ".vault.azure.cn",
+        ".vault.usgovcloudapi.net",
+        ".vault.microsoftazure.de",
+    ]
+    supported_vault_object_types = ["secrets"]
+
+    # https://learn.microsoft.com/en-us/azure/key-vault/general/about-keys-secrets-certificates has details on vault name structure
+    # Vault name and Managed HSM pool name must be a 3-24 character string, containing only 0-9, a-z, A-Z, and not consecutive -.
+    def _is_valid_vault_name(self, vault_name):
+        vault_name_pattern = r"^(?!.*--)[a-zA-Z0-9-]{3,24}$"
+        return re.match(vault_name_pattern, vault_name) is not None
+
+    # The type of the object can be, "keys", "secrets", or "certificates".
+    # Currently only secrets will be supported
+    def _is_valid_object_type(self, secret_type):
+        for type in self.supported_vault_object_types:
+            if secret_type == type:
+                return True
+        return False
+
+    # The secret name must be a 1-127 character string, starting with a letter and containing only 0-9, a-z, A-Z, and -.
+    def _is_valid_secret_name(self, secret_name):
+        secret_name_pattern = r"^[a-zA-Z][a-zA-Z0-9-]{0,126}$"
+        return re.match(secret_name_pattern, secret_name) is not None
+
+    # An object-version is a system-generated, 32 character string identifier that is optionally used to address a unique version of an object.
+    def _is_valid_object_version(self, secret_version):
+        object_version_pattern = r"^[a-zA-Z0-9]{32}$"
+        return re.match(object_version_pattern, secret_version) is not None
+
+    # This function will check if the secret_id is fully qualified url. It will return True iff the secret_id is of the form:
+    # https://myvault.vault.azure.net/secrets/mysecret/ec96f02080254f109c51a1f14cdb1931 OR
+    # https://myvault.vault.azure.net/secrets/mysecret/
+    # validating the above as per recommendations in https://devblogs.microsoft.com/azure-sdk/guidance-for-applications-using-the-key-vault-libraries/
+    def _is_secret_id_fully_qualified_url(self, secret_id):
+        # if the secret_id is None/empty/does not start with https then return false
+        if secret_id is None or secret_id == "" or not secret_id.startswith("https://"):
+            return False
+        try:
+            parsed_vault_url = urlparse(secret_id)
+        except ValueError:
+            print("invalid vault url", file=sys.stderr)
+            return False
+        hostname = parsed_vault_url.netloc
+
+        k_v_domain_found = False
+        actual_k_v_domain = ""
+        for k_v_domain in self.key_vault_domains:
+            if k_v_domain in hostname:
+                k_v_domain_found = True
+                actual_k_v_domain = k_v_domain
+                break
+        if not k_v_domain_found:
+            # the secret_id started with https:// however the key_vault_domains
+            # were not present in the secret_id which means
+            raise MetaflowAzureKeyVaultBadVault("bad key vault domain %s" % secret_id)
+
+        # given the secret_id seems to have a valid key vault domain
+        # lets verify that the vault name corresponds to its regex.
+        vault_name = hostname[: -len(actual_k_v_domain)]
+        # verify the vault name pattern
+        if not self._is_valid_vault_name(vault_name):
+            raise MetaflowAzureKeyVaultBadVault("bad key vault name %s" % vault_name)
+
+        path_parts = parsed_vault_url.path.strip("/").split("/")
+        total_path_parts = len(path_parts)
+        if total_path_parts < 2 or total_path_parts > 3:
+            raise MetaflowAzureKeyVaultBadSecretPath(
+                "bad secret uri path %s" % path_parts
+            )
+
+        object_type = path_parts[0]
+        if not self._is_valid_object_type(object_type):
+            raise MetaflowAzureKeyVaultBadSecretType("bad secret type %s" % object_type)
+
+        secret_name = path_parts[1]
+        if not self._is_valid_secret_name(secret_name=secret_name):
+            raise MetaflowAzureKeyVaultBadSecretName("bad secret name %s" % secret_name)
+
+        if total_path_parts == 3:
+            if not self._is_valid_object_version(path_parts[2]):
+                raise MetaflowAzureKeyVaultBadSecretVersion(
+                    "bad secret version %s" % path_parts[2]
+                )
+
+        return True
+
+    # This function will validate the correctness of the partial secret id.
+    # It will attempt to construct the fully qualified secret URL internally and
+    # call the _is_secret_id_fully_qualified_url to check validity
+    def _is_partial_secret_valid(self, secret_id):
+        secret_parts = secret_id.strip("/").split("/")
+        total_secret_parts = len(secret_parts)
+        if total_secret_parts < 1 or total_secret_parts > 2:
+            return False
+
+        # since the secret_id is supposedly a partial id, the AZURE_KEY_VAULT_PREFIX
+        # must be set.
+        if not AZURE_KEY_VAULT_PREFIX:
+            raise ValueError(
+                "cannot use simple secret id without setting METAFLOW_AZURE_KEY_VAULT_PREFIX. %s"
+                % AZURE_KEY_VAULT_PREFIX
+            )
+        domain = AZURE_KEY_VAULT_PREFIX.rstrip("/")
+        full_secret = "%s/secrets/%s" % (domain, secret_id)
+        if not self._is_secret_id_fully_qualified_url(full_secret):
+            return False
+
+        return True
+
+    def _sanitize_key_as_env_var(self, key):
+        """
+        Sanitize a key as an environment variable name.
+        This is purely a convenience trade-off to cover common cases well, vs. introducing
+        ambiguities (e.g. did the final '_' come from '.', or '-' or is original?).
+
+        1/27/2023(jackie):
+
+        We start with few rules and should *sparingly* add more over time.
+        Also, it's TBD whether all possible providers will share the same sanitization logic.
+        Therefore we will keep this function private for now
+        """
+        return key.replace("-", "_").replace(".", "_").replace("/", "_")
+
+    def get_secret_as_dict(self, secret_id, options={}, role=None):
+        # https://learn.microsoft.com/en-us/azure/app-service/app-service-key-vault-references?tabs=azure-cli has a lot of details on
+        # the patterns used in key vault
+        # Vault names and Managed HSM pool names are selected by the user and are globally unique.
+        # Vault name and Managed HSM pool name must be a 3-24 character string, containing only 0-9, a-z, A-Z, and not consecutive -.
+        # object-type	The type of the object. As of 05/08/24 only "secrets", are supported
+        # object-name	An object-name is a user provided name for and must be unique within a key vault. The name must be a 1-127 character string, starting with a letter and containing only 0-9, a-z, A-Z, and -.
+        # object-version	An object-version is a system-generated, 32 character string identifier that is optionally used to address a unique version of an object.
+
+        # We allow these forms of secret_id:
+        #
+        # 1. Full path like https://<key-vault-name><.vault-domain>/secrets/<secret-name>/<secret-version>. This is what you
+        # see in Azure portal and is easy to copy paste.
+        #
+        # 2. Full path but without the version like https://<key-vault-name><.vault-domain>/secrets/<secret-name>
+        #
+        # 3. Simple string like mysecret. This corresponds to the SecretName.
+        #
+        # 4. Simple string with <secret-name>/<secret-version> suffix like mysecret/123
+
+        # The latter two forms require METAFLOW_AZURE_KEY_VAULT_PREFIX to be set.
+
+        # if the secret_id is None/empty/does not start with https then return false
+        if secret_id is None or secret_id == "":
+            raise MetaflowAzureKeyVaultBadSecret("empty secret id is not supported")
+
+        # check if the passed in secret is a short-form ( #3/#4 in the above comment)
+        if not secret_id.startswith("https://"):
+            # check if the secret_id is of form `secret_name` OR `secret_name/secret_version`
+            if not self._is_partial_secret_valid(secret_id=secret_id):
+                raise MetaflowAzureKeyVaultBadSecret(
+                    "unsupported partial secret %s" % secret_id
+                )
+
+            domain = AZURE_KEY_VAULT_PREFIX.rstrip("/")
+            full_secret = "%s/secrets/%s" % (domain, secret_id)
+
+        # if the secret id is passed as a URL - then check if the url is fully qualified
+        if secret_id.startswith("https://"):
+            if not self._is_secret_id_fully_qualified_url(secret_id=secret_id):
+                raise MetaflowException("unsupported secret %s" % secret_id)
+            full_secret = secret_id
+
+        # at this point I know that the secret URL is good so we can start creating the Secret Client
+        az_credentials = create_cacheable_azure_credential()
+        res = urlparse(full_secret)
+        az_vault_url = "%s://%s" % (
+            res.scheme,
+            res.netloc,
+        )  # https://myvault.vault.azure.net
+        secret_data = res.path.strip("/").split("/")[1:]
+        secret_name = secret_data[0]
+        secret_version = None
+        if len(secret_data) > 1:
+            secret_version = secret_data[1]
+
+        from azure.keyvault.secrets import SecretClient
+
+        client = SecretClient(vault_url=az_vault_url, credential=az_credentials)
+
+        key_vault_secret_val = client.get_secret(
+            name=secret_name, version=secret_version
+        )
+
+        result = {}
+
+        if options.get("env_var_name") is not None:
+            env_var_name = options["env_var_name"]
+            sanitized_key = self._sanitize_key_as_env_var(env_var_name)
+        else:
+            sanitized_key = self._sanitize_key_as_env_var(key_vault_secret_val.name)
+
+        response_payload = key_vault_secret_val.value
+        result[sanitized_key] = response_payload
+        return result

metaflow/plugins/azure/azure_utils.py

@@ -7,6 +7,7 @@ from metaflow.plugins.azure.azure_exceptions import (
     MetaflowAzurePackageError,
 )
 from metaflow.exception import MetaflowInternalError, MetaflowException
+from metaflow.plugins.azure.azure_credential import create_cacheable_azure_credential
 
 
 def _check_and_init_azure_deps():
@@ -138,38 +139,6 @@ def handle_exceptions(func):
     return _inner_func
 
 
-@check_azure_deps
-def create_cacheable_default_azure_credentials(*args, **kwargs):
-    """azure.identity.DefaultAzureCredential is not readily cacheable in a dictionary
-    because it does not have a content based hash and equality implementations.
-
-    We implement a subclass CacheableDefaultAzureCredential to add them.
-
-    We need this because credentials will be part of the cache key in _ClientCache.
-    """
-    from azure.identity import DefaultAzureCredential
-
-    class CacheableDefaultAzureCredential(DefaultAzureCredential):
-        def __init__(self, *args, **kwargs):
-            super(CacheableDefaultAzureCredential, self).__init__(*args, **kwargs)
-            # Just hashing all the kwargs works because they are all individually
-            # hashable as of 7/15/2022.
-            #
-            # What if Azure adds unhashable things to kwargs?
-            # - We will have CI to catch this (it will always install the latest Azure SDKs)
-            # - In Metaflow usage today we never specify any kwargs anyway. (see last line
-            #   of the outer function.
-            self._hash_code = hash((args, tuple(sorted(kwargs.items()))))
-
-        def __hash__(self):
-            return self._hash_code
-
-        def __eq__(self, other):
-            return hash(self) == hash(other)
-
-    return CacheableDefaultAzureCredential(*args, **kwargs)
-
-
 @check_azure_deps
 def create_static_token_credential(token_):
     from azure.core.credentials import TokenCredential
@@ -200,9 +169,7 @@ def create_static_token_credential(token_):
         def get_token(self, *_scopes, **_kwargs):
 
             if (self._cached_token.expires_on - time.time()) < 300:
-                from azure.identity import DefaultAzureCredential
-
-                self._credential = DefaultAzureCredential()
+                self._credential = create_cacheable_azure_credential()
             if self._credential:
                 return self._credential.get_token(*_scopes, **_kwargs)
             return self._cached_token

metaflow/plugins/azure/blob_service_client_factory.py

@@ -1,9 +1,11 @@
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import AZURE_STORAGE_BLOB_SERVICE_ENDPOINT
 from metaflow.plugins.azure.azure_utils import (
-    create_cacheable_default_azure_credentials,
     check_azure_deps,
 )
+from metaflow.plugins.azure.azure_credential import (
+    create_cacheable_azure_credential,
+)
 
 import os
 import threading
@@ -125,7 +127,7 @@ def get_azure_blob_service_client(
     blob_service_endpoint = AZURE_STORAGE_BLOB_SERVICE_ENDPOINT
 
     if not credential:
-        credential = create_cacheable_default_azure_credentials()
+        credential = create_cacheable_azure_credential()
         credential_is_cacheable = True
 
     if not credential_is_cacheable:

metaflow/plugins/datastores/azure_storage.py

@@ -32,6 +32,8 @@ from metaflow.plugins.storage_executor import (
     handle_executor_exceptions,
 )
 
+from metaflow.plugins.azure.azure_credential import create_cacheable_azure_credential
+
 AZURE_STORAGE_DOWNLOAD_MAX_CONCURRENCY = 4
 AZURE_STORAGE_UPLOAD_MAX_CONCURRENCY = 16
 
@@ -272,12 +274,10 @@ class AzureStorage(DataStoreStorage):
         if not self._default_scope_token or (
             self._default_scope_token.expires_on - time.time() < 300
         ):
-            from azure.identity import DefaultAzureCredential
-
-            with DefaultAzureCredential() as credential:
-                self._default_scope_token = credential.get_token(
-                    AZURE_STORAGE_DEFAULT_SCOPE
-                )
+            credential = create_cacheable_azure_credential()
+            self._default_scope_token = credential.get_token(
+                AZURE_STORAGE_DEFAULT_SCOPE
+            )
         return self._default_scope_token
 
     @property

metaflow/plugins/datatools/s3/s3.py

@@ -21,7 +21,6 @@ from metaflow.metaflow_config import (
     TEMPDIR,
 )
 from metaflow.util import (
-    namedtuple_with_defaults,
     is_stringish,
     to_bytes,
     to_unicode,
@@ -29,6 +28,7 @@ from metaflow.util import (
     url_quote,
     url_unquote,
 )
+from metaflow.tuple_util import namedtuple_with_defaults
 from metaflow.exception import MetaflowException
 from metaflow.debug import debug
 import metaflow.tracing as tracing

metaflow/plugins/gcp/__init__.py

@@ -0,0 +1 @@
+from .gs_storage_client_factory import get_credentials

metaflow/plugins/gcp/gcp_secret_manager_secrets_provider.py

@@ -0,0 +1,169 @@
+import base64
+import json
+from json import JSONDecodeError
+
+
+from metaflow.exception import MetaflowException
+from metaflow.plugins.secrets import SecretsProvider
+import re
+from metaflow.plugins.gcp.gs_storage_client_factory import get_credentials
+from metaflow.metaflow_config import GCP_SECRET_MANAGER_PREFIX
+
+
+class MetaflowGcpSecretsManagerBadResponse(MetaflowException):
+    """Raised when the response from GCP Secrets Manager is not valid in some way"""
+
+
+class MetaflowGcpSecretsManagerDuplicateKey(MetaflowException):
+    """Raised when the response from GCP Secrets Manager contains duplicate keys"""
+
+
+class MetaflowGcpSecretsManagerJSONParseError(MetaflowException):
+    """Raised when the SecretString response from GCP Secrets Manager is not valid JSON"""
+
+
+class MetaflowGcpSecretsManagerNotJSONObject(MetaflowException):
+    """Raised when the SecretString response from GCP Secrets Manager is not valid JSON dictionary"""
+
+
+def _sanitize_key_as_env_var(key):
+    """
+    Sanitize a key as an environment variable name.
+    This is purely a convenience trade-off to cover common cases well, vs. introducing
+    ambiguities (e.g. did the final '_' come from '.', or '-' or is original?).
+
+    1/27/2023(jackie):
+
+    We start with few rules and should *sparingly* add more over time.
+    Also, it's TBD whether all possible providers will share the same sanitization logic.
+    Therefore we will keep this function private for now
+    """
+    return key.replace("-", "_").replace(".", "_").replace("/", "_")
+
+
+class GcpSecretManagerSecretsProvider(SecretsProvider):
+    TYPE = "gcp-secret-manager"
+
+    def get_secret_as_dict(self, secret_id, options={}, role=None):
+        """
+        Reads a secret from GCP Secrets Manager and returns it as a dictionary of environment variables.
+
+        If the secret contains a string payload ("SecretString"):
+        - if the `json` option is True:
+            Secret will be parsed as a JSON. If successfully parsed, AND the JSON contains a
+            top-level object, each entry K/V in the object will also be converted to an entry in the result. V will
+            always be casted to a string (if not already a string).
+        - If `json` option is False (default):
+            Will be returned as a single entry in the result, with the key being the last part after / in secret_id.
+
+        On GCP Secrets Manager, the secret payload is a binary blob. However, by default we interpret it as UTF8 encoded
+        string. To disable this, set the `binary` option to True, the binary will be base64 encoded in the result.
+
+        All keys in the result are sanitized to be more valid environment variable names. This is done on a best effort
+        basis. Further validation is expected to be done by the invoking @secrets decorator itself.
+
+        :param secret_id: GCP Secrets Manager secret ID
+        :param options: unused
+        :return: dict of environment variables. All keys and values are strings.
+        """
+        from google.cloud.secretmanager_v1.services.secret_manager_service import (
+            SecretManagerServiceClient,
+        )
+        from google.cloud.secretmanager_v1.services.secret_manager_service.transports import (
+            SecretManagerServiceTransport,
+        )
+
+        # Full secret id looks like projects/1234567890/secrets/mysecret/versions/latest
+        #
+        # We allow these forms of secret_id:
+        #
+        # 1. Full path like projects/1234567890/secrets/mysecret/versions/latest
+        #    This is what you'd specify if you used to GCP SDK.
+        #
+        # 2. Full path but without the version like projects/1234567890/secrets/mysecret.
+        #    This is what you see in the GCP console, makes it easier to copy & paste.
+        #
+        # 3. Simple string like mysecret
+        #
+        # 4. Simple string with /versions/<version> suffix like mysecret/versions/1
+
+        # The latter two forms require METAFLOW_GCP_SECRET_MANAGER_PREFIX to be set.
+
+        match_full = re.match(
+            r"^projects/\d+/secrets/([\w\-]+)(/versions/([\w\-]+))?$", secret_id
+        )
+        match_partial = re.match(r"^([\w\-]+)(/versions/[\w\-]+)?$", secret_id)
+        if match_full:
+            # Full path
+            env_var_name = match_full.group(1)
+            if match_full.group(3):
+                # With version specified
+                full_secret_name = secret_id
+            else:
+                # No version specified, use latest
+                full_secret_name = secret_id + "/versions/latest"
+        elif match_partial:
+            # Partial path, possibly with /versions/<version> suffix
+            env_var_name = secret_id
+            if not GCP_SECRET_MANAGER_PREFIX:
+                raise ValueError(
+                    "Cannot use simple secret_id without setting METAFLOW_GCP_SECRET_MANAGER_PREFIX. %s"
+                    % GCP_SECRET_MANAGER_PREFIX
+                )
+            if match_partial.group(2):
+                # With version specified
+                full_secret_name = "%s%s" % (GCP_SECRET_MANAGER_PREFIX, secret_id)
+                env_var_name = match_partial.group(1)
+            else:
+                # No version specified, use latest
+                full_secret_name = "%s%s/versions/latest" % (
+                    GCP_SECRET_MANAGER_PREFIX,
+                    secret_id,
+                )
+        else:
+            raise ValueError(
+                "Invalid secret_id: %s. Must be either a full path or a simple string."
+                % secret_id
+            )
+
+        result = {}
+
+        def _sanitize_and_add_entry_to_result(k, v):
+            # Two jobs - sanitize, and check for dupes
+            sanitized_k = _sanitize_key_as_env_var(k)
+            if sanitized_k in result:
+                raise MetaflowGcpSecretsManagerDuplicateKey(
+                    "Duplicate key in secret: '%s' (sanitizes to '%s')"
+                    % (k, sanitized_k)
+                )
+            result[sanitized_k] = v
+
+        credentials, _ = get_credentials(
+            scopes=SecretManagerServiceTransport.AUTH_SCOPES
+        )
+        client = SecretManagerServiceClient(credentials=credentials)
+        response = client.access_secret_version(request={"name": full_secret_name})
+        payload_str = response.payload.data.decode("UTF-8")
+        if options.get("json", False):
+            obj = json.loads(payload_str)
+            if type(obj) == dict:
+                for k, v in obj.items():
+                    # We try to make it work here - cast to string always
+                    _sanitize_and_add_entry_to_result(k, str(v))
+            else:
+                raise MetaflowGcpSecretsManagerNotJSONObject(
+                    "Secret string is a JSON, but not an object (dict-like) - actual type %s."
+                    % type(obj)
+                )
+        else:
+            if options.get("env_var_name"):
+                env_var_name = options["env_var_name"]
+
+            if options.get("binary", False):
+                _sanitize_and_add_entry_to_result(
+                    env_var_name, base64.b64encode(response.payload.data)
+                )
+            else:
+                _sanitize_and_add_entry_to_result(env_var_name, payload_str)
+
+        return result

metaflow/plugins/gcp/gs_storage_client_factory.py

@@ -8,7 +8,7 @@ def _get_cache_key():
     return os.getpid(), threading.get_ident()
 
 
-def get_gs_storage_client():
+def _get_gs_storage_client_default():
     cache_key = _get_cache_key()
     if cache_key not in _client_cache:
         from google.cloud import storage
@@ -19,3 +19,54 @@ def get_gs_storage_client():
             credentials=credentials, project=project_id
         )
     return _client_cache[cache_key]
+
+
+class GcpDefaultClientProvider(object):
+    name = "gcp-default"
+
+    @staticmethod
+    def get_gs_storage_client(*args, **kwargs):
+        return _get_gs_storage_client_default()
+
+    @staticmethod
+    def get_credentials(scopes, *args, **kwargs):
+        import google.auth
+
+        return google.auth.default(scopes=scopes)
+
+
+cached_provider_class = None
+
+
+def get_gs_storage_client():
+    global cached_provider_class
+    if cached_provider_class is None:
+        from metaflow.metaflow_config import DEFAULT_GCP_CLIENT_PROVIDER
+        from metaflow.plugins import GCP_CLIENT_PROVIDERS
+
+        for p in GCP_CLIENT_PROVIDERS:
+            if p.name == DEFAULT_GCP_CLIENT_PROVIDER:
+                cached_provider_class = p
+                break
+        else:
+            raise ValueError(
+                "Cannot find GCP Client provider %s" % DEFAULT_GCP_CLIENT_PROVIDER
+            )
+    return cached_provider_class.get_gs_storage_client()
+
+
+def get_credentials(scopes, *args, **kwargs):
+    global cached_provider_class
+    if cached_provider_class is None:
+        from metaflow.metaflow_config import DEFAULT_GCP_CLIENT_PROVIDER
+        from metaflow.plugins import GCP_CLIENT_PROVIDERS
+
+        for p in GCP_CLIENT_PROVIDERS:
+            if p.name == DEFAULT_GCP_CLIENT_PROVIDER:
+                cached_provider_class = p
+                break
+        else:
+            raise ValueError(
+                "Cannot find GCP Client provider %s" % DEFAULT_GCP_CLIENT_PROVIDER
+            )
+    return cached_provider_class.get_credentials(scopes, *args, **kwargs)