Secret

zenml.secret special

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

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."""

    schema_type: SecretSchemaType = SecretSchemaType.ARBITRARY

    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 that 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 that 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 that 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

BaseSecretSchema (BaseModel, ABC) pydantic-model

Source code in zenml/secret/base_secret.py

class BaseSecretSchema(BaseModel, ABC):
    name: str
    schema_type: SecretSchemaType

    @property
    def content(self) -> Dict[str, Any]:
        """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()
        fields_dict.pop("name")
        fields_dict.pop("schema_type")
        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", "schema_type"]
        return [
            schema_key
            for schema_key in cls.__fields__.keys()
            if schema_key not in ignored_keys
        ]

content: Dict[str, Any] property readonly

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.

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", "schema_type"]
    return [
        schema_key
        for schema_key in cls.__fields__.keys()
        if schema_key not in ignored_keys
    ]

secret_schema_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_schema: SecretSchemaType,
        secret: Type[BaseSecretSchema],
    ) -> None:
        """Registers a SecretSchema class.

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

        flavors[secret_schema_value] = secret
        logger.debug(
            "Registered stack component class for type '%s' and flavor '%s'.",
            secret_schema_value,
        )

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

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

        Returns:
            The SecretSchema class for the given type and flavor.

        Raises:
            KeyError: If no SecretSchema class is registered for the given type
                and flavor.
        """
        if isinstance(secret_schema, SecretSchemaType):
            secret_schema = secret_schema.value

        available_schemas = cls.secret_schema_classes
        try:
            return available_schemas[secret_schema]
        except KeyError:
            # TODO [ENG-722]: Probably remove this exception catch. Doesn't apply
            # to SecretSchema
            # 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 SecretSet  "
                    f"of flavor {secret_schema}. Registered flavors are:"
                    f" {set(available_schemas)}."
                ) from None

get_class(secret_schema) classmethod

Returns the SecretSchema class for the given type and flavor.

Parameters:

Name	Type	Description	Default
secret_schema	Union[zenml.enums.SecretSchemaType, 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 type and flavor.

Exceptions:

Type	Description
KeyError	If no SecretSchema class is registered for the given type and flavor.

Source code in zenml/secret/secret_schema_class_registry.py

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

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

    Returns:
        The SecretSchema class for the given type and flavor.

    Raises:
        KeyError: If no SecretSchema class is registered for the given type
            and flavor.
    """
    if isinstance(secret_schema, SecretSchemaType):
        secret_schema = secret_schema.value

    available_schemas = cls.secret_schema_classes
    try:
        return available_schemas[secret_schema]
    except KeyError:
        # TODO [ENG-722]: Probably remove this exception catch. Doesn't apply
        # to SecretSchema
        # 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 SecretSet  "
                f"of flavor {secret_schema}. Registered flavors are:"
                f" {set(available_schemas)}."
            ) from None

register_class(secret_schema, secret) classmethod

Registers a SecretSchema class.

Parameters:

Name	Type	Description	Default
secret_schema	SecretSchemaType	The flavor of the SecretSchema class to register.	required
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_schema: SecretSchemaType,
    secret: Type[BaseSecretSchema],
) -> None:
    """Registers a SecretSchema class.

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

    flavors[secret_schema_value] = secret
    logger.debug(
        "Registered stack component class for type '%s' and flavor '%s'.",
        secret_schema_value,
    )

register_secret_schema_class(secret_schema)

Parametrized decorator function to register SecretSchema classes.

Parameters:

Name	Type	Description	Default
secret_schema	SecretSchemaType	The flavor of the SecretSchema class to register.	required

Returns:

Type	Description
Callable[[Type[~C]], Type[~C]]	A decorator function that registers and returns the decorated SecretSchema class.

Source code in zenml/secret/secret_schema_class_registry.py

def register_secret_schema_class(
    secret_schema: SecretSchemaType,
) -> Callable[[Type[C]], Type[C]]:
    """Parametrized decorator function to register SecretSchema classes.

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

    Returns:
        A decorator function that registers and returns the decorated SecretSchema class.
    """

    def decorator_function(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_schema=secret_schema,
            secret=cls,
        )
        return cls

    return decorator_function