Skip to content

Secret

zenml.secret special

Initialization of the ZenML Secret module.

A ZenML Secret is a grouping of key-value pairs. These are accessed and administered via the ZenML Secret Manager (a stack component).

Secrets are distinguished by having different schemas. An AWS SecretSchema, for example, has key-value pairs for AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY as well as an optional AWS_SESSION_TOKEN. If you don't specify a schema at the point of registration, ZenML will set the schema as ArbitrarySecretSchema, a kind of default schema where things that aren't attached to a grouping can be stored.

arbitrary_secret_schema

Implementation of a SecretSchema for arbitrary key:value pairs.

ArbitrarySecretSchema (BaseSecretSchema) pydantic-model

Schema for arbitrary collections of key value pairs with no predefined schema.

Source code in zenml/secret/arbitrary_secret_schema.py
class ArbitrarySecretSchema(BaseSecretSchema):
    """Schema for arbitrary collections of key value pairs with no predefined schema."""

    TYPE: ClassVar[str] = ARBITRARY_SECRET_SCHEMA_TYPE

    arbitrary_kv_pairs: Dict[str, Any]

    @root_validator(pre=True)
    def build_arbitrary_kv_pairs(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        """Pydantic root_validator for the Secret Schemas.

        It takes all unused passed kwargs and passes them into the
        arbitrary_kv_pairs attribute.

        Args:
            values: Values passed to the object constructor

        Returns:
            Values passed to the object constructor
        """
        all_required_field_names = {
            field.alias
            for field in cls.__fields__.values()
            if field.alias != "arbitrary_kv_pairs"
        }

        arbitrary_kv_pairs: Dict[str, Any] = {
            field_name: values.pop(field_name)
            for field_name in list(values)
            if field_name not in all_required_field_names
        }

        values["arbitrary_kv_pairs"] = arbitrary_kv_pairs
        return values
build_arbitrary_kv_pairs(values) classmethod

Pydantic root_validator for the Secret Schemas.

It takes all unused passed kwargs and passes them into the arbitrary_kv_pairs attribute.

Parameters:

Name Type Description Default
values Dict[str, Any]

Values passed to the object constructor

required

Returns:

Type Description
Dict[str, Any]

Values passed to the object constructor

Source code in zenml/secret/arbitrary_secret_schema.py
@root_validator(pre=True)
def build_arbitrary_kv_pairs(cls, values: Dict[str, Any]) -> Dict[str, Any]:
    """Pydantic root_validator for the Secret Schemas.

    It takes all unused passed kwargs and passes them into the
    arbitrary_kv_pairs attribute.

    Args:
        values: Values passed to the object constructor

    Returns:
        Values passed to the object constructor
    """
    all_required_field_names = {
        field.alias
        for field in cls.__fields__.values()
        if field.alias != "arbitrary_kv_pairs"
    }

    arbitrary_kv_pairs: Dict[str, Any] = {
        field_name: values.pop(field_name)
        for field_name in list(values)
        if field_name not in all_required_field_names
    }

    values["arbitrary_kv_pairs"] = arbitrary_kv_pairs
    return values

base_secret

Implementation of the Base SecretSchema class.

BaseSecretSchema (BaseModel, ABC) pydantic-model

Base class for all Secret Schemas.

Source code in zenml/secret/base_secret.py
class BaseSecretSchema(BaseModel, ABC):
    """Base class for all Secret Schemas."""

    name: str
    TYPE: ClassVar[str]

    @property
    def content(self) -> Dict[str, Any]:
        """A dictionary for the content of the SecretSchema.

        The concept of SecretSchemas supports strongly typed secret schemas as
        well as arbitrary collections of key-value pairs. This property unifies
        all attributes into a content dictionary.

        Returns:
            A dictionary containing the content of the SecretSchema.
        """
        fields_dict = self.dict(exclude_none=True)
        fields_dict.pop("name")
        if "arbitrary_kv_pairs" in fields_dict:
            arbitrary_kv_pairs = fields_dict.pop("arbitrary_kv_pairs")
            fields_dict.update(arbitrary_kv_pairs)
        return fields_dict

    @classmethod
    def get_schema_keys(cls) -> List[str]:
        """Get all attribute keys that are not part of the ignored set.

        These schema keys can be used to define all required key-value pairs of
        a secret schema.

        Returns:
            A list of all attribute keys that are not part of the ignored set.
        """
        ignored_keys = ["name", "arbitrary_kv_pairs"]
        return [
            schema_key
            for schema_key in cls.__fields__.keys()
            if schema_key not in ignored_keys
        ]

    class Config:
        """Pydantic configuration class."""

        # validate attribute assignments
        validate_assignment = True
        # report extra attributes as validation failures
        extra = "forbid"
content: Dict[str, Any] property readonly

A dictionary for the content of the SecretSchema.

The concept of SecretSchemas supports strongly typed secret schemas as well as arbitrary collections of key-value pairs. This property unifies all attributes into a content dictionary.

Returns:

Type Description
Dict[str, Any]

A dictionary containing the content of the SecretSchema.

Config

Pydantic configuration class.

Source code in zenml/secret/base_secret.py
class Config:
    """Pydantic configuration class."""

    # validate attribute assignments
    validate_assignment = True
    # report extra attributes as validation failures
    extra = "forbid"
get_schema_keys() classmethod

Get all attribute keys that are not part of the ignored set.

These schema keys can be used to define all required key-value pairs of a secret schema.

Returns:

Type Description
List[str]

A list of all attribute keys that are not part of the ignored set.

Source code in zenml/secret/base_secret.py
@classmethod
def get_schema_keys(cls) -> List[str]:
    """Get all attribute keys that are not part of the ignored set.

    These schema keys can be used to define all required key-value pairs of
    a secret schema.

    Returns:
        A list of all attribute keys that are not part of the ignored set.
    """
    ignored_keys = ["name", "arbitrary_kv_pairs"]
    return [
        schema_key
        for schema_key in cls.__fields__.keys()
        if schema_key not in ignored_keys
    ]

schemas special

Initialization of secret schemas.

aws_secret_schema

AWS Authentication Secret Schema definition.

AWSSecretSchema (BaseSecretSchema) pydantic-model

AWS Authentication Secret Schema definition.

Source code in zenml/secret/schemas/aws_secret_schema.py
class AWSSecretSchema(BaseSecretSchema):
    """AWS Authentication Secret Schema definition."""

    TYPE: ClassVar[str] = AWS_SECRET_SCHEMA_TYPE

    aws_access_key_id: str
    aws_secret_access_key: str
    aws_session_token: Optional[str] = None

azure_secret_schema

Azure Authentication Secret Schema definition.

AzureSecretSchema (BaseSecretSchema) pydantic-model

Azure Authentication Secret Schema definition.

Source code in zenml/secret/schemas/azure_secret_schema.py
class AzureSecretSchema(BaseSecretSchema):
    """Azure Authentication Secret Schema definition."""

    TYPE: ClassVar[str] = AZURE_SECRET_SCHEMA_TYPE

    account_name: Optional[str] = None
    account_key: Optional[str] = None
    sas_token: Optional[str] = None
    connection_string: Optional[str] = None
    client_id: Optional[str] = None
    client_secret: Optional[str] = None
    tenant_id: Optional[str] = None

basic_auth_secret_schema

Basic Authentication Secret Schema definition.

BasicAuthSecretSchema (BaseSecretSchema) pydantic-model

Secret schema for basic authentication.

Attributes:

Name Type Description
username str

The username that should be used for authentication.

password str

The password that should be used for authentication.

Source code in zenml/secret/schemas/basic_auth_secret_schema.py
class BasicAuthSecretSchema(BaseSecretSchema):
    """Secret schema for basic authentication.

    Attributes:
        username: The username that should be used for authentication.
        password: The password that should be used for authentication.
    """

    username: str
    password: str

    # Class configuration
    TYPE: ClassVar[str] = BASIC_AUTH_SCHEMA_TYPE

gcp_secret_schema

GCP Authentication Secret Schema definition.

GCPSecretSchema (BaseSecretSchema) pydantic-model

GCP Authentication Secret Schema definition.

Source code in zenml/secret/schemas/gcp_secret_schema.py
class GCPSecretSchema(BaseSecretSchema):
    """GCP Authentication Secret Schema definition."""

    TYPE: ClassVar[str] = GCP_SECRET_SCHEMA_TYPE

    token: str

    def get_credential_dict(self) -> Dict[str, Any]:
        """Gets a dictionary of credentials for authenticating to GCP.

        Returns:
            A dictionary representing GCP credentials.

        Raises:
            ValueError: If the token value is not a JSON string of a dictionary.
        """
        try:
            dict_ = json.loads(self.token)
        except json.JSONDecodeError:
            raise ValueError(
                "Failed to parse GCP secret token. The token value is not a "
                "valid JSON string."
            )

        if not isinstance(dict_, Dict):
            raise ValueError(
                "Failed to parse GCP secret token. The token value does not "
                "represent a GCP credential dictionary."
            )

        return dict_
get_credential_dict(self)

Gets a dictionary of credentials for authenticating to GCP.

Returns:

Type Description
Dict[str, Any]

A dictionary representing GCP credentials.

Exceptions:

Type Description
ValueError

If the token value is not a JSON string of a dictionary.

Source code in zenml/secret/schemas/gcp_secret_schema.py
def get_credential_dict(self) -> Dict[str, Any]:
    """Gets a dictionary of credentials for authenticating to GCP.

    Returns:
        A dictionary representing GCP credentials.

    Raises:
        ValueError: If the token value is not a JSON string of a dictionary.
    """
    try:
        dict_ = json.loads(self.token)
    except json.JSONDecodeError:
        raise ValueError(
            "Failed to parse GCP secret token. The token value is not a "
            "valid JSON string."
        )

    if not isinstance(dict_, Dict):
        raise ValueError(
            "Failed to parse GCP secret token. The token value does not "
            "represent a GCP credential dictionary."
        )

    return dict_

secret_schema_class_registry

Implementation of the ZenML SecretSchema Class Registry.

SecretSchemaClassRegistry

Registry for SecretSchema classes.

All SecretSchema classes must be registered here so they can be instantiated from the component type and flavor specified inside the ZenML repository configuration.

Source code in zenml/secret/secret_schema_class_registry.py
class SecretSchemaClassRegistry:
    """Registry for SecretSchema classes.

    All SecretSchema classes must be registered here so they can be
    instantiated from the component type and flavor specified inside the
    ZenML repository configuration.
    """

    secret_schema_classes: ClassVar[Dict[str, Type[BaseSecretSchema]]] = dict()

    @classmethod
    def register_class(
        cls,
        secret: Type[BaseSecretSchema],
    ) -> None:
        """Registers a SecretSchema class.

        Args:
            secret: The SecretSchema class to register.
        """
        flavors = cls.secret_schema_classes
        if secret.TYPE in flavors:
            logger.warning(
                "Overwriting previously registered secret schema class `%s` "
                "for type '%s' and flavor '%s'.",
                flavors[secret.TYPE].__class__.__name__,
                secret.TYPE,
            )

        flavors[secret.TYPE] = secret
        logger.debug(
            "Registered secret schema class for type '%s' and flavor '%s'.",
            secret.__class__.__name__,
            secret.TYPE,
        )

    @classmethod
    def get_class(
        cls,
        secret_schema: str,
    ) -> Type[BaseSecretSchema]:
        """Returns the SecretSchema class for the given flavor.

        Args:
            secret_schema: The flavor of the SecretSchema class to return.

        Returns:
            The SecretSchema class for the given flavor.

        Raises:
            KeyError: If no SecretSchema class is registered for the given
                flavor.
        """
        available_schemas = cls.secret_schema_classes
        try:
            return available_schemas[secret_schema]
        except KeyError:
            # The SecretSchema might be part of an integration
            # -> Activate the integrations and try again
            from zenml.integrations.registry import integration_registry

            integration_registry.activate_integrations()

            try:
                return available_schemas[secret_schema]
            except KeyError:
                raise KeyError(
                    f"No SecretSchema class found for schema flavor "
                    f"`{secret_schema}`. Registered flavors are: "
                    f"{set(available_schemas)}. If your secret schema "
                    f"class is part of a ZenML integration, make "
                    f"sure the corresponding integration is installed by "
                    f"running `zenml integration install INTEGRATION_NAME`."
                ) from None
get_class(secret_schema) classmethod

Returns the SecretSchema class for the given flavor.

Parameters:

Name Type Description Default
secret_schema str

The flavor of the SecretSchema class to return.

required

Returns:

Type Description
Type[zenml.secret.base_secret.BaseSecretSchema]

The SecretSchema class for the given flavor.

Exceptions:

Type Description
KeyError

If no SecretSchema class is registered for the given flavor.

Source code in zenml/secret/secret_schema_class_registry.py
@classmethod
def get_class(
    cls,
    secret_schema: str,
) -> Type[BaseSecretSchema]:
    """Returns the SecretSchema class for the given flavor.

    Args:
        secret_schema: The flavor of the SecretSchema class to return.

    Returns:
        The SecretSchema class for the given flavor.

    Raises:
        KeyError: If no SecretSchema class is registered for the given
            flavor.
    """
    available_schemas = cls.secret_schema_classes
    try:
        return available_schemas[secret_schema]
    except KeyError:
        # The SecretSchema might be part of an integration
        # -> Activate the integrations and try again
        from zenml.integrations.registry import integration_registry

        integration_registry.activate_integrations()

        try:
            return available_schemas[secret_schema]
        except KeyError:
            raise KeyError(
                f"No SecretSchema class found for schema flavor "
                f"`{secret_schema}`. Registered flavors are: "
                f"{set(available_schemas)}. If your secret schema "
                f"class is part of a ZenML integration, make "
                f"sure the corresponding integration is installed by "
                f"running `zenml integration install INTEGRATION_NAME`."
            ) from None
register_class(secret) classmethod

Registers a SecretSchema class.

Parameters:

Name Type Description Default
secret Type[zenml.secret.base_secret.BaseSecretSchema]

The SecretSchema class to register.

required
Source code in zenml/secret/secret_schema_class_registry.py
@classmethod
def register_class(
    cls,
    secret: Type[BaseSecretSchema],
) -> None:
    """Registers a SecretSchema class.

    Args:
        secret: The SecretSchema class to register.
    """
    flavors = cls.secret_schema_classes
    if secret.TYPE in flavors:
        logger.warning(
            "Overwriting previously registered secret schema class `%s` "
            "for type '%s' and flavor '%s'.",
            flavors[secret.TYPE].__class__.__name__,
            secret.TYPE,
        )

    flavors[secret.TYPE] = secret
    logger.debug(
        "Registered secret schema class for type '%s' and flavor '%s'.",
        secret.__class__.__name__,
        secret.TYPE,
    )

register_secret_schema_class(cls)

Registers the SecretSchema class and returns it unmodified.

Parameters:

Name Type Description Default
cls Type[~C]

The SecretSchema class to register.

required

Returns:

Type Description
Type[~C]

The (unmodified) SecretSchema class to register.

Source code in zenml/secret/secret_schema_class_registry.py
def register_secret_schema_class(cls: Type[C]) -> Type[C]:
    """Registers the SecretSchema class and returns it unmodified.

    Args:
        cls: The SecretSchema class to register.

    Returns:
        The (unmodified) SecretSchema class to register.
    """
    SecretSchemaClassRegistry.register_class(secret=cls)
    return cls