Zen Stores

zenml.zen_stores

ZenStores define ways to store ZenML relevant data locally or remotely.

Modules

base_zen_store

Base Zen Store implementation.

Classes

BaseZenStore(skip_default_registrations: bool = False, **kwargs: Any)

Bases: BaseModel, ZenStoreInterface, ABC

Base class for accessing and persisting ZenML core objects.

Attributes:

Name	Type	Description
config	StoreConfiguration	The configuration of the store.

Create and initialize a store.

Parameters:

Name	Type	Description	Default
skip_default_registrations	bool	If True, the creation of the default stack and user in the store will be skipped.	False
**kwargs	Any	Additional keyword arguments to pass to the Pydantic constructor.	{}

Source code in src/zenml/zen_stores/base_zen_store.py

def __init__(
    self,
    skip_default_registrations: bool = False,
    **kwargs: Any,
) -> None:
    """Create and initialize a store.

    Args:
        skip_default_registrations: If `True`, the creation of the default
            stack and user in the store will be skipped.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.
    """
    super().__init__(**kwargs)

    self._initialize()

    if not skip_default_registrations:
        logger.debug("Initializing database")
        self._initialize_database()
    else:
        logger.debug("Skipping database initialization")

Attributes

type: StoreType property

The type of the store.

Returns:

Type	Description
StoreType	The type of the store.

url: str property

The URL of the store.

Returns:

Type	Description
str	The URL of the store.

Functions

convert_config(data: Dict[str, Any]) -> Dict[str, Any] classmethod

Method to infer the correct type of the config and convert.

Parameters:

Name	Type	Description	Default
data	Dict[str, Any]	The provided configuration object, can potentially be a generic object	required

Raises:

Type	Description
ValueError	If the provided config object's type does not match any of the current implementations.

Returns:

Type	Description
Dict[str, Any]	The converted configuration object.

Source code in src/zenml/zen_stores/base_zen_store.py

@model_validator(mode="before")
@classmethod
@before_validator_handler
def convert_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
    """Method to infer the correct type of the config and convert.

    Args:
        data: The provided configuration object, can potentially be a
            generic object

    Raises:
        ValueError: If the provided config object's type does not match
            any of the current implementations.

    Returns:
        The converted configuration object.
    """
    if data["config"].type == StoreType.SQL:
        from zenml.zen_stores.sql_zen_store import SqlZenStoreConfiguration

        data["config"] = SqlZenStoreConfiguration(
            **data["config"].model_dump()
        )

    elif data["config"].type == StoreType.REST:
        from zenml.zen_stores.rest_zen_store import (
            RestZenStoreConfiguration,
        )

        data["config"] = RestZenStoreConfiguration(
            **data["config"].model_dump()
        )
    else:
        raise ValueError(
            f"Unknown type '{data['config'].type}' for the configuration."
        )

    return data

create_store(config: StoreConfiguration, skip_default_registrations: bool = False, **kwargs: Any) -> BaseZenStore staticmethod

Create and initialize a store from a store configuration.

Parameters:

Name	Type	Description	Default
config	StoreConfiguration	The store configuration to use.	required
skip_default_registrations	bool	If True, the creation of the default stack and user in the store will be skipped.	False
**kwargs	Any	Additional keyword arguments to pass to the store class	{}

Returns:

Type	Description
BaseZenStore	The initialized store.

Source code in src/zenml/zen_stores/base_zen_store.py

@staticmethod
def create_store(
    config: StoreConfiguration,
    skip_default_registrations: bool = False,
    **kwargs: Any,
) -> "BaseZenStore":
    """Create and initialize a store from a store configuration.

    Args:
        config: The store configuration to use.
        skip_default_registrations: If `True`, the creation of the default
            stack and user in the store will be skipped.
        **kwargs: Additional keyword arguments to pass to the store class

    Returns:
        The initialized store.
    """
    store_class = BaseZenStore.get_store_class(config.type)
    store = store_class(
        config=config,
        skip_default_registrations=skip_default_registrations,
        **kwargs,
    )

    return store

get_default_store_config(path: str) -> StoreConfiguration staticmethod

Get the default store configuration.

The default store is a SQLite store that saves the DB contents on the local filesystem.

Parameters:

Name	Type	Description	Default
path	str	The local path where the store DB will be stored.	required

Returns:

Type	Description
StoreConfiguration	The default store configuration.

Source code in src/zenml/zen_stores/base_zen_store.py

@staticmethod
def get_default_store_config(path: str) -> StoreConfiguration:
    """Get the default store configuration.

    The default store is a SQLite store that saves the DB contents on the
    local filesystem.

    Args:
        path: The local path where the store DB will be stored.

    Returns:
        The default store configuration.
    """
    from zenml.zen_stores.secrets_stores.sql_secrets_store import (
        SqlSecretsStoreConfiguration,
    )
    from zenml.zen_stores.sql_zen_store import SqlZenStoreConfiguration

    config = SqlZenStoreConfiguration(
        type=StoreType.SQL,
        url=SqlZenStoreConfiguration.get_local_url(path),
        secrets_store=SqlSecretsStoreConfiguration(
            type=SecretsStoreType.SQL,
        ),
    )
    return config

get_store_class(store_type: StoreType) -> Type[BaseZenStore] staticmethod

Returns the class of the given store type.

Parameters:

Name	Type	Description	Default
store_type	StoreType	The type of the store to get the class for.	required

Returns:

Type	Description
Type[BaseZenStore]	The class of the given store type or None if the type is unknown.

Raises:

Type	Description
TypeError	If the store type is unsupported.

Source code in src/zenml/zen_stores/base_zen_store.py

@staticmethod
def get_store_class(store_type: StoreType) -> Type["BaseZenStore"]:
    """Returns the class of the given store type.

    Args:
        store_type: The type of the store to get the class for.

    Returns:
        The class of the given store type or None if the type is unknown.

    Raises:
        TypeError: If the store type is unsupported.
    """
    if store_type == StoreType.SQL:
        if os.environ.get(ENV_ZENML_SERVER):
            from zenml.zen_server.rbac.rbac_sql_zen_store import (
                RBACSqlZenStore,
            )

            return RBACSqlZenStore
        else:
            from zenml.zen_stores.sql_zen_store import SqlZenStore

            return SqlZenStore
    elif store_type == StoreType.REST:
        from zenml.zen_stores.rest_zen_store import RestZenStore

        return RestZenStore
    else:
        raise TypeError(
            f"No store implementation found for store type "
            f"`{store_type.value}`."
        )

get_store_config_class(store_type: StoreType) -> Type[StoreConfiguration] staticmethod

Returns the store config class of the given store type.

Parameters:

Name	Type	Description	Default
store_type	StoreType	The type of the store to get the class for.	required

Returns:

Type	Description
Type[StoreConfiguration]	The config class of the given store type.

Source code in src/zenml/zen_stores/base_zen_store.py

@staticmethod
def get_store_config_class(
    store_type: StoreType,
) -> Type["StoreConfiguration"]:
    """Returns the store config class of the given store type.

    Args:
        store_type: The type of the store to get the class for.

    Returns:
        The config class of the given store type.
    """
    store_class = BaseZenStore.get_store_class(store_type)
    return store_class.CONFIG_TYPE

get_store_info() -> ServerModel

Get information about the store.

Returns:

Type	Description
ServerModel	Information about the store.

Source code in src/zenml/zen_stores/base_zen_store.py

def get_store_info(self) -> ServerModel:
    """Get information about the store.

    Returns:
        Information about the store.
    """
    server_config = ServerConfiguration.get_server_config()
    deployment_type = server_config.deployment_type
    auth_scheme = server_config.auth_scheme
    metadata = server_config.metadata
    secrets_store_type = SecretsStoreType.NONE
    if self.config.type == StoreType.SQL and self.config.secrets_store:
        secrets_store_type = self.config.secrets_store.type
    store_info = ServerModel(
        id=GlobalConfiguration().user_id,
        active=True,
        version=zenml.__version__,
        deployment_type=deployment_type,
        database_type=ServerDatabaseType.OTHER,
        debug=IS_DEBUG_ENV,
        secrets_store_type=secrets_store_type,
        auth_scheme=auth_scheme,
        server_url=server_config.server_url or "",
        dashboard_url=server_config.dashboard_url or "",
        analytics_enabled=GlobalConfiguration().analytics_opt_in,
        metadata=metadata,
    )

    # Add ZenML Pro specific store information to the server model, if available.
    if store_info.deployment_type == ServerDeploymentType.CLOUD:
        from zenml.config.server_config import ServerProConfiguration

        pro_config = ServerProConfiguration.get_server_config()

        store_info.pro_api_url = pro_config.api_url
        store_info.pro_dashboard_url = pro_config.dashboard_url
        store_info.pro_organization_id = pro_config.organization_id
        store_info.pro_workspace_id = pro_config.workspace_id
        if pro_config.workspace_name:
            store_info.pro_workspace_name = pro_config.workspace_name
        if pro_config.organization_name:
            store_info.pro_organization_name = pro_config.organization_name

    return store_info

get_store_type(url: str) -> StoreType staticmethod

Returns the store type associated with a URL schema.

Parameters:

Name	Type	Description	Default
url	str	The store URL.	required

Returns:

Type	Description
StoreType	The store type associated with the supplied URL schema.

Raises:

Type	Description
TypeError	If no store type was found to support the supplied URL.

Source code in src/zenml/zen_stores/base_zen_store.py

@staticmethod
def get_store_type(url: str) -> StoreType:
    """Returns the store type associated with a URL schema.

    Args:
        url: The store URL.

    Returns:
        The store type associated with the supplied URL schema.

    Raises:
        TypeError: If no store type was found to support the supplied URL.
    """
    from zenml.zen_stores.rest_zen_store import RestZenStoreConfiguration

    if RestZenStoreConfiguration.supports_url_scheme(url):
        return StoreType.REST

    # Only import this once we've made sure it's not a REST URL, as the
    # zenml package without the local extra will fail this import due to
    # missing database dependencies.
    from zenml.zen_stores.sql_zen_store import SqlZenStoreConfiguration

    if SqlZenStoreConfiguration.supports_url_scheme(url):
        return StoreType.SQL

    raise TypeError(f"No store implementation found for URL: {url}.")

is_local_store() -> bool

Check if the store is local or connected to a local ZenML server.

Returns:

Type	Description
bool	True if the store is local, False otherwise.

Source code in src/zenml/zen_stores/base_zen_store.py

def is_local_store(self) -> bool:
    """Check if the store is local or connected to a local ZenML server.

    Returns:
        True if the store is local, False otherwise.
    """
    return self.get_store_info().is_local()

validate_active_config(active_project_id: Optional[UUID] = None, active_stack_id: Optional[UUID] = None, config_name: str = '') -> Tuple[Optional[ProjectResponse], StackResponse]

Validate the active configuration.

Call this method to validate the supplied active project and active stack values.

This method returns a valid project and stack values. If the supplied project and stack are not set or are not valid (e.g. they do not exist or are not accessible), the default project and default stack will be returned in their stead.

Parameters:

Name	Type	Description	Default
active_project_id	Optional[UUID]	The ID of the active project.	None
active_stack_id	Optional[UUID]	The ID of the active stack.	None
config_name	str	The name of the configuration to validate (used in the displayed logs/messages).	''

Returns:

Type	Description
Tuple[Optional[ProjectResponse], StackResponse]	A tuple containing the active project and active stack.

Source code in src/zenml/zen_stores/base_zen_store.py

def validate_active_config(
    self,
    active_project_id: Optional[UUID] = None,
    active_stack_id: Optional[UUID] = None,
    config_name: str = "",
) -> Tuple[Optional[ProjectResponse], StackResponse]:
    """Validate the active configuration.

    Call this method to validate the supplied active project and active
    stack values.

    This method returns a valid project and stack values. If the
    supplied project and stack are not set or are not valid (e.g. they
    do not exist or are not accessible), the default project and default
    stack will be returned in their stead.

    Args:
        active_project_id: The ID of the active project.
        active_stack_id: The ID of the active stack.
        config_name: The name of the configuration to validate (used in the
            displayed logs/messages).

    Returns:
        A tuple containing the active project and active stack.
    """
    active_project: Optional[ProjectResponse] = None

    if active_project_id:
        try:
            active_project = self.get_project(active_project_id)
        except (KeyError, IllegalOperationError):
            active_project_id = None
            logger.warning(
                f"The current {config_name} active project is no longer "
                f"available."
            )

    if active_project is None:
        user = self.get_user()
        if user.default_project_id:
            try:
                active_project = self.get_project(user.default_project_id)
            except (KeyError, IllegalOperationError):
                logger.warning(
                    "The default project %s for the active user is no "
                    "longer available.",
                    user.default_project_id,
                )
            else:
                logger.info(
                    f"Setting the {config_name} active project "
                    f"to '{active_project.name}'."
                )

    if active_project is None:
        try:
            projects = self.list_projects(
                project_filter_model=ProjectFilter()
            )
        except Exception:
            pass
        else:
            if len(projects) == 1:
                active_project = projects.items[0]
                logger.info(
                    f"Setting the {config_name} active project "
                    f"to '{active_project.name}'."
                )

    active_stack: StackResponse

    # Sanitize the active stack
    if active_stack_id:
        # Ensure that the active stack is still valid
        try:
            active_stack = self.get_stack(stack_id=active_stack_id)
        except (KeyError, IllegalOperationError):
            logger.warning(
                "The current %s active stack is no longer available. "
                "Resetting the active stack to default.",
                config_name,
            )
            active_stack = self._get_default_stack()

    else:
        logger.warning(
            "Setting the %s active stack to default.",
            config_name,
        )
        active_stack = self._get_default_stack()

    return active_project, active_stack

Functions

dag_generator

DAG generator helper.

Classes

DAGGeneratorHelper()

Helper class for generating pipeline run DAGs.

Initialize the DAG generator helper.

Source code in src/zenml/zen_stores/dag_generator.py

def __init__(self) -> None:
    """Initialize the DAG generator helper."""
    self.step_nodes: Dict[str, PipelineRunDAG.Node] = {}
    self.artifact_nodes: Dict[str, PipelineRunDAG.Node] = {}
    self.wait_condition_nodes: Dict[str, PipelineRunDAG.Node] = {}
    self.triggered_run_nodes: Dict[str, PipelineRunDAG.Node] = {}
    self.child_run_nodes: Dict[str, PipelineRunDAG.Node] = {}
    self.edges: List[PipelineRunDAG.Edge] = []

Functions

add_artifact_node(node_id: str, name: str, id: Optional[UUID] = None, **metadata: Any) -> PipelineRunDAG.Node

Add an artifact node to the DAG.

Parameters:

Name	Type	Description	Default
node_id	str	The ID of the node.	required
name	str	The name of the artifact.	required
id	Optional[UUID]	The ID of the artifact.	None
**metadata	Any	Additional node metadata.	{}

Returns:

Type	Description
Node	The added artifact node.

Source code in src/zenml/zen_stores/dag_generator.py

def add_artifact_node(
    self,
    node_id: str,
    name: str,
    id: Optional[UUID] = None,
    **metadata: Any,
) -> PipelineRunDAG.Node:
    """Add an artifact node to the DAG.

    Args:
        node_id: The ID of the node.
        name: The name of the artifact.
        id: The ID of the artifact.
        **metadata: Additional node metadata.

    Returns:
        The added artifact node.
    """
    artifact_node = PipelineRunDAG.Node(
        type="artifact",
        node_id=node_id,
        id=id,
        name=name,
        metadata=metadata,
    )
    self.artifact_nodes[artifact_node.node_id] = artifact_node
    return artifact_node

add_child_run_node(node_id: str, name: str, id: Optional[UUID] = None, **metadata: Any) -> PipelineRunDAG.Node

Add a child run node to the DAG.

Parameters:

Name	Type	Description	Default
node_id	str	The node ID.	required
name	str	The child run name.	required
id	Optional[UUID]	The child run ID.	None
**metadata	Any	Additional node metadata.	{}

Returns:

Type	Description
Node	The added child run node.

Source code in src/zenml/zen_stores/dag_generator.py

def add_child_run_node(
    self,
    node_id: str,
    name: str,
    id: Optional[UUID] = None,
    **metadata: Any,
) -> PipelineRunDAG.Node:
    """Add a child run node to the DAG.

    Args:
        node_id: The node ID.
        name: The child run name.
        id: The child run ID.
        **metadata: Additional node metadata.

    Returns:
        The added child run node.
    """
    child_run_node = PipelineRunDAG.Node(
        # TODO: change to child_run once the UI supports it
        type="triggered_run",
        id=id,
        node_id=node_id,
        name=name,
        metadata=metadata,
    )
    self.child_run_nodes[child_run_node.node_id] = child_run_node
    return child_run_node

add_edge(source: str, target: str, **metadata: Any) -> None

Add an edge to the DAG.

Parameters:

Name	Type	Description	Default
source	str	The source node ID.	required
target	str	The target node ID.	required
metadata	Any	Additional edge metadata.	{}

Source code in src/zenml/zen_stores/dag_generator.py

def add_edge(self, source: str, target: str, **metadata: Any) -> None:
    """Add an edge to the DAG.

    Args:
        source: The source node ID.
        target: The target node ID.
        metadata: Additional edge metadata.
    """
    self.edges.append(
        PipelineRunDAG.Edge(
            source=source, target=target, metadata=metadata
        )
    )

add_step_node(node_id: str, name: str, id: Optional[UUID] = None, **metadata: Any) -> PipelineRunDAG.Node

Add a step node to the DAG.

Parameters:

Name	Type	Description	Default
node_id	str	The ID of the node.	required
name	str	The name of the step.	required
id	Optional[UUID]	The ID of the step.	None
**metadata	Any	Additional node metadata.	{}

Returns:

Type	Description
Node	The added step node.

Source code in src/zenml/zen_stores/dag_generator.py

def add_step_node(
    self,
    node_id: str,
    name: str,
    id: Optional[UUID] = None,
    **metadata: Any,
) -> PipelineRunDAG.Node:
    """Add a step node to the DAG.

    Args:
        node_id: The ID of the node.
        name: The name of the step.
        id: The ID of the step.
        **metadata: Additional node metadata.

    Returns:
        The added step node.
    """
    step_node = PipelineRunDAG.Node(
        type="step",
        id=id,
        node_id=node_id,
        name=name,
        metadata=metadata,
    )
    self.step_nodes[step_node.node_id] = step_node
    return step_node

add_triggered_run_node(node_id: str, name: str, id: Optional[UUID] = None, **metadata: Any) -> PipelineRunDAG.Node

Add a triggered run node to the DAG.

Parameters:

Name	Type	Description	Default
node_id	str	The ID of the node.	required
name	str	The name of the triggered run.	required
id	Optional[UUID]	The ID of the triggered run.	None
**metadata	Any	Additional node metadata.	{}

Returns:

Type	Description
Node	The added triggered run node.

Source code in src/zenml/zen_stores/dag_generator.py

def add_triggered_run_node(
    self,
    node_id: str,
    name: str,
    id: Optional[UUID] = None,
    **metadata: Any,
) -> PipelineRunDAG.Node:
    """Add a triggered run node to the DAG.

    Args:
        node_id: The ID of the node.
        name: The name of the triggered run.
        id: The ID of the triggered run.
        **metadata: Additional node metadata.

    Returns:
        The added triggered run node.
    """
    triggered_run_node = PipelineRunDAG.Node(
        type="triggered_run",
        id=id,
        node_id=node_id,
        name=name,
        metadata=metadata,
    )
    self.triggered_run_nodes[triggered_run_node.node_id] = (
        triggered_run_node
    )
    return triggered_run_node

add_wait_condition_node(node_id: str, name: str, id: Optional[UUID] = None, **metadata: Any) -> PipelineRunDAG.Node

Add a wait condition node to the DAG.

Parameters:

Name	Type	Description	Default
node_id	str	The DAG node ID.	required
name	str	The wait condition display name.	required
id	Optional[UUID]	The wait condition ID.	None
**metadata	Any	Additional node metadata.	{}

Returns:

Type	Description
Node	The added wait condition node.

Source code in src/zenml/zen_stores/dag_generator.py

def add_wait_condition_node(
    self,
    node_id: str,
    name: str,
    id: Optional[UUID] = None,
    **metadata: Any,
) -> PipelineRunDAG.Node:
    """Add a wait condition node to the DAG.

    Args:
        node_id: The DAG node ID.
        name: The wait condition display name.
        id: The wait condition ID.
        **metadata: Additional node metadata.

    Returns:
        The added wait condition node.
    """
    wait_condition_node = PipelineRunDAG.Node(
        type="wait_condition",
        id=id,
        node_id=node_id,
        name=name,
        metadata=metadata,
    )
    self.wait_condition_nodes[wait_condition_node.node_id] = (
        wait_condition_node
    )
    return wait_condition_node

finalize_dag(pipeline_run_id: UUID, status: ExecutionStatus) -> PipelineRunDAG

Finalize the DAG.

Parameters:

Name	Type	Description	Default
pipeline_run_id	UUID	The ID of the pipeline run.	required
status	ExecutionStatus	The status of the pipeline run.	required

Returns:

Type	Description
PipelineRunDAG	The finalized DAG.

Source code in src/zenml/zen_stores/dag_generator.py

def finalize_dag(
    self, pipeline_run_id: UUID, status: ExecutionStatus
) -> PipelineRunDAG:
    """Finalize the DAG.

    Args:
        pipeline_run_id: The ID of the pipeline run.
        status: The status of the pipeline run.

    Returns:
        The finalized DAG.
    """
    return PipelineRunDAG(
        id=pipeline_run_id,
        status=status,
        nodes=list(self.step_nodes.values())
        + list(self.artifact_nodes.values())
        + list(self.wait_condition_nodes.values())
        + list(self.triggered_run_nodes.values())
        + list(self.child_run_nodes.values()),
        edges=self.edges,
    )

get_artifact_node_id(name: str, step_name: str, io_type: str, is_input: bool) -> str

Get the ID of an artifact node.

Parameters:

Name	Type	Description	Default
name	str	The name of the input or output artifact.	required
step_name	str	The name of the step.	required
io_type	str	The type of the input or output artifact.	required
is_input	bool	Whether the artifact is an input or output artifact.	required

Returns:

Type	Description
str	The ID of the artifact node.

Source code in src/zenml/zen_stores/dag_generator.py

def get_artifact_node_id(
    self, name: str, step_name: str, io_type: str, is_input: bool
) -> str:
    """Get the ID of an artifact node.

    Args:
        name: The name of the input or output artifact.
        step_name: The name of the step.
        io_type: The type of the input or output artifact.
        is_input: Whether the artifact is an input or output artifact.

    Returns:
        The ID of the artifact node.
    """
    # Make sure there is no slashes as we use them as delimiters
    name = name.replace("/", "-")
    step_name = step_name.replace("/", "-")
    io_str = "inputs" if is_input else "outputs"

    return f"{step_name}/{io_str}/{io_type}/{name}"

get_child_run_node_id(name: str) -> str

Get the ID of a child pipeline run node.

Parameters:

Name	Type	Description	Default
name	str	The child run name.	required

Returns:

Type	Description
str	The ID of the child run node.

Source code in src/zenml/zen_stores/dag_generator.py

def get_child_run_node_id(self, name: str) -> str:
    """Get the ID of a child pipeline run node.

    Args:
        name: The child run name.

    Returns:
        The ID of the child run node.
    """
    name = name.replace("/", "-")
    return f"child_run/{name}"

get_step_node_by_name(name: str) -> PipelineRunDAG.Node

Get a step node by name.

Parameters:

Name	Type	Description	Default
name	str	The name of the step.	required

Raises:

Type	Description
KeyError	If the step node with the given name is not found.

Returns:

Type	Description
Node	The step node.

Source code in src/zenml/zen_stores/dag_generator.py

def get_step_node_by_name(self, name: str) -> PipelineRunDAG.Node:
    """Get a step node by name.

    Args:
        name: The name of the step.

    Raises:
        KeyError: If the step node with the given name is not found.

    Returns:
        The step node.
    """
    for node in self.step_nodes.values():
        if node.name == name:
            return node
    raise KeyError(f"Step node with name {name} not found")

get_step_node_id(name: str) -> str

Get the ID of a step node.

Parameters:

Name	Type	Description	Default
name	str	The name of the step.	required

Returns:

Type	Description
str	The ID of the step node.

Source code in src/zenml/zen_stores/dag_generator.py

def get_step_node_id(self, name: str) -> str:
    """Get the ID of a step node.

    Args:
        name: The name of the step.

    Returns:
        The ID of the step node.
    """
    # Make sure there is no slashes as we use them as delimiters
    name = name.replace("/", "-")
    return f"step/{name}"

get_triggered_run_node_id(name: str) -> str

Get the ID of a triggered run node.

Parameters:

Name	Type	Description	Default
name	str	The name of the triggered run.	required

Returns:

Type	Description
str	The ID of the triggered run node.

Source code in src/zenml/zen_stores/dag_generator.py

def get_triggered_run_node_id(self, name: str) -> str:
    """Get the ID of a triggered run node.

    Args:
        name: The name of the triggered run.

    Returns:
        The ID of the triggered run node.
    """
    # Make sure there is no slashes as we use them as delimiters
    name = name.replace("/", "-")
    return f"run/{name}"

get_wait_condition_node_id(name: str) -> str

Get the ID of a wait condition node.

Parameters:

Name	Type	Description	Default
name	str	The wait condition name.	required

Returns:

Type	Description
str	The ID of the wait condition node.

Source code in src/zenml/zen_stores/dag_generator.py

def get_wait_condition_node_id(self, name: str) -> str:
    """Get the ID of a wait condition node.

    Args:
        name: The wait condition name.

    Returns:
        The ID of the wait condition node.
    """
    # Make sure there is no slashes as we use them as delimiters
    name = name.replace("/", "-")
    return f"wait_condition/{name}"

migrations

Alembic database migration utilities.

Modules

alembic

Alembic utilities wrapper.

The Alembic class defined here acts as a wrapper around the Alembic library that automatically configures Alembic to use the ZenML SQL store database connection.

Classes

Alembic(engine: Engine, metadata: MetaData = SQLModel.metadata, context: Optional[EnvironmentContext] = None, **kwargs: Any)

Alembic environment and migration API.

This class provides a wrapper around the Alembic library that automatically configures Alembic to use the ZenML SQL store database connection.

Initialize the Alembic wrapper.

Parameters:

Name	Type	Description	Default
engine	Engine	The SQLAlchemy engine to use.	required
metadata	MetaData	The SQLAlchemy metadata to use.	metadata
context	Optional[EnvironmentContext]	The Alembic environment context to use. If not set, a new context is created pointing to the ZenML migrations directory.	None
**kwargs	Any	Additional keyword arguments to pass to the Alembic environment context.	{}

Source code in src/zenml/zen_stores/migrations/alembic.py

def __init__(
    self,
    engine: Engine,
    metadata: MetaData = SQLModel.metadata,
    context: Optional[EnvironmentContext] = None,
    **kwargs: Any,
) -> None:
    """Initialize the Alembic wrapper.

    Args:
        engine: The SQLAlchemy engine to use.
        metadata: The SQLAlchemy metadata to use.
        context: The Alembic environment context to use. If not set, a new
            context is created pointing to the ZenML migrations directory.
        **kwargs: Additional keyword arguments to pass to the Alembic
            environment context.
    """
    self.engine = engine
    self.metadata = metadata
    self.context_kwargs = kwargs

    self.config = Config()
    self.config.set_main_option(
        "script_location", str(Path(__file__).parent)
    )

    self.script_directory = ScriptDirectory.from_config(self.config)
    if context is None:
        self.environment_context = EnvironmentContext(
            self.config, self.script_directory
        )
    else:
        self.environment_context = context

Functions

current_revisions() -> List[str]

Get the current database revisions.

Returns:

Type	Description
List[str]	List of head revisions.

Source code in src/zenml/zen_stores/migrations/alembic.py

def current_revisions(self) -> List[str]:
    """Get the current database revisions.

    Returns:
        List of head revisions.
    """
    current_revisions: List[str] = []

    def do_get_current_rev(rev: _RevIdType, context: Any) -> List[Any]:
        nonlocal current_revisions

        # Handle rev parameter in a way that's compatible with different alembic versions
        rev_input: Any
        if isinstance(rev, str):
            rev_input = rev
        else:
            rev_input = tuple(str(r) for r in rev)

        # Get current revision(s)
        for r in self.script_directory.get_all_current(rev_input):
            if r is None:
                continue
            current_revisions.append(r.revision)
        return []

    self.run_migrations(do_get_current_rev)

    return current_revisions

db_is_empty() -> bool

Check if the database is empty.

Returns:

Type	Description
bool	True if the database is empty, False otherwise.

Source code in src/zenml/zen_stores/migrations/alembic.py

def db_is_empty(self) -> bool:
    """Check if the database is empty.

    Returns:
        True if the database is empty, False otherwise.
    """
    # Check the existence of any of the SQLModel tables
    return not self.engine.dialect.has_table(
        self.engine.connect(), schemas.StackSchema.__tablename__
    )

downgrade(revision: str) -> None

Revert the database to a previous version.

Parameters:

Name	Type	Description	Default
revision	str	String revision target.	required

Source code in src/zenml/zen_stores/migrations/alembic.py

def downgrade(self, revision: str) -> None:
    """Revert the database to a previous version.

    Args:
        revision: String revision target.
    """

    def do_downgrade(rev: _RevIdType, context: Any) -> List[Any]:
        # Handle rev parameter in a way that's compatible with different alembic versions
        if isinstance(rev, str):
            return self.script_directory._downgrade_revs(revision, rev)
        else:
            if rev:
                return self.script_directory._downgrade_revs(
                    revision, str(rev[0])
                )
            return self.script_directory._downgrade_revs(revision, None)

    self.run_migrations(do_downgrade)

head_revisions() -> List[str]

Get the head database revisions.

Returns:

Type	Description
List[str]	List of head revisions.

Source code in src/zenml/zen_stores/migrations/alembic.py

def head_revisions(self) -> List[str]:
    """Get the head database revisions.

    Returns:
        List of head revisions.
    """
    head_revisions: List[str] = []

    def do_get_head_rev(rev: _RevIdType, context: Any) -> List[Any]:
        nonlocal head_revisions

        for r in self.script_directory.get_heads():
            if r is None:
                continue
            head_revisions.append(r)
        return []

    self.run_migrations(do_get_head_rev)

    return head_revisions

run_migrations(fn: Optional[Callable[[_RevIdType, MigrationContext], List[Any]]]) -> None

Run an online migration function in the current migration context.

Parameters:

Name	Type	Description	Default
fn	Optional[Callable[[_RevIdType, MigrationContext], List[Any]]]	Migration function to run. If not set, the function configured externally by the Alembic CLI command is used.	required

Source code in src/zenml/zen_stores/migrations/alembic.py

def run_migrations(
    self,
    fn: Optional[Callable[[_RevIdType, MigrationContext], List[Any]]],
) -> None:
    """Run an online migration function in the current migration context.

    Args:
        fn: Migration function to run. If not set, the function configured
            externally by the Alembic CLI command is used.
    """
    fn_context_args: Dict[Any, Any] = {}
    if fn is not None:
        fn_context_args["fn"] = fn

    with self.engine.connect() as connection:
        # Configure the context with our metadata
        self.environment_context.configure(
            connection=connection,
            target_metadata=self.metadata,
            include_object=include_object,
            compare_type=True,
            render_as_batch=True,
            **fn_context_args,
            **self.context_kwargs,
        )

        with self.environment_context.begin_transaction():
            self.environment_context.run_migrations()

stamp(revision: str) -> None

Stamp the revision table with the given revision without running any migrations.

Parameters:

Name	Type	Description	Default
revision	str	String revision target.	required

Source code in src/zenml/zen_stores/migrations/alembic.py

def stamp(self, revision: str) -> None:
    """Stamp the revision table with the given revision without running any migrations.

    Args:
        revision: String revision target.
    """

    def do_stamp(rev: _RevIdType, context: Any) -> List[Any]:
        # Handle rev parameter in a way that's compatible with different alembic versions
        if isinstance(rev, str):
            return self.script_directory._stamp_revs(revision, rev)
        else:
            # Convert to tuple for compatibility
            rev_tuple = tuple(str(r) for r in rev)
            return self.script_directory._stamp_revs(revision, rev_tuple)

    self.run_migrations(do_stamp)

upgrade(revision: str = 'heads') -> None

Upgrade the database to a later version.

Parameters:

Name	Type	Description	Default
revision	str	String revision target.	'heads'

Source code in src/zenml/zen_stores/migrations/alembic.py

def upgrade(self, revision: str = "heads") -> None:
    """Upgrade the database to a later version.

    Args:
        revision: String revision target.
    """

    def do_upgrade(rev: _RevIdType, context: Any) -> List[Any]:
        # Handle rev parameter in a way that's compatible with different alembic versions
        if isinstance(rev, str):
            return self.script_directory._upgrade_revs(revision, rev)
        else:
            if rev:
                # Use first element or revs for compatibility
                return self.script_directory._upgrade_revs(
                    revision, str(rev[0])
                )
            return []

    self.run_migrations(do_upgrade)

AlembicVersion

Bases: Base

Alembic version table.

Functions

include_object(object: Any, name: Optional[str], type_: str, *args: Any, **kwargs: Any) -> bool

Function used to exclude tables from the migration scripts.

Parameters:

Name	Type	Description	Default
object	Any	The schema item object to check.	required
name	Optional[str]	The name of the object to check.	required
type_	str	The type of the object to check.	required
*args	Any	Additional arguments.	()
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
bool	True if the object should be included, False otherwise.

Source code in src/zenml/zen_stores/migrations/alembic.py

def include_object(
    object: Any, name: Optional[str], type_: str, *args: Any, **kwargs: Any
) -> bool:
    """Function used to exclude tables from the migration scripts.

    Args:
        object: The schema item object to check.
        name: The name of the object to check.
        type_: The type of the object to check.
        *args: Additional arguments.
        **kwargs: Additional keyword arguments.

    Returns:
        True if the object should be included, False otherwise.
    """
    return not (type_ == "table" and name in exclude_tables)

Modules

backup

Database backup and restore utilities.

Modules

base

Base class for database backup engines.

Classes

BaseDatabaseBackupEngine(config: SqlZenStoreConfiguration, location: str | None = None)

Bases: ABC

Base class for database backup engines.

Initialize the backup engine.

Parameters:

Name	Type	Description	Default
config	SqlZenStoreConfiguration	The configuration of the store.	required
location	str | None	The custom location to store the backup.	None

Source code in src/zenml/zen_stores/migrations/backup/base.py

def __init__(
    self,
    config: "SqlZenStoreConfiguration",
    location: str | None = None,
) -> None:
    """Initialize the backup engine.

    Args:
        config: The configuration of the store.
        location: The custom location to store the backup.
    """
    self.config = config
    self._backup_location = location
    url, connect_args, engine_args = config.get_sqlalchemy_config()

    self.url = url
    self.connect_args = connect_args
    self.engine_args = engine_args
    self._engine: Engine | None = None
    self._master_engine: Engine | None = None

Attributes

backup_location: str property

The location where the database is backed up to.

Returns:

Type	Description
str	The location where the database is backed up to.

Raises:

Type	Description
RuntimeError	If the backup location is not set.

engine: Engine property

The SQLAlchemy engine.

Returns:

Type	Description
Engine	The SQLAlchemy engine.

master_engine: Engine property

The SQLAlchemy engine for the master database.

Returns:

Type	Description
Engine	The SQLAlchemy engine for the master database.

Functions

backup_database(overwrite: bool = False) -> None abstractmethod

Backup the database.

Parameters:

Name	Type	Description	Default
overwrite	bool	Whether to overwrite an existing backup if it exists. If set to False, the existing backup will be reused.	False

Source code in src/zenml/zen_stores/migrations/backup/base.py

@abstractmethod
def backup_database(
    self,
    overwrite: bool = False,
) -> None:
    """Backup the database.

    Args:
        overwrite: Whether to overwrite an existing backup if it exists.
            If set to False, the existing backup will be reused.
    """

backup_database_context() -> Generator[None, Any, Any]

Context manager for backing up and restoring the database.

Creates a backup before yielding to the caller. If an exception occurs during the wrapped operation, the database is restored from the backup. If the operation succeeds, the backup is cleaned up.

Yields:

Type	Description
None	None

Raises:

Type	Description
RuntimeError	If the wrapped operation fails. The original exception is chained.
Exception	If the wrapped operation fails. The original exception is chained.

Source code in src/zenml/zen_stores/migrations/backup/base.py

@contextmanager
def backup_database_context(
    self,
) -> Generator[None, Any, Any]:
    """Context manager for backing up and restoring the database.

    Creates a backup before yielding to the caller. If an exception
    occurs during the wrapped operation, the database is restored from
    the backup. If the operation succeeds, the backup is cleaned up.

    Yields:
        None

    Raises:
        RuntimeError: If the wrapped operation fails. The original
            exception is chained.
        Exception: If the wrapped operation fails. The original
            exception is chained.
    """
    logger.info(
        f"Backing up the database before migration to "
        f"`{self.backup_location}`."
    )

    try:
        self.backup_database(overwrite=False)
    except Exception as e:
        raise RuntimeError(
            f"Failed to backup the database: {str(e)}."
        ) from e

    logger.info(
        f"Database successfully backed up to `{self.backup_location}`. If "
        "something goes wrong with the upgrade, ZenML will attempt to "
        "restore the database from this backup automatically."
    )
    try:
        yield
    except Exception:
        logger.info(
            "The database operation failed. Attempting to restore the "
            f"database from `{self.backup_location}`."
        )
        try:
            self.restore_database(cleanup=True)
        except Exception:
            logger.exception(
                f"Failed to restore the database from "
                f"`{self.backup_location}`. You might need to restore the "
                "database manually."
            )
        else:
            logger.info(
                "The database was successfully restored from "
                f"`{self.backup_location}`."
            )

        raise

    try:
        self.cleanup_database_backup()
    except Exception:
        logger.exception("Failed to cleanup the database backup.")

cleanup_database_backup() -> None abstractmethod

Delete the database backup.

Source code in src/zenml/zen_stores/migrations/backup/base.py

@abstractmethod
def cleanup_database_backup(
    self,
) -> None:
    """Delete the database backup."""

create_database(database: str | None = None, drop: bool = False) -> None

Creates a mysql database.

Parameters:

Name	Type	Description	Default
database	str | None	The name of the database to create. If not set, the database name from the configuration will be used.	None
drop	bool	Whether to drop the database if it already exists.	False

Source code in src/zenml/zen_stores/migrations/backup/base.py

def create_database(
    self,
    database: str | None = None,
    drop: bool = False,
) -> None:
    """Creates a mysql database.

    Args:
        database: The name of the database to create. If not set, the
            database name from the configuration will be used.
        drop: Whether to drop the database if it already exists.
    """
    database = database or self.url.database
    if drop:
        self.drop_database(database=database)

    with self.master_engine.connect() as conn:
        logger.info(f"Creating database '{database}'")
        conn.execute(text(f"CREATE DATABASE IF NOT EXISTS `{database}`"))

create_engine(database: str | None = None) -> Engine

Get the SQLAlchemy engine for a database.

Parameters:

Name	Type	Description	Default
database	str | None	The name of the database. If not set, a master engine will be returned.	None

Returns:

Type	Description
Engine	The SQLAlchemy engine.

Source code in src/zenml/zen_stores/migrations/backup/base.py

def create_engine(self, database: str | None = None) -> Engine:
    """Get the SQLAlchemy engine for a database.

    Args:
        database: The name of the database. If not set, a master engine
            will be returned.

    Returns:
        The SQLAlchemy engine.
    """
    url = self.url._replace(database=database)
    return create_engine(
        url=url,
        connect_args=self.connect_args,
        **self.engine_args,
    )

database_exists(database: str | None = None) -> bool

Check if a database exists.

Parameters:

Name	Type	Description	Default
database	str | None	The name of the database to check. If not set, the database name from the configuration will be used.	None

Returns:

Type	Description
bool	Whether the database exists.

Raises:

Type	Description
OperationalError	If connecting to the database failed.

Source code in src/zenml/zen_stores/migrations/backup/base.py

def database_exists(
    self,
    database: str | None = None,
) -> bool:
    """Check if a database exists.

    Args:
        database: The name of the database to check. If not set, the
            database name from the configuration will be used.

    Returns:
        Whether the database exists.

    Raises:
        OperationalError: If connecting to the database failed.
    """
    database = database or self.url.database

    engine = self.create_engine(database=database)
    try:
        engine.connect()
    except OperationalError as e:
        if self.is_mysql_missing_database_error(e):
            return False
        else:
            logger.exception(
                f"Failed to connect to mysql database `{database}`.",
            )
            raise
    else:
        return True

drop_database(database: str | None = None) -> None

Drops a mysql database.

Parameters:

Name	Type	Description	Default
database	str | None	The name of the database to drop. If not set, the database name from the configuration will be used.	None

Source code in src/zenml/zen_stores/migrations/backup/base.py

def drop_database(
    self,
    database: str | None = None,
) -> None:
    """Drops a mysql database.

    Args:
        database: The name of the database to drop. If not set, the
            database name from the configuration will be used.
    """
    database = database or self.url.database
    with self.master_engine.connect() as conn:
        # drop the database if it exists
        logger.info(f"Dropping database '{database}'")
        conn.execute(text(f"DROP DATABASE IF EXISTS `{database}`"))

is_mysql_missing_database_error(error: OperationalError) -> bool classmethod

Checks if the given error is due to a missing database.

Parameters:

Name	Type	Description	Default
error	OperationalError	The error to check.	required

Returns:

Type	Description
bool	If the error because the MySQL database doesn't exist.

Source code in src/zenml/zen_stores/migrations/backup/base.py

@classmethod
def is_mysql_missing_database_error(cls, error: OperationalError) -> bool:
    """Checks if the given error is due to a missing database.

    Args:
        error: The error to check.

    Returns:
        If the error because the MySQL database doesn't exist.
    """
    from pymysql.constants.ER import BAD_DB_ERROR

    if not isinstance(error.orig, pymysql.err.OperationalError):
        return False

    error_code = cast(int, error.orig.args[0])
    return error_code == BAD_DB_ERROR

restore_database(cleanup: bool = False) -> None abstractmethod

Restore the database.

Parameters:

Name	Type	Description	Default
cleanup	bool	Whether to cleanup the backup after restoring the database.	False

Source code in src/zenml/zen_stores/migrations/backup/base.py

@abstractmethod
def restore_database(
    self,
    cleanup: bool = False,
) -> None:
    """Restore the database.

    Args:
        cleanup: Whether to cleanup the backup after restoring the database.
    """

DisabledDatabaseBackupEngine(config: SqlZenStoreConfiguration, location: str | None = None)

Bases: BaseDatabaseBackupEngine

Database backup engine that is disabled.

This is used when the backup strategy is set to disabled.

Source code in src/zenml/zen_stores/migrations/backup/base.py

def __init__(
    self,
    config: "SqlZenStoreConfiguration",
    location: str | None = None,
) -> None:
    """Initialize the backup engine.

    Args:
        config: The configuration of the store.
        location: The custom location to store the backup.
    """
    self.config = config
    self._backup_location = location
    url, connect_args, engine_args = config.get_sqlalchemy_config()

    self.url = url
    self.connect_args = connect_args
    self.engine_args = engine_args
    self._engine: Engine | None = None
    self._master_engine: Engine | None = None

Functions

backup_database(overwrite: bool = False) -> None

Backup the database.

Parameters:

Name	Type	Description	Default
overwrite	bool	Whether to overwrite an existing backup if it exists. If set to False, the existing backup will be reused.	False

Raises:

Type	Description
NotImplementedError	Database backup is not implemented.

Source code in src/zenml/zen_stores/migrations/backup/base.py

def backup_database(
    self,
    overwrite: bool = False,
) -> None:
    """Backup the database.

    Args:
        overwrite: Whether to overwrite an existing backup if it exists.
            If set to False, the existing backup will be reused.

    Raises:
        NotImplementedError: Database backup is not implemented.
    """
    raise NotImplementedError("Database backup is disabled.")

cleanup_database_backup() -> None

Delete the database backup.

Raises:

Type	Description
NotImplementedError	Database backup is not implemented.

Source code in src/zenml/zen_stores/migrations/backup/base.py

def cleanup_database_backup(
    self,
) -> None:
    """Delete the database backup.

    Raises:
        NotImplementedError: Database backup is not implemented.
    """
    raise NotImplementedError("Database backup is disabled.")

restore_database(cleanup: bool = False) -> None

Restore the database.

Parameters:

Name	Type	Description	Default
cleanup	bool	Whether to cleanup the backup after restoring the database.	False

Raises:

Type	Description
NotImplementedError	Database backup is not implemented.

Source code in src/zenml/zen_stores/migrations/backup/base.py

def restore_database(
    self,
    cleanup: bool = False,
) -> None:
    """Restore the database.

    Args:
        cleanup: Whether to cleanup the backup after restoring the database.

    Raises:
        NotImplementedError: Database backup is not implemented.
    """
    raise NotImplementedError("Database backup is disabled.")

Functions

mydumper

Base class for the mydumper/myloader database backup engine.

Classes

MyDumperDatabaseBackupEngine(config: SqlZenStoreConfiguration, location: str | None = None)

Bases: BaseDatabaseBackupEngine

Database backup engine that uses mydumper/myloader for parallel backup/restore.

Initialize the backup engine.

Parameters:

Name	Type	Description	Default
config	SqlZenStoreConfiguration	The configuration of the store.	required
location	str | None	The custom location to store the backup.	None

Source code in src/zenml/zen_stores/migrations/backup/mydumper.py

def __init__(
    self,
    config: "SqlZenStoreConfiguration",
    location: str | None = None,
) -> None:
    """Initialize the backup engine.

    Args:
        config: The configuration of the store.
        location: The custom location to store the backup.
    """
    self._check_mydumper_installed()
    self._check_myloader_installed()

    super().__init__(config, location)
    if self._backup_location is None:
        self._backup_location = os.path.join(
            self.config.backup_directory,
            f"{self.url.database}-backup",
        )

Functions

backup_database(overwrite: bool = False) -> None

Backup the database.

Parameters:

Name	Type	Description	Default
overwrite	bool	Whether to overwrite an existing backup if it exists. If set to False, the existing backup will be reused.	False

Raises:

Type	Description
RuntimeError	If the backup process fails.

Source code in src/zenml/zen_stores/migrations/backup/mydumper.py

def backup_database(
    self,
    overwrite: bool = False,
) -> None:
    """Backup the database.

    Args:
        overwrite: Whether to overwrite an existing backup if it exists.
            If set to False, the existing backup will be reused.

    Raises:
        RuntimeError: If the backup process fails.
    """
    if os.path.isdir(self.backup_location):
        if not overwrite:
            logger.warning(
                f"Backup directory `{self.backup_location}` already exists. "
                "Reusing the existing backup."
            )
            return

        self.cleanup_database_backup()

    os.makedirs(self.backup_location, exist_ok=True)

    cmd = self._build_mydumper_command()

    logger.info(
        f"Starting mydumper backup of database `{self.url.database}` "
        f"to `{self.backup_location}`"
    )
    logger.debug(f"mydumper command: {cmd}")

    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        env=self._get_mysql_env(),
    )

    assert process.stdout is not None
    for line in process.stdout:
        line = line.rstrip()
        if line:
            logger.info(f"[mydumper] {line}")

    return_code = process.wait()
    if return_code != 0:
        raise RuntimeError(
            f"mydumper backup failed with return code {return_code}"
        )

    logger.info(
        f"Database `{self.url.database}` successfully backed up "
        f"to `{self.backup_location}`"
    )

cleanup_database_backup() -> None

Delete the database backup.

Source code in src/zenml/zen_stores/migrations/backup/mydumper.py

def cleanup_database_backup(
    self,
) -> None:
    """Delete the database backup."""
    if os.path.isdir(self.backup_location):
        shutil.rmtree(self.backup_location)
        logger.info(
            f"Successfully cleaned up database backup `{self.backup_location}`."
        )

restore_database(cleanup: bool = False) -> None

Restore the database.

Parameters:

Name	Type	Description	Default
cleanup	bool	Whether to cleanup the backup after restoring the database.	False

Raises:

Type	Description
RuntimeError	If the backup directory does not exist or if the restore process fails.

Source code in src/zenml/zen_stores/migrations/backup/mydumper.py

def restore_database(
    self,
    cleanup: bool = False,
) -> None:
    """Restore the database.

    Args:
        cleanup: Whether to cleanup the backup after restoring the database.

    Raises:
        RuntimeError: If the backup directory does not exist or if the
            restore process fails.
    """
    if not os.path.isdir(self.backup_location):
        raise RuntimeError(
            f"Backup directory `{self.backup_location}` does not exist. "
            "Please backup the database first."
        )

    # Drop and re-create the primary database
    self.create_database(drop=True)

    cmd = self._build_myloader_command()

    logger.info(
        f"Starting myloader restore of database `{self.url.database}` "
        f"from `{self.backup_location}`"
    )
    logger.debug(f"myloader command: {cmd}")

    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        env=self._get_mysql_env(),
    )

    assert process.stdout is not None
    for line in process.stdout:
        line = line.rstrip()
        if line:
            logger.info(f"[myloader] {line}")

    return_code = process.wait()
    if return_code != 0:
        raise RuntimeError(
            f"myloader restore failed with return code {return_code}"
        )

    logger.info(
        f"Database `{self.url.database}` successfully restored "
        f"from `{self.backup_location}`"
    )

    if cleanup:
        self.cleanup_database_backup()

Functions

sqlalchemy

JSON database backup engine.

Classes

DBCloneDatabaseBackupEngine(config: SqlZenStoreConfiguration, location: str | None = None)

Bases: BaseDatabaseBackupEngine

Database backup engine that copies the database to a new database.

Initialize the database backup engine.

Parameters:

Name	Type	Description	Default
config	SqlZenStoreConfiguration	The configuration of the store.	required
location	str | None	The custom location to store the backup.	None

Raises:

Type	Description
ValueError	If the backup database name is not set in the store configuration.

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def __init__(
    self,
    config: "SqlZenStoreConfiguration",
    location: str | None = None,
) -> None:
    """Initialize the database backup engine.

    Args:
        config: The configuration of the store.
        location: The custom location to store the backup.

    Raises:
        ValueError: If the backup database name is not set in the store
            configuration.
    """
    super().__init__(config, location)
    if self._backup_location is None:
        if self.config.backup_database is None:
            raise ValueError(
                "The backup database name must be set in the store "
                "configuration to use the backup database strategy."
            )
        self._backup_location = self.config.backup_database

Functions

backup_database(overwrite: bool = False) -> None

Backup the database.

Parameters:

Name	Type	Description	Default
overwrite	bool	Whether to overwrite an existing backup if it exists. If set to False, the existing backup will be reused.	False

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def backup_database(
    self,
    overwrite: bool = False,
) -> None:
    """Backup the database.

    Args:
        overwrite: Whether to overwrite an existing backup if it exists.
            If set to False, the existing backup will be reused.
    """
    if self.database_exists(database=self.backup_location):
        if not overwrite:
            logger.warning(
                f"A previous backup database already exists at "
                f"`{self.backup_location}`. Reusing the existing backup."
            )
            return

        self.cleanup_database_backup()

    self.create_database(
        database=self.backup_location,
        drop=True,
    )

    backup_engine = self.create_engine(database=self.backup_location)

    self._copy_database(self.engine, backup_engine)

    logger.debug(
        f"Database backed up to the `{self.backup_location}` backup database."
    )

cleanup_database_backup() -> None

Delete the database backup.

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def cleanup_database_backup(
    self,
) -> None:
    """Delete the database backup."""
    if self.database_exists(database=self.backup_location):
        self.drop_database(
            database=self.backup_location,
        )
        logger.debug(
            f"Successfully cleaned up backup database "
            f"{self.backup_location}."
        )

restore_database(cleanup: bool = False) -> None

Restore the database.

Parameters:

Name	Type	Description	Default
cleanup	bool	Whether to cleanup the backup after restoring the database.	False

Raises:

Type	Description
RuntimeError	If the backup database does not exist.

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def restore_database(
    self,
    cleanup: bool = False,
) -> None:
    """Restore the database.

    Args:
        cleanup: Whether to cleanup the backup after restoring the database.

    Raises:
        RuntimeError: If the backup database does not exist.
    """
    if not self.database_exists(database=self.backup_location):
        raise RuntimeError(
            f"Backup database `{self.backup_location}` does not exist."
        )

    backup_engine = self.create_engine(database=self.backup_location)

    self.create_database(drop=True)

    self._copy_database(backup_engine, self.engine)

    logger.debug(
        f"Database restored from the `{self.backup_location}` backup database."
    )

    if cleanup:
        self.cleanup_database_backup()

FileDatabaseBackupEngine(config: SqlZenStoreConfiguration, location: str | None = None)

Bases: SQLAlchemyDatabaseBackupEngine

Database backup engine that stores the database data in a file.

Initialize the file database backup engine.

Parameters:

Name	Type	Description	Default
config	SqlZenStoreConfiguration	The configuration of the store.	required
location	str | None	The custom location to store the backup.	None

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def __init__(
    self,
    config: "SqlZenStoreConfiguration",
    location: str | None = None,
) -> None:
    """Initialize the file database backup engine.

    Args:
        config: The configuration of the store.
        location: The custom location to store the backup.
    """
    super().__init__(config, location)
    if self._backup_location is None:
        self._backup_location = self._get_db_backup_file_path()

Functions

backup_database(overwrite: bool = False) -> None

Backup the database.

Parameters:

Name	Type	Description	Default
overwrite	bool	Whether to overwrite an existing backup if it exists. If set to False, the existing backup will be reused.	False

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def backup_database(
    self,
    overwrite: bool = False,
) -> None:
    """Backup the database.

    Args:
        overwrite: Whether to overwrite an existing backup if it exists.
            If set to False, the existing backup will be reused.
    """
    if os.path.isfile(self.backup_location):
        if not overwrite:
            logger.warning(
                f"A previous backup file already exists at `{self.backup_location}`. "
                "Reusing the existing backup."
            )
            return

        self.cleanup_database_backup()

    if self.url.drivername == "sqlite":
        # For a sqlite database, we can just make a copy of the database
        # file
        assert self.url.database is not None
        shutil.copyfile(
            self.url.database,
            self.backup_location,
        )
        return

    with open(self.backup_location, "w") as f:
        self.backup_database_to_storage(dump_file=f)

    logger.debug(f"Database backed up to file `{self.backup_location}`.")

cleanup_database_backup() -> None

Delete the database backup.

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def cleanup_database_backup(
    self,
) -> None:
    """Delete the database backup."""
    if os.path.isfile(self.backup_location):
        try:
            os.remove(self.backup_location)
        except OSError:
            logger.warning(
                f"Failed to cleanup database dump file `{self.backup_location}`."
            )
        else:
            logger.info(
                f"Successfully cleaned up database dump file "
                f"`{self.backup_location}`."
            )

load_database_data(**kwargs: Any) -> Generator[dict[str, Any], None, None]

Generator that loads the database data.

Parameters:

Name	Type	Description	Default
**kwargs	Any	Must include dump_file (TextIO) - the file handle to load the database data from.	{}

Yields:

Type	Description
dict[str, Any]	The loaded database data.

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def load_database_data(
    self, **kwargs: Any
) -> Generator[dict[str, Any], None, None]:
    """Generator that loads the database data.

    Args:
        **kwargs: Must include `dump_file` (TextIO) - the file handle
            to load the database data from.

    Yields:
        The loaded database data.
    """
    dump_file: TextIO = kwargs["dump_file"]
    buffer = ""
    while True:
        chunk = dump_file.readline()
        if not chunk:
            break
        buffer += chunk
        if chunk.rstrip() == "}":
            yield json.loads(buffer)
            buffer = ""

restore_database(cleanup: bool = False) -> None

Restore the database.

Parameters:

Name	Type	Description	Default
cleanup	bool	Whether to cleanup the backup after restoring the database.	False

Raises:

Type	Description
RuntimeError	If the backup file does not exist or is not accessible.

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def restore_database(
    self,
    cleanup: bool = False,
) -> None:
    """Restore the database.

    Args:
        cleanup: Whether to cleanup the backup after restoring the database.

    Raises:
        RuntimeError: If the backup file does not exist or is not accessible.
    """
    if not os.path.isfile(self.backup_location):
        raise RuntimeError(
            f"Database backup file `{self.backup_location}` does not "
            "exist or is not accessible."
        )

    if self.url.drivername == "sqlite":
        # For a sqlite database, we just overwrite the database file
        # with the backup file
        assert self.url.database is not None
        shutil.copyfile(
            self.backup_location,
            self.url.database,
        )
        return

    with open(self.backup_location, "r") as f:
        self.restore_database_from_storage(dump_file=f)

    logger.debug(f"Database restored from file `{self.backup_location}`.")

    if cleanup:
        self.cleanup_database_backup()

store_database_data(data: dict[str, Any], **kwargs: Any) -> None

Store the database data.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	The database data to store.	required
**kwargs	Any	Must include dump_file (TextIO) - the file handle to store the database data.	{}

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def store_database_data(self, data: dict[str, Any], **kwargs: Any) -> None:
    """Store the database data.

    Args:
        data: The database data to store.
        **kwargs: Must include `dump_file` (TextIO) - the file handle
            to store the database data.
    """
    dump_file: TextIO = kwargs["dump_file"]
    # Write the data to the JSON file. Use an encoder that
    # can handle datetime, Decimal and other types.
    json.dump(
        data,
        dump_file,
        indent=4,
        default=pydantic_encoder,
    )
    dump_file.write("\n")

InMemoryDatabaseBackupEngine(config: SqlZenStoreConfiguration, location: str | None = None)

Bases: SQLAlchemyDatabaseBackupEngine

In-memory database backup engine.

Initialize the in-memory database backup engine.

Parameters:

Name	Type	Description	Default
config	SqlZenStoreConfiguration	The configuration of the store.	required
location	str | None	The custom location to store the backup.	None

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def __init__(
    self,
    config: "SqlZenStoreConfiguration",
    location: str | None = None,
) -> None:
    """Initialize the in-memory database backup engine.

    Args:
        config: The configuration of the store.
        location: The custom location to store the backup.
    """
    super().__init__(config, location or "memory")
    self.database_data: list[dict[str, Any]] = []

Functions

backup_database(overwrite: bool = False) -> None

Backup the database.

Parameters:

Name	Type	Description	Default
overwrite	bool	Whether to overwrite an existing backup if it exists. If set to False, the existing backup will be reused.	False

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def backup_database(
    self,
    overwrite: bool = False,
) -> None:
    """Backup the database.

    Args:
        overwrite: Whether to overwrite an existing backup if it exists.
            If set to False, the existing backup will be reused.
    """
    if len(self.database_data) > 0:
        if not overwrite:
            logger.warning(
                "An existing backup already exists. Reusing the existing backup."
            )
            return

        self.cleanup_database_backup()

    self.backup_database_to_storage()

    logger.debug("Database backed up to memory.")

cleanup_database_backup() -> None

Delete the database backup.

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def cleanup_database_backup(
    self,
) -> None:
    """Delete the database backup."""
    self.database_data = []

    logger.debug("In-memory database backup cleaned up.")

load_database_data(**kwargs: Any) -> Generator[dict[str, Any], None, None]

Generator that loads the database data.

Yields:

Type	Description
dict[str, Any]	The loaded database data.

Parameters:

Name	Type	Description	Default
kwargs	Any	Additional keyword arguments passed to the method.	{}

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def load_database_data(
    self, **kwargs: Any
) -> Generator[dict[str, Any], None, None]:
    """Generator that loads the database data.

    Yields:
        The loaded database data.

    Args:
        kwargs: Additional keyword arguments passed to the method.
    """
    for data in self.database_data:
        yield data

restore_database(cleanup: bool = False) -> None

Restore the database.

Parameters:

Name	Type	Description	Default
cleanup	bool	Whether to cleanup the backup after restoring the database.	False

Raises:

Type	Description
RuntimeError	If no in-memory backup exists.

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def restore_database(
    self,
    cleanup: bool = False,
) -> None:
    """Restore the database.

    Args:
        cleanup: Whether to cleanup the backup after restoring the database.

    Raises:
        RuntimeError: If no in-memory backup exists.
    """
    if len(self.database_data) == 0:
        raise RuntimeError(
            "No in-memory backup exists. Please backup the database first."
        )

    self.restore_database_from_storage()

    logger.debug("Database restored from memory.")

    if cleanup:
        self.cleanup_database_backup()

store_database_data(data: dict[str, Any], **kwargs: Any) -> None

Store the database data.

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	The database data to store.	required
kwargs	Any	Additional keyword arguments passed to the method.	{}

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def store_database_data(self, data: dict[str, Any], **kwargs: Any) -> None:
    """Store the database data.

    Args:
        data: The database data to store.
        kwargs: Additional keyword arguments passed to the method.
    """
    self.database_data.append(data)

SQLAlchemyDatabaseBackupEngine(config: SqlZenStoreConfiguration, location: str | None = None)

Bases: BaseDatabaseBackupEngine

Base class for SQLAlchemy-based database backup engines.

Source code in src/zenml/zen_stores/migrations/backup/base.py

def __init__(
    self,
    config: "SqlZenStoreConfiguration",
    location: str | None = None,
) -> None:
    """Initialize the backup engine.

    Args:
        config: The configuration of the store.
        location: The custom location to store the backup.
    """
    self.config = config
    self._backup_location = location
    url, connect_args, engine_args = config.get_sqlalchemy_config()

    self.url = url
    self.connect_args = connect_args
    self.engine_args = engine_args
    self._engine: Engine | None = None
    self._master_engine: Engine | None = None

Functions

backup_database_to_storage(**store_db_kwargs: Any) -> None

Backup the database to a storage location.

Backup the database to an abstract storage location. The storage location is implemented by the store_database_data method that is called repeatedly to store the database information.

Parameters:

Name	Type	Description	Default
store_db_kwargs	Any	Additional keyword arguments to pass to the store_database_data method.	{}

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def backup_database_to_storage(
    self,
    **store_db_kwargs: Any,
) -> None:
    """Backup the database to a storage location.

    Backup the database to an abstract storage location. The storage
    location is implemented by the `store_database_data` method that is
    called repeatedly to store the database information.

    Args:
        store_db_kwargs: Additional keyword arguments to pass to the
            `store_database_data` method.
    """
    metadata = MetaData()
    metadata.reflect(bind=self.engine)
    with self.engine.connect() as conn:
        for table in metadata.sorted_tables:
            # 1. extract the table creation statements

            create_table_construct = CreateTable(table)
            create_table_stmt = str(create_table_construct).strip()
            for column in create_table_construct.columns:
                # enclosing all column names in backticks. This is because
                # some column names are reserved keywords in MySQL. For
                # example, keys and values. So, instead of tracking all
                # keywords, we just enclose all column names in backticks.
                # enclose the first word in the column definition in
                # backticks
                words = str(column).split()
                words[0] = f"`{words[0]}`"
                create_table_stmt = create_table_stmt.replace(
                    f"\n\t{str(column)}", " ".join(words)
                )
            # if any double quotes are used for column names, replace them
            # with backticks
            create_table_stmt = create_table_stmt.replace('"', "") + ";"

            # enclose all table names in backticks. This is because some
            # table names are reserved keywords in MySQL (e.g key
            # and trigger).
            create_table_stmt = create_table_stmt.replace(
                f"CREATE TABLE {table.name}",
                f"CREATE TABLE `{table.name}`",
            )
            # do the same for references to other tables
            # (i.e. foreign key constraints) by replacing REFERENCES <word>
            # with REFERENCES `<word>`
            # use a regular expression for this
            create_table_stmt = re.sub(
                r"REFERENCES\s+(\w+)",
                r"REFERENCES `\1`",
                create_table_stmt,
            )

            # In SQLAlchemy, the CreateTable statement may not always
            # include unique constraints explicitly if they are implemented
            # as unique indexes instead. To make sure we get all unique
            # constraints, including those implemented as indexes, we
            # extract the unique constraints from the table schema and add
            # them to the create table statement.

            # Extract the unique constraints from the table schema
            index_create_statements = []
            unique_constraints = []
            for index in table.indexes:
                if index.unique:
                    unique_columns = [
                        f"`{column.name}`" for column in index.columns
                    ]
                    unique_constraints.append(
                        f"UNIQUE KEY `{index.name}` ({', '.join(unique_columns)})"
                    )
                else:
                    if index.name in {
                        fk.name for fk in table.foreign_key_constraints
                    }:
                        # Foreign key indices are already handled by the
                        # table creation statement.
                        continue

                    index_create = str(CreateIndex(index)).strip()
                    index_create = index_create.replace(
                        f"CREATE INDEX {index.name}",
                        f"CREATE INDEX `{index.name}`",
                    )
                    index_create = index_create.replace(
                        f"ON {table.name}", f"ON `{table.name}`"
                    )

                    for column_name in index.columns.keys():
                        # We need this logic here to avoid the column names
                        # inside the index name
                        index_create = index_create.replace(
                            f"({column_name}", f"(`{column_name}`"
                        )
                        index_create = index_create.replace(
                            f"{column_name},", f"`{column_name}`,"
                        )
                        index_create = index_create.replace(
                            f"{column_name})", f"`{column_name}`)"
                        )

                    index_create = index_create.replace('"', "") + ";"
                    index_create_statements.append(index_create)

            # Add the unique constraints to the create table statement
            if unique_constraints:
                # Remove the closing parenthesis, semicolon and any
                # whitespaces at the end of the create table statement
                create_table_stmt = re.sub(
                    r"\s*\)\s*;\s*$", "", create_table_stmt
                )
                create_table_stmt = (
                    create_table_stmt
                    + ", \n\t"
                    + ", \n\t".join(unique_constraints)
                    + "\n);"
                )

            # Detect self-referential foreign keys from the table schema
            has_self_referential_foreign_keys = False
            for fk in table.foreign_keys:
                # Check if the foreign key points to the same table
                if fk.column.table == table:
                    has_self_referential_foreign_keys = True
                    break

            # Store the table schema
            self.store_database_data(
                dict(
                    table=table.name,
                    create_stmt=create_table_stmt,
                    self_references=has_self_referential_foreign_keys,
                ),
                **store_db_kwargs,
            )

            for stmt in index_create_statements:
                self.store_database_data(
                    dict(
                        table=table.name,
                        index_create_stmt=stmt,
                    ),
                    **store_db_kwargs,
                )

            # 2. extract the table data in batches
            order_by = [col for col in table.primary_key]

            # Fetch the number of rows in the table
            row_count = conn.scalar(
                select(func.count()).select_from(table)
            )

            # Fetch the data from the table in batches
            if row_count is not None:
                batch_size = 100
                for i in range(0, row_count, batch_size):
                    rows = conn.execute(
                        table.select()
                        .order_by(*order_by)
                        .limit(batch_size)
                        .offset(i)
                    ).fetchall()

                    self.store_database_data(
                        dict(
                            table=table.name,
                            data=[row._asdict() for row in rows],
                        ),
                        **store_db_kwargs,
                    )

load_database_data(**kwargs: Any) -> Generator[dict[str, Any], None, None] abstractmethod

Generator that loads the database data.

This method is called repeatedly to load the database data. It must yield a dictionary containing either the table schema or table data, as documented in the store_database_data method.

Yields:

Type	Description
dict[str, Any]	The loaded database data.

Parameters:

Name	Type	Description	Default
kwargs	Any	Additional keyword arguments passed to the method.	{}

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

@abstractmethod
def load_database_data(
    self, **kwargs: Any
) -> Generator[dict[str, Any], None, None]:
    """Generator that loads the database data.

    This method is called repeatedly to load the database data. It must
    yield a dictionary containing either the table schema or table data,
    as documented in the `store_database_data` method.

    Yields:
        The loaded database data.

    Args:
        kwargs: Additional keyword arguments passed to the method.
    """

restore_database_from_storage(**load_db_kwargs: Any) -> None

Restore the database from a backup storage location.

Restores the database from an abstract storage location. The storage location is implemented by the load_database_data method that is called repeatedly to load the database information from the external storage chunk by chunk.

Parameters:

Name	Type	Description	Default
load_db_kwargs	Any	Additional keyword arguments to pass to the load_database_data method.	{}

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

def restore_database_from_storage(
    self,
    **load_db_kwargs: Any,
) -> None:
    """Restore the database from a backup storage location.

    Restores the database from an abstract storage location. The storage
    location is implemented by the `load_database_data` method that is
    called repeatedly to load the database information from the external
    storage chunk by chunk.

    Args:
        load_db_kwargs: Additional keyword arguments to pass to the
            `load_database_data` method.
    """
    # Drop and re-create the primary database
    self.create_database(drop=True)

    metadata = MetaData()

    with self.engine.begin() as connection:
        # read the DB information one JSON object at a time
        self_references: dict[str, bool] = {}
        for table_dump in self.load_database_data(**load_db_kwargs):
            table_name = table_dump["table"]
            if "create_stmt" in table_dump:
                # execute the table creation statement
                connection.execute(text(table_dump["create_stmt"]))
                # Reload the database metadata after creating the table
                metadata.reflect(bind=self.engine)
                self_references[table_name] = table_dump.get(
                    "self_references", False
                )

            if "index_create_stmt" in table_dump:
                # execute the index creation statement
                connection.execute(text(table_dump["index_create_stmt"]))
                # Reload the database metadata after creating the index
                metadata.reflect(bind=self.engine)

            if "data" in table_dump:
                # insert the data into the database
                table = metadata.tables[table_name]
                if self_references.get(table_name, False):
                    # If the table has self-referential foreign keys, we
                    # need to disable the foreign key checks before inserting
                    # the rows and re-enable them afterwards. This is because
                    # the rows need to be inserted in the correct order to
                    # satisfy the foreign key constraints and we don't sort
                    # the rows by creation time in the backup.
                    connection.execute(text("SET FOREIGN_KEY_CHECKS = 0"))

                for row in table_dump["data"]:
                    # Convert column values to the correct type
                    for column in table.columns:
                        # Blob columns are stored as binary strings
                        if column.type.python_type is bytes and isinstance(
                            row[column.name], str
                        ):
                            # Convert the string to bytes
                            row[column.name] = bytes(
                                row[column.name], "utf-8"
                            )

                # Insert the rows into the table in batches
                batch_size = 100
                for i in range(0, len(table_dump["data"]), batch_size):
                    connection.execute(
                        table.insert().values(
                            table_dump["data"][i : i + batch_size]
                        )
                    )

                if table_dump.get("self_references", False):
                    # Re-enable the foreign key checks after inserting the rows
                    connection.execute(text("SET FOREIGN_KEY_CHECKS = 1"))

store_database_data(data: dict[str, Any], **kwargs: Any) -> None abstractmethod

Store the database data.

This method is called repeatedly to store the database information. It is called with a single argument, which is a dictionary containing either the table schema or table data. The dictionary contains the following keys:

* `table`: The name of the table.
* `create_stmt`: The table creation statement.
* `data`: A list of rows in the table.

The format of the dump is as depicted in the following example:

{
    "table": "table1",
    "create_stmt": "CREATE TABLE table1 (id INTEGER NOT NULL, "
        "name VARCHAR(255), PRIMARY KEY (id))"
}
{
    "table": "table1",
    "data": [
    {
        "id": 1,
        "name": "foo"
    },
    {
        "id": 1,
        "name": "bar"
    },
    ...
    ]
}
{
    "table": "table1",
    "data": [
    {
        "id": 101,
        "name": "fee"
    },
    {
        "id": 102,
        "name": "bee"
    },
    ...
    ]
}

Parameters:

Name	Type	Description	Default
data	dict[str, Any]	The database data to store.	required
kwargs	Any	Additional keyword arguments passed to the method.	{}

Source code in src/zenml/zen_stores/migrations/backup/sqlalchemy.py

@abstractmethod
def store_database_data(self, data: dict[str, Any], **kwargs: Any) -> None:
    """Store the database data.

    This method is called repeatedly to store the database information. It
    is called with a single argument, which is a dictionary containing
    either the table schema or table data. The dictionary contains the
    following keys:

        * `table`: The name of the table.
        * `create_stmt`: The table creation statement.
        * `data`: A list of rows in the table.

    The format of the dump is as depicted in the following example:

    ```json
    {
        "table": "table1",
        "create_stmt": "CREATE TABLE table1 (id INTEGER NOT NULL, "
            "name VARCHAR(255), PRIMARY KEY (id))"
    }
    {
        "table": "table1",
        "data": [
        {
            "id": 1,
            "name": "foo"
        },
        {
            "id": 1,
            "name": "bar"
        },
        ...
        ]
    }
    {
        "table": "table1",
        "data": [
        {
            "id": 101,
            "name": "fee"
        },
        {
            "id": 102,
            "name": "bee"
        },
        ...
        ]
    }
    ```

    Args:
        data: The database data to store.
        kwargs: Additional keyword arguments passed to the method.
    """

Functions

resource_pools

Resource pools management.

Modules

store_interface

Resource pools store interface.

Classes

ResourcePoolsSQLStoreInterface(store: SqlZenStore)

Bases: ResourcePoolsStoreInterface

Resource pools SQL store interface.

Initialize the resource pools SQL store.

Parameters:

Name	Type	Description	Default
store	SqlZenStore	The store to use.	required

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

def __init__(self, store: "SqlZenStore") -> None:
    """Initialize the resource pools SQL store.

    Args:
        store: The store to use.
    """
    super().__init__()
    self.store = store
    store.resource_pools = self

Functions

create_resource_request(session: Session, resource_request: ResourceRequestRequest) -> ResourceRequestResponse | None abstractmethod

Create a resource request.

Parameters:

Name	Type	Description	Default
session	Session	DB session.	required
resource_request	ResourceRequestRequest	The resource request to create.	required

Returns:

Type	Description
ResourceRequestResponse | None	The created resource request or None if the feature is not enabled.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def create_resource_request(
    self, session: "Session", resource_request: ResourceRequestRequest
) -> ResourceRequestResponse | None:
    """Create a resource request.

    Args:
        session: DB session.
        resource_request: The resource request to create.

    Returns:
        The created resource request or None if the feature is not enabled.
    """

release_step_run_resources(session: Session, step_run_id: UUID) -> None abstractmethod

Release potentially acquired resources for a step run.

Parameters:

Name	Type	Description	Default
session	Session	DB session.	required
step_run_id	UUID	The ID of the step run to release resources for.	required

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def release_step_run_resources(
    self, session: "Session", step_run_id: UUID
) -> None:
    """Release potentially acquired resources for a step run.

    Args:
        session: DB session.
        step_run_id: The ID of the step run to release resources for.
    """

ResourcePoolsStoreInterface

Bases: ABC

Resource pools store interface.

Functions

create_resource_pool(resource_pool: ResourcePoolRequest) -> ResourcePoolResponse abstractmethod

Create a resource pool.

Parameters:

Name	Type	Description	Default
resource_pool	ResourcePoolRequest	The resource pool to create.	required

Returns:

Type	Description
ResourcePoolResponse	The created resource pool.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def create_resource_pool(
    self, resource_pool: ResourcePoolRequest
) -> ResourcePoolResponse:
    """Create a resource pool.

    Args:
        resource_pool: The resource pool to create.

    Returns:
        The created resource pool.
    """

create_resource_pool_subject_policy(policy: ResourcePoolSubjectPolicyRequest) -> ResourcePoolSubjectPolicyResponse abstractmethod

Create a resource pool subject policy.

Parameters:

Name	Type	Description	Default
policy	ResourcePoolSubjectPolicyRequest	The policy to create.	required

Returns:

Type	Description
ResourcePoolSubjectPolicyResponse	The created policy.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def create_resource_pool_subject_policy(
    self, policy: ResourcePoolSubjectPolicyRequest
) -> ResourcePoolSubjectPolicyResponse:
    """Create a resource pool subject policy.

    Args:
        policy: The policy to create.

    Returns:
        The created policy.
    """

delete_resource_pool(resource_pool_id: UUID) -> None abstractmethod

Delete a resource pool.

Parameters:

Name	Type	Description	Default
resource_pool_id	UUID	The ID of the resource pool to delete.	required

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def delete_resource_pool(self, resource_pool_id: UUID) -> None:
    """Delete a resource pool.

    Args:
        resource_pool_id: The ID of the resource pool to delete.
    """

delete_resource_pool_subject_policy(policy_id: UUID) -> None abstractmethod

Delete a resource pool subject policy.

Parameters:

Name	Type	Description	Default
policy_id	UUID	The ID of the policy to delete.	required

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def delete_resource_pool_subject_policy(self, policy_id: UUID) -> None:
    """Delete a resource pool subject policy.

    Args:
        policy_id: The ID of the policy to delete.
    """

delete_resource_request(resource_request_id: UUID) -> None abstractmethod

Delete a resource request.

Parameters:

Name	Type	Description	Default
resource_request_id	UUID	The ID of the resource request to delete.	required

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def delete_resource_request(self, resource_request_id: UUID) -> None:
    """Delete a resource request.

    Args:
        resource_request_id: The ID of the resource request to delete.
    """

get_resource_pool(resource_pool_id: UUID, hydrate: bool = True) -> ResourcePoolResponse abstractmethod

Get a resource pool by ID.

Parameters:

Name	Type	Description	Default
resource_pool_id	UUID	The ID of the resource pool to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ResourcePoolResponse	The resource pool.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def get_resource_pool(
    self, resource_pool_id: UUID, hydrate: bool = True
) -> ResourcePoolResponse:
    """Get a resource pool by ID.

    Args:
        resource_pool_id: The ID of the resource pool to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The resource pool.
    """

get_resource_pool_subject_policy(policy_id: UUID, hydrate: bool = True) -> ResourcePoolSubjectPolicyResponse abstractmethod

Get a resource pool subject policy by ID.

Parameters:

Name	Type	Description	Default
policy_id	UUID	The ID of the policy to get.	required
hydrate	bool	Whether to include metadata fields.	True

Returns:

Type	Description
ResourcePoolSubjectPolicyResponse	The requested policy.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def get_resource_pool_subject_policy(
    self, policy_id: UUID, hydrate: bool = True
) -> ResourcePoolSubjectPolicyResponse:
    """Get a resource pool subject policy by ID.

    Args:
        policy_id: The ID of the policy to get.
        hydrate: Whether to include metadata fields.

    Returns:
        The requested policy.
    """

get_resource_request(resource_request_id: UUID, hydrate: bool = True) -> ResourceRequestResponse abstractmethod

Get a resource request by ID.

Parameters:

Name	Type	Description	Default
resource_request_id	UUID	The ID of the resource request to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ResourceRequestResponse	The resource request.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def get_resource_request(
    self, resource_request_id: UUID, hydrate: bool = True
) -> ResourceRequestResponse:
    """Get a resource request by ID.

    Args:
        resource_request_id: The ID of the resource request to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The resource request.
    """

list_resource_pool_subject_policies(filter_model: ResourcePoolSubjectPolicyFilter, hydrate: bool = False) -> Page[ResourcePoolSubjectPolicyResponse] abstractmethod

List resource pool subject policies.

Parameters:

Name	Type	Description	Default
filter_model	ResourcePoolSubjectPolicyFilter	All filter parameters including pagination params.	required
hydrate	bool	Whether to include metadata fields.	False

Returns:

Type	Description
Page[ResourcePoolSubjectPolicyResponse]	Matching policies.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def list_resource_pool_subject_policies(
    self,
    filter_model: ResourcePoolSubjectPolicyFilter,
    hydrate: bool = False,
) -> Page[ResourcePoolSubjectPolicyResponse]:
    """List resource pool subject policies.

    Args:
        filter_model: All filter parameters including pagination params.
        hydrate: Whether to include metadata fields.

    Returns:
        Matching policies.
    """

list_resource_pools(filter_model: ResourcePoolFilter, hydrate: bool = False) -> Page[ResourcePoolResponse] abstractmethod

List all resource pools matching the given filter criteria.

Parameters:

Name	Type	Description	Default
filter_model	ResourcePoolFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ResourcePoolResponse]	A list of all resource pools matching the filter criteria.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def list_resource_pools(
    self, filter_model: ResourcePoolFilter, hydrate: bool = False
) -> Page[ResourcePoolResponse]:
    """List all resource pools matching the given filter criteria.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all resource pools matching the filter criteria.
    """

list_resource_requests(filter_model: ResourceRequestFilter, hydrate: bool = False) -> Page[ResourceRequestResponse] abstractmethod

List all resource requests matching the given filter criteria.

Parameters:

Name	Type	Description	Default
filter_model	ResourceRequestFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ResourceRequestResponse]	A list of all resource requests matching the filter criteria.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def list_resource_requests(
    self, filter_model: ResourceRequestFilter, hydrate: bool = False
) -> Page[ResourceRequestResponse]:
    """List all resource requests matching the given filter criteria.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all resource requests matching the filter criteria.
    """

update_resource_pool(resource_pool_id: UUID, update: ResourcePoolUpdate) -> ResourcePoolResponse abstractmethod

Update an existing resource pool.

Parameters:

Name	Type	Description	Default
resource_pool_id	UUID	The ID of the resource pool to update.	required
update	ResourcePoolUpdate	The update to be applied to the resource pool.	required

Returns:

Type	Description
ResourcePoolResponse	The updated resource pool.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def update_resource_pool(
    self, resource_pool_id: UUID, update: ResourcePoolUpdate
) -> ResourcePoolResponse:
    """Update an existing resource pool.

    Args:
        resource_pool_id: The ID of the resource pool to update.
        update: The update to be applied to the resource pool.

    Returns:
        The updated resource pool.
    """

update_resource_pool_subject_policy(policy_id: UUID, update: ResourcePoolSubjectPolicyUpdate) -> ResourcePoolSubjectPolicyResponse abstractmethod

Update an existing resource pool subject policy.

Parameters:

Name	Type	Description	Default
policy_id	UUID	The ID of the policy to update.	required
update	ResourcePoolSubjectPolicyUpdate	The update model.	required

Returns:

Type	Description
ResourcePoolSubjectPolicyResponse	The updated policy.

Source code in src/zenml/zen_stores/resource_pools/store_interface.py

@abstractmethod
def update_resource_pool_subject_policy(
    self, policy_id: UUID, update: ResourcePoolSubjectPolicyUpdate
) -> ResourcePoolSubjectPolicyResponse:
    """Update an existing resource pool subject policy.

    Args:
        policy_id: The ID of the policy to update.
        update: The update model.

    Returns:
        The updated policy.
    """

rest_zen_store

REST Zen Store implementation.

Classes

RestZenStore(skip_default_registrations: bool = False, **kwargs: Any)

Bases: BaseZenStore

Store implementation for accessing data from a REST API.

Source code in src/zenml/zen_stores/base_zen_store.py

def __init__(
    self,
    skip_default_registrations: bool = False,
    **kwargs: Any,
) -> None:
    """Create and initialize a store.

    Args:
        skip_default_registrations: If `True`, the creation of the default
            stack and user in the store will be skipped.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.
    """
    super().__init__(**kwargs)

    self._initialize()

    if not skip_default_registrations:
        logger.debug("Initializing database")
        self._initialize_database()
    else:
        logger.debug("Skipping database initialization")

Attributes

server_info: ServerModel property

Get cached information about the server.

Returns:

Type	Description
ServerModel	Cached information about the server.

session: requests.Session property

Initialize and return a requests session.

Returns:

Type	Description
Session	A requests session.

Functions

attach_trigger_to_snapshot(trigger_id: UUID, snapshot_id: UUID, run_configuration: PipelineRunConfiguration | None = None, allow_replace: bool = False) -> None

Attaches (links) a trigger to a snapshot.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	The ID of the trigger.	required
snapshot_id	UUID	The ID of the snapshot.	required
run_configuration	PipelineRunConfiguration | None	The configuration applied to subsequent runs of this trigger & snapshot.	None
allow_replace	bool	Allow replacement if attachment already exists.	False

Source code in src/zenml/zen_stores/rest_zen_store.py

def attach_trigger_to_snapshot(
    self,
    trigger_id: UUID,
    snapshot_id: UUID,
    run_configuration: PipelineRunConfiguration | None = None,
    allow_replace: bool = False,
) -> None:
    """Attaches (links) a trigger to a snapshot.

    Args:
        trigger_id: The ID of the trigger.
        snapshot_id: The ID of the snapshot.
        run_configuration: The configuration applied to subsequent runs of this trigger & snapshot.
        allow_replace: Allow replacement if attachment already exists.
    """
    self.put(
        path=f"{TRIGGERS}/{trigger_id}{PIPELINE_SNAPSHOTS}/{snapshot_id}",
        body=run_configuration,
        params={"allow_replace": allow_replace},
    )

authenticate(force: bool = False) -> None

Authenticate or re-authenticate to the ZenML server.

Parameters:

Name	Type	Description	Default
force	bool	If True, force a re-authentication even if a valid API token is currently cached. This is useful when the current API token is known to be invalid or expired.	False

Source code in src/zenml/zen_stores/rest_zen_store.py

def authenticate(self, force: bool = False) -> None:
    """Authenticate or re-authenticate to the ZenML server.

    Args:
        force: If True, force a re-authentication even if a valid API token
            is currently cached. This is useful when the current API token
            is known to be invalid or expired.
    """
    # This is called to trigger an authentication flow, either because
    # the current API token is expired or no longer valid, or because
    # a configuration change has happened or merely because an
    # authentication was never attempted before.
    #
    # 1. Drop the API token currently being used, if any.
    # 2. If force=True, clear the current API token from the credentials
    # store, if any, otherwise it will just be re-used on the next call.
    # 3. Get a new API token

    # The authentication token could have expired or invalidated through
    # other means; refresh it and try again. This will clear any cached
    # token and trigger a new authentication flow.
    if self._api_token and not force:
        if self._api_token.expired:
            logger.info(
                "Authentication session expired; attempting to "
                "re-authenticate."
            )
        else:
            logger.info(
                "Authentication session was invalidated by the server; "
                "This can happen for example if the user's permissions "
                "have been revoked or if the server has been restarted "
                "and lost its session state. Attempting to "
                "re-authenticate."
            )
    else:
        if force:
            # Clear the current API token from the credentials store, if
            # any, to force a new authentication flow.
            get_credentials_store().clear_token(self.url)
        # Never authenticated since the client was created or the API token
        # was explicitly cleared.
        logger.debug(f"Authenticating to {self.url}...")

    self._api_token = None

    new_api_token = self.get_or_generate_api_token()

    # Set or refresh the authentication token
    self.session.headers.update(
        {"Authorization": "Bearer " + new_api_token}
    )
    logger.debug(f"Authenticated to {self.url}")
    self._last_authenticated = utc_now()

backup_secrets(ignore_errors: bool = True, delete_secrets: bool = False) -> None

Backs up all secrets to the configured backup secrets store.

Parameters:

Name	Type	Description	Default
ignore_errors	bool	Whether to ignore individual errors during the backup process and attempt to backup all secrets.	True
delete_secrets	bool	Whether to delete the secrets that have been successfully backed up from the primary secrets store. Setting this flag effectively moves all secrets from the primary secrets store to the backup secrets store.	False

Source code in src/zenml/zen_stores/rest_zen_store.py

def backup_secrets(
    self, ignore_errors: bool = True, delete_secrets: bool = False
) -> None:
    """Backs up all secrets to the configured backup secrets store.

    Args:
        ignore_errors: Whether to ignore individual errors during the backup
            process and attempt to backup all secrets.
        delete_secrets: Whether to delete the secrets that have been
            successfully backed up from the primary secrets store. Setting
            this flag effectively moves all secrets from the primary secrets
            store to the backup secrets store.
    """
    params: Dict[str, Any] = {
        "ignore_errors": ignore_errors,
        "delete_secrets": delete_secrets,
    }
    self.put(
        f"{SECRETS_OPERATIONS}{SECRETS_BACKUP}",
        params=params,
    )

batch_create_artifact_versions(artifact_versions: List[ArtifactVersionRequest]) -> List[ArtifactVersionResponse]

Creates a batch of artifact versions.

Parameters:

Name	Type	Description	Default
artifact_versions	List[ArtifactVersionRequest]	The artifact versions to create.	required

Returns:

Type	Description
List[ArtifactVersionResponse]	The created artifact versions.

Source code in src/zenml/zen_stores/rest_zen_store.py

def batch_create_artifact_versions(
    self, artifact_versions: List[ArtifactVersionRequest]
) -> List[ArtifactVersionResponse]:
    """Creates a batch of artifact versions.

    Args:
        artifact_versions: The artifact versions to create.

    Returns:
        The created artifact versions.
    """
    return self._batch_create_resources(
        resources=artifact_versions,
        response_model=ArtifactVersionResponse,
        route=ARTIFACT_VERSIONS,
    )

batch_create_tag_resource(tag_resources: List[TagResourceRequest]) -> List[TagResourceResponse]

Create a batch of tag resource relationships.

Parameters:

Name	Type	Description	Default
tag_resources	List[TagResourceRequest]	The tag resource relationships to be created.	required

Returns:

Type	Description
List[TagResourceResponse]	The newly created tag resource relationships.

Source code in src/zenml/zen_stores/rest_zen_store.py

def batch_create_tag_resource(
    self, tag_resources: List[TagResourceRequest]
) -> List[TagResourceResponse]:
    """Create a batch of tag resource relationships.

    Args:
        tag_resources: The tag resource relationships to be created.

    Returns:
        The newly created tag resource relationships.
    """
    return self._batch_create_resources(
        resources=tag_resources,
        response_model=TagResourceResponse,
        route=TAG_RESOURCES,
    )

batch_delete_tag_resource(tag_resources: List[TagResourceRequest]) -> None

Delete a batch of tag resources.

Parameters:

Name	Type	Description	Default
tag_resources	List[TagResourceRequest]	The tag resource relationships to be deleted.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def batch_delete_tag_resource(
    self, tag_resources: List[TagResourceRequest]
) -> None:
    """Delete a batch of tag resources.

    Args:
        tag_resources: The tag resource relationships to be deleted.
    """
    self._batch_delete_resources(
        resources=tag_resources,
        route=TAG_RESOURCES,
    )

clear_trigger_dispatch_error(trigger_id: UUID, snapshot_id: UUID | None = None) -> None

Clear dispatch error details for trigger dispatch associations.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	The ID of the trigger.	required
snapshot_id	UUID | None	Optional snapshot ID filter.	None

Source code in src/zenml/zen_stores/rest_zen_store.py

def clear_trigger_dispatch_error(
    self,
    trigger_id: UUID,
    snapshot_id: UUID | None = None,
) -> None:
    """Clear dispatch error details for trigger dispatch associations.

    Args:
        trigger_id: The ID of the trigger.
        snapshot_id: Optional snapshot ID filter.
    """
    params = {"snapshot_id": str(snapshot_id)} if snapshot_id else None
    self.delete(
        path=f"{TRIGGERS}/{trigger_id}{TRIGGER_SNAPSHOT_DISPATCH_STATE}",
        params=params,
    )

create_api_key(service_account_id: UUID, api_key: APIKeyRequest) -> APIKeyResponse

Create a new API key for a service account.

Parameters:

Name	Type	Description	Default
service_account_id	UUID	The ID of the service account for which to create the API key.	required
api_key	APIKeyRequest	The API key to create.	required

Returns:

Type	Description
APIKeyResponse	The created API key.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_api_key(
    self, service_account_id: UUID, api_key: APIKeyRequest
) -> APIKeyResponse:
    """Create a new API key for a service account.

    Args:
        service_account_id: The ID of the service account for which to
            create the API key.
        api_key: The API key to create.

    Returns:
        The created API key.
    """
    return self._create_resource(
        resource=api_key,
        route=f"{SERVICE_ACCOUNTS}/{str(service_account_id)}{API_KEYS}",
        response_model=APIKeyResponse,
    )

create_artifact(artifact: ArtifactRequest) -> ArtifactResponse

Creates a new artifact.

Parameters:

Name	Type	Description	Default
artifact	ArtifactRequest	The artifact to create.	required

Returns:

Type	Description
ArtifactResponse	The newly created artifact.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_artifact(self, artifact: ArtifactRequest) -> ArtifactResponse:
    """Creates a new artifact.

    Args:
        artifact: The artifact to create.

    Returns:
        The newly created artifact.
    """
    return self._create_resource(
        resource=artifact,
        response_model=ArtifactResponse,
        route=ARTIFACTS,
    )

create_artifact_version(artifact_version: ArtifactVersionRequest) -> ArtifactVersionResponse

Creates an artifact version.

Parameters:

Name	Type	Description	Default
artifact_version	ArtifactVersionRequest	The artifact version to create.	required

Returns:

Type	Description
ArtifactVersionResponse	The created artifact version.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_artifact_version(
    self, artifact_version: ArtifactVersionRequest
) -> ArtifactVersionResponse:
    """Creates an artifact version.

    Args:
        artifact_version: The artifact version to create.

    Returns:
        The created artifact version.
    """
    return self._create_resource(
        resource=artifact_version,
        response_model=ArtifactVersionResponse,
        route=ARTIFACT_VERSIONS,
    )

create_build(build: PipelineBuildRequest) -> PipelineBuildResponse

Creates a new build.

Parameters:

Name	Type	Description	Default
build	PipelineBuildRequest	The build to create.	required

Returns:

Type	Description
PipelineBuildResponse	The newly created build.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_build(
    self,
    build: PipelineBuildRequest,
) -> PipelineBuildResponse:
    """Creates a new build.

    Args:
        build: The build to create.

    Returns:
        The newly created build.
    """
    return self._create_resource(
        resource=build,
        route=PIPELINE_BUILDS,
        response_model=PipelineBuildResponse,
    )

create_code_repository(code_repository: CodeRepositoryRequest) -> CodeRepositoryResponse

Creates a new code repository.

Parameters:

Name	Type	Description	Default
code_repository	CodeRepositoryRequest	Code repository to be created.	required

Returns:

Type	Description
CodeRepositoryResponse	The newly created code repository.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_code_repository(
    self, code_repository: CodeRepositoryRequest
) -> CodeRepositoryResponse:
    """Creates a new code repository.

    Args:
        code_repository: Code repository to be created.

    Returns:
        The newly created code repository.
    """
    return self._create_resource(
        resource=code_repository,
        response_model=CodeRepositoryResponse,
        route=CODE_REPOSITORIES,
    )

create_curated_visualization(visualization: CuratedVisualizationRequest) -> CuratedVisualizationResponse

Create a curated visualization via REST API.

Parameters:

Name	Type	Description	Default
visualization	CuratedVisualizationRequest	The curated visualization to create.	required

Returns:

Type	Description
CuratedVisualizationResponse	The created curated visualization.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_curated_visualization(
    self, visualization: CuratedVisualizationRequest
) -> CuratedVisualizationResponse:
    """Create a curated visualization via REST API.

    Args:
        visualization: The curated visualization to create.

    Returns:
        The created curated visualization.
    """
    return self._create_resource(
        resource=visualization,
        response_model=CuratedVisualizationResponse,
        route=CURATED_VISUALIZATIONS,
        params={"hydrate": True},
    )

create_deployment(deployment: DeploymentRequest) -> DeploymentResponse

Create a new deployment.

Parameters:

Name	Type	Description	Default
deployment	DeploymentRequest	The deployment to create.	required

Returns:

Type	Description
DeploymentResponse	The newly created deployment.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_deployment(
    self, deployment: DeploymentRequest
) -> DeploymentResponse:
    """Create a new deployment.

    Args:
        deployment: The deployment to create.

    Returns:
        The newly created deployment.
    """
    return self._create_resource(
        resource=deployment,
        route=DEPLOYMENTS,
        response_model=DeploymentResponse,
    )

create_flavor(flavor: FlavorRequest) -> FlavorResponse

Creates a new stack component flavor.

Parameters:

Name	Type	Description	Default
flavor	FlavorRequest	The stack component flavor to create.	required

Returns:

Type	Description
FlavorResponse	The newly created flavor.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_flavor(self, flavor: FlavorRequest) -> FlavorResponse:
    """Creates a new stack component flavor.

    Args:
        flavor: The stack component flavor to create.

    Returns:
        The newly created flavor.
    """
    return self._create_resource(
        resource=flavor,
        route=FLAVORS,
        response_model=FlavorResponse,
    )

create_logs(logs: LogsRequest) -> LogsResponse

Create a logs entry.

Parameters:

Name	Type	Description	Default
logs	LogsRequest	The logs entry to create.	required

Returns:

Type	Description
LogsResponse	The created logs entry.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_logs(self, logs: LogsRequest) -> LogsResponse:
    """Create a logs entry.

    Args:
        logs: The logs entry to create.

    Returns:
        The created logs entry.
    """
    return self._create_resource(
        resource=logs,
        route=LOGS,
        response_model=LogsResponse,
    )

create_model(model: ModelRequest) -> ModelResponse

Creates a new model.

Parameters:

Name	Type	Description	Default
model	ModelRequest	the Model to be created.	required

Returns:

Type	Description
ModelResponse	The newly created model.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_model(self, model: ModelRequest) -> ModelResponse:
    """Creates a new model.

    Args:
        model: the Model to be created.

    Returns:
        The newly created model.
    """
    return self._create_resource(
        resource=model,
        response_model=ModelResponse,
        route=MODELS,
    )

create_model_version(model_version: ModelVersionRequest) -> ModelVersionResponse

Creates a new model version.

Parameters:

Name	Type	Description	Default
model_version	ModelVersionRequest	the Model Version to be created.	required

Returns:

Type	Description
ModelVersionResponse	The newly created model version.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_model_version(
    self, model_version: ModelVersionRequest
) -> ModelVersionResponse:
    """Creates a new model version.

    Args:
        model_version: the Model Version to be created.

    Returns:
        The newly created model version.
    """
    return self._create_resource(
        resource=model_version,
        response_model=ModelVersionResponse,
        route=MODEL_VERSIONS,
    )

create_model_version_artifact_link(model_version_artifact_link: ModelVersionArtifactRequest) -> ModelVersionArtifactResponse

Creates a new model version link.

Parameters:

Name	Type	Description	Default
model_version_artifact_link	ModelVersionArtifactRequest	the Model Version to Artifact Link to be created.	required

Returns:

Type	Description
ModelVersionArtifactResponse	The newly created model version to artifact link.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_model_version_artifact_link(
    self, model_version_artifact_link: ModelVersionArtifactRequest
) -> ModelVersionArtifactResponse:
    """Creates a new model version link.

    Args:
        model_version_artifact_link: the Model Version to Artifact Link
            to be created.

    Returns:
        The newly created model version to artifact link.
    """
    return self._create_resource(
        resource=model_version_artifact_link,
        response_model=ModelVersionArtifactResponse,
        route=MODEL_VERSION_ARTIFACTS,
    )

create_model_version_pipeline_run_link(model_version_pipeline_run_link: ModelVersionPipelineRunRequest) -> ModelVersionPipelineRunResponse

Creates a new model version to pipeline run link.

Parameters:

Name	Type	Description	Default
model_version_pipeline_run_link	ModelVersionPipelineRunRequest	the Model Version to Pipeline Run Link to be created.	required

Returns:

Type	Description
ModelVersionPipelineRunResponse	If Model Version to Pipeline Run Link already exists - returns the existing link.
ModelVersionPipelineRunResponse	Otherwise, returns the newly created model version to pipeline run link.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_model_version_pipeline_run_link(
    self,
    model_version_pipeline_run_link: ModelVersionPipelineRunRequest,
) -> ModelVersionPipelineRunResponse:
    """Creates a new model version to pipeline run link.

    Args:
        model_version_pipeline_run_link: the Model Version to Pipeline Run
            Link to be created.

    Returns:
        - If Model Version to Pipeline Run Link already exists - returns
            the existing link.
        - Otherwise, returns the newly created model version to pipeline
            run link.
    """
    return self._create_resource(
        resource=model_version_pipeline_run_link,
        response_model=ModelVersionPipelineRunResponse,
        route=MODEL_VERSION_PIPELINE_RUNS,
    )

create_pipeline(pipeline: PipelineRequest) -> PipelineResponse

Creates a new pipeline.

Parameters:

Name	Type	Description	Default
pipeline	PipelineRequest	The pipeline to create.	required

Returns:

Type	Description
PipelineResponse	The newly created pipeline.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_pipeline(self, pipeline: PipelineRequest) -> PipelineResponse:
    """Creates a new pipeline.

    Args:
        pipeline: The pipeline to create.

    Returns:
        The newly created pipeline.
    """
    return self._create_resource(
        resource=pipeline,
        route=PIPELINES,
        response_model=PipelineResponse,
    )

create_project(project: ProjectRequest) -> ProjectResponse

Creates a new project.

Parameters:

Name	Type	Description	Default
project	ProjectRequest	The project to create.	required

Returns:

Type	Description
ProjectResponse	The newly created project.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_project(self, project: ProjectRequest) -> ProjectResponse:
    """Creates a new project.

    Args:
        project: The project to create.

    Returns:
        The newly created project.
    """
    return self._create_resource(
        resource=project,
        route=PROJECTS,
        response_model=ProjectResponse,
    )

create_resource_pool(resource_pool: ResourcePoolRequest) -> ResourcePoolResponse

Create a resource pool.

Parameters:

Name	Type	Description	Default
resource_pool	ResourcePoolRequest	The resource pool to create.	required

Returns:

Type	Description
ResourcePoolResponse	The created resource pool.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_resource_pool(
    self, resource_pool: ResourcePoolRequest
) -> ResourcePoolResponse:
    """Create a resource pool.

    Args:
        resource_pool: The resource pool to create.

    Returns:
        The created resource pool.
    """
    return self._create_resource(
        resource=resource_pool,
        route=RESOURCE_POOLS,
        response_model=ResourcePoolResponse,
    )

create_resource_pool_subject_policy(policy: ResourcePoolSubjectPolicyRequest) -> ResourcePoolSubjectPolicyResponse

Create a resource pool subject policy.

Parameters:

Name	Type	Description	Default
policy	ResourcePoolSubjectPolicyRequest	The policy to create.	required

Returns:

Type	Description
ResourcePoolSubjectPolicyResponse	The created policy.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_resource_pool_subject_policy(
    self, policy: ResourcePoolSubjectPolicyRequest
) -> ResourcePoolSubjectPolicyResponse:
    """Create a resource pool subject policy.

    Args:
        policy: The policy to create.

    Returns:
        The created policy.
    """
    return self._create_resource(
        resource=policy,
        route=RESOURCE_POOL_SUBJECT_POLICIES,
        response_model=ResourcePoolSubjectPolicyResponse,
    )

create_run_metadata(run_metadata: RunMetadataRequest) -> None

Creates run metadata.

Parameters:

Name	Type	Description	Default
run_metadata	RunMetadataRequest	The run metadata to create.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_run_metadata(self, run_metadata: RunMetadataRequest) -> None:
    """Creates run metadata.

    Args:
        run_metadata: The run metadata to create.
    """
    self.post(RUN_METADATA, body=run_metadata)

create_run_step(step_run: StepRunRequest) -> StepRunResponse

Creates a step run.

Parameters:

Name	Type	Description	Default
step_run	StepRunRequest	The step run to create.	required

Returns:

Type	Description
StepRunResponse	The created step run.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_run_step(self, step_run: StepRunRequest) -> StepRunResponse:
    """Creates a step run.

    Args:
        step_run: The step run to create.

    Returns:
        The created step run.
    """
    return self._create_resource(
        resource=step_run,
        response_model=StepRunResponse,
        route=STEPS,
    )

create_run_template(template: RunTemplateRequest) -> RunTemplateResponse

Create a new run template.

Parameters:

Name	Type	Description	Default
template	RunTemplateRequest	The template to create.	required

Returns:

Type	Description
RunTemplateResponse	The newly created template.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_run_template(
    self,
    template: RunTemplateRequest,
) -> RunTemplateResponse:
    """Create a new run template.

    Args:
        template: The template to create.

    Returns:
        The newly created template.
    """
    return self._create_resource(
        resource=template,
        route=RUN_TEMPLATES,
        response_model=RunTemplateResponse,
    )

create_run_wait_condition(run_wait_condition: RunWaitConditionRequest) -> RunWaitConditionResponse

Create a run wait condition.

Parameters:

Name	Type	Description	Default
run_wait_condition	RunWaitConditionRequest	Wait condition creation payload.	required

Returns:

Type	Description
RunWaitConditionResponse	The created wait condition.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_run_wait_condition(
    self, run_wait_condition: RunWaitConditionRequest
) -> RunWaitConditionResponse:
    """Create a run wait condition.

    Args:
        run_wait_condition: Wait condition creation payload.

    Returns:
        The created wait condition.
    """
    response_body = self.post(
        RUN_WAIT_CONDITIONS,
        body=run_wait_condition,
    )
    return RunWaitConditionResponse.model_validate(response_body)

create_schedule(schedule: ScheduleRequest) -> ScheduleResponse

Creates a new schedule.

Parameters:

Name	Type	Description	Default
schedule	ScheduleRequest	The schedule to create.	required

Returns:

Type	Description
ScheduleResponse	The newly created schedule.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_schedule(self, schedule: ScheduleRequest) -> ScheduleResponse:
    """Creates a new schedule.

    Args:
        schedule: The schedule to create.

    Returns:
        The newly created schedule.
    """
    return self._create_resource(
        resource=schedule,
        route=SCHEDULES,
        response_model=ScheduleResponse,
    )

create_secret(secret: SecretRequest) -> SecretResponse

Creates a new secret.

The new secret is also validated against the scoping rules enforced in the secrets store:

• only one private secret with the given name can exist.
• only one public secret with the given name can exist.

Parameters:

Name	Type	Description	Default
secret	SecretRequest	The secret to create.	required

Returns:

Type	Description
SecretResponse	The newly created secret.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_secret(self, secret: SecretRequest) -> SecretResponse:
    """Creates a new secret.

    The new secret is also validated against the scoping rules enforced in
    the secrets store:

      - only one private secret with the given name can exist.
      - only one public secret with the given name can exist.

    Args:
        secret: The secret to create.

    Returns:
        The newly created secret.
    """
    return self._create_resource(
        resource=secret,
        route=SECRETS,
        response_model=SecretResponse,
    )

create_service(service_request: ServiceRequest) -> ServiceResponse

Create a new service.

Parameters:

Name	Type	Description	Default
service_request	ServiceRequest	The service to create.	required

Returns:

Type	Description
ServiceResponse	The created service.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_service(
    self, service_request: ServiceRequest
) -> ServiceResponse:
    """Create a new service.

    Args:
        service_request: The service to create.

    Returns:
        The created service.
    """
    return self._create_resource(
        resource=service_request,
        response_model=ServiceResponse,
        route=SERVICES,
    )

create_service_account(service_account: ServiceAccountRequest) -> ServiceAccountResponse

Creates a new service account.

Parameters:

Name	Type	Description	Default
service_account	ServiceAccountRequest	Service account to be created.	required

Returns:

Type	Description
ServiceAccountResponse	The newly created service account.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_service_account(
    self, service_account: ServiceAccountRequest
) -> ServiceAccountResponse:
    """Creates a new service account.

    Args:
        service_account: Service account to be created.

    Returns:
        The newly created service account.
    """
    return self._create_resource(
        resource=service_account,
        route=SERVICE_ACCOUNTS,
        response_model=ServiceAccountResponse,
    )

create_service_connector(service_connector: ServiceConnectorRequest) -> ServiceConnectorResponse

Creates a new service connector.

Parameters:

Name	Type	Description	Default
service_connector	ServiceConnectorRequest	Service connector to be created.	required

Returns:

Type	Description
ServiceConnectorResponse	The newly created service connector.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_service_connector(
    self, service_connector: ServiceConnectorRequest
) -> ServiceConnectorResponse:
    """Creates a new service connector.

    Args:
        service_connector: Service connector to be created.

    Returns:
        The newly created service connector.
    """
    connector_model = self._create_resource(
        resource=service_connector,
        route=SERVICE_CONNECTORS,
        response_model=ServiceConnectorResponse,
    )
    self._populate_connector_type(connector_model)
    # Call this to properly split the secrets from the configuration
    try:
        connector_model.validate_configuration()
    except ValueError as e:
        logger.error(
            f"Error validating connector configuration for "
            f"{connector_model.name}: {e}"
        )
    return connector_model

create_snapshot(snapshot: PipelineSnapshotRequest) -> PipelineSnapshotResponse

Creates a new snapshot.

Parameters:

Name	Type	Description	Default
snapshot	PipelineSnapshotRequest	The snapshot to create.	required

Returns:

Type	Description
PipelineSnapshotResponse	The newly created snapshot.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_snapshot(
    self,
    snapshot: PipelineSnapshotRequest,
) -> PipelineSnapshotResponse:
    """Creates a new snapshot.

    Args:
        snapshot: The snapshot to create.

    Returns:
        The newly created snapshot.
    """
    return self._create_resource(
        resource=snapshot,
        route=PIPELINE_SNAPSHOTS,
        response_model=PipelineSnapshotResponse,
    )

create_stack(stack: StackRequest) -> StackResponse

Register a new stack.

Parameters:

Name	Type	Description	Default
stack	StackRequest	The stack to register.	required

Returns:

Type	Description
StackResponse	The registered stack.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_stack(self, stack: StackRequest) -> StackResponse:
    """Register a new stack.

    Args:
        stack: The stack to register.

    Returns:
        The registered stack.
    """
    return self._create_resource(
        resource=stack,
        response_model=StackResponse,
        route=STACKS,
    )

create_stack_component(component: ComponentRequest) -> ComponentResponse

Create a stack component.

Parameters:

Name	Type	Description	Default
component	ComponentRequest	The stack component to create.	required

Returns:

Type	Description
ComponentResponse	The created stack component.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_stack_component(
    self,
    component: ComponentRequest,
) -> ComponentResponse:
    """Create a stack component.

    Args:
        component: The stack component to create.

    Returns:
        The created stack component.
    """
    return self._create_resource(
        resource=component,
        route=STACK_COMPONENTS,
        response_model=ComponentResponse,
    )

create_tag(tag: TagRequest) -> TagResponse

Creates a new tag.

Parameters:

Name	Type	Description	Default
tag	TagRequest	the tag to be created.	required

Returns:

Type	Description
TagResponse	The newly created tag.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_tag(self, tag: TagRequest) -> TagResponse:
    """Creates a new tag.

    Args:
        tag: the tag to be created.

    Returns:
        The newly created tag.
    """
    return self._create_resource(
        resource=tag,
        response_model=TagResponse,
        route=TAGS,
    )

create_tag_resource(tag_resource: TagResourceRequest) -> TagResourceResponse

Create a new tag resource.

Parameters:

Name	Type	Description	Default
tag_resource	TagResourceRequest	The tag resource to be created.	required

Returns:

Type	Description
TagResourceResponse	The newly created tag resource.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_tag_resource(
    self,
    tag_resource: TagResourceRequest,
) -> TagResourceResponse:
    """Create a new tag resource.

    Args:
        tag_resource: The tag resource to be created.

    Returns:
        The newly created tag resource.
    """
    return self._create_resource(
        resource=tag_resource,
        response_model=TagResourceResponse,
        route=TAG_RESOURCES,
    )

create_trigger(trigger: TriggerRequest) -> TRIGGER_RETURN_TYPE_UNION

Creates a new trigger.

Parameters:

Name	Type	Description	Default
trigger	TriggerRequest	The trigger to create.	required

Returns:

Type	Description
TRIGGER_RETURN_TYPE_UNION	The created trigger.

Raises:

Type	Description
ValueError	If an unexpected payload is retrieved.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_trigger(
    self, trigger: TriggerRequest
) -> TRIGGER_RETURN_TYPE_UNION:
    """Creates a new trigger.

    Args:
        trigger: The trigger to create.

    Returns:
        The created trigger.

    Raises:
        ValueError: If an unexpected payload is retrieved.
    """
    body: dict[str, Any] = self.post(TRIGGERS, body=trigger)  # type: ignore[assignment]

    try:
        response_model = TYPE_TO_RESPONSE_MAPPING[body["body"]["type"]]
        return response_model.model_validate(body)
    except (KeyError, TypeError):
        raise ValueError("Bad response, expected a trigger type object.")

create_user(user: UserRequest) -> UserResponse

Creates a new user.

Parameters:

Name	Type	Description	Default
user	UserRequest	User to be created.	required

Returns:

Type	Description
UserResponse	The newly created user.

Source code in src/zenml/zen_stores/rest_zen_store.py

def create_user(self, user: UserRequest) -> UserResponse:
    """Creates a new user.

    Args:
        user: User to be created.

    Returns:
        The newly created user.
    """
    return self._create_resource(
        resource=user,
        route=USERS,
        response_model=UserResponse,
    )

deactivate_user(user_name_or_id: Union[str, UUID]) -> UserResponse

Deactivates a user.

Parameters:

Name	Type	Description	Default
user_name_or_id	Union[str, UUID]	The name or ID of the user to delete.	required

Returns:

Type	Description
UserResponse	The deactivated user containing the activation token.

Source code in src/zenml/zen_stores/rest_zen_store.py

def deactivate_user(
    self, user_name_or_id: Union[str, UUID]
) -> UserResponse:
    """Deactivates a user.

    Args:
        user_name_or_id: The name or ID of the user to delete.

    Returns:
        The deactivated user containing the activation token.
    """
    response_body = self.put(
        f"{USERS}/{str(user_name_or_id)}{DEACTIVATE}",
    )

    return UserResponse.model_validate(response_body)

delete(path: str, body: Optional[BaseModel] = None, params: Optional[Dict[str, Any]] = None, timeout: Optional[int] = None, **kwargs: Any) -> Json

Make a DELETE request to the given endpoint path.

Parameters:

Name	Type	Description	Default
path	str	The path to the endpoint.	required
body	Optional[BaseModel]	The body to send.	None
params	Optional[Dict[str, Any]]	The query parameters to pass to the endpoint.	None
timeout	Optional[int]	The request timeout in seconds.	None
kwargs	Any	Additional keyword arguments to pass to the request.	{}

Returns:

Type	Description
Json	The response body.

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete(
    self,
    path: str,
    body: Optional[BaseModel] = None,
    params: Optional[Dict[str, Any]] = None,
    timeout: Optional[int] = None,
    **kwargs: Any,
) -> Json:
    """Make a DELETE request to the given endpoint path.

    Args:
        path: The path to the endpoint.
        body: The body to send.
        params: The query parameters to pass to the endpoint.
        timeout: The request timeout in seconds.
        kwargs: Additional keyword arguments to pass to the request.

    Returns:
        The response body.
    """
    return self._request(
        "DELETE",
        self.url + API + VERSION_1 + path,
        json=body.model_dump(mode="json") if body else None,
        params=params,
        timeout=timeout,
        **kwargs,
    )

delete_all_model_version_artifact_links(model_version_id: UUID, only_links: bool = True) -> None

Deletes all links between model version and an artifact.

Parameters:

Name	Type	Description	Default
model_version_id	UUID	ID of the model version containing the link.	required
only_links	bool	Flag deciding whether to delete only links or all.	True

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_all_model_version_artifact_links(
    self,
    model_version_id: UUID,
    only_links: bool = True,
) -> None:
    """Deletes all links between model version and an artifact.

    Args:
        model_version_id: ID of the model version containing the link.
        only_links: Flag deciding whether to delete only links or all.
    """
    self.delete(
        f"{MODEL_VERSIONS}/{model_version_id}{ARTIFACTS}",
        params={"only_links": only_links},
    )

delete_api_key(service_account_id: UUID, api_key_name_or_id: Union[str, UUID]) -> None

Delete an API key for a service account.

Parameters:

Name	Type	Description	Default
service_account_id	UUID	The ID of the service account for which to delete the API key.	required
api_key_name_or_id	Union[str, UUID]	The name or ID of the API key to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_api_key(
    self,
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
) -> None:
    """Delete an API key for a service account.

    Args:
        service_account_id: The ID of the service account for which to
            delete the API key.
        api_key_name_or_id: The name or ID of the API key to delete.
    """
    self._delete_resource(
        resource_id=api_key_name_or_id,
        route=f"{SERVICE_ACCOUNTS}/{str(service_account_id)}{API_KEYS}",
    )

delete_artifact(artifact_id: UUID) -> None

Deletes an artifact.

Parameters:

Name	Type	Description	Default
artifact_id	UUID	The ID of the artifact to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_artifact(self, artifact_id: UUID) -> None:
    """Deletes an artifact.

    Args:
        artifact_id: The ID of the artifact to delete.
    """
    self._delete_resource(resource_id=artifact_id, route=ARTIFACTS)

delete_artifact_version(artifact_version_id: UUID) -> None

Deletes an artifact version.

Parameters:

Name	Type	Description	Default
artifact_version_id	UUID	The ID of the artifact version to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_artifact_version(self, artifact_version_id: UUID) -> None:
    """Deletes an artifact version.

    Args:
        artifact_version_id: The ID of the artifact version to delete.
    """
    self._delete_resource(
        resource_id=artifact_version_id, route=ARTIFACT_VERSIONS
    )

delete_authorized_device(device_id: UUID) -> None

Deletes an OAuth 2.0 authorized device.

Parameters:

Name	Type	Description	Default
device_id	UUID	The ID of the device to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_authorized_device(self, device_id: UUID) -> None:
    """Deletes an OAuth 2.0 authorized device.

    Args:
        device_id: The ID of the device to delete.
    """
    self._delete_resource(resource_id=device_id, route=DEVICES)

delete_build(build_id: UUID) -> None

Deletes a build.

Parameters:

Name	Type	Description	Default
build_id	UUID	The ID of the build to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_build(self, build_id: UUID) -> None:
    """Deletes a build.

    Args:
        build_id: The ID of the build to delete.
    """
    self._delete_resource(
        resource_id=build_id,
        route=PIPELINE_BUILDS,
    )

delete_code_repository(code_repository_id: UUID) -> None

Deletes a code repository.

Parameters:

Name	Type	Description	Default
code_repository_id	UUID	The ID of the code repository to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_code_repository(self, code_repository_id: UUID) -> None:
    """Deletes a code repository.

    Args:
        code_repository_id: The ID of the code repository to delete.
    """
    self._delete_resource(
        resource_id=code_repository_id, route=CODE_REPOSITORIES
    )

delete_curated_visualization(visualization_id: UUID) -> None

Delete a curated visualization via REST API.

Parameters:

Name	Type	Description	Default
visualization_id	UUID	The ID of the curated visualization to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_curated_visualization(self, visualization_id: UUID) -> None:
    """Delete a curated visualization via REST API.

    Args:
        visualization_id: The ID of the curated visualization to delete.
    """
    self._delete_resource(
        resource_id=visualization_id,
        route=CURATED_VISUALIZATIONS,
    )

delete_deployment(deployment_id: UUID) -> None

Delete a deployment.

Parameters:

Name	Type	Description	Default
deployment_id	UUID	The ID of the deployment to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_deployment(self, deployment_id: UUID) -> None:
    """Delete a deployment.

    Args:
        deployment_id: The ID of the deployment to delete.
    """
    self._delete_resource(
        resource_id=deployment_id,
        route=DEPLOYMENTS,
    )

delete_flavor(flavor_id: UUID) -> None

Delete a stack component flavor.

Parameters:

Name	Type	Description	Default
flavor_id	UUID	The ID of the stack component flavor to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_flavor(self, flavor_id: UUID) -> None:
    """Delete a stack component flavor.

    Args:
        flavor_id: The ID of the stack component flavor to delete.
    """
    self._delete_resource(
        resource_id=flavor_id,
        route=FLAVORS,
    )

delete_model(model_id: UUID) -> None

Deletes a model.

Parameters:

Name	Type	Description	Default
model_id	UUID	id of the model to be deleted.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_model(self, model_id: UUID) -> None:
    """Deletes a model.

    Args:
        model_id: id of the model to be deleted.
    """
    self._delete_resource(resource_id=model_id, route=MODELS)

delete_model_version(model_version_id: UUID) -> None

Deletes a model version.

Parameters:

Name	Type	Description	Default
model_version_id	UUID	name or id of the model version to be deleted.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_model_version(
    self,
    model_version_id: UUID,
) -> None:
    """Deletes a model version.

    Args:
        model_version_id: name or id of the model version to be deleted.
    """
    self._delete_resource(
        resource_id=model_version_id,
        route=MODEL_VERSIONS,
    )

delete_model_version_artifact_link(model_version_id: UUID, model_version_artifact_link_name_or_id: Union[str, UUID]) -> None

Deletes a model version to artifact link.

Parameters:

Name	Type	Description	Default
model_version_id	UUID	ID of the model version containing the link.	required
model_version_artifact_link_name_or_id	Union[str, UUID]	name or ID of the model version to artifact link to be deleted.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_model_version_artifact_link(
    self,
    model_version_id: UUID,
    model_version_artifact_link_name_or_id: Union[str, UUID],
) -> None:
    """Deletes a model version to artifact link.

    Args:
        model_version_id: ID of the model version containing the link.
        model_version_artifact_link_name_or_id: name or ID of the model
            version to artifact link to be deleted.
    """
    self._delete_resource(
        resource_id=model_version_artifact_link_name_or_id,
        route=f"{MODEL_VERSIONS}/{model_version_id}{ARTIFACTS}",
    )

delete_model_version_pipeline_run_link(model_version_id: UUID, model_version_pipeline_run_link_name_or_id: Union[str, UUID]) -> None

Deletes a model version to pipeline run link.

Parameters:

Name	Type	Description	Default
model_version_id	UUID	ID of the model version containing the link.	required
model_version_pipeline_run_link_name_or_id	Union[str, UUID]	name or ID of the model version to pipeline run link to be deleted.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_model_version_pipeline_run_link(
    self,
    model_version_id: UUID,
    model_version_pipeline_run_link_name_or_id: Union[str, UUID],
) -> None:
    """Deletes a model version to pipeline run link.

    Args:
        model_version_id: ID of the model version containing the link.
        model_version_pipeline_run_link_name_or_id: name or ID of the model version to pipeline run link to be deleted.
    """
    self._delete_resource(
        resource_id=model_version_pipeline_run_link_name_or_id,
        route=f"{MODEL_VERSIONS}/{model_version_id}{RUNS}",
    )

delete_pipeline(pipeline_id: UUID) -> None

Deletes a pipeline.

Parameters:

Name	Type	Description	Default
pipeline_id	UUID	The ID of the pipeline to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_pipeline(self, pipeline_id: UUID) -> None:
    """Deletes a pipeline.

    Args:
        pipeline_id: The ID of the pipeline to delete.
    """
    self._delete_resource(
        resource_id=pipeline_id,
        route=PIPELINES,
    )

delete_project(project_name_or_id: Union[str, UUID]) -> None

Deletes a project.

Parameters:

Name	Type	Description	Default
project_name_or_id	Union[str, UUID]	Name or ID of the project to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_project(self, project_name_or_id: Union[str, UUID]) -> None:
    """Deletes a project.

    Args:
        project_name_or_id: Name or ID of the project to delete.
    """
    self._delete_resource(
        resource_id=project_name_or_id,
        route=PROJECTS,
    )

delete_resource_pool(resource_pool_id: UUID) -> None

Delete a resource pool.

Parameters:

Name	Type	Description	Default
resource_pool_id	UUID	The ID of the resource pool to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_resource_pool(self, resource_pool_id: UUID) -> None:
    """Delete a resource pool.

    Args:
        resource_pool_id: The ID of the resource pool to delete.
    """
    self._delete_resource(
        resource_id=resource_pool_id,
        route=RESOURCE_POOLS,
    )

delete_resource_pool_subject_policy(policy_id: UUID) -> None

Delete a resource pool subject policy.

Parameters:

Name	Type	Description	Default
policy_id	UUID	The ID of the policy to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_resource_pool_subject_policy(self, policy_id: UUID) -> None:
    """Delete a resource pool subject policy.

    Args:
        policy_id: The ID of the policy to delete.
    """
    self._delete_resource(
        resource_id=policy_id,
        route=RESOURCE_POOL_SUBJECT_POLICIES,
    )

delete_resource_request(resource_request_id: UUID) -> None

Delete a resource request.

Parameters:

Name	Type	Description	Default
resource_request_id	UUID	The ID of the resource request to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_resource_request(self, resource_request_id: UUID) -> None:
    """Delete a resource request.

    Args:
        resource_request_id: The ID of the resource request to delete.
    """
    self._delete_resource(
        resource_id=resource_request_id,
        route=RESOURCE_REQUESTS,
    )

delete_run(run_id: UUID) -> None

Deletes a pipeline run.

Parameters:

Name	Type	Description	Default
run_id	UUID	The ID of the pipeline run to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_run(self, run_id: UUID) -> None:
    """Deletes a pipeline run.

    Args:
        run_id: The ID of the pipeline run to delete.
    """
    self._delete_resource(
        resource_id=run_id,
        route=RUNS,
    )

delete_run_template(template_id: UUID) -> None

Delete a run template.

Parameters:

Name	Type	Description	Default
template_id	UUID	The ID of the template to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_run_template(self, template_id: UUID) -> None:
    """Delete a run template.

    Args:
        template_id: The ID of the template to delete.
    """
    self._delete_resource(
        resource_id=template_id,
        route=RUN_TEMPLATES,
    )

delete_schedule(schedule_id: UUID, soft: bool = False) -> None

Deletes a schedule.

Parameters:

Name	Type	Description	Default
schedule_id	UUID	The ID of the schedule to delete.	required
soft	bool	Soft deletion will archive the schedule.	False

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_schedule(self, schedule_id: UUID, soft: bool = False) -> None:
    """Deletes a schedule.

    Args:
        schedule_id: The ID of the schedule to delete.
        soft: Soft deletion will archive the schedule.
    """
    self._delete_resource(
        resource_id=schedule_id,
        route=SCHEDULES,
        params={"soft": soft},
    )

delete_secret(secret_id: UUID) -> None

Delete a secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The id of the secret to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_secret(self, secret_id: UUID) -> None:
    """Delete a secret.

    Args:
        secret_id: The id of the secret to delete.
    """
    self._delete_resource(
        resource_id=secret_id,
        route=SECRETS,
    )

delete_service(service_id: UUID) -> None

Delete a service.

Parameters:

Name	Type	Description	Default
service_id	UUID	The ID of the service to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_service(self, service_id: UUID) -> None:
    """Delete a service.

    Args:
        service_id: The ID of the service to delete.
    """
    self._delete_resource(resource_id=service_id, route=SERVICES)

delete_service_account(service_account_name_or_id: Union[str, UUID]) -> None

Delete a service account.

Parameters:

Name	Type	Description	Default
service_account_name_or_id	Union[str, UUID]	The name or the ID of the service account to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_service_account(
    self,
    service_account_name_or_id: Union[str, UUID],
) -> None:
    """Delete a service account.

    Args:
        service_account_name_or_id: The name or the ID of the service
            account to delete.
    """
    self._delete_resource(
        resource_id=service_account_name_or_id,
        route=SERVICE_ACCOUNTS,
    )

delete_service_connector(service_connector_id: UUID) -> None

Deletes a service connector.

Parameters:

Name	Type	Description	Default
service_connector_id	UUID	The ID of the service connector to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_service_connector(self, service_connector_id: UUID) -> None:
    """Deletes a service connector.

    Args:
        service_connector_id: The ID of the service connector to delete.
    """
    self._delete_resource(
        resource_id=service_connector_id, route=SERVICE_CONNECTORS
    )

delete_snapshot(snapshot_id: UUID) -> None

Deletes a snapshot.

Parameters:

Name	Type	Description	Default
snapshot_id	UUID	The ID of the snapshot to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_snapshot(self, snapshot_id: UUID) -> None:
    """Deletes a snapshot.

    Args:
        snapshot_id: The ID of the snapshot to delete.
    """
    self._delete_resource(
        resource_id=snapshot_id,
        route=PIPELINE_SNAPSHOTS,
    )

delete_stack(stack_id: UUID) -> None

Delete a stack.

Parameters:

Name	Type	Description	Default
stack_id	UUID	The ID of the stack to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_stack(self, stack_id: UUID) -> None:
    """Delete a stack.

    Args:
        stack_id: The ID of the stack to delete.
    """
    self._delete_resource(
        resource_id=stack_id,
        route=STACKS,
    )

delete_stack_component(component_id: UUID) -> None

Delete a stack component.

Parameters:

Name	Type	Description	Default
component_id	UUID	The ID of the stack component to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_stack_component(self, component_id: UUID) -> None:
    """Delete a stack component.

    Args:
        component_id: The ID of the stack component to delete.
    """
    self._delete_resource(
        resource_id=component_id,
        route=STACK_COMPONENTS,
    )

delete_tag(tag_id: UUID) -> None

Deletes a tag.

Parameters:

Name	Type	Description	Default
tag_id	UUID	id of the tag to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_tag(
    self,
    tag_id: UUID,
) -> None:
    """Deletes a tag.

    Args:
        tag_id: id of the tag to delete.
    """
    self._delete_resource(
        resource_id=tag_id,
        route=TAGS,
    )

delete_tag_resource(tag_resource: TagResourceRequest) -> None

Delete a tag resource.

Parameters:

Name	Type	Description	Default
tag_resource	TagResourceRequest	The tag resource relationship to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_tag_resource(
    self,
    tag_resource: TagResourceRequest,
) -> None:
    """Delete a tag resource.

    Args:
        tag_resource: The tag resource relationship to delete.
    """
    self.delete(path=TAG_RESOURCES, body=tag_resource)

delete_trigger(trigger_id: UUID, soft: bool = True) -> None

Deletes a trigger.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	The ID of the trigger.	required
soft	bool	Flag deciding whether to soft-delete the trigger.	True

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_trigger(self, trigger_id: UUID, soft: bool = True) -> None:
    """Deletes a trigger.

    Args:
        trigger_id: The ID of the trigger.
        soft: Flag deciding whether to soft-delete the trigger.
    """
    self._delete_resource(
        resource_id=trigger_id,
        route=TRIGGERS,
        params={"soft": soft},
    )

delete_user(user_name_or_id: Union[str, UUID]) -> None

Deletes a user.

Parameters:

Name	Type	Description	Default
user_name_or_id	Union[str, UUID]	The name or ID of the user to delete.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def delete_user(self, user_name_or_id: Union[str, UUID]) -> None:
    """Deletes a user.

    Args:
        user_name_or_id: The name or ID of the user to delete.
    """
    self._delete_resource(
        resource_id=user_name_or_id,
        route=USERS,
    )

detach_trigger_from_snapshot(trigger_id: UUID, snapshot_id: UUID) -> None

Detaches (unlinks) a trigger from a snapshot.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	The ID of the trigger.	required
snapshot_id	UUID	The ID of the snapshot.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def detach_trigger_from_snapshot(
    self, trigger_id: UUID, snapshot_id: UUID
) -> None:
    """Detaches (unlinks) a trigger from a snapshot.

    Args:
        trigger_id: The ID of the trigger.
        snapshot_id: The ID of the snapshot.
    """
    self.delete(
        path=f"{TRIGGERS}/{trigger_id}{PIPELINE_SNAPSHOTS}/{snapshot_id}",
    )

disable_run_heartbeat(run_id: UUID) -> None

Disables heartbeat for a pipeline run.

Parameters:

Name	Type	Description	Default
run_id	UUID	The ID of the pipeline run.	required

Source code in src/zenml/zen_stores/rest_zen_store.py

def disable_run_heartbeat(self, run_id: UUID) -> None:
    """Disables heartbeat for a pipeline run.

    Args:
        run_id: The ID of the pipeline run.
    """
    self.put(
        path=f"{RUNS}/{str(run_id)}{DISABLE_HEARTBEAT}",
    )

get(path: str, params: Optional[Dict[str, Any]] = None, timeout: Optional[int] = None, **kwargs: Any) -> Json

Make a GET request to the given endpoint path.

Parameters:

Name	Type	Description	Default
path	str	The path to the endpoint.	required
params	Optional[Dict[str, Any]]	The query parameters to pass to the endpoint.	None
timeout	Optional[int]	The request timeout in seconds.	None
kwargs	Any	Additional keyword arguments to pass to the request.	{}

Returns:

Type	Description
Json	The response body.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get(
    self,
    path: str,
    params: Optional[Dict[str, Any]] = None,
    timeout: Optional[int] = None,
    **kwargs: Any,
) -> Json:
    """Make a GET request to the given endpoint path.

    Args:
        path: The path to the endpoint.
        params: The query parameters to pass to the endpoint.
        timeout: The request timeout in seconds.
        kwargs: Additional keyword arguments to pass to the request.

    Returns:
        The response body.
    """
    return self._request(
        "GET",
        self.url + API + VERSION_1 + path,
        params=params,
        timeout=timeout,
        **kwargs,
    )

get_api_key(service_account_id: UUID, api_key_name_or_id: Union[str, UUID], hydrate: bool = True) -> APIKeyResponse

Get an API key for a service account.

Parameters:

Name	Type	Description	Default
service_account_id	UUID	The ID of the service account for which to fetch the API key.	required
api_key_name_or_id	Union[str, UUID]	The name or ID of the API key to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
APIKeyResponse	The API key with the given ID.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_api_key(
    self,
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    hydrate: bool = True,
) -> APIKeyResponse:
    """Get an API key for a service account.

    Args:
        service_account_id: The ID of the service account for which to fetch
            the API key.
        api_key_name_or_id: The name or ID of the API key to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The API key with the given ID.
    """
    return self._get_resource(
        resource_id=api_key_name_or_id,
        route=f"{SERVICE_ACCOUNTS}/{str(service_account_id)}{API_KEYS}",
        response_model=APIKeyResponse,
        params={"hydrate": hydrate},
    )

get_api_token(token_type: APITokenType = APITokenType.WORKLOAD, expires_in: Optional[int] = None, schedule_id: Optional[UUID] = None, pipeline_run_id: Optional[UUID] = None, deployment_id: Optional[UUID] = None) -> str

Get an API token.

Parameters:

Name	Type	Description	Default
token_type	APITokenType	The type of the token to get.	WORKLOAD
expires_in	Optional[int]	The time in seconds until the token expires.	None
schedule_id	Optional[UUID]	The ID of the schedule to get a token for.	None
pipeline_run_id	Optional[UUID]	The ID of the pipeline run to get a token for.	None
deployment_id	Optional[UUID]	The ID of the deployment to get a token for.	None

Returns:

Type	Description
str	The API token.

Raises:

Type	Description
ValueError	if the server response is not valid.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_api_token(
    self,
    token_type: APITokenType = APITokenType.WORKLOAD,
    expires_in: Optional[int] = None,
    schedule_id: Optional[UUID] = None,
    pipeline_run_id: Optional[UUID] = None,
    deployment_id: Optional[UUID] = None,
) -> str:
    """Get an API token.

    Args:
        token_type: The type of the token to get.
        expires_in: The time in seconds until the token expires.
        schedule_id: The ID of the schedule to get a token for.
        pipeline_run_id: The ID of the pipeline run to get a token for.
        deployment_id: The ID of the deployment to get a token for.

    Returns:
        The API token.

    Raises:
        ValueError: if the server response is not valid.
    """
    params: Dict[str, Any] = {
        "token_type": token_type.value,
    }
    if expires_in:
        params["expires_in"] = expires_in
    if schedule_id:
        params["schedule_id"] = schedule_id
    if pipeline_run_id:
        params["pipeline_run_id"] = pipeline_run_id
    if deployment_id:
        params["deployment_id"] = deployment_id
    response_body = self.get(API_TOKEN, params=params)
    if not isinstance(response_body, str):
        raise ValueError(
            f"Bad API Response. Expected API token, got "
            f"{type(response_body)}"
        )
    return response_body

get_artifact(artifact_id: UUID, hydrate: bool = True) -> ArtifactResponse

Gets an artifact.

Parameters:

Name	Type	Description	Default
artifact_id	UUID	The ID of the artifact to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ArtifactResponse	The artifact.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_artifact(
    self, artifact_id: UUID, hydrate: bool = True
) -> ArtifactResponse:
    """Gets an artifact.

    Args:
        artifact_id: The ID of the artifact to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The artifact.
    """
    return self._get_resource(
        resource_id=artifact_id,
        route=ARTIFACTS,
        response_model=ArtifactResponse,
        params={"hydrate": hydrate},
    )

get_artifact_version(artifact_version_id: UUID, hydrate: bool = True) -> ArtifactVersionResponse

Gets an artifact.

Parameters:

Name	Type	Description	Default
artifact_version_id	UUID	The ID of the artifact version to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ArtifactVersionResponse	The artifact version.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_artifact_version(
    self, artifact_version_id: UUID, hydrate: bool = True
) -> ArtifactVersionResponse:
    """Gets an artifact.

    Args:
        artifact_version_id: The ID of the artifact version to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The artifact version.
    """
    return self._get_resource(
        resource_id=artifact_version_id,
        route=ARTIFACT_VERSIONS,
        response_model=ArtifactVersionResponse,
        params={"hydrate": hydrate},
    )

get_artifact_visualization(artifact_visualization_id: UUID, hydrate: bool = True) -> ArtifactVisualizationResponse

Gets an artifact visualization.

Parameters:

Name	Type	Description	Default
artifact_visualization_id	UUID	The ID of the artifact visualization to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ArtifactVisualizationResponse	The artifact visualization.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_artifact_visualization(
    self, artifact_visualization_id: UUID, hydrate: bool = True
) -> ArtifactVisualizationResponse:
    """Gets an artifact visualization.

    Args:
        artifact_visualization_id: The ID of the artifact visualization to
            get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The artifact visualization.
    """
    return self._get_resource(
        resource_id=artifact_visualization_id,
        route=ARTIFACT_VISUALIZATIONS,
        response_model=ArtifactVisualizationResponse,
        params={"hydrate": hydrate},
    )

get_authorized_device(device_id: UUID, hydrate: bool = True) -> OAuthDeviceResponse

Gets a specific OAuth 2.0 authorized device.

Parameters:

Name	Type	Description	Default
device_id	UUID	The ID of the device to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
OAuthDeviceResponse	The requested device, if it was found.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_authorized_device(
    self, device_id: UUID, hydrate: bool = True
) -> OAuthDeviceResponse:
    """Gets a specific OAuth 2.0 authorized device.

    Args:
        device_id: The ID of the device to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested device, if it was found.
    """
    return self._get_resource(
        resource_id=device_id,
        route=DEVICES,
        response_model=OAuthDeviceResponse,
        params={"hydrate": hydrate},
    )

get_build(build_id: UUID, hydrate: bool = True) -> PipelineBuildResponse

Get a build with a given ID.

Parameters:

Name	Type	Description	Default
build_id	UUID	ID of the build.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
PipelineBuildResponse	The build.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_build(
    self, build_id: UUID, hydrate: bool = True
) -> PipelineBuildResponse:
    """Get a build with a given ID.

    Args:
        build_id: ID of the build.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The build.
    """
    return self._get_resource(
        resource_id=build_id,
        route=PIPELINE_BUILDS,
        response_model=PipelineBuildResponse,
        params={"hydrate": hydrate},
    )

get_code_reference(code_reference_id: UUID, hydrate: bool = True) -> CodeReferenceResponse

Gets a code reference.

Parameters:

Name	Type	Description	Default
code_reference_id	UUID	The ID of the code reference to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
CodeReferenceResponse	The code reference.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_code_reference(
    self, code_reference_id: UUID, hydrate: bool = True
) -> CodeReferenceResponse:
    """Gets a code reference.

    Args:
        code_reference_id: The ID of the code reference to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The code reference.
    """
    return self._get_resource(
        resource_id=code_reference_id,
        route=CODE_REFERENCES,
        response_model=CodeReferenceResponse,
        params={"hydrate": hydrate},
    )

get_code_repository(code_repository_id: UUID, hydrate: bool = True) -> CodeRepositoryResponse

Gets a specific code repository.

Parameters:

Name	Type	Description	Default
code_repository_id	UUID	The ID of the code repository to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
CodeRepositoryResponse	The requested code repository, if it was found.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_code_repository(
    self, code_repository_id: UUID, hydrate: bool = True
) -> CodeRepositoryResponse:
    """Gets a specific code repository.

    Args:
        code_repository_id: The ID of the code repository to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested code repository, if it was found.
    """
    return self._get_resource(
        resource_id=code_repository_id,
        route=CODE_REPOSITORIES,
        response_model=CodeRepositoryResponse,
        params={"hydrate": hydrate},
    )

get_curated_visualization(visualization_id: UUID, hydrate: bool = True) -> CuratedVisualizationResponse

Get a curated visualization by ID.

Parameters:

Name	Type	Description	Default
visualization_id	UUID	The ID of the curated visualization to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
CuratedVisualizationResponse	The curated visualization with the given ID.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_curated_visualization(
    self, visualization_id: UUID, hydrate: bool = True
) -> CuratedVisualizationResponse:
    """Get a curated visualization by ID.

    Args:
        visualization_id: The ID of the curated visualization to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The curated visualization with the given ID.
    """
    return self._get_resource(
        resource_id=visualization_id,
        route=CURATED_VISUALIZATIONS,
        response_model=CuratedVisualizationResponse,
        params={"hydrate": hydrate},
    )

get_deployment(deployment_id: UUID, hydrate: bool = True) -> DeploymentResponse

Get a deployment with a given ID.

Parameters:

Name	Type	Description	Default
deployment_id	UUID	ID of the deployment.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
DeploymentResponse	The deployment.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_deployment(
    self, deployment_id: UUID, hydrate: bool = True
) -> DeploymentResponse:
    """Get a deployment with a given ID.

    Args:
        deployment_id: ID of the deployment.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The deployment.
    """
    return self._get_resource(
        resource_id=deployment_id,
        route=DEPLOYMENTS,
        response_model=DeploymentResponse,
        params={"hydrate": hydrate},
    )

get_deployment_id() -> UUID

Get the ID of the deployment.

Returns:

Type	Description
UUID	The ID of the deployment.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_deployment_id(self) -> UUID:
    """Get the ID of the deployment.

    Returns:
        The ID of the deployment.
    """
    return self.server_info.id

get_flavor(flavor_id: UUID, hydrate: bool = True) -> FlavorResponse

Get a stack component flavor by ID.

Parameters:

Name	Type	Description	Default
flavor_id	UUID	The ID of the stack component flavor to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
FlavorResponse	The stack component flavor.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_flavor(
    self, flavor_id: UUID, hydrate: bool = True
) -> FlavorResponse:
    """Get a stack component flavor by ID.

    Args:
        flavor_id: The ID of the stack component flavor to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The stack component flavor.
    """
    return self._get_resource(
        resource_id=flavor_id,
        route=FLAVORS,
        response_model=FlavorResponse,
        params={"hydrate": hydrate},
    )

get_logs(logs_id: UUID, hydrate: bool = True) -> LogsResponse

Gets logs with the given ID.

Parameters:

Name	Type	Description	Default
logs_id	UUID	The ID of the logs to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
LogsResponse	The logs.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_logs(self, logs_id: UUID, hydrate: bool = True) -> LogsResponse:
    """Gets logs with the given ID.

    Args:
        logs_id: The ID of the logs to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The logs.
    """
    return self._get_resource(
        resource_id=logs_id,
        route=LOGS,
        response_model=LogsResponse,
        params={"hydrate": hydrate},
    )

get_model(model_id: UUID, hydrate: bool = True) -> ModelResponse

Get an existing model.

Parameters:

Name	Type	Description	Default
model_id	UUID	id of the model to be retrieved.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ModelResponse	The model of interest.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_model(self, model_id: UUID, hydrate: bool = True) -> ModelResponse:
    """Get an existing model.

    Args:
        model_id: id of the model to be retrieved.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The model of interest.
    """
    return self._get_resource(
        resource_id=model_id,
        route=MODELS,
        response_model=ModelResponse,
        params={"hydrate": hydrate},
    )

get_model_version(model_version_id: UUID, hydrate: bool = True) -> ModelVersionResponse

Get an existing model version.

Parameters:

Name	Type	Description	Default
model_version_id	UUID	name, id, stage or number of the model version to be retrieved. If skipped - latest is retrieved.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ModelVersionResponse	The model version of interest.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_model_version(
    self, model_version_id: UUID, hydrate: bool = True
) -> ModelVersionResponse:
    """Get an existing model version.

    Args:
        model_version_id: name, id, stage or number of the model version to
            be retrieved. If skipped - latest is retrieved.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The model version of interest.
    """
    return self._get_resource(
        resource_id=model_version_id,
        route=MODEL_VERSIONS,
        response_model=ModelVersionResponse,
        params={"hydrate": hydrate},
    )

get_or_create_run(pipeline_run: PipelineRunRequest) -> Tuple[PipelineRunResponse, bool]

Gets or creates a pipeline run.

If a run with the same ID or name already exists, it is returned. Otherwise, a new run is created.

Parameters:

Name	Type	Description	Default
pipeline_run	PipelineRunRequest	The pipeline run to get or create.	required

Returns:

Type	Description
PipelineRunResponse	The pipeline run, and a boolean indicating whether the run was
bool	created or not.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_or_create_run(
    self, pipeline_run: PipelineRunRequest
) -> Tuple[PipelineRunResponse, bool]:
    """Gets or creates a pipeline run.

    If a run with the same ID or name already exists, it is returned.
    Otherwise, a new run is created.

    Args:
        pipeline_run: The pipeline run to get or create.

    Returns:
        The pipeline run, and a boolean indicating whether the run was
        created or not.
    """
    return self._get_or_create_resource(
        resource=pipeline_run,
        route=RUNS,
        response_model=PipelineRunResponse,
    )

get_or_generate_api_token() -> str

Get or generate an API token.

Returns:

Type	Description
str	The API token.

Raises:

Type	Description
CredentialsNotValid	if an API token cannot be fetched or generated because the client credentials are not valid.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_or_generate_api_token(self) -> str:
    """Get or generate an API token.

    Returns:
        The API token.

    Raises:
        CredentialsNotValid: if an API token cannot be fetched or
            generated because the client credentials are not valid.
    """
    if self._api_token is None or self._api_token.expired:
        # Check if a valid API token is already in the cache
        credentials_store = get_credentials_store()
        credentials = credentials_store.get_credentials(self.url)
        token = credentials.api_token if credentials else None
        if credentials and token and not token.expired:
            self._api_token = token
            return self._api_token.access_token

        # Token is expired or not found in the cache. Time to get a new one.

        if not token:
            logger.debug(f"Authenticating to {self.url}")
        else:
            logger.debug(
                f"Authentication token for {self.url} expired; refreshing..."
            )

        data: Optional[Dict[str, str]] = None

        # Use a custom user agent to identify the ZenML client in the server
        # logs.
        headers: Dict[str, str] = {
            "User-Agent": "zenml/" + zenml.__version__,
        }

        # Check if an API key is configured
        api_key = credentials_store.get_api_key(self.url)

        # Check if username and password are configured
        username, password = credentials_store.get_password(self.url)

        if api_key is not None:
            # An API key is configured. Use it as a password to
            # authenticate.
            data = {
                "grant_type": OAuthGrantTypes.ZENML_API_KEY.value,
                "password": api_key,
            }
        elif username is not None and password is not None:
            # Username and password are configured. Use them to authenticate.
            data = {
                "grant_type": OAuthGrantTypes.OAUTH_PASSWORD.value,
                "username": username,
                "password": password,
            }
        elif self.server_info.is_pro_server():
            # ZenML Pro workspaces use a proprietary authorization grant
            # where the ZenML Pro API session token is exchanged for a
            # regular ZenML server access token.

            # Get the ZenML Pro API session token, if cached and valid

            # We need to determine the right ZenML Pro API URL to use
            pro_api_url = self.server_info.pro_api_url
            if not pro_api_url and credentials and credentials.pro_api_url:
                pro_api_url = credentials.pro_api_url
            if not pro_api_url:
                pro_api_url = ZENML_PRO_API_URL

            pro_credentials = credentials_store.get_pro_credentials(
                pro_api_url
            )
            if not pro_credentials:
                raise CredentialsNotValid(
                    "You need to be logged in to ZenML Pro in order to "
                    f"access the ZenML Pro server '{self.url}'. Please run "
                    "'zenml login' to log in or choose a different server."
                )

            elif pro_credentials.has_valid_token:
                assert pro_credentials.api_token is not None
                pro_token = pro_credentials.api_token
            elif pro_credentials.can_refresh_token:
                pro_token = ZenMLProClient(pro_api_url).authenticate()
            else:
                raise CredentialsNotValid(
                    "Your ZenML Pro login session has expired. "
                    "Please log in again using 'zenml login'."
                )

            data = {
                "grant_type": OAuthGrantTypes.ZENML_EXTERNAL.value,
            }
            headers.update(
                {"Authorization": "Bearer " + pro_token.access_token}
            )
        else:
            if not token:
                raise CredentialsNotValid(
                    "No valid credentials found. Please run 'zenml login "
                    f"{self.url}' to connect to the current server."
                )
            elif token.expired:
                raise CredentialsNotValid(
                    "Your authentication to the current server has expired. "
                    "Please log in again using 'zenml login "
                    f"{self.url}'."
                )

        response = self._handle_response(
            requests.post(
                self.url + API + VERSION_1 + LOGIN,
                data=data,
                verify=self.config.verify_ssl,
                timeout=self.config.http_timeout,
                headers=headers,
            )
        )
        try:
            token_response = OAuthTokenResponse.model_validate(response)
        except ValidationError as e:
            raise CredentialsNotValid(
                "Unexpected response received while authenticating to "
                f"the server {e}"
            ) from e

        # Cache the token
        self._api_token = credentials_store.set_token(
            self.url, token_response
        )

        # Update the server info in the credentials store with the latest
        # information from the server.
        # NOTE: this is the best place to do this because we know that
        # the token is valid and the server is reachable.
        try:
            server_info = self.get_store_info()
        except Exception as e:
            logger.warning(f"Failed to get server info: {e}.")
        else:
            credentials_store.update_server_info(self.url, server_info)

    return self._api_token.access_token

get_pipeline(pipeline_id: UUID, hydrate: bool = True) -> PipelineResponse

Get a pipeline with a given ID.

Parameters:

Name	Type	Description	Default
pipeline_id	UUID	ID of the pipeline.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
PipelineResponse	The pipeline.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_pipeline(
    self, pipeline_id: UUID, hydrate: bool = True
) -> PipelineResponse:
    """Get a pipeline with a given ID.

    Args:
        pipeline_id: ID of the pipeline.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The pipeline.
    """
    return self._get_resource(
        resource_id=pipeline_id,
        route=PIPELINES,
        response_model=PipelineResponse,
        params={"hydrate": hydrate},
    )

get_project(project_name_or_id: Union[UUID, str], hydrate: bool = True) -> ProjectResponse

Get an existing project by name or ID.

Parameters:

Name	Type	Description	Default
project_name_or_id	Union[UUID, str]	Name or ID of the project to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ProjectResponse	The requested project.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_project(
    self, project_name_or_id: Union[UUID, str], hydrate: bool = True
) -> ProjectResponse:
    """Get an existing project by name or ID.

    Args:
        project_name_or_id: Name or ID of the project to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested project.
    """
    return self._get_resource(
        resource_id=project_name_or_id,
        route=PROJECTS,
        response_model=ProjectResponse,
        params={"hydrate": hydrate},
    )

get_resource_pool(resource_pool_id: UUID, hydrate: bool = True) -> ResourcePoolResponse

Get a resource pool by ID.

Parameters:

Name	Type	Description	Default
resource_pool_id	UUID	The ID of the resource pool to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ResourcePoolResponse	The resource pool.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_resource_pool(
    self, resource_pool_id: UUID, hydrate: bool = True
) -> ResourcePoolResponse:
    """Get a resource pool by ID.

    Args:
        resource_pool_id: The ID of the resource pool to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The resource pool.
    """
    return self._get_resource(
        resource_id=resource_pool_id,
        route=RESOURCE_POOLS,
        response_model=ResourcePoolResponse,
        params={"hydrate": hydrate},
    )

get_resource_pool_subject_policy(policy_id: UUID, hydrate: bool = True) -> ResourcePoolSubjectPolicyResponse

Get a resource pool subject policy by ID.

Parameters:

Name	Type	Description	Default
policy_id	UUID	The ID of the policy to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s).	True

Returns:

Type	Description
ResourcePoolSubjectPolicyResponse	The policy.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_resource_pool_subject_policy(
    self, policy_id: UUID, hydrate: bool = True
) -> ResourcePoolSubjectPolicyResponse:
    """Get a resource pool subject policy by ID.

    Args:
        policy_id: The ID of the policy to get.
        hydrate: Flag deciding whether to hydrate the output model(s).

    Returns:
        The policy.
    """
    return self._get_resource(
        resource_id=policy_id,
        route=RESOURCE_POOL_SUBJECT_POLICIES,
        response_model=ResourcePoolSubjectPolicyResponse,
        params={"hydrate": hydrate},
    )

get_resource_request(resource_request_id: UUID, hydrate: bool = True) -> ResourceRequestResponse

Get a resource request by ID.

Parameters:

Name	Type	Description	Default
resource_request_id	UUID	The ID of the resource request to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ResourceRequestResponse	The resource request.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_resource_request(
    self, resource_request_id: UUID, hydrate: bool = True
) -> ResourceRequestResponse:
    """Get a resource request by ID.

    Args:
        resource_request_id: The ID of the resource request to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The resource request.
    """
    return self._get_resource(
        resource_id=resource_request_id,
        route=RESOURCE_REQUESTS,
        response_model=ResourceRequestResponse,
        params={"hydrate": hydrate},
    )

get_run(run_id: UUID, hydrate: bool = True, include_full_metadata: bool = False) -> PipelineRunResponse

Gets a pipeline run.

Parameters:

Name	Type	Description	Default
run_id	UUID	The ID of the pipeline run to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True
include_full_metadata	bool	If True, include metadata of all steps in the response.	False

Returns:

Type	Description
PipelineRunResponse	The pipeline run.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_run(
    self,
    run_id: UUID,
    hydrate: bool = True,
    include_full_metadata: bool = False,
) -> PipelineRunResponse:
    """Gets a pipeline run.

    Args:
        run_id: The ID of the pipeline run to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        include_full_metadata: If True, include metadata of all steps in
            the response.

    Returns:
        The pipeline run.
    """
    return self._get_resource(
        resource_id=run_id,
        route=RUNS,
        response_model=PipelineRunResponse,
        params={
            "hydrate": hydrate,
            "include_full_metadata": include_full_metadata,
        },
    )

get_run_statistics(request: RunStatisticsRequest) -> RunStatisticsResponse

Compute grouped statistics over pipeline runs.

Parameters:

Name	Type	Description	Default
request	RunStatisticsRequest	Statistics request.	required

Returns:

Type	Description
RunStatisticsResponse	Grouped statistics.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_run_statistics(
    self, request: RunStatisticsRequest
) -> RunStatisticsResponse:
    """Compute grouped statistics over pipeline runs.

    Args:
        request: Statistics request.

    Returns:
        Grouped statistics.
    """
    response_body = self.post(f"{RUNS}{STATISTICS}", body=request)
    return RunStatisticsResponse.model_validate(response_body)

get_run_step(step_run_id: UUID, hydrate: bool = True) -> StepRunResponse

Get a step run by ID.

Parameters:

Name	Type	Description	Default
step_run_id	UUID	The ID of the step run to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
StepRunResponse	The step run.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_run_step(
    self, step_run_id: UUID, hydrate: bool = True
) -> StepRunResponse:
    """Get a step run by ID.

    Args:
        step_run_id: The ID of the step run to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The step run.
    """
    return self._get_resource(
        resource_id=step_run_id,
        route=STEPS,
        response_model=StepRunResponse,
        params={"hydrate": hydrate},
    )

get_run_template(template_id: UUID, hydrate: bool = True) -> RunTemplateResponse

Get a run template with a given ID.

Parameters:

Name	Type	Description	Default
template_id	UUID	ID of the template.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
RunTemplateResponse	The template.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_run_template(
    self, template_id: UUID, hydrate: bool = True
) -> RunTemplateResponse:
    """Get a run template with a given ID.

    Args:
        template_id: ID of the template.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The template.
    """
    return self._get_resource(
        resource_id=template_id,
        route=RUN_TEMPLATES,
        response_model=RunTemplateResponse,
        params={"hydrate": hydrate},
    )

get_run_wait_condition(run_wait_condition_id: UUID, hydrate: bool = True) -> RunWaitConditionResponse

Get a run wait condition.

Parameters:

Name	Type	Description	Default
run_wait_condition_id	UUID	Wait condition ID.	required
hydrate	bool	Whether to hydrate metadata/resources.	True

Returns:

Type	Description
RunWaitConditionResponse	The requested wait condition.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_run_wait_condition(
    self, run_wait_condition_id: UUID, hydrate: bool = True
) -> RunWaitConditionResponse:
    """Get a run wait condition.

    Args:
        run_wait_condition_id: Wait condition ID.
        hydrate: Whether to hydrate metadata/resources.

    Returns:
        The requested wait condition.
    """
    return self._get_resource(
        resource_id=run_wait_condition_id,
        route=RUN_WAIT_CONDITIONS,
        response_model=RunWaitConditionResponse,
        params={"hydrate": hydrate},
    )

get_schedule(schedule_id: UUID, hydrate: bool = True) -> ScheduleResponse

Get a schedule with a given ID.

Parameters:

Name	Type	Description	Default
schedule_id	UUID	ID of the schedule.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ScheduleResponse	The schedule.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_schedule(
    self, schedule_id: UUID, hydrate: bool = True
) -> ScheduleResponse:
    """Get a schedule with a given ID.

    Args:
        schedule_id: ID of the schedule.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The schedule.
    """
    return self._get_resource(
        resource_id=schedule_id,
        route=SCHEDULES,
        response_model=ScheduleResponse,
        params={"hydrate": hydrate},
    )

get_secret(secret_id: UUID, hydrate: bool = True) -> SecretResponse

Get a secret by ID.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret to fetch.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
SecretResponse	The secret.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_secret(
    self, secret_id: UUID, hydrate: bool = True
) -> SecretResponse:
    """Get a secret by ID.

    Args:
        secret_id: The ID of the secret to fetch.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The secret.
    """
    return self._get_resource(
        resource_id=secret_id,
        route=SECRETS,
        response_model=SecretResponse,
        params={"hydrate": hydrate},
    )

get_server_settings(hydrate: bool = True) -> ServerSettingsResponse

Get the server settings.

Parameters:

Name	Type	Description	Default
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ServerSettingsResponse	The server settings.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_server_settings(
    self, hydrate: bool = True
) -> ServerSettingsResponse:
    """Get the server settings.

    Args:
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The server settings.
    """
    response_body = self.get(SERVER_SETTINGS, params={"hydrate": hydrate})
    return ServerSettingsResponse.model_validate(response_body)

get_service(service_id: UUID, hydrate: bool = True) -> ServiceResponse

Get a service.

Parameters:

Name	Type	Description	Default
service_id	UUID	The ID of the service to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ServiceResponse	The service.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_service(
    self, service_id: UUID, hydrate: bool = True
) -> ServiceResponse:
    """Get a service.

    Args:
        service_id: The ID of the service to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The service.
    """
    return self._get_resource(
        resource_id=service_id,
        route=SERVICES,
        response_model=ServiceResponse,
        params={"hydrate": hydrate},
    )

get_service_account(service_account_name_or_id: Union[str, UUID], hydrate: bool = True) -> ServiceAccountResponse

Gets a specific service account.

Parameters:

Name	Type	Description	Default
service_account_name_or_id	Union[str, UUID]	The name or ID of the service account to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ServiceAccountResponse	The requested service account, if it was found.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_service_account(
    self,
    service_account_name_or_id: Union[str, UUID],
    hydrate: bool = True,
) -> ServiceAccountResponse:
    """Gets a specific service account.

    Args:
        service_account_name_or_id: The name or ID of the service account to
            get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested service account, if it was found.
    """
    return self._get_resource(
        resource_id=service_account_name_or_id,
        route=SERVICE_ACCOUNTS,
        response_model=ServiceAccountResponse,
        params={"hydrate": hydrate},
    )

get_service_connector(service_connector_id: UUID, hydrate: bool = True, expand_secrets: bool = False) -> ServiceConnectorResponse

Gets a specific service connector.

Parameters:

Name	Type	Description	Default
service_connector_id	UUID	The ID of the service connector to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True
expand_secrets	bool	Flag deciding whether to include the secrets associated with the service connector.	False

Returns:

Type	Description
ServiceConnectorResponse	The requested service connector, if it was found.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_service_connector(
    self,
    service_connector_id: UUID,
    hydrate: bool = True,
    expand_secrets: bool = False,
) -> ServiceConnectorResponse:
    """Gets a specific service connector.

    Args:
        service_connector_id: The ID of the service connector to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        expand_secrets: Flag deciding whether to include the secrets
            associated with the service connector.

    Returns:
        The requested service connector, if it was found.
    """
    connector_model = self._get_resource(
        resource_id=service_connector_id,
        route=SERVICE_CONNECTORS,
        response_model=ServiceConnectorResponse,
        params={"hydrate": hydrate, "expand_secrets": expand_secrets},
    )
    self._populate_connector_type(connector_model)
    if expand_secrets:
        try:
            # Call this to properly split the secrets from the configuration
            connector_model.validate_configuration()
        except ValueError as e:
            logger.error(
                f"Error validating connector configuration for "
                f"{connector_model.name}: {e}"
            )
    return connector_model

get_service_connector_client(service_connector_id: UUID, resource_type: Optional[str] = None, resource_id: Optional[str] = None) -> ServiceConnectorResponse

Get a service connector client for a service connector and given resource.

Parameters:

Name	Type	Description	Default
service_connector_id	UUID	The ID of the base service connector to use.	required
resource_type	Optional[str]	The type of resource to get a client for.	None
resource_id	Optional[str]	The ID of the resource to get a client for.	None

Returns:

Type	Description
ServiceConnectorResponse	A service connector client that can be used to access the given
ServiceConnectorResponse	resource.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_service_connector_client(
    self,
    service_connector_id: UUID,
    resource_type: Optional[str] = None,
    resource_id: Optional[str] = None,
) -> ServiceConnectorResponse:
    """Get a service connector client for a service connector and given resource.

    Args:
        service_connector_id: The ID of the base service connector to use.
        resource_type: The type of resource to get a client for.
        resource_id: The ID of the resource to get a client for.

    Returns:
        A service connector client that can be used to access the given
        resource.
    """
    params = {}
    if resource_type:
        params["resource_type"] = resource_type
    if resource_id:
        params["resource_id"] = resource_id
    response_body = self.get(
        f"{SERVICE_CONNECTORS}/{str(service_connector_id)}{SERVICE_CONNECTOR_CLIENT}",
        params=params,
    )

    connector = ServiceConnectorResponse.model_validate(response_body)
    self._populate_connector_type(connector)
    # Call this to properly split the secrets from the configuration
    try:
        connector.validate_configuration()
    except ValueError as e:
        logger.error(
            f"Error validating connector configuration for connector client "
            f"{connector.name}: {e}"
        )
    return connector

get_service_connector_type(connector_type: str) -> ServiceConnectorTypeModel

Returns the requested service connector type.

Parameters:

Name	Type	Description	Default
connector_type	str	the service connector type identifier.	required

Returns:

Type	Description
ServiceConnectorTypeModel	The requested service connector type.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_service_connector_type(
    self,
    connector_type: str,
) -> ServiceConnectorTypeModel:
    """Returns the requested service connector type.

    Args:
        connector_type: the service connector type identifier.

    Returns:
        The requested service connector type.
    """
    # Use the local registry to get the service connector type, if it
    # exists.
    local_connector_type: Optional[ServiceConnectorTypeModel] = None
    if service_connector_registry.is_registered(connector_type):
        local_connector_type = (
            service_connector_registry.get_service_connector_type(
                connector_type
            )
        )
    try:
        response_body = self.get(
            f"{SERVICE_CONNECTOR_TYPES}/{connector_type}",
        )
        remote_connector_type = ServiceConnectorTypeModel.model_validate(
            response_body
        )
        if local_connector_type:
            # If locally available, return the local connector type but
            # mark it as being remotely available.
            local_connector_type.remote = True
            return local_connector_type

        # Mark the remote connector type as being only remotely available
        remote_connector_type.local = False
        remote_connector_type.remote = True

        return remote_connector_type
    except KeyError:
        # If the service connector type is not found, check the local
        # registry.
        return service_connector_registry.get_service_connector_type(
            connector_type
        )

get_snapshot(snapshot_id: UUID, hydrate: bool = True, step_configuration_filter: Optional[List[str]] = None, include_config_schema: Optional[bool] = None) -> PipelineSnapshotResponse

Get a snapshot with a given ID.

Parameters:

Name	Type	Description	Default
snapshot_id	UUID	ID of the snapshot.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True
step_configuration_filter	Optional[List[str]]	List of step configurations to include in the response. If not given, all step configurations will be included.	None
include_config_schema	Optional[bool]	Whether the config schema will be filled.	None

Returns:

Type	Description
PipelineSnapshotResponse	The snapshot.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_snapshot(
    self,
    snapshot_id: UUID,
    hydrate: bool = True,
    step_configuration_filter: Optional[List[str]] = None,
    include_config_schema: Optional[bool] = None,
) -> PipelineSnapshotResponse:
    """Get a snapshot with a given ID.

    Args:
        snapshot_id: ID of the snapshot.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        step_configuration_filter: List of step configurations to include in
            the response. If not given, all step configurations will be
            included.
        include_config_schema: Whether the config schema will be filled.

    Returns:
        The snapshot.
    """
    return self._get_resource(
        resource_id=snapshot_id,
        route=PIPELINE_SNAPSHOTS,
        response_model=PipelineSnapshotResponse,
        params={
            "hydrate": hydrate,
            "step_configuration_filter": step_configuration_filter,
            "include_config_schema": include_config_schema,
        },
    )

get_stack(stack_id: UUID, hydrate: bool = True) -> StackResponse

Get a stack by its unique ID.

Parameters:

Name	Type	Description	Default
stack_id	UUID	The ID of the stack to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
StackResponse	The stack with the given ID.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_stack(self, stack_id: UUID, hydrate: bool = True) -> StackResponse:
    """Get a stack by its unique ID.

    Args:
        stack_id: The ID of the stack to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The stack with the given ID.
    """
    return self._get_resource(
        resource_id=stack_id,
        route=STACKS,
        response_model=StackResponse,
        params={"hydrate": hydrate},
    )

get_stack_component(component_id: UUID, hydrate: bool = True) -> ComponentResponse

Get a stack component by ID.

Parameters:

Name	Type	Description	Default
component_id	UUID	The ID of the stack component to get.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
ComponentResponse	The stack component.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_stack_component(
    self, component_id: UUID, hydrate: bool = True
) -> ComponentResponse:
    """Get a stack component by ID.

    Args:
        component_id: The ID of the stack component to get.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The stack component.
    """
    return self._get_resource(
        resource_id=component_id,
        route=STACK_COMPONENTS,
        response_model=ComponentResponse,
        params={"hydrate": hydrate},
    )

get_stack_deployment_config(provider: StackDeploymentProvider, stack_name: str, location: Optional[str] = None) -> StackDeploymentConfig

Return the cloud provider console URL and configuration needed to deploy the ZenML stack.

Parameters:

Name	Type	Description	Default
provider	StackDeploymentProvider	The stack deployment provider.	required
stack_name	str	The name of the stack.	required
location	Optional[str]	The location where the stack should be deployed.	None

Returns:

Type	Description
StackDeploymentConfig	The cloud provider console URL and configuration needed to deploy
StackDeploymentConfig	the ZenML stack to the specified cloud provider.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_stack_deployment_config(
    self,
    provider: StackDeploymentProvider,
    stack_name: str,
    location: Optional[str] = None,
) -> StackDeploymentConfig:
    """Return the cloud provider console URL and configuration needed to deploy the ZenML stack.

    Args:
        provider: The stack deployment provider.
        stack_name: The name of the stack.
        location: The location where the stack should be deployed.

    Returns:
        The cloud provider console URL and configuration needed to deploy
        the ZenML stack to the specified cloud provider.
    """
    params = {
        "provider": provider.value,
        "stack_name": stack_name,
    }
    if location:
        params["location"] = location
    body = self.get(f"{STACK_DEPLOYMENT}{CONFIG}", params=params)
    return StackDeploymentConfig.model_validate(body)

get_stack_deployment_info(provider: StackDeploymentProvider) -> StackDeploymentInfo

Get information about a stack deployment provider.

Parameters:

Name	Type	Description	Default
provider	StackDeploymentProvider	The stack deployment provider.	required

Returns:

Type	Description
StackDeploymentInfo	Information about the stack deployment provider.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_stack_deployment_info(
    self,
    provider: StackDeploymentProvider,
) -> StackDeploymentInfo:
    """Get information about a stack deployment provider.

    Args:
        provider: The stack deployment provider.

    Returns:
        Information about the stack deployment provider.
    """
    body = self.get(
        f"{STACK_DEPLOYMENT}{INFO}",
        params={"provider": provider.value},
    )
    return StackDeploymentInfo.model_validate(body)

get_stack_deployment_stack(provider: StackDeploymentProvider, stack_name: str, location: Optional[str] = None, date_start: Optional[datetime] = None) -> Optional[DeployedStack]

Return a matching ZenML stack that was deployed and registered.

Parameters:

Name	Type	Description	Default
provider	StackDeploymentProvider	The stack deployment provider.	required
stack_name	str	The name of the stack.	required
location	Optional[str]	The location where the stack should be deployed.	None
date_start	Optional[datetime]	The date when the deployment started.	None

Returns:

Type	Description
Optional[DeployedStack]	The ZenML stack that was deployed and registered or None if the
Optional[DeployedStack]	stack was not found.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_stack_deployment_stack(
    self,
    provider: StackDeploymentProvider,
    stack_name: str,
    location: Optional[str] = None,
    date_start: Optional[datetime] = None,
) -> Optional[DeployedStack]:
    """Return a matching ZenML stack that was deployed and registered.

    Args:
        provider: The stack deployment provider.
        stack_name: The name of the stack.
        location: The location where the stack should be deployed.
        date_start: The date when the deployment started.

    Returns:
        The ZenML stack that was deployed and registered or None if the
        stack was not found.
    """
    params = {
        "provider": provider.value,
        "stack_name": stack_name,
    }
    if location:
        params["location"] = location
    if date_start:
        params["date_start"] = str(date_start)
    body = self.get(
        f"{STACK_DEPLOYMENT}{STACK}",
        params=params,
    )
    if body:
        return DeployedStack.model_validate(body)

    return None

get_store_info() -> ServerModel

Get information about the server.

Returns:

Type	Description
ServerModel	Information about the server.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_store_info(self) -> ServerModel:
    """Get information about the server.

    Returns:
        Information about the server.
    """
    body = self.get(INFO)
    self._server_info = ServerModel.model_validate(body)
    return self._server_info

get_tag(tag_id: UUID, hydrate: bool = True) -> TagResponse

Get an existing tag.

Parameters:

Name	Type	Description	Default
tag_id	UUID	id of the tag to be retrieved.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
TagResponse	The tag of interest.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_tag(
    self,
    tag_id: UUID,
    hydrate: bool = True,
) -> TagResponse:
    """Get an existing tag.

    Args:
        tag_id: id of the tag to be retrieved.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The tag of interest.
    """
    return self._get_resource(
        resource_id=tag_id,
        route=TAGS,
        response_model=TagResponse,
        params={"hydrate": hydrate},
    )

get_trigger(trigger_id: UUID, hydrate: bool = True) -> TRIGGER_RETURN_TYPE_UNION

Retrieves a trigger.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	The ID of the trigger.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s)	True

Returns:

Type	Description
TRIGGER_RETURN_TYPE_UNION	The trigger.

Raises:

Type	Description
ValueError	In case of bad response.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_trigger(
    self, trigger_id: UUID, hydrate: bool = True
) -> TRIGGER_RETURN_TYPE_UNION:
    """Retrieves a trigger.

    Args:
        trigger_id: The ID of the trigger.
        hydrate: Flag deciding whether to hydrate the output model(s)

    Returns:
        The trigger.

    Raises:
        ValueError: In case of bad response.
    """
    body: dict[str, Any] = self.get(  # type: ignore[assignment]
        f"{TRIGGERS}/{str(trigger_id)}", params={"hydrate": hydrate}
    )

    try:
        response_model = TYPE_TO_RESPONSE_MAPPING[body["body"]["type"]]
        return response_model.model_validate(body)
    except (KeyError, TypeError):
        raise ValueError("Bad response, expected a trigger type object.")

get_user(user_name_or_id: Optional[Union[str, UUID]] = None, include_private: bool = False, hydrate: bool = True) -> UserResponse

Gets a specific user, when no id is specified get the active user.

The include_private parameter is ignored here as it is handled implicitly by the /current-user endpoint that is queried when no user_name_or_id is set. Raises a KeyError in case a user with that id does not exist.

Parameters:

Name	Type	Description	Default
user_name_or_id	Optional[Union[str, UUID]]	The name or ID of the user to get.	None
include_private	bool	Whether to include private user information.	False
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	True

Returns:

Type	Description
UserResponse	The requested user, if it was found.

Source code in src/zenml/zen_stores/rest_zen_store.py

def get_user(
    self,
    user_name_or_id: Optional[Union[str, UUID]] = None,
    include_private: bool = False,
    hydrate: bool = True,
) -> UserResponse:
    """Gets a specific user, when no id is specified get the active user.

    The `include_private` parameter is ignored here as it is handled
    implicitly by the /current-user endpoint that is queried when no
    user_name_or_id is set. Raises a KeyError in case a user with that id
    does not exist.

    Args:
        user_name_or_id: The name or ID of the user to get.
        include_private: Whether to include private user information.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The requested user, if it was found.
    """
    if user_name_or_id:
        return self._get_resource(
            resource_id=user_name_or_id,
            route=USERS,
            response_model=UserResponse,
            params={"hydrate": hydrate},
        )
    else:
        body = self.get(CURRENT_USER, params={"hydrate": hydrate})
        return UserResponse.model_validate(body)

list_api_keys(service_account_id: UUID, filter_model: APIKeyFilter, hydrate: bool = False) -> Page[APIKeyResponse]

List all API keys for a service account matching the given filter criteria.

Parameters:

Name	Type	Description	Default
service_account_id	UUID	The ID of the service account for which to list the API keys.	required
filter_model	APIKeyFilter	All filter parameters including pagination params	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[APIKeyResponse]	A list of all API keys matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_api_keys(
    self,
    service_account_id: UUID,
    filter_model: APIKeyFilter,
    hydrate: bool = False,
) -> Page[APIKeyResponse]:
    """List all API keys for a service account matching the given filter criteria.

    Args:
        service_account_id: The ID of the service account for which to list
            the API keys.
        filter_model: All filter parameters including pagination
            params
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all API keys matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=f"{SERVICE_ACCOUNTS}/{str(service_account_id)}{API_KEYS}",
        response_model=APIKeyResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate},
    )

list_artifact_versions(artifact_version_filter_model: ArtifactVersionFilter, hydrate: bool = False) -> Page[ArtifactVersionResponse]

List all artifact versions matching the given filter criteria.

Parameters:

Name	Type	Description	Default
artifact_version_filter_model	ArtifactVersionFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ArtifactVersionResponse]	A list of all artifact versions matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_artifact_versions(
    self,
    artifact_version_filter_model: ArtifactVersionFilter,
    hydrate: bool = False,
) -> Page[ArtifactVersionResponse]:
    """List all artifact versions matching the given filter criteria.

    Args:
        artifact_version_filter_model: All filter parameters including
            pagination params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all artifact versions matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=ARTIFACT_VERSIONS,
        response_model=ArtifactVersionResponse,
        filter_model=artifact_version_filter_model,
        params={"hydrate": hydrate},
    )

list_artifacts(filter_model: ArtifactFilter, hydrate: bool = False) -> Page[ArtifactResponse]

List all artifacts matching the given filter criteria.

Parameters:

Name	Type	Description	Default
filter_model	ArtifactFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ArtifactResponse]	A list of all artifacts matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_artifacts(
    self, filter_model: ArtifactFilter, hydrate: bool = False
) -> Page[ArtifactResponse]:
    """List all artifacts matching the given filter criteria.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all artifacts matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=ARTIFACTS,
        response_model=ArtifactResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate},
    )

list_authorized_devices(filter_model: OAuthDeviceFilter, hydrate: bool = False) -> Page[OAuthDeviceResponse]

List all OAuth 2.0 authorized devices for a user.

Parameters:

Name	Type	Description	Default
filter_model	OAuthDeviceFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[OAuthDeviceResponse]	A page of all matching OAuth 2.0 authorized devices.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_authorized_devices(
    self, filter_model: OAuthDeviceFilter, hydrate: bool = False
) -> Page[OAuthDeviceResponse]:
    """List all OAuth 2.0 authorized devices for a user.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all matching OAuth 2.0 authorized devices.
    """
    return self._list_paginated_resources(
        route=DEVICES,
        response_model=OAuthDeviceResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate},
    )

list_builds(build_filter_model: PipelineBuildFilter, hydrate: bool = False) -> Page[PipelineBuildResponse]

List all builds matching the given filter criteria.

Parameters:

Name	Type	Description	Default
build_filter_model	PipelineBuildFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[PipelineBuildResponse]	A page of all builds matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_builds(
    self,
    build_filter_model: PipelineBuildFilter,
    hydrate: bool = False,
) -> Page[PipelineBuildResponse]:
    """List all builds matching the given filter criteria.

    Args:
        build_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all builds matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=PIPELINE_BUILDS,
        response_model=PipelineBuildResponse,
        filter_model=build_filter_model,
        params={"hydrate": hydrate},
    )

list_code_repositories(filter_model: CodeRepositoryFilter, hydrate: bool = False) -> Page[CodeRepositoryResponse]

List all code repositories.

Parameters:

Name	Type	Description	Default
filter_model	CodeRepositoryFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[CodeRepositoryResponse]	A page of all code repositories.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_code_repositories(
    self,
    filter_model: CodeRepositoryFilter,
    hydrate: bool = False,
) -> Page[CodeRepositoryResponse]:
    """List all code repositories.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all code repositories.
    """
    return self._list_paginated_resources(
        route=CODE_REPOSITORIES,
        response_model=CodeRepositoryResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate},
    )

list_deployments(deployment_filter_model: DeploymentFilter, hydrate: bool = False) -> Page[DeploymentResponse]

List all deployments matching the given filter criteria.

Parameters:

Name	Type	Description	Default
deployment_filter_model	DeploymentFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[DeploymentResponse]	A page of all deployments matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_deployments(
    self,
    deployment_filter_model: DeploymentFilter,
    hydrate: bool = False,
) -> Page[DeploymentResponse]:
    """List all deployments matching the given filter criteria.

    Args:
        deployment_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all deployments matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=DEPLOYMENTS,
        response_model=DeploymentResponse,
        filter_model=deployment_filter_model,
        params={"hydrate": hydrate},
    )

list_flavors(flavor_filter_model: FlavorFilter, hydrate: bool = False) -> Page[FlavorResponse]

List all stack component flavors matching the given filter criteria.

Parameters:

Name	Type	Description	Default
flavor_filter_model	FlavorFilter	All filter parameters including pagination params	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[FlavorResponse]	List of all the stack component flavors matching the given criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_flavors(
    self,
    flavor_filter_model: FlavorFilter,
    hydrate: bool = False,
) -> Page[FlavorResponse]:
    """List all stack component flavors matching the given filter criteria.

    Args:
        flavor_filter_model: All filter parameters including pagination
            params
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        List of all the stack component flavors matching the given criteria.
    """
    return self._list_paginated_resources(
        route=FLAVORS,
        response_model=FlavorResponse,
        filter_model=flavor_filter_model,
        params={"hydrate": hydrate},
    )

list_model_version_artifact_links(model_version_artifact_link_filter_model: ModelVersionArtifactFilter, hydrate: bool = False) -> Page[ModelVersionArtifactResponse]

Get all model version to artifact links by filter.

Parameters:

Name	Type	Description	Default
model_version_artifact_link_filter_model	ModelVersionArtifactFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ModelVersionArtifactResponse]	A page of all model version to artifact links.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_model_version_artifact_links(
    self,
    model_version_artifact_link_filter_model: ModelVersionArtifactFilter,
    hydrate: bool = False,
) -> Page[ModelVersionArtifactResponse]:
    """Get all model version to artifact links by filter.

    Args:
        model_version_artifact_link_filter_model: All filter parameters
            including pagination params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all model version to artifact links.
    """
    return self._list_paginated_resources(
        route=MODEL_VERSION_ARTIFACTS,
        response_model=ModelVersionArtifactResponse,
        filter_model=model_version_artifact_link_filter_model,
        params={"hydrate": hydrate},
    )

list_model_version_pipeline_run_links(model_version_pipeline_run_link_filter_model: ModelVersionPipelineRunFilter, hydrate: bool = False) -> Page[ModelVersionPipelineRunResponse]

Get all model version to pipeline run links by filter.

Parameters:

Name	Type	Description	Default
model_version_pipeline_run_link_filter_model	ModelVersionPipelineRunFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ModelVersionPipelineRunResponse]	A page of all model version to pipeline run links.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_model_version_pipeline_run_links(
    self,
    model_version_pipeline_run_link_filter_model: ModelVersionPipelineRunFilter,
    hydrate: bool = False,
) -> Page[ModelVersionPipelineRunResponse]:
    """Get all model version to pipeline run links by filter.

    Args:
        model_version_pipeline_run_link_filter_model: All filter parameters
            including pagination params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all model version to pipeline run links.
    """
    return self._list_paginated_resources(
        route=MODEL_VERSION_PIPELINE_RUNS,
        response_model=ModelVersionPipelineRunResponse,
        filter_model=model_version_pipeline_run_link_filter_model,
        params={"hydrate": hydrate},
    )

list_model_versions(model_version_filter_model: ModelVersionFilter, hydrate: bool = False) -> Page[ModelVersionResponse]

Get all model versions by filter.

Parameters:

Name	Type	Description	Default
model_version_filter_model	ModelVersionFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ModelVersionResponse]	A page of all model versions.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_model_versions(
    self,
    model_version_filter_model: ModelVersionFilter,
    hydrate: bool = False,
) -> Page[ModelVersionResponse]:
    """Get all model versions by filter.

    Args:
        model_version_filter_model: All filter parameters including
            pagination params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all model versions.
    """
    return self._list_paginated_resources(
        route=MODEL_VERSIONS,
        response_model=ModelVersionResponse,
        filter_model=model_version_filter_model,
        params={"hydrate": hydrate},
    )

list_models(model_filter_model: ModelFilter, hydrate: bool = False) -> Page[ModelResponse]

Get all models by filter.

Parameters:

Name	Type	Description	Default
model_filter_model	ModelFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ModelResponse]	A page of all models.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_models(
    self,
    model_filter_model: ModelFilter,
    hydrate: bool = False,
) -> Page[ModelResponse]:
    """Get all models by filter.

    Args:
        model_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all models.
    """
    return self._list_paginated_resources(
        route=MODELS,
        response_model=ModelResponse,
        filter_model=model_filter_model,
        params={"hydrate": hydrate},
    )

list_pipelines(pipeline_filter_model: PipelineFilter, hydrate: bool = False) -> Page[PipelineResponse]

List all pipelines matching the given filter criteria.

Parameters:

Name	Type	Description	Default
pipeline_filter_model	PipelineFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[PipelineResponse]	A list of all pipelines matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_pipelines(
    self,
    pipeline_filter_model: PipelineFilter,
    hydrate: bool = False,
) -> Page[PipelineResponse]:
    """List all pipelines matching the given filter criteria.

    Args:
        pipeline_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all pipelines matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=PIPELINES,
        response_model=PipelineResponse,
        filter_model=pipeline_filter_model,
        params={"hydrate": hydrate},
    )

list_projects(project_filter_model: ProjectFilter, hydrate: bool = False) -> Page[ProjectResponse]

List all projects matching the given filter criteria.

Parameters:

Name	Type	Description	Default
project_filter_model	ProjectFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ProjectResponse]	A list of all projects matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_projects(
    self,
    project_filter_model: ProjectFilter,
    hydrate: bool = False,
) -> Page[ProjectResponse]:
    """List all projects matching the given filter criteria.

    Args:
        project_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all projects matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=PROJECTS,
        response_model=ProjectResponse,
        filter_model=project_filter_model,
        params={"hydrate": hydrate},
    )

list_resource_pool_subject_policies(filter_model: ResourcePoolSubjectPolicyFilter, hydrate: bool = False) -> Page[ResourcePoolSubjectPolicyResponse]

List resource pool subject policies.

Parameters:

Name	Type	Description	Default
filter_model	ResourcePoolSubjectPolicyFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s).	False

Returns:

Type	Description
Page[ResourcePoolSubjectPolicyResponse]	Matching policies.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_resource_pool_subject_policies(
    self,
    filter_model: ResourcePoolSubjectPolicyFilter,
    hydrate: bool = False,
) -> Page[ResourcePoolSubjectPolicyResponse]:
    """List resource pool subject policies.

    Args:
        filter_model: All filter parameters including pagination params.
        hydrate: Flag deciding whether to hydrate the output model(s).

    Returns:
        Matching policies.
    """
    return self._list_paginated_resources(
        route=RESOURCE_POOL_SUBJECT_POLICIES,
        response_model=ResourcePoolSubjectPolicyResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate},
    )

list_resource_pools(filter_model: ResourcePoolFilter, hydrate: bool = False) -> Page[ResourcePoolResponse]

List all resource pools matching the given filter criteria.

Parameters:

Name	Type	Description	Default
filter_model	ResourcePoolFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ResourcePoolResponse]	A list of all resource pools matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_resource_pools(
    self, filter_model: ResourcePoolFilter, hydrate: bool = False
) -> Page[ResourcePoolResponse]:
    """List all resource pools matching the given filter criteria.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all resource pools matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=RESOURCE_POOLS,
        response_model=ResourcePoolResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate},
    )

list_resource_requests(filter_model: ResourceRequestFilter, hydrate: bool = False) -> Page[ResourceRequestResponse]

List all resource requests matching the given filter criteria.

Parameters:

Name	Type	Description	Default
filter_model	ResourceRequestFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ResourceRequestResponse]	A list of all resource requests matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_resource_requests(
    self, filter_model: ResourceRequestFilter, hydrate: bool = False
) -> Page[ResourceRequestResponse]:
    """List all resource requests matching the given filter criteria.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all resource requests matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=RESOURCE_REQUESTS,
        response_model=ResourceRequestResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate},
    )

list_run_steps(step_run_filter_model: StepRunFilter, hydrate: bool = False) -> Page[StepRunResponse]

List all step runs matching the given filter criteria.

Parameters:

Name	Type	Description	Default
step_run_filter_model	StepRunFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[StepRunResponse]	A list of all step runs matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_run_steps(
    self,
    step_run_filter_model: StepRunFilter,
    hydrate: bool = False,
) -> Page[StepRunResponse]:
    """List all step runs matching the given filter criteria.

    Args:
        step_run_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all step runs matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=STEPS,
        response_model=StepRunResponse,
        filter_model=step_run_filter_model,
        params={"hydrate": hydrate},
    )

list_run_templates(template_filter_model: RunTemplateFilter, hydrate: bool = False) -> Page[RunTemplateResponse]

List all run templates matching the given filter criteria.

Parameters:

Name	Type	Description	Default
template_filter_model	RunTemplateFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[RunTemplateResponse]	A list of all templates matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_run_templates(
    self,
    template_filter_model: RunTemplateFilter,
    hydrate: bool = False,
) -> Page[RunTemplateResponse]:
    """List all run templates matching the given filter criteria.

    Args:
        template_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all templates matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=RUN_TEMPLATES,
        response_model=RunTemplateResponse,
        filter_model=template_filter_model,
        params={"hydrate": hydrate},
    )

list_run_wait_conditions(run_wait_condition_filter_model: RunWaitConditionFilter, hydrate: bool = False) -> Page[RunWaitConditionResponse]

List run wait conditions.

Parameters:

Name	Type	Description	Default
run_wait_condition_filter_model	RunWaitConditionFilter	Wait condition filter model.	required
hydrate	bool	Whether to hydrate metadata/resources.	False

Returns:

Type	Description
Page[RunWaitConditionResponse]	A page of wait conditions.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_run_wait_conditions(
    self,
    run_wait_condition_filter_model: RunWaitConditionFilter,
    hydrate: bool = False,
) -> Page[RunWaitConditionResponse]:
    """List run wait conditions.

    Args:
        run_wait_condition_filter_model: Wait condition filter model.
        hydrate: Whether to hydrate metadata/resources.

    Returns:
        A page of wait conditions.
    """
    return self._list_paginated_resources(
        route=RUN_WAIT_CONDITIONS,
        response_model=RunWaitConditionResponse,
        filter_model=run_wait_condition_filter_model,
        params={"hydrate": hydrate},
    )

list_runs(runs_filter_model: PipelineRunFilter, hydrate: bool = False, include_full_metadata: bool = False) -> Page[PipelineRunResponse]

List all pipeline runs matching the given filter criteria.

Parameters:

Name	Type	Description	Default
runs_filter_model	PipelineRunFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False
include_full_metadata	bool	If True, include metadata of all steps in the response.	False

Returns:

Type	Description
Page[PipelineRunResponse]	A list of all pipeline runs matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_runs(
    self,
    runs_filter_model: PipelineRunFilter,
    hydrate: bool = False,
    include_full_metadata: bool = False,
) -> Page[PipelineRunResponse]:
    """List all pipeline runs matching the given filter criteria.

    Args:
        runs_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        include_full_metadata: If True, include metadata of all steps in
            the response.

    Returns:
        A list of all pipeline runs matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=RUNS,
        response_model=PipelineRunResponse,
        filter_model=runs_filter_model,
        params={
            "hydrate": hydrate,
            "include_full_metadata": include_full_metadata,
        },
    )

list_schedules(schedule_filter_model: ScheduleFilter, hydrate: bool = False) -> Page[ScheduleResponse]

List all schedules.

Parameters:

Name	Type	Description	Default
schedule_filter_model	ScheduleFilter	All filter parameters including pagination params	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ScheduleResponse]	A list of schedules.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_schedules(
    self,
    schedule_filter_model: ScheduleFilter,
    hydrate: bool = False,
) -> Page[ScheduleResponse]:
    """List all schedules.

    Args:
        schedule_filter_model: All filter parameters including pagination
            params
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of schedules.
    """
    return self._list_paginated_resources(
        route=SCHEDULES,
        response_model=ScheduleResponse,
        filter_model=schedule_filter_model,
        params={"hydrate": hydrate},
    )

list_secrets(secret_filter_model: SecretFilter, hydrate: bool = False) -> Page[SecretResponse]

List all secrets matching the given filter criteria.

Note that returned secrets do not include any secret values. To fetch the secret values, use get_secret.

Parameters:

Name	Type	Description	Default
secret_filter_model	SecretFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[SecretResponse]	A list of all secrets matching the filter criteria, with pagination
Page[SecretResponse]	information and sorted according to the filter criteria. The
Page[SecretResponse]	returned secrets do not include any secret values, only metadata. To
Page[SecretResponse]	fetch the secret values, use get_secret individually with each
Page[SecretResponse]	secret.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_secrets(
    self, secret_filter_model: SecretFilter, hydrate: bool = False
) -> Page[SecretResponse]:
    """List all secrets matching the given filter criteria.

    Note that returned secrets do not include any secret values. To fetch
    the secret values, use `get_secret`.

    Args:
        secret_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all secrets matching the filter criteria, with pagination
        information and sorted according to the filter criteria. The
        returned secrets do not include any secret values, only metadata. To
        fetch the secret values, use `get_secret` individually with each
        secret.
    """
    return self._list_paginated_resources(
        route=SECRETS,
        response_model=SecretResponse,
        filter_model=secret_filter_model,
        params={"hydrate": hydrate},
    )

list_service_accounts(filter_model: ServiceAccountFilter, hydrate: bool = False) -> Page[ServiceAccountResponse]

List all service accounts.

Parameters:

Name	Type	Description	Default
filter_model	ServiceAccountFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ServiceAccountResponse]	A list of filtered service accounts.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_service_accounts(
    self, filter_model: ServiceAccountFilter, hydrate: bool = False
) -> Page[ServiceAccountResponse]:
    """List all service accounts.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of filtered service accounts.
    """
    return self._list_paginated_resources(
        route=SERVICE_ACCOUNTS,
        response_model=ServiceAccountResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate},
    )

list_service_connector_resources(filter_model: ServiceConnectorFilter) -> List[ServiceConnectorResourcesModel]

List resources that can be accessed by service connectors.

Parameters:

Name	Type	Description	Default
filter_model	ServiceConnectorFilter	The filter model to use when fetching service connectors.	required

Returns:

Type	Description
List[ServiceConnectorResourcesModel]	The matching list of resources that available service
List[ServiceConnectorResourcesModel]	connectors have access to.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_service_connector_resources(
    self,
    filter_model: ServiceConnectorFilter,
) -> List[ServiceConnectorResourcesModel]:
    """List resources that can be accessed by service connectors.

    Args:
        filter_model: The filter model to use when fetching service
            connectors.

    Returns:
        The matching list of resources that available service
        connectors have access to.
    """
    response_body = self.get(
        SERVICE_CONNECTORS + SERVICE_CONNECTOR_RESOURCES,
        params=filter_model.model_dump(exclude_none=True),
        timeout=max(
            self.config.http_timeout,
            SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
        ),
    )

    assert isinstance(response_body, list)
    resource_list = [
        ServiceConnectorResourcesModel.model_validate(item)
        for item in response_body
    ]

    self._populate_connector_type(*resource_list)

    # For service connectors with types that are only locally available,
    # we need to retrieve the resource list locally
    for idx, resources in enumerate(resource_list):
        if isinstance(resources.connector_type, str):
            # Skip connector types that are neither locally nor remotely
            # available
            continue
        if resources.connector_type.remote:
            # Skip connector types that are remotely available
            continue

        # Retrieve the resource list locally
        assert resources.id is not None
        connector = self.get_service_connector(
            resources.id, expand_secrets=True
        )
        connector_instance = (
            service_connector_registry.instantiate_connector(
                model=connector
            )
        )

        try:
            local_resources = connector_instance.verify(
                resource_type=filter_model.resource_type,
                resource_id=filter_model.resource_id,
            )
        except (ValueError, AuthorizationException) as e:
            logger.error(
                f"Failed to fetch {filter_model.resource_type or 'available'} "
                f"resources from service connector {connector.name}/"
                f"{connector.id}: {e}"
            )
            continue

        resource_list[idx] = local_resources

    return resource_list

list_service_connector_types(connector_type: Optional[str] = None, resource_type: Optional[str] = None, auth_method: Optional[str] = None) -> List[ServiceConnectorTypeModel]

Get a list of service connector types.

Parameters:

Name	Type	Description	Default
connector_type	Optional[str]	Filter by connector type.	None
resource_type	Optional[str]	Filter by resource type.	None
auth_method	Optional[str]	Filter by authentication method.	None

Returns:

Type	Description
List[ServiceConnectorTypeModel]	List of service connector types.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_service_connector_types(
    self,
    connector_type: Optional[str] = None,
    resource_type: Optional[str] = None,
    auth_method: Optional[str] = None,
) -> List[ServiceConnectorTypeModel]:
    """Get a list of service connector types.

    Args:
        connector_type: Filter by connector type.
        resource_type: Filter by resource type.
        auth_method: Filter by authentication method.

    Returns:
        List of service connector types.
    """
    params = {}
    if connector_type:
        params["connector_type"] = connector_type
    if resource_type:
        params["resource_type"] = resource_type
    if auth_method:
        params["auth_method"] = auth_method
    response_body = self.get(
        SERVICE_CONNECTOR_TYPES,
        params=params,
    )

    assert isinstance(response_body, list)
    remote_connector_types = [
        ServiceConnectorTypeModel.model_validate(item)
        for item in response_body
    ]

    # Mark the remote connector types as being only remotely available
    for c in remote_connector_types:
        c.local = False
        c.remote = True

    local_connector_types = (
        service_connector_registry.list_service_connector_types(
            connector_type=connector_type,
            resource_type=resource_type,
            auth_method=auth_method,
        )
    )

    # Add the connector types in the local registry to the list of
    # connector types available remotely. Overwrite those that have
    # the same connector type but mark them as being remotely available.
    connector_types_map = {
        connector_type.connector_type: connector_type
        for connector_type in remote_connector_types
    }

    for connector in local_connector_types:
        if connector.connector_type in connector_types_map:
            connector.remote = True
        connector_types_map[connector.connector_type] = connector

    return list(connector_types_map.values())

list_service_connectors(filter_model: ServiceConnectorFilter, hydrate: bool = False, expand_secrets: bool = False) -> Page[ServiceConnectorResponse]

List all service connectors.

Parameters:

Name	Type	Description	Default
filter_model	ServiceConnectorFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False
expand_secrets	bool	Flag deciding whether to include the secrets associated with the service connector.	False

Returns:

Type	Description
Page[ServiceConnectorResponse]	A page of all service connectors.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_service_connectors(
    self,
    filter_model: ServiceConnectorFilter,
    hydrate: bool = False,
    expand_secrets: bool = False,
) -> Page[ServiceConnectorResponse]:
    """List all service connectors.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        expand_secrets: Flag deciding whether to include the secrets
            associated with the service connector.

    Returns:
        A page of all service connectors.
    """
    connector_models = self._list_paginated_resources(
        route=SERVICE_CONNECTORS,
        response_model=ServiceConnectorResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate, "expand_secrets": expand_secrets},
    )
    self._populate_connector_type(*connector_models.items)
    if expand_secrets:
        # Call this to properly split the secrets from the configuration
        for connector_model in connector_models.items:
            try:
                connector_model.validate_configuration()
            except ValueError as e:
                logger.error(
                    f"Error validating connector configuration for "
                    f"{connector_model.name}: {e}"
                )
    return connector_models

list_services(filter_model: ServiceFilter, hydrate: bool = False) -> Page[ServiceResponse]

List all services matching the given filter criteria.

Parameters:

Name	Type	Description	Default
filter_model	ServiceFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ServiceResponse]	A list of all services matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_services(
    self, filter_model: ServiceFilter, hydrate: bool = False
) -> Page[ServiceResponse]:
    """List all services matching the given filter criteria.

    Args:
        filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all services matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=SERVICES,
        response_model=ServiceResponse,
        filter_model=filter_model,
        params={"hydrate": hydrate},
    )

list_snapshots(snapshot_filter_model: PipelineSnapshotFilter, hydrate: bool = False) -> Page[PipelineSnapshotResponse]

List all snapshots matching the given filter criteria.

Parameters:

Name	Type	Description	Default
snapshot_filter_model	PipelineSnapshotFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[PipelineSnapshotResponse]	A page of all snapshots matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_snapshots(
    self,
    snapshot_filter_model: PipelineSnapshotFilter,
    hydrate: bool = False,
) -> Page[PipelineSnapshotResponse]:
    """List all snapshots matching the given filter criteria.

    Args:
        snapshot_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all snapshots matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=PIPELINE_SNAPSHOTS,
        response_model=PipelineSnapshotResponse,
        filter_model=snapshot_filter_model,
        params={"hydrate": hydrate},
    )

list_stack_components(component_filter_model: ComponentFilter, hydrate: bool = False) -> Page[ComponentResponse]

List all stack components matching the given filter criteria.

Parameters:

Name	Type	Description	Default
component_filter_model	ComponentFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[ComponentResponse]	A list of all stack components matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_stack_components(
    self,
    component_filter_model: ComponentFilter,
    hydrate: bool = False,
) -> Page[ComponentResponse]:
    """List all stack components matching the given filter criteria.

    Args:
        component_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all stack components matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=STACK_COMPONENTS,
        response_model=ComponentResponse,
        filter_model=component_filter_model,
        params={"hydrate": hydrate},
    )

list_stacks(stack_filter_model: StackFilter, hydrate: bool = False) -> Page[StackResponse]

List all stacks matching the given filter criteria.

Parameters:

Name	Type	Description	Default
stack_filter_model	StackFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[StackResponse]	A list of all stacks matching the filter criteria.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_stacks(
    self, stack_filter_model: StackFilter, hydrate: bool = False
) -> Page[StackResponse]:
    """List all stacks matching the given filter criteria.

    Args:
        stack_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all stacks matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=STACKS,
        response_model=StackResponse,
        filter_model=stack_filter_model,
        params={"hydrate": hydrate},
    )

list_tags(tag_filter_model: TagFilter, hydrate: bool = False) -> Page[TagResponse]

Get all tags by filter.

Parameters:

Name	Type	Description	Default
tag_filter_model	TagFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[TagResponse]	A page of all tags.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_tags(
    self,
    tag_filter_model: TagFilter,
    hydrate: bool = False,
) -> Page[TagResponse]:
    """Get all tags by filter.

    Args:
        tag_filter_model: All filter parameters including pagination params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A page of all tags.
    """
    return self._list_paginated_resources(
        route=TAGS,
        response_model=TagResponse,
        filter_model=tag_filter_model,
        params={"hydrate": hydrate},
    )

list_triggers(triggers_filter_model: TriggerFilter, hydrate: bool = False) -> Page[TRIGGER_RETURN_TYPE_UNION]

List all triggers.

Parameters:

Name	Type	Description	Default
triggers_filter_model	TriggerFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[TRIGGER_RETURN_TYPE_UNION]	A list of triggers matching the filter criteria.

Raises:

Type	Description
ValueError	In case of bad response.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_triggers(
    self, triggers_filter_model: TriggerFilter, hydrate: bool = False
) -> Page[TRIGGER_RETURN_TYPE_UNION]:
    """List all triggers.

    Args:
        triggers_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of triggers matching the filter criteria.

    Raises:
        ValueError: In case of bad response.
    """
    body: dict[str, Any] = self.get(  # type: ignore[assignment]
        TRIGGERS,
        params={
            "hydrate": hydrate,
            **triggers_filter_model.model_dump(exclude_none=True),
        },
    )

    page_of_items: Page[AnyResponse] = Page.model_validate(body)  # type: ignore[valid-type]

    if not page_of_items.items:
        return page_of_items

    try:
        page_of_items.items = [
            TYPE_TO_RESPONSE_MAPPING[
                generic_item["body"]["type"]
            ].model_validate(generic_item)
            for generic_item in body["items"]
        ]
    except (KeyError, TypeError):
        raise ValueError("Bad response, expected a trigger type object.")

    return page_of_items

list_users(user_filter_model: UserFilter, hydrate: bool = False) -> Page[UserResponse]

List all users.

Parameters:

Name	Type	Description	Default
user_filter_model	UserFilter	All filter parameters including pagination params.	required
hydrate	bool	Flag deciding whether to hydrate the output model(s) by including metadata fields in the response.	False

Returns:

Type	Description
Page[UserResponse]	A list of all users.

Source code in src/zenml/zen_stores/rest_zen_store.py

def list_users(
    self,
    user_filter_model: UserFilter,
    hydrate: bool = False,
) -> Page[UserResponse]:
    """List all users.

    Args:
        user_filter_model: All filter parameters including pagination
            params.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of all users.
    """
    return self._list_paginated_resources(
        route=USERS,
        response_model=UserResponse,
        filter_model=user_filter_model,
        params={"hydrate": hydrate},
    )

post(path: str, body: BaseModel, params: Optional[Dict[str, Any]] = None, timeout: Optional[int] = None, **kwargs: Any) -> Json

Make a POST request to the given endpoint path.

Parameters:

Name	Type	Description	Default
path	str	The path to the endpoint.	required
body	BaseModel	The body to send.	required
params	Optional[Dict[str, Any]]	The query parameters to pass to the endpoint.	None
timeout	Optional[int]	The request timeout in seconds.	None
kwargs	Any	Additional keyword arguments to pass to the request.	{}

Returns:

Type	Description
Json	The response body.

Source code in src/zenml/zen_stores/rest_zen_store.py

def post(
    self,
    path: str,
    body: BaseModel,
    params: Optional[Dict[str, Any]] = None,
    timeout: Optional[int] = None,
    **kwargs: Any,
) -> Json:
    """Make a POST request to the given endpoint path.

    Args:
        path: The path to the endpoint.
        body: The body to send.
        params: The query parameters to pass to the endpoint.
        timeout: The request timeout in seconds.
        kwargs: Additional keyword arguments to pass to the request.

    Returns:
        The response body.
    """
    return self._request(
        "POST",
        self.url + API + VERSION_1 + path,
        json=body.model_dump(mode="json"),
        params=params,
        timeout=timeout,
        **kwargs,
    )

prune_artifact_versions(project_name_or_id: Union[str, UUID], only_versions: bool = True) -> None

Prunes unused artifact versions and their artifacts.

Parameters:

Name	Type	Description	Default
project_name_or_id	Union[str, UUID]	The project name or ID to prune artifact versions for.	required
only_versions	bool	Only delete artifact versions, keeping artifacts	True

Source code in src/zenml/zen_stores/rest_zen_store.py

def prune_artifact_versions(
    self,
    project_name_or_id: Union[str, UUID],
    only_versions: bool = True,
) -> None:
    """Prunes unused artifact versions and their artifacts.

    Args:
        project_name_or_id: The project name or ID to prune artifact
            versions for.
        only_versions: Only delete artifact versions, keeping artifacts
    """
    self.delete(
        path=ARTIFACT_VERSIONS,
        params={
            "only_versions": only_versions,
            "project_name_or_id": project_name_or_id,
        },
    )

publish_run_events(pipeline_run_id: UUID, batch: StreamBatchRequest) -> StreamBatchResponse

Publish a batch of live events to a pipeline run's stream.

Parameters:

Name	Type	Description	Default
pipeline_run_id	UUID	The ID of the run the events belong to.	required
batch	StreamBatchRequest	The batch of events to publish.	required

Returns:

Type	Description
StreamBatchResponse	The server-side ingest response.

Source code in src/zenml/zen_stores/rest_zen_store.py

def publish_run_events(
    self, pipeline_run_id: UUID, batch: StreamBatchRequest
) -> StreamBatchResponse:
    """Publish a batch of live events to a pipeline run's stream.

    Args:
        pipeline_run_id: The ID of the run the events belong to.
        batch: The batch of events to publish.

    Returns:
        The server-side ingest response.
    """
    response_body = self.post(
        f"{RUNS}/{pipeline_run_id}{EVENTS}", body=batch
    )
    return StreamBatchResponse.model_validate(response_body)

put(path: str, body: Optional[BaseModel] = None, params: Optional[Dict[str, Any]] = None, timeout: Optional[int] = None, exclude_unset: bool = True, **kwargs: Any) -> Json

Make a PUT request to the given endpoint path.

Parameters:

Name	Type	Description	Default
path	str	The path to the endpoint.	required
body	Optional[BaseModel]	The body to send.	None
params	Optional[Dict[str, Any]]	The query parameters to pass to the endpoint.	None
timeout	Optional[int]	The request timeout in seconds.	None
exclude_unset	bool	Exclude unset fields, defaults to True.	True
kwargs	Any	Additional keyword arguments to pass to the request.	{}

Returns:

Type	Description
Json	The response body.

Source code in src/zenml/zen_stores/rest_zen_store.py

def put(
    self,
    path: str,
    body: Optional[BaseModel] = None,
    params: Optional[Dict[str, Any]] = None,
    timeout: Optional[int] = None,
    exclude_unset: bool = True,
    **kwargs: Any,
) -> Json:
    """Make a PUT request to the given endpoint path.

    Args:
        path: The path to the endpoint.
        body: The body to send.
        params: The query parameters to pass to the endpoint.
        timeout: The request timeout in seconds.
        exclude_unset: Exclude unset fields, defaults to True.
        kwargs: Additional keyword arguments to pass to the request.

    Returns:
        The response body.
    """
    json = (
        body.model_dump(mode="json", exclude_unset=exclude_unset)
        if body
        else None
    )
    return self._request(
        "PUT",
        self.url + API + VERSION_1 + path,
        json=json,
        params=params,
        timeout=timeout,
        **kwargs,
    )

reinitialize_session() -> None

Reinitialize the session.

This is used to reset the session to a new one with a new connection pool.

Source code in src/zenml/zen_stores/rest_zen_store.py

def reinitialize_session(self) -> None:
    """Reinitialize the session.

    This is used to reset the session to a new one with a new connection pool.
    """
    with self._session_lock:
        if self._session is not None:
            headers = dict(self._session.headers.items())
            self._session.close()
            self._session = None
            self.session.headers.update(headers)

replay_run(run_id: UUID, run_configuration: ReplayRunConfiguration) -> PipelineRunResponse

Replay a pipeline run.

Parameters:

Name	Type	Description	Default
run_id	UUID	The ID of the pipeline run to replay.	required
run_configuration	ReplayRunConfiguration	Replay configuration.	required

Raises:

Type	Description
RuntimeError	If the server does not support replaying a run.

Returns:

Type	Description
PipelineRunResponse	The replayed pipeline run.

Source code in src/zenml/zen_stores/rest_zen_store.py

def replay_run(
    self, run_id: UUID, run_configuration: ReplayRunConfiguration
) -> PipelineRunResponse:
    """Replay a pipeline run.

    Args:
        run_id: The ID of the pipeline run to replay.
        run_configuration: Replay configuration.

    Raises:
        RuntimeError: If the server does not support replaying a run.

    Returns:
        The replayed pipeline run.
    """
    try:
        response_body = self.post(
            f"{RUNS}/{run_id}{REPLAY}", body=run_configuration
        )
    except MethodNotAllowedError as e:
        raise RuntimeError(
            "Replaying a run is not supported for this server."
        ) from e

    return PipelineRunResponse.model_validate(response_body)

resolve_run_wait_condition(run_wait_condition_id: UUID, resolve_request: RunWaitConditionResolveRequest) -> RunWaitConditionResponse

Resolve a run wait condition.

Parameters:

Name	Type	Description	Default
run_wait_condition_id	UUID	Wait condition ID.	required
resolve_request	RunWaitConditionResolveRequest	Resolution payload.	required

Returns:

Type	Description
RunWaitConditionResponse	The resolved wait condition.

Source code in src/zenml/zen_stores/rest_zen_store.py

def resolve_run_wait_condition(
    self,
    run_wait_condition_id: UUID,
    resolve_request: RunWaitConditionResolveRequest,
) -> RunWaitConditionResponse:
    """Resolve a run wait condition.

    Args:
        run_wait_condition_id: Wait condition ID.
        resolve_request: Resolution payload.

    Returns:
        The resolved wait condition.
    """
    response_body = self.put(
        path=f"{RUN_WAIT_CONDITIONS}/{run_wait_condition_id}{RESOLVE}",
        body=resolve_request,
    )
    return RunWaitConditionResponse.model_validate(response_body)

restore_secrets(ignore_errors: bool = False, delete_secrets: bool = False) -> None

Restore all secrets from the configured backup secrets store.

Parameters:

Name	Type	Description	Default
ignore_errors	bool	Whether to ignore individual errors during the restore process and attempt to restore all secrets.	False
delete_secrets	bool	Whether to delete the secrets that have been successfully restored from the backup secrets store. Setting this flag effectively moves all secrets from the backup secrets store to the primary secrets store.	False

Source code in src/zenml/zen_stores/rest_zen_store.py

def restore_secrets(
    self, ignore_errors: bool = False, delete_secrets: bool = False
) -> None:
    """Restore all secrets from the configured backup secrets store.

    Args:
        ignore_errors: Whether to ignore individual errors during the
            restore process and attempt to restore all secrets.
        delete_secrets: Whether to delete the secrets that have been
            successfully restored from the backup secrets store. Setting
            this flag effectively moves all secrets from the backup secrets
            store to the primary secrets store.
    """
    params: Dict[str, Any] = {
        "ignore_errors": ignore_errors,
        "delete_secrets": delete_secrets,
    }
    self.put(
        f"{SECRETS_OPERATIONS}{SECRETS_RESTORE}",
        params=params,
    )

rotate_api_key(service_account_id: UUID, api_key_name_or_id: Union[str, UUID], rotate_request: APIKeyRotateRequest) -> APIKeyResponse

Rotate an API key for a service account.

Parameters:

Name	Type	Description	Default
service_account_id	UUID	The ID of the service account for which to rotate the API key.	required
api_key_name_or_id	Union[str, UUID]	The name or ID of the API key to rotate.	required
rotate_request	APIKeyRotateRequest	The rotate request on the API key.	required

Returns:

Type	Description
APIKeyResponse	The updated API key.

Source code in src/zenml/zen_stores/rest_zen_store.py

def rotate_api_key(
    self,
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    rotate_request: APIKeyRotateRequest,
) -> APIKeyResponse:
    """Rotate an API key for a service account.

    Args:
        service_account_id: The ID of the service account for which to
            rotate the API key.
        api_key_name_or_id: The name or ID of the API key to rotate.
        rotate_request: The rotate request on the API key.

    Returns:
        The updated API key.
    """
    response_body = self.put(
        f"{SERVICE_ACCOUNTS}/{str(service_account_id)}{API_KEYS}/{str(api_key_name_or_id)}{API_KEY_ROTATE}",
        body=rotate_request,
    )
    return APIKeyResponse.model_validate(response_body)

run_snapshot(snapshot_id: UUID, run_request: PipelineSnapshotRunRequest) -> PipelineRunResponse

Run a snapshot.

Parameters:

Name	Type	Description	Default
snapshot_id	UUID	The ID of the snapshot to run.	required
run_request	PipelineSnapshotRunRequest	Configuration for the run.	required

Raises:

Type	Description
RuntimeError	If the server does not support running a snapshot.

Returns:

Type	Description
PipelineRunResponse	The created pipeline run.

Source code in src/zenml/zen_stores/rest_zen_store.py

def run_snapshot(
    self,
    snapshot_id: UUID,
    run_request: PipelineSnapshotRunRequest,
) -> PipelineRunResponse:
    """Run a snapshot.

    Args:
        snapshot_id: The ID of the snapshot to run.
        run_request: Configuration for the run.

    Raises:
        RuntimeError: If the server does not support running a snapshot.

    Returns:
        The created pipeline run.
    """
    try:
        response_body = self.post(
            f"{PIPELINE_SNAPSHOTS}/{snapshot_id}/runs",
            body=run_request,
        )
    except MethodNotAllowedError as e:
        raise RuntimeError(
            "Running a snapshot is not supported for this server."
        ) from e

    return PipelineRunResponse.model_validate(response_body)

run_template(template_id: UUID, run_configuration: Optional[PipelineRunConfiguration] = None) -> PipelineRunResponse

Run a template.

Parameters:

Name	Type	Description	Default
template_id	UUID	The ID of the template to run.	required
run_configuration	Optional[PipelineRunConfiguration]	Configuration for the run.	None

Raises:

Type	Description
RuntimeError	If the server does not support running a template.

Returns:

Type	Description
PipelineRunResponse	Model of the pipeline run.

Source code in src/zenml/zen_stores/rest_zen_store.py

def run_template(
    self,
    template_id: UUID,
    run_configuration: Optional[PipelineRunConfiguration] = None,
) -> PipelineRunResponse:
    """Run a template.

    Args:
        template_id: The ID of the template to run.
        run_configuration: Configuration for the run.

    Raises:
        RuntimeError: If the server does not support running a template.

    Returns:
        Model of the pipeline run.
    """
    run_configuration = run_configuration or PipelineRunConfiguration()

    try:
        response_body = self.post(
            f"{RUN_TEMPLATES}/{template_id}/runs",
            body=run_configuration,
        )
    except MethodNotAllowedError as e:
        raise RuntimeError(
            "Running a template is not supported for this server."
        ) from e

    return PipelineRunResponse.model_validate(response_body)

update_api_key(service_account_id: UUID, api_key_name_or_id: Union[str, UUID], api_key_update: APIKeyUpdate) -> APIKeyResponse

Update an API key for a service account.

Parameters:

Name	Type	Description	Default
service_account_id	UUID	The ID of the service account for which to update the API key.	required
api_key_name_or_id	Union[str, UUID]	The name or ID of the API key to update.	required
api_key_update	APIKeyUpdate	The update request on the API key.	required

Returns:

Type	Description
APIKeyResponse	The updated API key.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_api_key(
    self,
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    api_key_update: APIKeyUpdate,
) -> APIKeyResponse:
    """Update an API key for a service account.

    Args:
        service_account_id: The ID of the service account for which to
            update the API key.
        api_key_name_or_id: The name or ID of the API key to update.
        api_key_update: The update request on the API key.

    Returns:
        The updated API key.
    """
    return self._update_resource(
        resource_id=api_key_name_or_id,
        resource_update=api_key_update,
        route=f"{SERVICE_ACCOUNTS}/{str(service_account_id)}{API_KEYS}",
        response_model=APIKeyResponse,
    )

update_artifact(artifact_id: UUID, artifact_update: ArtifactUpdate) -> ArtifactResponse

Updates an artifact.

Parameters:

Name	Type	Description	Default
artifact_id	UUID	The ID of the artifact to update.	required
artifact_update	ArtifactUpdate	The update to be applied to the artifact.	required

Returns:

Type	Description
ArtifactResponse	The updated artifact.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_artifact(
    self, artifact_id: UUID, artifact_update: ArtifactUpdate
) -> ArtifactResponse:
    """Updates an artifact.

    Args:
        artifact_id: The ID of the artifact to update.
        artifact_update: The update to be applied to the artifact.

    Returns:
        The updated artifact.
    """
    return self._update_resource(
        resource_id=artifact_id,
        resource_update=artifact_update,
        response_model=ArtifactResponse,
        route=ARTIFACTS,
    )

update_artifact_version(artifact_version_id: UUID, artifact_version_update: ArtifactVersionUpdate) -> ArtifactVersionResponse

Updates an artifact version.

Parameters:

Name	Type	Description	Default
artifact_version_id	UUID	The ID of the artifact version to update.	required
artifact_version_update	ArtifactVersionUpdate	The update to be applied to the artifact version.	required

Returns:

Type	Description
ArtifactVersionResponse	The updated artifact version.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_artifact_version(
    self,
    artifact_version_id: UUID,
    artifact_version_update: ArtifactVersionUpdate,
) -> ArtifactVersionResponse:
    """Updates an artifact version.

    Args:
        artifact_version_id: The ID of the artifact version to update.
        artifact_version_update: The update to be applied to the artifact
            version.

    Returns:
        The updated artifact version.
    """
    return self._update_resource(
        resource_id=artifact_version_id,
        resource_update=artifact_version_update,
        response_model=ArtifactVersionResponse,
        route=ARTIFACT_VERSIONS,
    )

update_authorized_device(device_id: UUID, update: OAuthDeviceUpdate) -> OAuthDeviceResponse

Updates an existing OAuth 2.0 authorized device for internal use.

Parameters:

Name	Type	Description	Default
device_id	UUID	The ID of the device to update.	required
update	OAuthDeviceUpdate	The update to be applied to the device.	required

Returns:

Type	Description
OAuthDeviceResponse	The updated OAuth 2.0 authorized device.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_authorized_device(
    self, device_id: UUID, update: OAuthDeviceUpdate
) -> OAuthDeviceResponse:
    """Updates an existing OAuth 2.0 authorized device for internal use.

    Args:
        device_id: The ID of the device to update.
        update: The update to be applied to the device.

    Returns:
        The updated OAuth 2.0 authorized device.
    """
    return self._update_resource(
        resource_id=device_id,
        resource_update=update,
        response_model=OAuthDeviceResponse,
        route=DEVICES,
    )

update_code_repository(code_repository_id: UUID, update: CodeRepositoryUpdate) -> CodeRepositoryResponse

Updates an existing code repository.

Parameters:

Name	Type	Description	Default
code_repository_id	UUID	The ID of the code repository to update.	required
update	CodeRepositoryUpdate	The update to be applied to the code repository.	required

Returns:

Type	Description
CodeRepositoryResponse	The updated code repository.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_code_repository(
    self, code_repository_id: UUID, update: CodeRepositoryUpdate
) -> CodeRepositoryResponse:
    """Updates an existing code repository.

    Args:
        code_repository_id: The ID of the code repository to update.
        update: The update to be applied to the code repository.

    Returns:
        The updated code repository.
    """
    return self._update_resource(
        resource_id=code_repository_id,
        resource_update=update,
        response_model=CodeRepositoryResponse,
        route=CODE_REPOSITORIES,
    )

update_curated_visualization(visualization_id: UUID, visualization_update: CuratedVisualizationUpdate) -> CuratedVisualizationResponse

Update a curated visualization via REST API.

Parameters:

Name	Type	Description	Default
visualization_id	UUID	The ID of the curated visualization to update.	required
visualization_update	CuratedVisualizationUpdate	The update to apply to the curated visualization.	required

Returns:

Type	Description
CuratedVisualizationResponse	The updated curated visualization.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_curated_visualization(
    self,
    visualization_id: UUID,
    visualization_update: CuratedVisualizationUpdate,
) -> CuratedVisualizationResponse:
    """Update a curated visualization via REST API.

    Args:
        visualization_id: The ID of the curated visualization to update.
        visualization_update: The update to apply to the curated
            visualization.

    Returns:
        The updated curated visualization.
    """
    return self._update_resource(
        resource_id=visualization_id,
        resource_update=visualization_update,
        response_model=CuratedVisualizationResponse,
        route=CURATED_VISUALIZATIONS,
    )

update_deployment(deployment_id: UUID, deployment_update: DeploymentUpdate) -> DeploymentResponse

Update a deployment.

Parameters:

Name	Type	Description	Default
deployment_id	UUID	The ID of the deployment to update.	required
deployment_update	DeploymentUpdate	The update to apply.	required

Returns:

Type	Description
DeploymentResponse	The updated deployment.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_deployment(
    self, deployment_id: UUID, deployment_update: DeploymentUpdate
) -> DeploymentResponse:
    """Update a deployment.

    Args:
        deployment_id: The ID of the deployment to update.
        deployment_update: The update to apply.

    Returns:
        The updated deployment.
    """
    return self._update_resource(
        resource_id=deployment_id,
        resource_update=deployment_update,
        route=DEPLOYMENTS,
        response_model=DeploymentResponse,
    )

update_flavor(flavor_id: UUID, flavor_update: FlavorUpdate) -> FlavorResponse

Updates an existing user.

Parameters:

Name	Type	Description	Default
flavor_id	UUID	The id of the flavor to update.	required
flavor_update	FlavorUpdate	The update to be applied to the flavor.	required

Returns:

Type	Description
FlavorResponse	The updated flavor.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_flavor(
    self, flavor_id: UUID, flavor_update: FlavorUpdate
) -> FlavorResponse:
    """Updates an existing user.

    Args:
        flavor_id: The id of the flavor to update.
        flavor_update: The update to be applied to the flavor.

    Returns:
        The updated flavor.
    """
    return self._update_resource(
        resource_id=flavor_id,
        resource_update=flavor_update,
        route=FLAVORS,
        response_model=FlavorResponse,
    )

update_logs(logs_id: UUID, logs_update: LogsUpdate) -> LogsResponse

Update an existing logs entry.

Parameters:

Name	Type	Description	Default
logs_id	UUID	The ID of the logs entry to update.	required
logs_update	LogsUpdate	The update to be applied to the logs entry.	required

Returns:

Type	Description
LogsResponse	The updated logs entry.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_logs(
    self, logs_id: UUID, logs_update: LogsUpdate
) -> LogsResponse:
    """Update an existing logs entry.

    Args:
        logs_id: The ID of the logs entry to update.
        logs_update: The update to be applied to the logs entry.

    Returns:
        The updated logs entry.
    """
    return self._update_resource(
        resource_id=logs_id,
        resource_update=logs_update,
        route=LOGS,
        response_model=LogsResponse,
    )

update_model(model_id: UUID, model_update: ModelUpdate) -> ModelResponse

Updates an existing model.

Parameters:

Name	Type	Description	Default
model_id	UUID	UUID of the model to be updated.	required
model_update	ModelUpdate	the Model to be updated.	required

Returns:

Type	Description
ModelResponse	The updated model.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_model(
    self,
    model_id: UUID,
    model_update: ModelUpdate,
) -> ModelResponse:
    """Updates an existing model.

    Args:
        model_id: UUID of the model to be updated.
        model_update: the Model to be updated.

    Returns:
        The updated model.
    """
    return self._update_resource(
        resource_id=model_id,
        resource_update=model_update,
        route=MODELS,
        response_model=ModelResponse,
    )

update_model_version(model_version_id: UUID, model_version_update_model: ModelVersionUpdate) -> ModelVersionResponse

Get all model versions by filter.

Parameters:

Name	Type	Description	Default
model_version_id	UUID	The ID of model version to be updated.	required
model_version_update_model	ModelVersionUpdate	The model version to be updated.	required

Returns:

Type	Description
ModelVersionResponse	An updated model version.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_model_version(
    self,
    model_version_id: UUID,
    model_version_update_model: ModelVersionUpdate,
) -> ModelVersionResponse:
    """Get all model versions by filter.

    Args:
        model_version_id: The ID of model version to be updated.
        model_version_update_model: The model version to be updated.

    Returns:
        An updated model version.

    """
    return self._update_resource(
        resource_id=model_version_id,
        resource_update=model_version_update_model,
        route=MODEL_VERSIONS,
        response_model=ModelVersionResponse,
    )

update_pipeline(pipeline_id: UUID, pipeline_update: PipelineUpdate) -> PipelineResponse

Updates a pipeline.

Parameters:

Name	Type	Description	Default
pipeline_id	UUID	The ID of the pipeline to be updated.	required
pipeline_update	PipelineUpdate	The update to be applied.	required

Returns:

Type	Description
PipelineResponse	The updated pipeline.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_pipeline(
    self, pipeline_id: UUID, pipeline_update: PipelineUpdate
) -> PipelineResponse:
    """Updates a pipeline.

    Args:
        pipeline_id: The ID of the pipeline to be updated.
        pipeline_update: The update to be applied.

    Returns:
        The updated pipeline.
    """
    return self._update_resource(
        resource_id=pipeline_id,
        resource_update=pipeline_update,
        route=PIPELINES,
        response_model=PipelineResponse,
    )

update_project(project_id: UUID, project_update: ProjectUpdate) -> ProjectResponse

Update an existing project.

Parameters:

Name	Type	Description	Default
project_id	UUID	The ID of the project to be updated.	required
project_update	ProjectUpdate	The update to be applied to the project.	required

Returns:

Type	Description
ProjectResponse	The updated project.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_project(
    self, project_id: UUID, project_update: ProjectUpdate
) -> ProjectResponse:
    """Update an existing project.

    Args:
        project_id: The ID of the project to be updated.
        project_update: The update to be applied to the project.

    Returns:
        The updated project.
    """
    return self._update_resource(
        resource_id=project_id,
        resource_update=project_update,
        route=PROJECTS,
        response_model=ProjectResponse,
    )

update_resource_pool(resource_pool_id: UUID, update: ResourcePoolUpdate) -> ResourcePoolResponse

Update an existing resource pool.

Parameters:

Name	Type	Description	Default
resource_pool_id	UUID	The ID of the resource pool to update.	required
update	ResourcePoolUpdate	The update to be applied to the resource pool.	required

Returns:

Type	Description
ResourcePoolResponse	The updated resource pool.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_resource_pool(
    self, resource_pool_id: UUID, update: ResourcePoolUpdate
) -> ResourcePoolResponse:
    """Update an existing resource pool.

    Args:
        resource_pool_id: The ID of the resource pool to update.
        update: The update to be applied to the resource pool.

    Returns:
        The updated resource pool.
    """
    return self._update_resource(
        resource_id=resource_pool_id,
        resource_update=update,
        route=RESOURCE_POOLS,
        response_model=ResourcePoolResponse,
    )

update_resource_pool_subject_policy(policy_id: UUID, update: ResourcePoolSubjectPolicyUpdate) -> ResourcePoolSubjectPolicyResponse

Update an existing resource pool subject policy.

Parameters:

Name	Type	Description	Default
policy_id	UUID	The ID of the policy to update.	required
update	ResourcePoolSubjectPolicyUpdate	The update to be applied to the policy.	required

Returns:

Type	Description
ResourcePoolSubjectPolicyResponse	The updated policy.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_resource_pool_subject_policy(
    self, policy_id: UUID, update: ResourcePoolSubjectPolicyUpdate
) -> ResourcePoolSubjectPolicyResponse:
    """Update an existing resource pool subject policy.

    Args:
        policy_id: The ID of the policy to update.
        update: The update to be applied to the policy.

    Returns:
        The updated policy.
    """
    return self._update_resource(
        resource_id=policy_id,
        resource_update=update,
        route=RESOURCE_POOL_SUBJECT_POLICIES,
        response_model=ResourcePoolSubjectPolicyResponse,
    )

update_run(run_id: UUID, run_update: PipelineRunUpdate) -> PipelineRunResponse

Updates a pipeline run.

Parameters:

Name	Type	Description	Default
run_id	UUID	The ID of the pipeline run to update.	required
run_update	PipelineRunUpdate	The update to be applied to the pipeline run.	required

Returns:

Type	Description
PipelineRunResponse	The updated pipeline run.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_run(
    self, run_id: UUID, run_update: PipelineRunUpdate
) -> PipelineRunResponse:
    """Updates a pipeline run.

    Args:
        run_id: The ID of the pipeline run to update.
        run_update: The update to be applied to the pipeline run.


    Returns:
        The updated pipeline run.
    """
    return self._update_resource(
        resource_id=run_id,
        resource_update=run_update,
        response_model=PipelineRunResponse,
        route=RUNS,
    )

update_run_step(step_run_id: UUID, step_run_update: StepRunUpdate) -> StepRunResponse

Updates a step run.

Parameters:

Name	Type	Description	Default
step_run_id	UUID	The ID of the step to update.	required
step_run_update	StepRunUpdate	The update to be applied to the step.	required

Returns:

Type	Description
StepRunResponse	The updated step run.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_run_step(
    self,
    step_run_id: UUID,
    step_run_update: StepRunUpdate,
) -> StepRunResponse:
    """Updates a step run.

    Args:
        step_run_id: The ID of the step to update.
        step_run_update: The update to be applied to the step.

    Returns:
        The updated step run.
    """
    return self._update_resource(
        resource_id=step_run_id,
        resource_update=step_run_update,
        response_model=StepRunResponse,
        route=STEPS,
    )

update_run_template(template_id: UUID, template_update: RunTemplateUpdate) -> RunTemplateResponse

Updates a run template.

Parameters:

Name	Type	Description	Default
template_id	UUID	The ID of the template to update.	required
template_update	RunTemplateUpdate	The update to apply.	required

Returns:

Type	Description
RunTemplateResponse	The updated template.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_run_template(
    self,
    template_id: UUID,
    template_update: RunTemplateUpdate,
) -> RunTemplateResponse:
    """Updates a run template.

    Args:
        template_id: The ID of the template to update.
        template_update: The update to apply.

    Returns:
        The updated template.
    """
    return self._update_resource(
        resource_id=template_id,
        resource_update=template_update,
        route=RUN_TEMPLATES,
        response_model=RunTemplateResponse,
    )

update_run_wait_condition_lease(run_wait_condition_id: UUID, lease_update: RunWaitConditionLeaseUpdate) -> RunWaitConditionStatus

Update a run wait condition polling lease.

Parameters:

Name	Type	Description	Default
run_wait_condition_id	UUID	Wait condition ID.	required
lease_update	RunWaitConditionLeaseUpdate	Lease refresh payload.	required

Returns:

Type	Description
RunWaitConditionStatus	The current wait condition status after attempting the lease update.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_run_wait_condition_lease(
    self,
    run_wait_condition_id: UUID,
    lease_update: RunWaitConditionLeaseUpdate,
) -> RunWaitConditionStatus:
    """Update a run wait condition polling lease.

    Args:
        run_wait_condition_id: Wait condition ID.
        lease_update: Lease refresh payload.

    Returns:
        The current wait condition status after attempting the lease update.
    """
    response_body = self.put(
        path=f"{RUN_WAIT_CONDITIONS}/{run_wait_condition_id}",
        body=lease_update,
    )
    return RunWaitConditionStatus(response_body)

update_schedule(schedule_id: UUID, schedule_update: ScheduleUpdate) -> ScheduleResponse

Updates a schedule.

Parameters:

Name	Type	Description	Default
schedule_id	UUID	The ID of the schedule to be updated.	required
schedule_update	ScheduleUpdate	The update to be applied.	required

Returns:

Type	Description
ScheduleResponse	The updated schedule.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_schedule(
    self,
    schedule_id: UUID,
    schedule_update: ScheduleUpdate,
) -> ScheduleResponse:
    """Updates a schedule.

    Args:
        schedule_id: The ID of the schedule to be updated.
        schedule_update: The update to be applied.

    Returns:
        The updated schedule.
    """
    return self._update_resource(
        resource_id=schedule_id,
        resource_update=schedule_update,
        route=SCHEDULES,
        response_model=ScheduleResponse,
    )

update_secret(secret_id: UUID, secret_update: SecretUpdate) -> SecretResponse

Updates a secret.

Secret values that are specified as None in the update that are present in the existing secret are removed from the existing secret. Values that are present in both secrets are overwritten. All other values in both the existing secret and the update are kept (merged).

If the update includes a change of name or scope, the scoping rules enforced in the secrets store are used to validate the update:

• only one private secret with the given name can exist.
• only one public secret with the given name can exist.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret to be updated.	required
secret_update	SecretUpdate	The update to be applied.	required

Returns:

Type	Description
SecretResponse	The updated secret.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_secret(
    self, secret_id: UUID, secret_update: SecretUpdate
) -> SecretResponse:
    """Updates a secret.

    Secret values that are specified as `None` in the update that are
    present in the existing secret are removed from the existing secret.
    Values that are present in both secrets are overwritten. All other
    values in both the existing secret and the update are kept (merged).

    If the update includes a change of name or scope, the scoping rules
    enforced in the secrets store are used to validate the update:

      - only one private secret with the given name can exist.
      - only one public secret with the given name can exist.

    Args:
        secret_id: The ID of the secret to be updated.
        secret_update: The update to be applied.

    Returns:
        The updated secret.
    """
    return self._update_resource(
        resource_id=secret_id,
        resource_update=secret_update,
        route=SECRETS,
        response_model=SecretResponse,
        # The default endpoint behavior is to replace all secret values
        # with the values in the update. We want to merge the values
        # instead.
        params=dict(patch_values=True),
    )

update_server_settings(settings_update: ServerSettingsUpdate) -> ServerSettingsResponse

Update the server settings.

Parameters:

Name	Type	Description	Default
settings_update	ServerSettingsUpdate	The server settings update.	required

Returns:

Type	Description
ServerSettingsResponse	The updated server settings.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_server_settings(
    self, settings_update: ServerSettingsUpdate
) -> ServerSettingsResponse:
    """Update the server settings.

    Args:
        settings_update: The server settings update.

    Returns:
        The updated server settings.
    """
    response_body = self.put(SERVER_SETTINGS, body=settings_update)
    return ServerSettingsResponse.model_validate(response_body)

update_service(service_id: UUID, update: ServiceUpdate) -> ServiceResponse

Update a service.

Parameters:

Name	Type	Description	Default
service_id	UUID	The ID of the service to update.	required
update	ServiceUpdate	The update to be applied to the service.	required

Returns:

Type	Description
ServiceResponse	The updated service.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_service(
    self, service_id: UUID, update: ServiceUpdate
) -> ServiceResponse:
    """Update a service.

    Args:
        service_id: The ID of the service to update.
        update: The update to be applied to the service.

    Returns:
        The updated service.
    """
    return self._update_resource(
        resource_id=service_id,
        resource_update=update,
        response_model=ServiceResponse,
        route=SERVICES,
    )

update_service_account(service_account_name_or_id: Union[str, UUID], service_account_update: ServiceAccountUpdate) -> ServiceAccountResponse

Updates an existing service account.

Parameters:

Name	Type	Description	Default
service_account_name_or_id	Union[str, UUID]	The name or the ID of the service account to update.	required
service_account_update	ServiceAccountUpdate	The update to be applied to the service account.	required

Returns:

Type	Description
ServiceAccountResponse	The updated service account.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_service_account(
    self,
    service_account_name_or_id: Union[str, UUID],
    service_account_update: ServiceAccountUpdate,
) -> ServiceAccountResponse:
    """Updates an existing service account.

    Args:
        service_account_name_or_id: The name or the ID of the service
            account to update.
        service_account_update: The update to be applied to the service
            account.

    Returns:
        The updated service account.
    """
    return self._update_resource(
        resource_id=service_account_name_or_id,
        resource_update=service_account_update,
        route=SERVICE_ACCOUNTS,
        response_model=ServiceAccountResponse,
    )

update_service_connector(service_connector_id: UUID, update: ServiceConnectorUpdate) -> ServiceConnectorResponse

Updates an existing service connector.

The update model contains the fields to be updated. If a field value is set to None in the model, the field is not updated, but there are special rules concerning some fields:

• the configuration and secrets fields together represent a full valid configuration update, not just a partial update. If either is set (i.e. not None) in the update, their values are merged together and will replace the existing configuration and secrets values.
• the resource_id field value is also a full replacement value: if set to None, the resource ID is removed from the service connector.
• the expiration_seconds field value is also a full replacement value: if set to None, the expiration is removed from the service connector.
• the secret_id field value in the update is ignored, given that secrets are managed internally by the ZenML store.
• the labels field is also a full labels update: if set (i.e. not None), all existing labels are removed and replaced by the new labels in the update.

Parameters:

Name	Type	Description	Default
service_connector_id	UUID	The ID of the service connector to update.	required
update	ServiceConnectorUpdate	The update to be applied to the service connector.	required

Returns:

Type	Description
ServiceConnectorResponse	The updated service connector.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_service_connector(
    self, service_connector_id: UUID, update: ServiceConnectorUpdate
) -> ServiceConnectorResponse:
    """Updates an existing service connector.

    The update model contains the fields to be updated. If a field value is
    set to None in the model, the field is not updated, but there are
    special rules concerning some fields:

    * the `configuration` and `secrets` fields together represent a full
    valid configuration update, not just a partial update. If either is
    set (i.e. not None) in the update, their values are merged together and
    will replace the existing configuration and secrets values.
    * the `resource_id` field value is also a full replacement value: if set
    to `None`, the resource ID is removed from the service connector.
    * the `expiration_seconds` field value is also a full replacement value:
    if set to `None`, the expiration is removed from the service connector.
    * the `secret_id` field value in the update is ignored, given that
    secrets are managed internally by the ZenML store.
    * the `labels` field is also a full labels update: if set (i.e. not
    `None`), all existing labels are removed and replaced by the new labels
    in the update.

    Args:
        service_connector_id: The ID of the service connector to update.
        update: The update to be applied to the service connector.

    Returns:
        The updated service connector.
    """
    connector_model = self._update_resource(
        resource_id=service_connector_id,
        resource_update=update,
        response_model=ServiceConnectorResponse,
        route=SERVICE_CONNECTORS,
    )
    self._populate_connector_type(connector_model)
    # Call this to properly split the secrets from the configuration
    try:
        connector_model.validate_configuration()
    except ValueError as e:
        logger.error(
            f"Error validating connector configuration for "
            f"{connector_model.name}: {e}"
        )
    return connector_model

update_snapshot(snapshot_id: UUID, snapshot_update: PipelineSnapshotUpdate) -> PipelineSnapshotResponse

Update a snapshot.

Parameters:

Name	Type	Description	Default
snapshot_id	UUID	The ID of the snapshot to update.	required
snapshot_update	PipelineSnapshotUpdate	The update to apply.	required

Returns:

Type	Description
PipelineSnapshotResponse	The updated snapshot.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_snapshot(
    self,
    snapshot_id: UUID,
    snapshot_update: PipelineSnapshotUpdate,
) -> PipelineSnapshotResponse:
    """Update a snapshot.

    Args:
        snapshot_id: The ID of the snapshot to update.
        snapshot_update: The update to apply.

    Returns:
        The updated snapshot.
    """
    return self._update_resource(
        resource_id=snapshot_id,
        resource_update=snapshot_update,
        route=PIPELINE_SNAPSHOTS,
        response_model=PipelineSnapshotResponse,
    )

update_stack(stack_id: UUID, stack_update: StackUpdate) -> StackResponse

Update a stack.

Parameters:

Name	Type	Description	Default
stack_id	UUID	The ID of the stack update.	required
stack_update	StackUpdate	The update request on the stack.	required

Returns:

Type	Description
StackResponse	The updated stack.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_stack(
    self, stack_id: UUID, stack_update: StackUpdate
) -> StackResponse:
    """Update a stack.

    Args:
        stack_id: The ID of the stack update.
        stack_update: The update request on the stack.

    Returns:
        The updated stack.
    """
    return self._update_resource(
        resource_id=stack_id,
        resource_update=stack_update,
        route=STACKS,
        response_model=StackResponse,
    )

update_stack_component(component_id: UUID, component_update: ComponentUpdate) -> ComponentResponse

Update an existing stack component.

Parameters:

Name	Type	Description	Default
component_id	UUID	The ID of the stack component to update.	required
component_update	ComponentUpdate	The update to be applied to the stack component.	required

Returns:

Type	Description
ComponentResponse	The updated stack component.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_stack_component(
    self,
    component_id: UUID,
    component_update: ComponentUpdate,
) -> ComponentResponse:
    """Update an existing stack component.

    Args:
        component_id: The ID of the stack component to update.
        component_update: The update to be applied to the stack component.

    Returns:
        The updated stack component.
    """
    return self._update_resource(
        resource_id=component_id,
        resource_update=component_update,
        route=STACK_COMPONENTS,
        response_model=ComponentResponse,
    )

update_step_heartbeat(step_run_id: UUID) -> StepHeartbeatResponse

Updates a step run heartbeat.

Parameters:

Name	Type	Description	Default
step_run_id	UUID	The ID of the step to update.	required

Returns:

Type	Description
StepHeartbeatResponse	The step heartbeat response.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_step_heartbeat(
    self, step_run_id: UUID
) -> StepHeartbeatResponse:
    """Updates a step run heartbeat.

    Args:
        step_run_id: The ID of the step to update.

    Returns:
        The step heartbeat response.
    """
    response_body = self.put(
        path=f"{STEPS}/{str(step_run_id)}{HEARTBEAT}",
        timeout=5,
    )

    return StepHeartbeatResponse.model_validate(response_body)

update_tag(tag_id: UUID, tag_update_model: TagUpdate) -> TagResponse

Update tag.

Parameters:

Name	Type	Description	Default
tag_id	UUID	id of the tag to be updated.	required
tag_update_model	TagUpdate	Tag to use for the update.	required

Returns:

Type	Description
TagResponse	An updated tag.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_tag(
    self,
    tag_id: UUID,
    tag_update_model: TagUpdate,
) -> TagResponse:
    """Update tag.

    Args:
        tag_id: id of the tag to be updated.
        tag_update_model: Tag to use for the update.

    Returns:
        An updated tag.
    """
    return self._update_resource(
        resource_id=tag_id,
        resource_update=tag_update_model,
        route=TAGS,
        response_model=TagResponse,
    )

update_trigger(trigger_id: UUID, trigger_update: TriggerUpdate) -> TRIGGER_RETURN_TYPE_UNION

Updates a trigger.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	The ID of the trigger to update.	required
trigger_update	TriggerUpdate	The update to be applied to the trigger.	required

Returns:

Type	Description
TRIGGER_RETURN_TYPE_UNION	The updated trigger.

Raises:

Type	Description
ValueError	In case of bad response.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_trigger(
    self, trigger_id: UUID, trigger_update: TriggerUpdate
) -> TRIGGER_RETURN_TYPE_UNION:
    """Updates a trigger.

    Args:
        trigger_id: The ID of the trigger to update.
        trigger_update: The update to be applied to the trigger.

    Returns:
        The updated trigger.

    Raises:
        ValueError: In case of bad response.
    """
    body: dict[str, Any] = self.put(  # type: ignore[assignment]
        f"{TRIGGERS}/{trigger_id}",
        body=trigger_update,
        params=None,
        exclude_unset=False,
    )
    try:
        response_model = TYPE_TO_RESPONSE_MAPPING[body["body"]["type"]]
        return response_model.model_validate(body)
    except (KeyError, TypeError):
        raise ValueError("Bad response, expected a trigger type object.")

update_user(user_id: UUID, user_update: UserUpdate) -> UserResponse

Updates an existing user.

Parameters:

Name	Type	Description	Default
user_id	UUID	The id of the user to update.	required
user_update	UserUpdate	The update to be applied to the user.	required

Returns:

Type	Description
UserResponse	The updated user.

Source code in src/zenml/zen_stores/rest_zen_store.py

def update_user(
    self, user_id: UUID, user_update: UserUpdate
) -> UserResponse:
    """Updates an existing user.

    Args:
        user_id: The id of the user to update.
        user_update: The update to be applied to the user.

    Returns:
        The updated user.
    """
    return self._update_resource(
        resource_id=user_id,
        resource_update=user_update,
        route=USERS,
        response_model=UserResponse,
    )

verify_service_connector(service_connector_id: UUID, resource_type: Optional[str] = None, resource_id: Optional[str] = None, list_resources: bool = True) -> ServiceConnectorResourcesModel

Verifies if a service connector instance has access to one or more resources.

Parameters:

Name	Type	Description	Default
service_connector_id	UUID	The ID of the service connector to verify.	required
resource_type	Optional[str]	The type of resource to verify access to.	None
resource_id	Optional[str]	The ID of the resource to verify access to.	None
list_resources	bool	If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned.	True

Returns:

Type	Description
ServiceConnectorResourcesModel	The list of resources that the service connector has access to,
ServiceConnectorResourcesModel	scoped to the supplied resource type and ID, if provided.

Source code in src/zenml/zen_stores/rest_zen_store.py

def verify_service_connector(
    self,
    service_connector_id: UUID,
    resource_type: Optional[str] = None,
    resource_id: Optional[str] = None,
    list_resources: bool = True,
) -> ServiceConnectorResourcesModel:
    """Verifies if a service connector instance has access to one or more resources.

    Args:
        service_connector_id: The ID of the service connector to verify.
        resource_type: The type of resource to verify access to.
        resource_id: The ID of the resource to verify access to.
        list_resources: If True, the list of all resources accessible
            through the service connector and matching the supplied resource
            type and ID are returned.

    Returns:
        The list of resources that the service connector has access to,
        scoped to the supplied resource type and ID, if provided.
    """
    params: Dict[str, Any] = {"list_resources": list_resources}
    if resource_type:
        params["resource_type"] = resource_type
    if resource_id:
        params["resource_id"] = resource_id
    response_body = self.put(
        f"{SERVICE_CONNECTORS}/{str(service_connector_id)}{SERVICE_CONNECTOR_VERIFY}",
        params=params,
        timeout=max(
            self.config.http_timeout,
            SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
        ),
    )

    resources = ServiceConnectorResourcesModel.model_validate(
        response_body
    )
    self._populate_connector_type(resources)
    return resources

verify_service_connector_config(service_connector: ServiceConnectorRequest, list_resources: bool = True) -> ServiceConnectorResourcesModel

Verifies if a service connector configuration has access to resources.

Parameters:

Name	Type	Description	Default
service_connector	ServiceConnectorRequest	The service connector configuration to verify.	required
list_resources	bool	If True, the list of all resources accessible through the service connector and matching the supplied resource type and ID are returned.	True

Returns:

Type	Description
ServiceConnectorResourcesModel	The list of resources that the service connector configuration has
ServiceConnectorResourcesModel	access to.

Source code in src/zenml/zen_stores/rest_zen_store.py

def verify_service_connector_config(
    self,
    service_connector: ServiceConnectorRequest,
    list_resources: bool = True,
) -> ServiceConnectorResourcesModel:
    """Verifies if a service connector configuration has access to resources.

    Args:
        service_connector: The service connector configuration to verify.
        list_resources: If True, the list of all resources accessible
            through the service connector and matching the supplied resource
            type and ID are returned.

    Returns:
        The list of resources that the service connector configuration has
        access to.
    """
    response_body = self.post(
        f"{SERVICE_CONNECTORS}{SERVICE_CONNECTOR_VERIFY}",
        body=service_connector,
        params={"list_resources": list_resources},
        timeout=max(
            self.config.http_timeout,
            SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT,
        ),
    )

    resources = ServiceConnectorResourcesModel.model_validate(
        response_body
    )
    self._populate_connector_type(resources)
    return resources

RestZenStoreConfiguration

Bases: StoreConfiguration

REST ZenML store configuration.

Attributes:

Name	Type	Description
type	StoreType	The type of the store.
username	StoreType	The username to use to connect to the Zen server.
password	StoreType	The password to use to connect to the Zen server.
api_key	StoreType	The service account API key to use to connect to the Zen server. This is only set if the API key is configured explicitly via environment variables or the ZenML global configuration file. API keys configured via the CLI are stored in the credentials store instead.
verify_ssl	Union[bool, str]	Either a boolean, in which case it controls whether we verify the server's TLS certificate, or a string, in which case it must be a path to a CA bundle to use or the CA bundle value itself.
http_timeout	int	The timeout to use for all requests.
connection_pool_size	int	The size of the connection pool to use for all requests.

Functions

supports_url_scheme(url: str) -> bool classmethod

Check if a URL scheme is supported by this store.

Parameters:

Name	Type	Description	Default
url	str	The URL to check.	required

Returns:

Type	Description
bool	True if the URL scheme is supported, False otherwise.

Source code in src/zenml/zen_stores/rest_zen_store.py

@classmethod
def supports_url_scheme(cls, url: str) -> bool:
    """Check if a URL scheme is supported by this store.

    Args:
        url: The URL to check.

    Returns:
        True if the URL scheme is supported, False otherwise.
    """
    return urlparse(url).scheme in ("http", "https")

validate_url(url: str) -> str classmethod

Validates that the URL is a well-formed REST store URL.

Parameters:

Name	Type	Description	Default
url	str	The URL to be validated.	required

Returns:

Type	Description
str	The validated URL without trailing slashes.

Raises:

Type	Description
ValueError	If the URL is not a well-formed REST store URL.

Source code in src/zenml/zen_stores/rest_zen_store.py

@field_validator("url")
@classmethod
def validate_url(cls, url: str) -> str:
    """Validates that the URL is a well-formed REST store URL.

    Args:
        url: The URL to be validated.

    Returns:
        The validated URL without trailing slashes.

    Raises:
        ValueError: If the URL is not a well-formed REST store URL.
    """
    url = url.rstrip("/")
    scheme = re.search("^([a-z0-9]+://)", url)
    if scheme is None or scheme.group() not in ("https://", "http://"):
        raise ValueError(
            "Invalid URL for REST store: {url}. Should be in the form "
            "https://hostname[:port] or http://hostname[:port]."
        )

    # When running inside a container, if the URL uses localhost, the
    # target service will not be available. We try to replace localhost
    # with one of the special Docker or K3D internal hostnames.
    url = replace_localhost_with_internal_hostname(url)

    return url

validate_verify_ssl(verify_ssl: Union[bool, str]) -> Union[bool, str] classmethod

Validates that the verify_ssl either points to a file or is a bool.

Parameters:

Name	Type	Description	Default
verify_ssl	Union[bool, str]	The verify_ssl value to be validated.	required

Returns:

Type	Description
Union[bool, str]	The validated verify_ssl value.

Source code in src/zenml/zen_stores/rest_zen_store.py

@field_validator("verify_ssl")
@classmethod
def validate_verify_ssl(
    cls, verify_ssl: Union[bool, str]
) -> Union[bool, str]:
    """Validates that the verify_ssl either points to a file or is a bool.

    Args:
        verify_ssl: The verify_ssl value to be validated.

    Returns:
        The validated verify_ssl value.
    """
    secret_folder = Path(
        GlobalConfiguration().local_stores_path,
        "certificates",
    )
    if isinstance(verify_ssl, bool) or verify_ssl.startswith(
        str(secret_folder)
    ):
        return verify_ssl

    if os.path.isfile(verify_ssl):
        with open(verify_ssl, "r") as f:
            cert_content = f.read()
    else:
        cert_content = verify_ssl

    fileio.makedirs(str(secret_folder))
    file_path = Path(secret_folder, "ca_bundle.pem")
    with os.fdopen(
        os.open(file_path, flags=os.O_RDWR | os.O_CREAT, mode=0o600), "w"
    ) as f:
        f.write(cert_content)

    return str(file_path)

Functions

Modules

schemas

SQL Model Implementations.

Classes

APIKeySchema

Bases: NamedSchema

SQL Model for API keys.

Functions

from_request(service_account_id: UUID, request: APIKeyRequest) -> Tuple[APIKeySchema, str] classmethod

Convert a APIKeyRequest to a APIKeySchema.

Parameters:

Name	Type	Description	Default
service_account_id	UUID	The service account id to associate the key with.	required
request	APIKeyRequest	The request model to convert.	required

Returns:

Type	Description
Tuple[APIKeySchema, str]	The converted schema and the un-hashed API key.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

@classmethod
def from_request(
    cls,
    service_account_id: UUID,
    request: APIKeyRequest,
) -> Tuple["APIKeySchema", str]:
    """Convert a `APIKeyRequest` to a `APIKeySchema`.

    Args:
        service_account_id: The service account id to associate the key
            with.
        request: The request model to convert.

    Returns:
        The converted schema and the un-hashed API key.
    """
    key = cls._generate_jwt_secret_key()
    hashed_key = cls._get_hashed_key(key)
    now = utc_now()
    return (
        cls(
            name=request.name,
            description=request.description or "",
            key=hashed_key,
            service_account_id=service_account_id,
            created=now,
            updated=now,
        ),
        key,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = [
        joinedload(jl_arg(APIKeySchema.service_account), innerjoin=True),
    ]

    return options

internal_update(update: APIKeyInternalUpdate) -> APIKeySchema

Update an APIKeySchema with an APIKeyInternalUpdate.

The internal update can also update the last used timestamp.

Parameters:

Name	Type	Description	Default
update	APIKeyInternalUpdate	The update model.	required

Returns:

Type	Description
APIKeySchema	The updated APIKeySchema.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def internal_update(self, update: APIKeyInternalUpdate) -> "APIKeySchema":
    """Update an `APIKeySchema` with an `APIKeyInternalUpdate`.

    The internal update can also update the last used timestamp.

    Args:
        update: The update model.

    Returns:
        The updated `APIKeySchema`.
    """
    self.update(update)

    if update.update_last_login:
        self.last_login = self.updated

    return self

rotate(rotate_request: APIKeyRotateRequest) -> Tuple[APIKeySchema, str]

Rotate the key for an APIKeySchema.

Parameters:

Name	Type	Description	Default
rotate_request	APIKeyRotateRequest	The rotate request model.	required

Returns:

Type	Description
Tuple[APIKeySchema, str]	The updated APIKeySchema and the new un-hashed key.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def rotate(
    self,
    rotate_request: APIKeyRotateRequest,
) -> Tuple["APIKeySchema", str]:
    """Rotate the key for an `APIKeySchema`.

    Args:
        rotate_request: The rotate request model.

    Returns:
        The updated `APIKeySchema` and the new un-hashed key.
    """
    self.updated = utc_now()
    self.previous_key = self.key
    self.retain_period = rotate_request.retain_period_minutes
    new_key = self._generate_jwt_secret_key()
    self.key = self._get_hashed_key(new_key)
    self.last_rotated = self.updated

    return self, new_key

to_internal_model(include_metadata: bool = False, include_resources: bool = False) -> APIKeyInternalResponse

Convert a APIKeySchema to an APIKeyInternalResponse.

The internal response model includes the hashed key values.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False

Returns:

Type	Description
APIKeyInternalResponse	The created APIKeyInternalResponse.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def to_internal_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
) -> APIKeyInternalResponse:
    """Convert a `APIKeySchema` to an `APIKeyInternalResponse`.

    The internal response model includes the hashed key values.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.

    Returns:
        The created APIKeyInternalResponse.
    """
    model = self.to_model(
        include_metadata=include_metadata,
        include_resources=include_resources,
    )
    model.get_body().key = self.key

    return APIKeyInternalResponse(
        id=self.id,
        name=self.name,
        previous_key=self.previous_key,
        body=model.body,
        metadata=model.metadata,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> APIKeyResponse

Convert a APIKeySchema to an APIKeyResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic.	{}

Returns:

Type	Description
APIKeyResponse	The created APIKeyResponse.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> APIKeyResponse:
    """Convert a `APIKeySchema` to an `APIKeyResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic.

    Returns:
        The created APIKeyResponse.
    """
    metadata = None
    if include_metadata:
        metadata = APIKeyResponseMetadata(
            description=self.description,
            retain_period_minutes=self.retain_period,
            last_login=self.last_login,
            last_rotated=self.last_rotated,
        )

    body = APIKeyResponseBody(
        created=self.created,
        updated=self.updated,
        active=self.active,
        service_account=self.service_account.to_service_account_model(),
    )

    return APIKeyResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
    )

update(update: APIKeyUpdate) -> APIKeySchema

Update an APIKeySchema with an APIKeyUpdate.

Parameters:

Name	Type	Description	Default
update	APIKeyUpdate	The update model.	required

Returns:

Type	Description
APIKeySchema	The updated APIKeySchema.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def update(self, update: APIKeyUpdate) -> "APIKeySchema":
    """Update an `APIKeySchema` with an `APIKeyUpdate`.

    Args:
        update: The update model.

    Returns:
        The updated `APIKeySchema`.
    """
    for field, value in update.model_dump(exclude_none=True).items():
        if hasattr(self, field):
            setattr(self, field, value)

    self.updated = utc_now()
    return self

ApiTransactionResultSchema

Bases: SQLModel

SQL Model for API transaction results.

ApiTransactionSchema

Bases: BaseSchema

SQL Model for API transactions.

The result payload is stored in a separate table to keep this table small and fast for cleanup operations. Deleting rows with large blobs is expensive because the entire row must be copied to the undo log.

Functions

from_request(request: ApiTransactionRequest) -> ApiTransactionSchema classmethod

Create a new API transaction from a request.

Parameters:

Name	Type	Description	Default
request	ApiTransactionRequest	The API transaction request.	required

Returns:

Type	Description
ApiTransactionSchema	The API transaction schema.

Source code in src/zenml/zen_stores/schemas/api_transaction_schemas.py

@classmethod
def from_request(
    cls, request: ApiTransactionRequest
) -> "ApiTransactionSchema":
    """Create a new API transaction from a request.

    Args:
        request: The API transaction request.

    Returns:
        The API transaction schema.
    """
    assert request.user is not None, "User must be set."
    return cls(
        id=request.transaction_id,
        user_id=request.user,
        method=request.method,
        url=request.url,
        completed=False,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ApiTransactionResponse

Convert the SQL model to a ZenML model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
ApiTransactionResponse	The API transaction response.

Source code in src/zenml/zen_stores/schemas/api_transaction_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ApiTransactionResponse:
    """Convert the SQL model to a ZenML model.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        **kwargs: Additional keyword arguments.

    Returns:
        The API transaction response.
    """
    response = ApiTransactionResponse(
        id=self.id,
        body=ApiTransactionResponseBody(
            method=self.method,
            url=self.url,
            created=self.created,
            updated=self.updated,
            user_id=self.user_id,
            completed=self.completed,
        ),
    )
    return response

update(update: ApiTransactionUpdate) -> ApiTransactionSchema

Update the API transaction.

Parameters:

Name	Type	Description	Default
update	ApiTransactionUpdate	The API transaction update.	required

Returns:

Type	Description
ApiTransactionSchema	The API transaction schema.

Source code in src/zenml/zen_stores/schemas/api_transaction_schemas.py

def update(self, update: ApiTransactionUpdate) -> "ApiTransactionSchema":
    """Update the API transaction.

    Args:
        update: The API transaction update.

    Returns:
        The API transaction schema.
    """
    self.updated = utc_now()
    self.expired = self.updated + timedelta(seconds=update.cache_time)
    return self

ArtifactSchema

Bases: NamedSchema

SQL Model for artifacts.

Attributes

latest_version: Optional[ArtifactVersionSchema] property

Fetch the latest version for this artifact.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[ArtifactVersionSchema]	The latest version for this artifact.

Functions

from_request(artifact_request: ArtifactRequest) -> ArtifactSchema classmethod

Convert an ArtifactRequest to an ArtifactSchema.

Parameters:

Name	Type	Description	Default
artifact_request	ArtifactRequest	The request model to convert.	required

Returns:

Type	Description
ArtifactSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

@classmethod
def from_request(
    cls,
    artifact_request: ArtifactRequest,
) -> "ArtifactSchema":
    """Convert an `ArtifactRequest` to an `ArtifactSchema`.

    Args:
        artifact_request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=artifact_request.name,
        has_custom_name=artifact_request.has_custom_name,
        project_id=artifact_request.project,
        user_id=artifact_request.user,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ArtifactSchema.user)),
                # joinedload(jl_arg(ArtifactSchema.tags)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ArtifactResponse

Convert an ArtifactSchema to an ArtifactResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ArtifactResponse	The created ArtifactResponse.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ArtifactResponse:
    """Convert an `ArtifactSchema` to an `ArtifactResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic



    Returns:
        The created `ArtifactResponse`.
    """
    # Create the body of the model
    body = ArtifactResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
    )

    # Create the metadata of the model
    metadata = None
    if include_metadata:
        metadata = ArtifactResponseMetadata(
            has_custom_name=self.has_custom_name,
        )

    resources = None
    if include_resources:
        latest_id, latest_name = None, None
        if latest_version := self.latest_version:
            latest_id = latest_version.id
            latest_name = latest_version.version

        resources = ArtifactResponseResources(
            user=self.user.to_model() if self.user else None,
            tags=[tag.to_model() for tag in self.tags],
            latest_version_id=latest_id,
            latest_version_name=latest_name,
        )

    return ArtifactResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(artifact_update: ArtifactUpdate) -> ArtifactSchema

Update an ArtifactSchema with an ArtifactUpdate.

Parameters:

Name	Type	Description	Default
artifact_update	ArtifactUpdate	The update model to apply.	required

Returns:

Type	Description
ArtifactSchema	The updated ArtifactSchema.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

def update(self, artifact_update: ArtifactUpdate) -> "ArtifactSchema":
    """Update an `ArtifactSchema` with an `ArtifactUpdate`.

    Args:
        artifact_update: The update model to apply.

    Returns:
        The updated `ArtifactSchema`.
    """
    self.updated = utc_now()
    if artifact_update.name:
        self.name = artifact_update.name
        self.has_custom_name = True
    if artifact_update.has_custom_name is not None:
        self.has_custom_name = artifact_update.has_custom_name
    return self

ArtifactVersionSchema

Bases: BaseSchema, RunMetadataInterface

SQL Model for artifact versions.

Attributes

producer_run_ids: Optional[Tuple[UUID, UUID]] property

Fetch the producer run IDs for this artifact version.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[Tuple[UUID, UUID]]	The producer step run ID and pipeline run ID for this artifact
Optional[Tuple[UUID, UUID]]	version.

Functions

from_request(artifact_version_request: ArtifactVersionRequest) -> ArtifactVersionSchema classmethod

Convert an ArtifactVersionRequest to an ArtifactVersionSchema.

Parameters:

Name	Type	Description	Default
artifact_version_request	ArtifactVersionRequest	The request model to convert.	required

Raises:

Type	Description
ValueError	If the request does not specify a version number.

Returns:

Type	Description
ArtifactVersionSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

@classmethod
def from_request(
    cls,
    artifact_version_request: ArtifactVersionRequest,
) -> "ArtifactVersionSchema":
    """Convert an `ArtifactVersionRequest` to an `ArtifactVersionSchema`.

    Args:
        artifact_version_request: The request model to convert.

    Raises:
        ValueError: If the request does not specify a version number.

    Returns:
        The converted schema.
    """
    if not artifact_version_request.version:
        raise ValueError("Missing version for artifact version request.")

    try:
        version_number = int(artifact_version_request.version)
    except ValueError:
        version_number = None
    return cls(
        artifact_id=artifact_version_request.artifact_id,
        version=str(artifact_version_request.version),
        version_number=version_number,
        artifact_store_id=artifact_version_request.artifact_store_id,
        project_id=artifact_version_request.project,
        user_id=artifact_version_request.user,
        type=artifact_version_request.type.value,
        uri=artifact_version_request.uri,
        materializer=artifact_version_request.materializer.model_dump_json(),
        data_type=artifact_version_request.data_type.model_dump_json(),
        save_type=artifact_version_request.save_type.value,
        content_hash=artifact_version_request.content_hash,
        item_count=artifact_version_request.item_count,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    # if include_metadata:
    #     options.extend(
    #         [
    #             joinedload(jl_arg(ArtifactVersionSchema.visualizations)),
    #             joinedload(jl_arg(ArtifactVersionSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ArtifactVersionSchema.user)),
                # joinedload(jl_arg(ArtifactVersionSchema.tags)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ArtifactVersionResponse

Convert an ArtifactVersionSchema to an ArtifactVersionResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ArtifactVersionResponse	The created ArtifactVersionResponse.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ArtifactVersionResponse:
    """Convert an `ArtifactVersionSchema` to an `ArtifactVersionResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic



    Returns:
        The created `ArtifactVersionResponse`.
    """
    try:
        materializer = Source.model_validate_json(self.materializer)
    except ValidationError:
        # This is an old source which was an importable source path
        materializer = Source.from_import_path(self.materializer)

    try:
        data_type = Source.model_validate_json(self.data_type)
    except ValidationError:
        # This is an old source which was an importable source path
        data_type = Source.from_import_path(self.data_type)

    # Create the body of the model
    artifact = self.artifact.to_model()
    body = ArtifactVersionResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        artifact=artifact,
        version=self.version or str(self.version_number),
        uri=self.uri,
        type=ArtifactType(self.type),
        materializer=materializer,
        data_type=data_type,
        created=self.created,
        updated=self.updated,
        save_type=ArtifactSaveType(self.save_type),
        artifact_store_id=self.artifact_store_id,
        content_hash=self.content_hash,
        item_count=self.item_count,
    )

    # Create the metadata of the model
    metadata = None
    if include_metadata:
        metadata = ArtifactVersionResponseMetadata(
            visualizations=[v.to_model() for v in self.visualizations],
            run_metadata=self.fetch_metadata(),
        )

    resources = None
    if include_resources:
        producer_step_run_id, producer_pipeline_run_id = None, None
        if producer_run_ids := self.producer_run_ids:
            # TODO: Why was the producer_pipeline_run_id only set for one
            # of the cases before?
            producer_step_run_id, producer_pipeline_run_id = (
                producer_run_ids
            )

        resources = ArtifactVersionResponseResources(
            user=self.user.to_model() if self.user else None,
            tags=[tag.to_model() for tag in self.tags],
            producer_step_run_id=producer_step_run_id,
            producer_pipeline_run_id=producer_pipeline_run_id,
        )

    return ArtifactVersionResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(artifact_version_update: ArtifactVersionUpdate) -> ArtifactVersionSchema

Update an ArtifactVersionSchema with an ArtifactVersionUpdate.

Parameters:

Name	Type	Description	Default
artifact_version_update	ArtifactVersionUpdate	The update model to apply.	required

Returns:

Type	Description
ArtifactVersionSchema	The updated ArtifactVersionSchema.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

def update(
    self, artifact_version_update: ArtifactVersionUpdate
) -> "ArtifactVersionSchema":
    """Update an `ArtifactVersionSchema` with an `ArtifactVersionUpdate`.

    Args:
        artifact_version_update: The update model to apply.

    Returns:
        The updated `ArtifactVersionSchema`.
    """
    self.updated = utc_now()
    return self

ArtifactVisualizationSchema

Bases: BaseSchema

SQL Model for visualizations of artifacts.

Functions

from_model(artifact_visualization_request: ArtifactVisualizationRequest, artifact_version_id: UUID) -> ArtifactVisualizationSchema classmethod

Convert a ArtifactVisualizationRequest to a ArtifactVisualizationSchema.

Parameters:

Name	Type	Description	Default
artifact_visualization_request	ArtifactVisualizationRequest	The visualization.	required
artifact_version_id	UUID	The UUID of the artifact version.	required

Returns:

Type	Description
ArtifactVisualizationSchema	The ArtifactVisualizationSchema.

Source code in src/zenml/zen_stores/schemas/artifact_visualization_schemas.py

@classmethod
def from_model(
    cls,
    artifact_visualization_request: ArtifactVisualizationRequest,
    artifact_version_id: UUID,
) -> "ArtifactVisualizationSchema":
    """Convert a `ArtifactVisualizationRequest` to a `ArtifactVisualizationSchema`.

    Args:
        artifact_visualization_request: The visualization.
        artifact_version_id: The UUID of the artifact version.

    Returns:
        The `ArtifactVisualizationSchema`.
    """
    return cls(
        type=artifact_visualization_request.type.value,
        uri=artifact_visualization_request.uri,
        artifact_version_id=artifact_version_id,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ArtifactVisualizationResponse

Convert an ArtifactVisualizationSchema to a Visualization.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ArtifactVisualizationResponse	The Visualization.

Source code in src/zenml/zen_stores/schemas/artifact_visualization_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ArtifactVisualizationResponse:
    """Convert an `ArtifactVisualizationSchema` to a `Visualization`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic



    Returns:
        The `Visualization`.
    """
    body = ArtifactVisualizationResponseBody(
        type=VisualizationType(self.type),
        uri=self.uri,
        created=self.created,
        updated=self.updated,
    )

    metadata = None
    if include_metadata:
        metadata = ArtifactVisualizationResponseMetadata(
            artifact_version_id=self.artifact_version_id,
        )

    resources = None
    if include_resources:
        if self.artifact_version is not None:
            artifact_version = self.artifact_version.to_model(
                include_metadata=False,
                include_resources=False,
            )
        else:
            artifact_version = None
        resources = ArtifactVisualizationResponseResources(
            artifact_version=artifact_version,
        )

    return ArtifactVisualizationResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

BaseSchema

Bases: SQLModel

Base SQL Model for ZenML entities.

Functions

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

This method should return query options that improve the performance when trying to later on converting that schema to a model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/base_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    This method should return query options that improve the performance
    when trying to later on converting that schema to a model.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    return []

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Any

In case the Schema has a corresponding Model, this allows conversion to that model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Any	The model corresponding to this schema.

Raises:

Type	Description
NotImplementedError	When the base class fails to implement this.

Source code in src/zenml/zen_stores/schemas/base_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Any:
    """In case the Schema has a corresponding Model, this allows conversion to that model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The model corresponding to this schema.

    Raises:
        NotImplementedError: When the base class fails to implement this.
    """
    raise NotImplementedError(
        "No 'to_model()' method implemented for this"
        f"schema: '{self.__class__.__name__}'."
    )

CodeReferenceSchema

Bases: BaseSchema

SQL Model for code references.

Functions

from_request(request: CodeReferenceRequest, project_id: UUID) -> CodeReferenceSchema classmethod

Convert a CodeReferenceRequest to a CodeReferenceSchema.

Parameters:

Name	Type	Description	Default
request	CodeReferenceRequest	The request model to convert.	required
project_id	UUID	The project ID.	required

Returns:

Type	Description
CodeReferenceSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

@classmethod
def from_request(
    cls, request: "CodeReferenceRequest", project_id: UUID
) -> "CodeReferenceSchema":
    """Convert a `CodeReferenceRequest` to a `CodeReferenceSchema`.

    Args:
        request: The request model to convert.
        project_id: The project ID.

    Returns:
        The converted schema.
    """
    return cls(
        project_id=project_id,
        commit=request.commit,
        subdirectory=request.subdirectory,
        code_repository_id=request.code_repository,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> CodeReferenceResponse

Convert a CodeReferenceSchema to a CodeReferenceResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic.	{}

Returns:

Type	Description
CodeReferenceResponse	The converted model.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "CodeReferenceResponse":
    """Convert a `CodeReferenceSchema` to a `CodeReferenceResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic.

    Returns:
        The converted model.
    """
    body = CodeReferenceResponseBody(
        commit=self.commit,
        subdirectory=self.subdirectory,
        code_repository=self.code_repository.to_model(),
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = CodeReferenceResponseMetadata()

    return CodeReferenceResponse(
        id=self.id,
        body=body,
        metadata=metadata,
    )

CodeRepositorySchema

Bases: NamedSchema

SQL Model for code repositories.

Functions

from_request(request: CodeRepositoryRequest) -> CodeRepositorySchema classmethod

Convert a CodeRepositoryRequest to a CodeRepositorySchema.

Parameters:

Name	Type	Description	Default
request	CodeRepositoryRequest	The request model to convert.	required

Returns:

Type	Description
CodeRepositorySchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

@classmethod
def from_request(
    cls, request: "CodeRepositoryRequest"
) -> "CodeRepositorySchema":
    """Convert a `CodeRepositoryRequest` to a `CodeRepositorySchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        config=json.dumps(request.config),
        source=request.source.model_dump_json(),
        description=request.description,
        logo_url=request.logo_url,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(CodeRepositorySchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> CodeRepositoryResponse

Convert a CodeRepositorySchema to a CodeRepositoryResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic.	{}

Returns:

Type	Description
CodeRepositoryResponse	The created CodeRepositoryResponse.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "CodeRepositoryResponse":
    """Convert a `CodeRepositorySchema` to a `CodeRepositoryResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic.

    Returns:
        The created CodeRepositoryResponse.
    """
    body = CodeRepositoryResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        source=json.loads(self.source),
        logo_url=self.logo_url,
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = CodeRepositoryResponseMetadata(
            config=json.loads(self.config),
            description=self.description,
        )

    resources = None
    if include_resources:
        resources = CodeRepositoryResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return CodeRepositoryResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: CodeRepositoryUpdate) -> CodeRepositorySchema

Update a CodeRepositorySchema with a CodeRepositoryUpdate.

Parameters:

Name	Type	Description	Default
update	CodeRepositoryUpdate	The update model.	required

Returns:

Type	Description
CodeRepositorySchema	The updated CodeRepositorySchema.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

def update(self, update: "CodeRepositoryUpdate") -> "CodeRepositorySchema":
    """Update a `CodeRepositorySchema` with a `CodeRepositoryUpdate`.

    Args:
        update: The update model.

    Returns:
        The updated `CodeRepositorySchema`.
    """
    if update.name:
        self.name = update.name

    if update.description:
        self.description = update.description

    if update.logo_url:
        self.logo_url = update.logo_url

    if update.config:
        self.config = json.dumps(update.config)

    self.updated = utc_now()
    return self

CuratedVisualizationSchema

Bases: BaseSchema

SQL Model for curated visualizations.

Functions

from_request(request: CuratedVisualizationRequest) -> CuratedVisualizationSchema classmethod

Convert a request into a schema instance.

Parameters:

Name	Type	Description	Default
request	CuratedVisualizationRequest	The request to convert.	required

Returns:

Type	Description
CuratedVisualizationSchema	The created schema.

Source code in src/zenml/zen_stores/schemas/curated_visualization_schemas.py

@classmethod
def from_request(
    cls, request: CuratedVisualizationRequest
) -> "CuratedVisualizationSchema":
    """Convert a request into a schema instance.

    Args:
        request: The request to convert.

    Returns:
        The created schema.
    """
    return cls(
        project_id=request.project,
        artifact_visualization_id=request.artifact_visualization_id,
        display_name=request.display_name,
        display_order=request.display_order,
        layout_size=request.layout_size.value,
        resource_id=request.resource_id,
        resource_type=request.resource_type.value,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/curated_visualization_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options: List[ExecutableOption] = []

    if include_resources:
        options.append(selectinload(jl_arg(cls.artifact_visualization)))

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> CuratedVisualizationResponse

Convert schema into response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
CuratedVisualizationResponse	The created response model.

Source code in src/zenml/zen_stores/schemas/curated_visualization_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> CuratedVisualizationResponse:
    """Convert schema into response model.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        **kwargs: Additional keyword arguments.

    Returns:
        The created response model.
    """
    try:
        layout_size_enum = CuratedVisualizationSize(self.layout_size)
    except ValueError:
        layout_size_enum = CuratedVisualizationSize.FULL_WIDTH

    try:
        resource_type_enum = VisualizationResourceTypes(self.resource_type)
    except ValueError:
        resource_type_enum = VisualizationResourceTypes.PROJECT

    artifact_version_id = self.artifact_visualization.artifact_version_id

    body = CuratedVisualizationResponseBody(
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        artifact_visualization_id=self.artifact_visualization_id,
        artifact_version_id=artifact_version_id,
        display_name=self.display_name,
        display_order=self.display_order,
        layout_size=layout_size_enum,
        resource_id=self.resource_id,
        resource_type=resource_type_enum,
    )

    metadata = None
    if include_metadata:
        metadata = CuratedVisualizationResponseMetadata()

    resources = None
    if include_resources:
        artifact_visualization = self.artifact_visualization.to_model(
            include_metadata=False,
            include_resources=False,
        )
        resources = CuratedVisualizationResponseResources(
            artifact_visualization=artifact_visualization,
        )

    return CuratedVisualizationResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: CuratedVisualizationUpdate) -> CuratedVisualizationSchema

Update a schema instance from an update model.

Parameters:

Name	Type	Description	Default
update	CuratedVisualizationUpdate	The update definition.	required

Returns:

Type	Description
CuratedVisualizationSchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/curated_visualization_schemas.py

def update(
    self,
    update: CuratedVisualizationUpdate,
) -> "CuratedVisualizationSchema":
    """Update a schema instance from an update model.

    Args:
        update: The update definition.

    Returns:
        The updated schema.
    """
    changes = update.model_dump(exclude_unset=True)
    layout_size_update = changes.pop("layout_size", None)
    if layout_size_update is not None:
        self.layout_size = layout_size_update.value

    for field, value in changes.items():
        if hasattr(self, field):
            setattr(self, field, value)

    from zenml.utils.time_utils import utc_now

    self.updated = utc_now()
    return self

DeploymentSchema

Bases: NamedSchema

SQL Model for pipeline deployment.

Functions

from_request(request: DeploymentRequest) -> DeploymentSchema classmethod

Convert a DeploymentRequest to a DeploymentSchema.

Parameters:

Name	Type	Description	Default
request	DeploymentRequest	The request model to convert.	required

Returns:

Type	Description
DeploymentSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/deployment_schemas.py

@classmethod
def from_request(cls, request: DeploymentRequest) -> "DeploymentSchema":
    """Convert a `DeploymentRequest` to a `DeploymentSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        status=DeploymentStatus.UNKNOWN.value,
        snapshot_id=request.snapshot_id,
        deployer_id=request.deployer_id,
        auth_key=request.auth_key,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/deployment_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(DeploymentSchema.user)),
                selectinload(jl_arg(DeploymentSchema.deployer)),
                selectinload(jl_arg(DeploymentSchema.snapshot)).joinedload(
                    jl_arg(PipelineSnapshotSchema.pipeline)
                ),
                selectinload(jl_arg(DeploymentSchema.snapshot)).joinedload(
                    jl_arg(PipelineSnapshotSchema.stack)
                ),
                selectinload(jl_arg(DeploymentSchema.visualizations)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> DeploymentResponse

Convert a DeploymentSchema to a DeploymentResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
DeploymentResponse	The created DeploymentResponse.

Source code in src/zenml/zen_stores/schemas/deployment_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> DeploymentResponse:
    """Convert a `DeploymentSchema` to a `DeploymentResponse`.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        kwargs: Additional keyword arguments.

    Returns:
        The created `DeploymentResponse`.
    """
    status: Optional[DeploymentStatus] = None
    if self.status in DeploymentStatus.values():
        status = DeploymentStatus(self.status)
    elif self.status is not None:
        status = DeploymentStatus.UNKNOWN
        logger.warning(
            f"Deployment status '{self.status}' used for deployment "
            f"{self.name} is not a valid DeploymentStatus value. "
            "Using UNKNOWN instead."
        )

    body = DeploymentResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        url=self.url,
        status=status,
    )

    metadata = None
    if include_metadata:
        metadata = DeploymentResponseMetadata(
            deployment_metadata=json.loads(self.deployment_metadata),
            auth_key=self.auth_key,
        )

    resources = None
    if include_resources:
        resources = DeploymentResponseResources(
            user=self.user.to_model() if self.user else None,
            tags=[tag.to_model() for tag in self.tags],
            snapshot=self.snapshot.to_model() if self.snapshot else None,
            deployer=self.deployer.to_model() if self.deployer else None,
            pipeline=self.snapshot.pipeline.to_model()
            if self.snapshot and self.snapshot.pipeline
            else None,
            stack=self.snapshot.stack.to_model()
            if self.snapshot and self.snapshot.stack
            else None,
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    return DeploymentResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: DeploymentUpdate) -> DeploymentSchema

Updates a DeploymentSchema from a DeploymentUpdate.

Parameters:

Name	Type	Description	Default
update	DeploymentUpdate	The DeploymentUpdate to update from.	required

Returns:

Type	Description
DeploymentSchema	The updated DeploymentSchema.

Source code in src/zenml/zen_stores/schemas/deployment_schemas.py

def update(
    self,
    update: DeploymentUpdate,
) -> "DeploymentSchema":
    """Updates a `DeploymentSchema` from a `DeploymentUpdate`.

    Args:
        update: The `DeploymentUpdate` to update from.

    Returns:
        The updated `DeploymentSchema`.
    """
    for field, value in update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field == "deployment_metadata":
            setattr(self, field, json.dumps(value))
        elif hasattr(self, field):
            setattr(self, field, value)

    self.updated = utc_now()
    return self

FlavorSchema

Bases: NamedSchema

SQL Model for flavors.

Attributes:

Name	Type	Description
type	str	The type of the flavor.
source	str	The source of the flavor.
config_schema	str	The config schema of the flavor.
integration	Optional[str]	The integration associated with the flavor.

Functions

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/flavor_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(FlavorSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> FlavorResponse

Converts a flavor schema to a flavor model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
FlavorResponse	The flavor model.

Source code in src/zenml/zen_stores/schemas/flavor_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "FlavorResponse":
    """Converts a flavor schema to a flavor model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The flavor model.
    """
    body = FlavorResponseBody(
        user_id=self.user_id,
        type=StackComponentType(self.type),
        display_name=self.display_name
        or self.name.replace("_", " ").title(),
        integration=self.integration,
        source=self.source,
        logo_url=self.logo_url,
        is_custom=self.is_custom,
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = FlavorResponseMetadata(
            config_schema=json.loads(self.config_schema),
            connector_type=self.connector_type,
            connector_resource_type=self.connector_resource_type,
            connector_resource_id_attr=self.connector_resource_id_attr,
            docs_url=self.docs_url,
            sdk_docs_url=self.sdk_docs_url,
        )
    resources = None
    if include_resources:
        resources = FlavorResponseResources(
            user=self.user.to_model() if self.user else None,
        )
    return FlavorResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(flavor_update: FlavorUpdate) -> FlavorSchema

Update a FlavorSchema from a FlavorUpdate.

Parameters:

Name	Type	Description	Default
flavor_update	FlavorUpdate	The FlavorUpdate from which to update the schema.	required

Returns:

Type	Description
FlavorSchema	The updated FlavorSchema.

Source code in src/zenml/zen_stores/schemas/flavor_schemas.py

def update(
    self,
    flavor_update: "FlavorUpdate",
) -> "FlavorSchema":
    """Update a `FlavorSchema` from a `FlavorUpdate`.

    Args:
        flavor_update: The `FlavorUpdate` from which to update the schema.

    Returns:
        The updated `FlavorSchema`.
    """
    for field, value in flavor_update.model_dump(
        exclude_unset=True, exclude={"user"}
    ).items():
        if field == "config_schema":
            setattr(self, field, json.dumps(value))
        elif field == "type":
            setattr(self, field, value.value)
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

LogsSchema

Bases: BaseSchema

SQL Model for logs.

Functions

from_request(request: LogsRequest) -> LogsSchema classmethod

Create a LogsSchema from a LogsRequest.

Parameters:

Name	Type	Description	Default
request	LogsRequest	The LogsRequest to create the LogsSchema from.	required

Returns:

Type	Description
LogsSchema	The created LogsSchema.

Source code in src/zenml/zen_stores/schemas/logs_schemas.py

@classmethod
def from_request(cls, request: LogsRequest) -> "LogsSchema":
    """Create a `LogsSchema` from a `LogsRequest`.

    Args:
        request: The `LogsRequest` to create the `LogsSchema` from.

    Returns:
        The created `LogsSchema`.
    """
    return LogsSchema(
        id=request.id,
        uri=request.uri,
        source=request.source,
        project_id=request.project,
        user_id=request.user,
        pipeline_run_id=request.pipeline_run_id,
        step_run_id=request.step_run_id,
        artifact_store_id=request.artifact_store_id,
        log_store_id=request.log_store_id,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> LogsResponse

Convert a LogsSchema to a LogsResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
LogsResponse	The created LogsResponse.

Source code in src/zenml/zen_stores/schemas/logs_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "LogsResponse":
    """Convert a `LogsSchema` to a `LogsResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The created `LogsResponse`.
    """
    body = LogsResponseBody(
        uri=self.uri,
        source=self.source,
        created=self.created,
        updated=self.updated,
        project_id=self.project_id,
        user_id=self.user_id,
    )

    metadata = None
    if include_metadata:
        metadata = LogsResponseMetadata(
            step_run_id=self.step_run_id,
            pipeline_run_id=self.pipeline_run_id,
            artifact_store_id=self.artifact_store_id,
            log_store_id=self.log_store_id,
        )

    resources = None
    if include_resources:
        resources = LogsResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return LogsResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

ModelSchema

Bases: NamedSchema

SQL Model for model.

Attributes

latest_version: Optional[ModelVersionSchema] property

Fetch the latest version for this model.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[ModelVersionSchema]	The latest version for this model.

Functions

from_request(model_request: ModelRequest) -> ModelSchema classmethod

Convert an ModelRequest to an ModelSchema.

Parameters:

Name	Type	Description	Default
model_request	ModelRequest	The request model to convert.	required

Returns:

Type	Description
ModelSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def from_request(cls, model_request: ModelRequest) -> "ModelSchema":
    """Convert an `ModelRequest` to an `ModelSchema`.

    Args:
        model_request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=model_request.name,
        project_id=model_request.project,
        user_id=model_request.user,
        license=model_request.license,
        description=model_request.description,
        audience=model_request.audience,
        use_cases=model_request.use_cases,
        limitations=model_request.limitations,
        trade_offs=model_request.trade_offs,
        ethics=model_request.ethics,
        save_models_to_registry=model_request.save_models_to_registry,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ModelSchema.user)),
                # joinedload(jl_arg(ModelSchema.tags)),
                selectinload(jl_arg(ModelSchema.visualizations)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ModelResponse

Convert an ModelSchema to an ModelResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ModelResponse	The created ModelResponse.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ModelResponse:
    """Convert an `ModelSchema` to an `ModelResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ModelResponse`.
    """
    metadata = None
    if include_metadata:
        metadata = ModelResponseMetadata(
            license=self.license,
            description=self.description,
            audience=self.audience,
            use_cases=self.use_cases,
            limitations=self.limitations,
            trade_offs=self.trade_offs,
            ethics=self.ethics,
            save_models_to_registry=self.save_models_to_registry,
        )

    resources = None
    if include_resources:
        if latest_version := self.latest_version:
            latest_version_name = latest_version.name
            latest_version_id = latest_version.id
        else:
            latest_version_name = None
            latest_version_id = None

        resources = ModelResponseResources(
            user=self.user.to_model() if self.user else None,
            tags=[tag.to_model() for tag in self.tags],
            latest_version_name=latest_version_name,
            latest_version_id=latest_version_id,
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    body = ModelResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
    )

    return ModelResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(model_update: ModelUpdate) -> ModelSchema

Updates a ModelSchema from a ModelUpdate.

Parameters:

Name	Type	Description	Default
model_update	ModelUpdate	The ModelUpdate to update from.	required

Returns:

Type	Description
ModelSchema	The updated ModelSchema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def update(
    self,
    model_update: ModelUpdate,
) -> "ModelSchema":
    """Updates a `ModelSchema` from a `ModelUpdate`.

    Args:
        model_update: The `ModelUpdate` to update from.

    Returns:
        The updated `ModelSchema`.
    """
    for field, value in model_update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field in ["add_tags", "remove_tags"]:
            # Tags are handled separately
            continue
        setattr(self, field, value)
    self.updated = utc_now()
    return self

ModelVersionArtifactSchema

Bases: BaseSchema

SQL Model for linking of Model Versions and Artifacts M:M.

Functions

from_request(model_version_artifact_request: ModelVersionArtifactRequest) -> ModelVersionArtifactSchema classmethod

Convert an ModelVersionArtifactRequest to a ModelVersionArtifactSchema.

Parameters:

Name	Type	Description	Default
model_version_artifact_request	ModelVersionArtifactRequest	The request link to convert.	required

Returns:

Type	Description
ModelVersionArtifactSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def from_request(
    cls,
    model_version_artifact_request: ModelVersionArtifactRequest,
) -> "ModelVersionArtifactSchema":
    """Convert an `ModelVersionArtifactRequest` to a `ModelVersionArtifactSchema`.

    Args:
        model_version_artifact_request: The request link to convert.

    Returns:
        The converted schema.
    """
    return cls(
        model_version_id=model_version_artifact_request.model_version,
        artifact_version_id=model_version_artifact_request.artifact_version,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ModelVersionArtifactResponse

Convert an ModelVersionArtifactSchema to an ModelVersionArtifactResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ModelVersionArtifactResponse	The created ModelVersionArtifactResponseModel.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ModelVersionArtifactResponse:
    """Convert an `ModelVersionArtifactSchema` to an `ModelVersionArtifactResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ModelVersionArtifactResponseModel`.
    """
    return ModelVersionArtifactResponse(
        id=self.id,
        body=ModelVersionArtifactResponseBody(
            created=self.created,
            updated=self.updated,
            model_version=self.model_version_id,
            artifact_version=self.artifact_version.to_model(),
        ),
        metadata=BaseResponseMetadata() if include_metadata else None,
    )

ModelVersionPipelineRunSchema

Bases: BaseSchema

SQL Model for linking of Model Versions and Pipeline Runs M:M.

Functions

from_request(model_version_pipeline_run_request: ModelVersionPipelineRunRequest) -> ModelVersionPipelineRunSchema classmethod

Convert an ModelVersionPipelineRunRequest to an ModelVersionPipelineRunSchema.

Parameters:

Name	Type	Description	Default
model_version_pipeline_run_request	ModelVersionPipelineRunRequest	The request link to convert.	required

Returns:

Type	Description
ModelVersionPipelineRunSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def from_request(
    cls,
    model_version_pipeline_run_request: ModelVersionPipelineRunRequest,
) -> "ModelVersionPipelineRunSchema":
    """Convert an `ModelVersionPipelineRunRequest` to an `ModelVersionPipelineRunSchema`.

    Args:
        model_version_pipeline_run_request: The request link to convert.

    Returns:
        The converted schema.
    """
    return cls(
        model_version_id=model_version_pipeline_run_request.model_version,
        pipeline_run_id=model_version_pipeline_run_request.pipeline_run,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ModelVersionPipelineRunResponse

Convert an ModelVersionPipelineRunSchema to an ModelVersionPipelineRunResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ModelVersionPipelineRunResponse	The created ModelVersionPipelineRunResponse.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ModelVersionPipelineRunResponse:
    """Convert an `ModelVersionPipelineRunSchema` to an `ModelVersionPipelineRunResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ModelVersionPipelineRunResponse`.
    """
    return ModelVersionPipelineRunResponse(
        id=self.id,
        body=ModelVersionPipelineRunResponseBody(
            created=self.created,
            updated=self.updated,
            model_version=self.model_version_id,
            pipeline_run=self.pipeline_run.to_model(),
        ),
        metadata=BaseResponseMetadata() if include_metadata else None,
    )

ModelVersionSchema

Bases: NamedSchema, RunMetadataInterface

SQL Model for model version.

Functions

from_request(model_version_request: ModelVersionRequest, model_version_number: int, producer_run_id: Optional[UUID] = None) -> ModelVersionSchema classmethod

Convert an ModelVersionRequest to an ModelVersionSchema.

Parameters:

Name	Type	Description	Default
model_version_request	ModelVersionRequest	The request model version to convert.	required
model_version_number	int	The model version number.	required
producer_run_id	Optional[UUID]	The ID of the producer run.	None

Returns:

Type	Description
ModelVersionSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def from_request(
    cls,
    model_version_request: ModelVersionRequest,
    model_version_number: int,
    producer_run_id: Optional[UUID] = None,
) -> "ModelVersionSchema":
    """Convert an `ModelVersionRequest` to an `ModelVersionSchema`.

    Args:
        model_version_request: The request model version to convert.
        model_version_number: The model version number.
        producer_run_id: The ID of the producer run.

    Returns:
        The converted schema.
    """
    id_ = uuid4()
    is_numeric = str(model_version_number) == model_version_request.name

    return cls(
        id=id_,
        project_id=model_version_request.project,
        user_id=model_version_request.user,
        model_id=model_version_request.model,
        name=model_version_request.name,
        number=model_version_number,
        description=model_version_request.description,
        stage=model_version_request.stage,
        producer_run_id_if_numeric=producer_run_id
        if (producer_run_id and is_numeric)
        else id_,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = [
        joinedload(jl_arg(ModelVersionSchema.model), innerjoin=True),
    ]

    # if include_metadata:
    #     options.extend(
    #         [
    #             joinedload(jl_arg(ModelVersionSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ModelVersionSchema.user)),
                # joinedload(jl_arg(ModelVersionSchema.services)),
                # joinedload(jl_arg(ModelVersionSchema.tags)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ModelVersionResponse

Convert an ModelVersionSchema to an ModelVersionResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ModelVersionResponse	The created ModelVersionResponse.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ModelVersionResponse:
    """Convert an `ModelVersionSchema` to an `ModelVersionResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ModelVersionResponse`.
    """
    from zenml.models import ServiceResponse

    metadata = None
    if include_metadata:
        metadata = ModelVersionResponseMetadata(
            description=self.description,
            run_metadata=self.fetch_metadata(),
        )

    resources = None
    if include_resources:
        services = cast(
            Page[ServiceResponse],
            get_page_from_list(
                items_list=self.services,
                response_model=ServiceResponse,
                include_resources=include_resources,
                include_metadata=include_metadata,
            ),
        )
        resources = ModelVersionResponseResources(
            user=self.user.to_model() if self.user else None,
            services=services,
            tags=[tag.to_model() for tag in self.tags],
        )

    body = ModelVersionResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        stage=self.stage,
        number=self.number,
        model=self.model.to_model(),
    )

    return ModelVersionResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(target_stage: Optional[str] = None, target_name: Optional[str] = None, target_description: Optional[str] = None) -> ModelVersionSchema

Updates a ModelVersionSchema to a target stage.

Parameters:

Name	Type	Description	Default
target_stage	Optional[str]	The stage to be updated.	None
target_name	Optional[str]	The version name to be updated.	None
target_description	Optional[str]	The version description to be updated.	None

Returns:

Type	Description
ModelVersionSchema	The updated ModelVersionSchema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def update(
    self,
    target_stage: Optional[str] = None,
    target_name: Optional[str] = None,
    target_description: Optional[str] = None,
) -> "ModelVersionSchema":
    """Updates a `ModelVersionSchema` to a target stage.

    Args:
        target_stage: The stage to be updated.
        target_name: The version name to be updated.
        target_description: The version description to be updated.

    Returns:
        The updated `ModelVersionSchema`.
    """
    if target_stage is not None:
        self.stage = target_stage
    if target_name is not None:
        self.name = target_name
    if target_description is not None:
        self.description = target_description
    self.updated = utc_now()
    return self

NamedSchema

Bases: BaseSchema

Base Named SQL Model.

OAuthDeviceSchema

Bases: BaseSchema

SQL Model for authorized OAuth2 devices.

Functions

from_request(request: OAuthDeviceInternalRequest) -> Tuple[OAuthDeviceSchema, str, str] classmethod

Create an authorized device DB entry from a device authorization request.

Parameters:

Name	Type	Description	Default
request	OAuthDeviceInternalRequest	The device authorization request.	required

Returns:

Type	Description
Tuple[OAuthDeviceSchema, str, str]	The created OAuthDeviceSchema, the user code and the device code.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

@classmethod
def from_request(
    cls, request: OAuthDeviceInternalRequest
) -> Tuple["OAuthDeviceSchema", str, str]:
    """Create an authorized device DB entry from a device authorization request.

    Args:
        request: The device authorization request.

    Returns:
        The created `OAuthDeviceSchema`, the user code and the device code.
    """
    user_code = cls._generate_user_code()
    device_code = cls._generate_device_code()
    hashed_user_code = cls._get_hashed_code(user_code)
    hashed_device_code = cls._get_hashed_code(device_code)
    now = utc_now()
    return (
        cls(
            client_id=request.client_id,
            user_code=hashed_user_code,
            device_code=hashed_device_code,
            status=OAuthDeviceStatus.PENDING.value,
            failed_auth_attempts=0,
            expires=now + timedelta(seconds=request.expires_in),
            os=request.os,
            ip_address=request.ip_address,
            hostname=request.hostname,
            python_version=request.python_version,
            zenml_version=request.zenml_version,
            city=request.city,
            region=request.region,
            country=request.country,
            created=now,
            updated=now,
        ),
        user_code,
        device_code,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(OAuthDeviceSchema.user)),
            ]
        )

    return options

internal_update(device_update: OAuthDeviceInternalUpdate) -> Tuple[OAuthDeviceSchema, Optional[str], Optional[str]]

Update an authorized device from an internal device update model.

Parameters:

Name	Type	Description	Default
device_update	OAuthDeviceInternalUpdate	The internal device update model.	required

Returns:

Type	Description
OAuthDeviceSchema	The updated OAuthDeviceSchema and the new user code and device
Optional[str]	code, if they were generated.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

def internal_update(
    self, device_update: OAuthDeviceInternalUpdate
) -> Tuple["OAuthDeviceSchema", Optional[str], Optional[str]]:
    """Update an authorized device from an internal device update model.

    Args:
        device_update: The internal device update model.

    Returns:
        The updated `OAuthDeviceSchema` and the new user code and device
        code, if they were generated.
    """
    now = utc_now()
    user_code: Optional[str] = None
    device_code: Optional[str] = None

    # This call also takes care of setting fields that have the same
    # name in the internal model and the schema.
    self.update(device_update)

    if device_update.expires_in is not None:
        if device_update.expires_in <= 0:
            self.expires = None
        else:
            self.expires = now + timedelta(
                seconds=device_update.expires_in
            )
    if device_update.update_last_login:
        self.last_login = now
    if device_update.generate_new_codes:
        user_code = self._generate_user_code()
        device_code = self._generate_device_code()
        self.user_code = self._get_hashed_code(user_code)
        self.device_code = self._get_hashed_code(device_code)
    self.updated = now
    return self, user_code, device_code

to_internal_model(include_metadata: bool = False, include_resources: bool = False) -> OAuthDeviceInternalResponse

Convert a device schema to an internal device response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False

Returns:

Type	Description
OAuthDeviceInternalResponse	The converted internal device response model.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

def to_internal_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
) -> OAuthDeviceInternalResponse:
    """Convert a device schema to an internal device response model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.

    Returns:
        The converted internal device response model.
    """
    device_model = self.to_model(
        include_metadata=include_metadata,
        include_resources=include_resources,
    )
    return OAuthDeviceInternalResponse(
        id=device_model.id,
        body=device_model.body,
        metadata=device_model.metadata,
        resources=device_model.resources,
        user_code=self.user_code,
        device_code=self.device_code,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> OAuthDeviceResponse

Convert a device schema to a device response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
OAuthDeviceResponse	The converted device response model.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> OAuthDeviceResponse:
    """Convert a device schema to a device response model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The converted device response model.
    """
    metadata = None
    if include_metadata:
        metadata = OAuthDeviceResponseMetadata(
            python_version=self.python_version,
            zenml_version=self.zenml_version,
            city=self.city,
            region=self.region,
            country=self.country,
            failed_auth_attempts=self.failed_auth_attempts,
            last_login=self.last_login,
        )

    body = OAuthDeviceResponseBody(
        user_id=self.user_id,
        created=self.created,
        updated=self.updated,
        client_id=self.client_id,
        expires=self.expires,
        trusted_device=self.trusted_device,
        status=OAuthDeviceStatus(self.status),
        os=self.os,
        ip_address=self.ip_address,
        hostname=self.hostname,
    )
    resources = None
    if include_resources:
        resources = OAuthDeviceResponseResources(
            user=self.user.to_model() if self.user else None,
        )
    return OAuthDeviceResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(device_update: OAuthDeviceUpdate) -> OAuthDeviceSchema

Update an authorized device from a device update model.

Parameters:

Name	Type	Description	Default
device_update	OAuthDeviceUpdate	The device update model.	required

Returns:

Type	Description
OAuthDeviceSchema	The updated OAuthDeviceSchema.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

def update(self, device_update: OAuthDeviceUpdate) -> "OAuthDeviceSchema":
    """Update an authorized device from a device update model.

    Args:
        device_update: The device update model.

    Returns:
        The updated `OAuthDeviceSchema`.
    """
    for field, value in device_update.model_dump(
        exclude_none=True
    ).items():
        if hasattr(self, field):
            setattr(self, field, value)

    if device_update.locked is True:
        self.status = OAuthDeviceStatus.LOCKED.value
    elif device_update.locked is False:
        self.status = OAuthDeviceStatus.ACTIVE.value

    self.updated = utc_now()
    return self

PipelineBuildSchema

Bases: BaseSchema

SQL Model for pipeline builds.

Functions

from_request(request: PipelineBuildRequest) -> PipelineBuildSchema classmethod

Convert a PipelineBuildRequest to a PipelineBuildSchema.

Parameters:

Name	Type	Description	Default
request	PipelineBuildRequest	The request to convert.	required

Returns:

Type	Description
PipelineBuildSchema	The created PipelineBuildSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_build_schemas.py

@classmethod
def from_request(
    cls, request: PipelineBuildRequest
) -> "PipelineBuildSchema":
    """Convert a `PipelineBuildRequest` to a `PipelineBuildSchema`.

    Args:
        request: The request to convert.

    Returns:
        The created `PipelineBuildSchema`.
    """
    return cls(
        stack_id=request.stack,
        project_id=request.project,
        user_id=request.user,
        pipeline_id=request.pipeline,
        images=json.dumps(request.images, default=pydantic_encoder),
        is_local=request.is_local,
        contains_code=request.contains_code,
        zenml_version=request.zenml_version,
        python_version=request.python_version,
        checksum=request.checksum,
        stack_checksum=request.stack_checksum,
        duration=request.duration,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/pipeline_build_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_metadata:
        options.extend(
            [
                joinedload(jl_arg(PipelineBuildSchema.pipeline)),
                joinedload(jl_arg(PipelineBuildSchema.stack)),
            ]
        )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(PipelineBuildSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> PipelineBuildResponse

Convert a PipelineBuildSchema to a PipelineBuildResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
PipelineBuildResponse	The created PipelineBuildResponse.

Source code in src/zenml/zen_stores/schemas/pipeline_build_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> PipelineBuildResponse:
    """Convert a `PipelineBuildSchema` to a `PipelineBuildResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `PipelineBuildResponse`.
    """
    body = PipelineBuildResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = PipelineBuildResponseMetadata(
            pipeline=self.pipeline.to_model() if self.pipeline else None,
            stack=self.stack.to_model() if self.stack else None,
            images=json.loads(self.images),
            zenml_version=self.zenml_version,
            python_version=self.python_version,
            checksum=self.checksum,
            stack_checksum=self.stack_checksum,
            is_local=self.is_local,
            contains_code=self.contains_code,
            duration=self.duration,
        )

    resources = None
    if include_resources:
        resources = PipelineBuildResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return PipelineBuildResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

PipelineRunOutputSchema

Bases: BaseSchema

SQL model defining pipeline run outputs.

PipelineRunSchema

Bases: NamedSchema, RunMetadataInterface

SQL Model for pipeline runs.

Functions

fetch_metadata_collection(include_full_metadata: bool = False, **kwargs: Any) -> Dict[str, List[RunMetadataEntry]]

Fetches all the metadata entries related to the pipeline run.

Parameters:

Name	Type	Description	Default
include_full_metadata	bool	Whether the full metadata will be included.	False
**kwargs	Any	Keyword arguments.	{}

Returns:

Type	Description
Dict[str, List[RunMetadataEntry]]	a dictionary, where the key is the key of the metadata entry and the values represent the list of entries with this key.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def fetch_metadata_collection(
    self, include_full_metadata: bool = False, **kwargs: Any
) -> Dict[str, List[RunMetadataEntry]]:
    """Fetches all the metadata entries related to the pipeline run.

    Args:
        include_full_metadata: Whether the full metadata will be included.
        **kwargs: Keyword arguments.

    Returns:
        a dictionary, where the key is the key of the metadata entry
            and the values represent the list of entries with this key.
    """
    # Fetch the metadata related to this run
    metadata_collection = super().fetch_metadata_collection(**kwargs)

    if include_full_metadata:
        # Fetch the metadata related to the steps of this run
        for s in self.step_runs:
            step_metadata = s.fetch_metadata_collection()
            for k, v in step_metadata.items():
                metadata_collection[f"{s.name}::{k}"] = v

        # Fetch the metadata related to the schedule of this run
        if self.snapshot is not None:
            if schedule := self.snapshot.schedule:
                schedule_metadata = schedule.fetch_metadata_collection()
                for k, v in schedule_metadata.items():
                    metadata_collection[f"schedule:{k}"] = v

    return metadata_collection

from_request(request: PipelineRunRequest, pipeline_id: UUID, index: int, enable_heartbeat: bool, root_run_id: Optional[UUID] = None) -> PipelineRunSchema classmethod

Convert a PipelineRunRequest to a PipelineRunSchema.

Parameters:

Name	Type	Description	Default
request	PipelineRunRequest	The request to convert.	required
pipeline_id	UUID	The ID of the pipeline.	required
index	int	The index of the pipeline run.	required
enable_heartbeat	bool	Whether the heartbeat should be enabled.	required
root_run_id	Optional[UUID]	The root_run_id of the parent run, if this run is a child run.	None

Returns:

Type	Description
PipelineRunSchema	The created PipelineRunSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

@classmethod
def from_request(
    cls,
    request: "PipelineRunRequest",
    pipeline_id: UUID,
    index: int,
    enable_heartbeat: bool,
    root_run_id: Optional[UUID] = None,
) -> "PipelineRunSchema":
    """Convert a `PipelineRunRequest` to a `PipelineRunSchema`.

    Args:
        request: The request to convert.
        pipeline_id: The ID of the pipeline.
        index: The index of the pipeline run.
        enable_heartbeat: Whether the heartbeat should be enabled.
        root_run_id: The `root_run_id` of the parent run, if this
            run is a child run.

    Returns:
        The created `PipelineRunSchema`.
    """
    orchestrator_environment = json.dumps(request.orchestrator_environment)
    if len(orchestrator_environment) > TEXT_FIELD_MAX_LENGTH:
        logger.warning(
            "Orchestrator environment is too large to be stored in the "
            "database. Skipping."
        )
        orchestrator_environment = "{}"

    triggered_by = None
    triggered_by_type = None
    if request.trigger_info:
        if request.trigger_info.step_run_id:
            triggered_by = request.trigger_info.step_run_id
            triggered_by_type = PipelineRunTriggeredByType.STEP_RUN.value
        elif request.trigger_info.deployment_id:
            triggered_by = request.trigger_info.deployment_id
            triggered_by_type = PipelineRunTriggeredByType.DEPLOYMENT.value

    return cls(
        project_id=request.project,
        user_id=request.user,
        name=request.name,
        orchestrator_run_id=request.orchestrator_run_id,
        orchestrator_environment=orchestrator_environment,
        start_time=request.start_time,
        end_time=request.end_time,
        status=request.status.value,
        index=index,
        in_progress=not request.status.is_finished,
        status_reason=request.status_reason,
        pipeline_id=pipeline_id,
        snapshot_id=request.snapshot,
        triggered_by=triggered_by,
        triggered_by_type=triggered_by_type,
        enable_heartbeat=enable_heartbeat,
        exception_info=request.exception_info.model_dump_json()
        if request.exception_info
        else None,
        original_run_id=request.original_run_id,
        parent_run_id=request.parent_run_id,
        child_key=request.child_key,
        root_run_id=root_run_id,
    )

get_pipeline_configuration() -> PipelineConfiguration

Get the pipeline configuration for the pipeline run.

Raises:

Type	Description
RuntimeError	if the pipeline run has no snapshot and no pipeline configuration.

Returns:

Type	Description
PipelineConfiguration	The pipeline configuration.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def get_pipeline_configuration(self) -> PipelineConfiguration:
    """Get the pipeline configuration for the pipeline run.

    Raises:
        RuntimeError: if the pipeline run has no snapshot and no pipeline
            configuration.

    Returns:
        The pipeline configuration.
    """
    if self.snapshot:
        pipeline_config = PipelineConfiguration.model_validate_json(
            self.snapshot.pipeline_configuration
        )
    elif self.pipeline_configuration:
        pipeline_config = PipelineConfiguration.model_validate_json(
            self.pipeline_configuration
        )
    else:
        raise RuntimeError(
            "Pipeline run has no snapshot and no pipeline configuration."
        )

    pipeline_config.finalize_substitutions(
        start_time=self.start_time, inplace=True
    )
    return pipeline_config

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    from zenml.zen_stores.schemas import ModelVersionSchema

    options = []

    if include_metadata:
        options.extend(
            [
                selectinload(jl_arg(PipelineRunSchema.trigger_execution)),
            ]
        )

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(PipelineRunSchema.outputs)),
                selectinload(jl_arg(PipelineRunSchema.parent_run)),
                selectinload(
                    jl_arg(PipelineRunSchema.model_version)
                ).joinedload(
                    jl_arg(ModelVersionSchema.model), innerjoin=True
                ),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(
                    jl_arg(PipelineSnapshotSchema.source_snapshot)
                ),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.pipeline)),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.stack)),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.build)),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.schedule)),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(
                    jl_arg(PipelineSnapshotSchema.code_reference)
                ),
                selectinload(jl_arg(PipelineRunSchema.logs)),
                selectinload(jl_arg(PipelineRunSchema.wait_conditions)),
                selectinload(jl_arg(PipelineRunSchema.user)),
                selectinload(jl_arg(PipelineRunSchema.tags)),
                selectinload(jl_arg(PipelineRunSchema.visualizations)),
                joinedload(jl_arg(PipelineRunSchema.trigger)),
            ]
        )

    return options

get_step_configuration(step_name: str) -> Step

Get the step configuration for the pipeline run.

Parameters:

Name	Type	Description	Default
step_name	str	The name of the step to get the configuration for.	required

Raises:

Type	Description
RuntimeError	If the pipeline run has no snapshot.

Returns:

Type	Description
Step	The step configuration.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def get_step_configuration(self, step_name: str) -> Step:
    """Get the step configuration for the pipeline run.

    Args:
        step_name: The name of the step to get the configuration for.

    Raises:
        RuntimeError: If the pipeline run has no snapshot.

    Returns:
        The step configuration.
    """
    if self.snapshot:
        pipeline_configuration = self.get_pipeline_configuration()
        return Step.from_dict(
            data=json.loads(
                self.snapshot.get_step_configuration(step_name).config
            ),
            pipeline_configuration=pipeline_configuration,
        )
    else:
        raise RuntimeError("Pipeline run has no snapshot.")

get_upstream_steps() -> Dict[str, List[str]]

Get the list of all the upstream steps for each step.

Returns:

Type	Description
Dict[str, List[str]]	The list of upstream steps for each step.

Raises:

Type	Description
RuntimeError	If the pipeline run has no snapshot or the snapshot has no pipeline spec.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def get_upstream_steps(self) -> Dict[str, List[str]]:
    """Get the list of all the upstream steps for each step.

    Returns:
        The list of upstream steps for each step.

    Raises:
        RuntimeError: If the pipeline run has no snapshot or
            the snapshot has no pipeline spec.
    """
    if self.snapshot and self.snapshot.pipeline_spec:
        pipeline_spec = PipelineSpec.model_validate_json(
            self.snapshot.pipeline_spec
        )
        steps = {}
        for step_spec in pipeline_spec.steps:
            steps[step_spec.invocation_id] = step_spec.upstream_steps
        return steps
    else:
        raise RuntimeError("Pipeline run has no snapshot.")

is_placeholder_run() -> bool

Whether the pipeline run is a placeholder run.

Returns:

Type	Description
bool	Whether the pipeline run is a placeholder run.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def is_placeholder_run(self) -> bool:
    """Whether the pipeline run is a placeholder run.

    Returns:
        Whether the pipeline run is a placeholder run.
    """
    return self.status in {
        ExecutionStatus.INITIALIZING.value,
        ExecutionStatus.PROVISIONING.value,
    }

to_model(include_metadata: bool = False, include_resources: bool = False, include_python_packages: bool = False, include_full_metadata: bool = False, **kwargs: Any) -> PipelineRunResponse

Convert a PipelineRunSchema to a PipelineRunResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
include_python_packages	bool	Whether the python packages will be filled.	False
include_full_metadata	bool	Whether the full metadata will be included.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
PipelineRunResponse	The created PipelineRunResponse.

Raises:

Type	Description
RuntimeError	if the model creation fails.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    include_python_packages: bool = False,
    include_full_metadata: bool = False,
    **kwargs: Any,
) -> "PipelineRunResponse":
    """Convert a `PipelineRunSchema` to a `PipelineRunResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        include_python_packages: Whether the python packages will be filled.
        include_full_metadata: Whether the full metadata will be included.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `PipelineRunResponse`.

    Raises:
        RuntimeError: if the model creation fails.
    """
    if self.snapshot is not None:
        config = PipelineConfiguration.model_validate_json(
            self.snapshot.pipeline_configuration
        )
        client_environment = json.loads(self.snapshot.client_environment)
    elif self.pipeline_configuration is not None:
        config = PipelineConfiguration.model_validate_json(
            self.pipeline_configuration
        )
        client_environment = (
            json.loads(self.client_environment)
            if self.client_environment
            else {}
        )
    else:
        raise RuntimeError(
            "Pipeline run model creation has failed. Each pipeline run "
            "entry should either have a snapshot_id or "
            "pipeline_configuration."
        )

    config.finalize_substitutions(start_time=self.start_time, inplace=True)

    body = PipelineRunResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        status=ExecutionStatus(self.status),
        status_reason=self.status_reason,
        created=self.created,
        updated=self.updated,
        in_progress=self.in_progress,
        index=self.index,
        pipeline_id=self.pipeline_id,
        child_key=self.child_key,
        root_run_id=self.root_run_id,
    )
    metadata = None
    if include_metadata:
        is_templatable = False
        if (
            self.snapshot
            and self.snapshot.build
            and not self.snapshot.build.is_local
            and self.snapshot.build.stack_id
        ):
            is_templatable = True

        orchestrator_environment = (
            json.loads(self.orchestrator_environment)
            if self.orchestrator_environment
            else {}
        )

        if not include_python_packages:
            client_environment.pop("python_packages", None)
            orchestrator_environment.pop("python_packages", None)

        trigger_info: Optional[PipelineRunTriggerInfo] = None
        if self.triggered_by and self.triggered_by_type:
            if (
                self.triggered_by_type
                == PipelineRunTriggeredByType.STEP_RUN.value
            ):
                trigger_info = PipelineRunTriggerInfo(
                    step_run_id=self.triggered_by,
                )
            elif (
                self.triggered_by_type
                == PipelineRunTriggeredByType.DEPLOYMENT.value
            ):
                trigger_info = PipelineRunTriggerInfo(
                    deployment_id=self.triggered_by,
                )

        metadata = PipelineRunResponseMetadata(
            run_metadata=self.fetch_metadata(
                include_full_metadata=include_full_metadata
            ),
            config=config,
            start_time=self.start_time,
            end_time=self.end_time,
            client_environment=client_environment,
            orchestrator_environment=orchestrator_environment,
            orchestrator_run_id=self.orchestrator_run_id,
            code_path=self.snapshot.code_path if self.snapshot else None,
            template_id=self.snapshot.template_id
            if self.snapshot
            else None,
            is_templatable=is_templatable,
            trigger_info=trigger_info,
            enable_heartbeat=self.enable_heartbeat,
            exception_info=json.loads(self.exception_info)
            if self.exception_info
            else None,
            trigger_execution_info=json.loads(self.trigger_execution.info)
            if self.trigger_execution and self.trigger_execution.info
            else None,
        )

    resources = None
    if include_resources:
        if self.snapshot:
            source_snapshot = (
                self.snapshot.source_snapshot.to_model()
                if self.snapshot.source_snapshot
                else None
            )
            stack = (
                self.snapshot.stack.to_model()
                if self.snapshot.stack
                else None
            )
            pipeline: Optional["PipelineResponse"] = (
                self.snapshot.pipeline.to_model()
            )
            build = (
                self.snapshot.build.to_model()
                if self.snapshot.build
                else None
            )
            schedule = (
                self.snapshot.schedule.to_model()
                if self.snapshot.schedule
                else None
            )
            code_reference = (
                self.snapshot.code_reference.to_model()
                if self.snapshot.code_reference
                else None
            )
        else:
            source_snapshot = None
            stack = self.stack.to_model() if self.stack else None
            pipeline = self.pipeline.to_model() if self.pipeline else None
            build = self.build.to_model() if self.build else None
            schedule = self.schedule.to_model() if self.schedule else None
            code_reference = None

        resources = PipelineRunResponseResources(
            user=self.user.to_model() if self.user else None,
            snapshot=self.snapshot.to_model() if self.snapshot else None,
            source_snapshot=source_snapshot,
            stack=stack,
            pipeline=pipeline,
            build=build,
            schedule=schedule,
            code_reference=code_reference,
            model_version=self.model_version.to_model()
            if self.model_version
            else None,
            tags=[tag.to_model() for tag in self.tags],
            log_collection=[
                log.to_model()
                for log in sorted(self.logs, key=lambda log: log.created)
            ],
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
            trigger=self.trigger.to_model() if self.trigger else None,
            original_run=self.original_run.to_model()
            if self.original_run
            else None,
            parent_run=self.parent_run.to_model()
            if self.parent_run
            else None,
            active_wait_condition=next(
                (
                    condition.to_model()
                    for condition in self.wait_conditions
                    if condition.status
                    == RunWaitConditionStatus.PENDING.value
                ),
                None,
            ),
            outputs={
                output.name: output.artifact_version.to_model()
                for output in sorted(
                    self.outputs, key=lambda output: output.output_index
                )
            },
        )

    return PipelineRunResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(run_update: PipelineRunUpdate) -> PipelineRunSchema

Update a PipelineRunSchema with a PipelineRunUpdate.

Parameters:

Name	Type	Description	Default
run_update	PipelineRunUpdate	The PipelineRunUpdate to update with.	required

Raises:

Type	Description
ValueError	When trying to update the orchestrator run ID of a run that already has a different one.

Returns:

Type	Description
PipelineRunSchema	The updated PipelineRunSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def update(self, run_update: "PipelineRunUpdate") -> "PipelineRunSchema":
    """Update a `PipelineRunSchema` with a `PipelineRunUpdate`.

    Args:
        run_update: The `PipelineRunUpdate` to update with.

    Raises:
        ValueError: When trying to update the orchestrator run ID of a
            run that already has a different one.

    Returns:
        The updated `PipelineRunSchema`.
    """
    if run_update.orchestrator_run_id:
        if (
            self.orchestrator_run_id
            and self.orchestrator_run_id != run_update.orchestrator_run_id
        ):
            raise ValueError(
                "Updating the orchestrator run ID of a run with an "
                "existing orchestrator run ID "
                f"({self.orchestrator_run_id}) is not allowed."
            )
        self.orchestrator_run_id = run_update.orchestrator_run_id

    if run_update.exception_info:
        self.exception_info = run_update.exception_info.model_dump_json()

    self.updated = utc_now()
    return self

update_placeholder(request: PipelineRunRequest) -> PipelineRunSchema

Update a placeholder run.

Parameters:

Name	Type	Description	Default
request	PipelineRunRequest	The pipeline run request which should replace the placeholder.	required

Raises:

Type	Description
RuntimeError	If the DB entry does not represent a placeholder run.
ValueError	If the run request is not a valid request to replace the placeholder run.

Returns:

Type	Description
PipelineRunSchema	The updated PipelineRunSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def update_placeholder(
    self, request: "PipelineRunRequest"
) -> "PipelineRunSchema":
    """Update a placeholder run.

    Args:
        request: The pipeline run request which should replace the
            placeholder.

    Raises:
        RuntimeError: If the DB entry does not represent a placeholder run.
        ValueError: If the run request is not a valid request to replace the
            placeholder run.

    Returns:
        The updated `PipelineRunSchema`.
    """
    if not self.is_placeholder_run():
        raise RuntimeError(
            f"Unable to replace pipeline run {self.id} which is not a "
            "placeholder run."
        )

    if request.is_placeholder_request:
        raise ValueError(
            "Cannot replace a placeholder run with another placeholder run."
        )

    if (
        self.snapshot_id != request.snapshot
        or self.project_id != request.project
    ):
        raise ValueError(
            "Snapshot or project ID of placeholder run "
            "do not match the IDs of the run request."
        )

    if not request.orchestrator_run_id:
        raise ValueError(
            "Orchestrator run ID is required to replace a placeholder run."
        )

    if (
        self.orchestrator_run_id
        and self.orchestrator_run_id != request.orchestrator_run_id
    ):
        raise ValueError(
            "Orchestrator run ID of placeholder run does not match the "
            "ID of the run request."
        )

    orchestrator_environment = json.dumps(request.orchestrator_environment)

    self.orchestrator_run_id = request.orchestrator_run_id
    self.orchestrator_environment = orchestrator_environment
    self.status = request.status.value
    self.in_progress = not request.status.is_finished

    self.updated = utc_now()

    return self

update_status(requested_status: Optional[ExecutionStatus], status_reason: Optional[str] = None) -> bool

Update the status of the pipeline run.

Parameters:

Name	Type	Description	Default
requested_status	Optional[ExecutionStatus]	The requested status of the pipeline run.	required
status_reason	Optional[str]	The reason for the status of the pipeline run.	None

Raises:

Type	Description
IllegalOperationError	If the requested status transition is invalid.

Returns:

Type	Description
bool	Whether the status was updated.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def update_status(
    self,
    requested_status: Optional[ExecutionStatus],
    status_reason: Optional[str] = None,
) -> bool:
    """Update the status of the pipeline run.

    Args:
        requested_status: The requested status of the pipeline run.
        status_reason: The reason for the status of the pipeline run.

    Raises:
        IllegalOperationError: If the requested status transition is
            invalid.

    Returns:
        Whether the status was updated.
    """
    if requested_status in {
        ExecutionStatus.CACHED,
        ExecutionStatus.SKIPPED,
        ExecutionStatus.RETRYING,
        ExecutionStatus.RETRIED,
    }:
        raise IllegalOperationError(
            f"Execution status `{requested_status}` is not valid for "
            "pipeline runs."
        )

    current_status = ExecutionStatus(self.status)

    if (
        requested_status == ExecutionStatus.RESUMING
        and current_status
        not in {
            ExecutionStatus.PAUSED,
            ExecutionStatus.FAILED,
        }
    ):
        raise IllegalOperationError(
            "Only failed or paused runs can be resumed."
        )

    if (
        requested_status == ExecutionStatus.RESUMING
        and self.parent_run_id is not None
    ):
        raise IllegalOperationError(
            "Cannot resume a child run. Resume the parent run instead."
        )

    if (
        requested_status == ExecutionStatus.PROVISIONING
        and current_status != ExecutionStatus.INITIALIZING
    ):
        # Ignore transitions to provisioning from non-initializing states.
        # This could happen if the orchestrator starts running the pipeline
        # before the client environment can update the status to provisioning.
        return False

    # Snapshot always exists for pipeline runs of newer versions
    assert self.snapshot
    is_dynamic_pipeline = self.snapshot.is_dynamic

    if is_dynamic_pipeline:
        # In dynamic pipelines, the run status is only updated on manual
        # status updates and does not depend on step statuses.
        if requested_status is None and status_reason is None:
            return False

        new_status = requested_status or current_status
    else:
        # For static pipelines we compute the run status based on the step
        # statuses.
        new_status = _compute_static_pipeline_run_status(
            run_status=requested_status or current_status,
            step_statuses=self._get_step_run_statuses(),
            num_steps=self.snapshot.step_count,
        )

    if current_status.is_finished:
        if (
            current_status == ExecutionStatus.FAILED
            and new_status == ExecutionStatus.RESUMING
        ):
            # Allow failed -> resuming transition for retries.
            pass
        elif current_status != new_status:
            raise IllegalOperationError(
                "Cannot update the status of a finished run."
            )

    self.status = new_status.value

    if is_dynamic_pipeline:
        self.in_progress = not new_status.is_finished
    else:
        self.in_progress = self._check_if_run_in_progress()

    now = utc_now()
    if not self.in_progress and self.end_time is None:
        self.end_time = now
    elif new_status == ExecutionStatus.RESUMING:
        self.end_time = None

    if status_reason:
        self.status_reason = status_reason
    elif (
        current_status == ExecutionStatus.STOPPING
        and new_status == ExecutionStatus.STOPPED
    ):
        # Don't clear status reason for stopping -> stopped transition
        pass
    elif current_status != new_status:
        # Clear status reason when the status changes.
        self.status_reason = None

    self.updated = now
    return True

PipelineSchema

Bases: NamedSchema

SQL Model for pipelines.

Attributes

latest_run: Optional[PipelineRunSchema] property

Fetch the latest run for this pipeline.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[PipelineRunSchema]	The latest run for this pipeline.

Functions

from_request(pipeline_request: PipelineRequest) -> PipelineSchema classmethod

Convert a PipelineRequest to a PipelineSchema.

Parameters:

Name	Type	Description	Default
pipeline_request	PipelineRequest	The request model to convert.	required

Returns:

Type	Description
PipelineSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/pipeline_schemas.py

@classmethod
def from_request(
    cls,
    pipeline_request: "PipelineRequest",
) -> "PipelineSchema":
    """Convert a `PipelineRequest` to a `PipelineSchema`.

    Args:
        pipeline_request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=pipeline_request.name,
        description=pipeline_request.description,
        project_id=pipeline_request.project,
        user_id=pipeline_request.user,
        run_count=0,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/pipeline_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(PipelineSchema.user)),
                # joinedload(jl_arg(PipelineSchema.tags)),
                selectinload(jl_arg(PipelineSchema.visualizations)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> PipelineResponse

Convert a PipelineSchema to a PipelineResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
PipelineResponse	The created PipelineResponse.

Source code in src/zenml/zen_stores/schemas/pipeline_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "PipelineResponse":
    """Convert a `PipelineSchema` to a `PipelineResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The created PipelineResponse.
    """
    body = PipelineResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
    )

    metadata = None
    if include_metadata:
        metadata = PipelineResponseMetadata(
            description=self.description,
        )

    resources = None
    if include_resources:
        latest_run = self.latest_run
        latest_run_user = latest_run.user if latest_run else None

        resources = PipelineResponseResources(
            user=self.user.to_model() if self.user else None,
            latest_run_user=latest_run_user.to_model()
            if latest_run_user
            else None,
            latest_run_id=latest_run.id if latest_run else None,
            latest_run_status=latest_run.status if latest_run else None,
            tags=[tag.to_model() for tag in self.tags],
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    return PipelineResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(pipeline_update: PipelineUpdate) -> PipelineSchema

Update a PipelineSchema with a PipelineUpdate.

Parameters:

Name	Type	Description	Default
pipeline_update	PipelineUpdate	The update model.	required

Returns:

Type	Description
PipelineSchema	The updated PipelineSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_schemas.py

def update(self, pipeline_update: "PipelineUpdate") -> "PipelineSchema":
    """Update a `PipelineSchema` with a `PipelineUpdate`.

    Args:
        pipeline_update: The update model.

    Returns:
        The updated `PipelineSchema`.
    """
    self.description = pipeline_update.description
    self.updated = utc_now()
    return self

PipelineSnapshotSchema

Bases: BaseSchema

SQL Model for pipeline snapshots.

Attributes

is_runnable: bool property

Implements the is_runnable property.

Returns:

Type	Description
bool	True if the snapshot is runnable from server.

latest_run: Optional[PipelineRunSchema] property

Fetch the latest run for this snapshot.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[PipelineRunSchema]	The latest run for this snapshot.

Functions

from_request(request: PipelineSnapshotRequest, code_reference_id: Optional[UUID]) -> PipelineSnapshotSchema classmethod

Create schema from request.

Parameters:

Name	Type	Description	Default
request	PipelineSnapshotRequest	The request to convert.	required
code_reference_id	Optional[UUID]	Optional ID of the code reference for the snapshot.	required

Returns:

Type	Description
PipelineSnapshotSchema	The created schema.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

@classmethod
def from_request(
    cls,
    request: PipelineSnapshotRequest,
    code_reference_id: Optional[UUID],
) -> "PipelineSnapshotSchema":
    """Create schema from request.

    Args:
        request: The request to convert.
        code_reference_id: Optional ID of the code reference for the
            snapshot.

    Returns:
        The created schema.
    """
    client_env = json.dumps(request.client_environment)
    if len(client_env) > TEXT_FIELD_MAX_LENGTH:
        logger.warning(
            "Client environment is too large to be stored in the database. "
            "Skipping."
        )
        client_env = "{}"

    name = None
    if isinstance(request.name, str):
        name = request.name

    return cls(
        name=name,
        description=request.description,
        source_code=request.source_code,
        is_dynamic=request.is_dynamic,
        stack_id=request.stack,
        project_id=request.project,
        pipeline_id=request.pipeline,
        build_id=request.build,
        user_id=request.user,
        schedule_id=request.schedule,
        template_id=request.template,
        source_snapshot_id=request.source_snapshot,
        code_reference_id=code_reference_id,
        run_name_template=request.run_name_template,
        pipeline_configuration=request.pipeline_configuration.model_dump_json(),
        step_count=len(request.step_configurations),
        client_environment=client_env,
        client_version=request.client_version,
        server_version=request.server_version,
        pipeline_version_hash=request.pipeline_version_hash,
        pipeline_spec=json.dumps(
            request.pipeline_spec.model_dump(mode="json"), sort_keys=True
        )
        if request.pipeline_spec
        else None,
        code_path=request.code_path,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_metadata:
        options.extend(
            [
                joinedload(jl_arg(PipelineSnapshotSchema.stack)),
                joinedload(jl_arg(PipelineSnapshotSchema.build)),
                joinedload(jl_arg(PipelineSnapshotSchema.pipeline)),
                joinedload(jl_arg(PipelineSnapshotSchema.schedule)),
                joinedload(jl_arg(PipelineSnapshotSchema.code_reference)),
            ]
        )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(PipelineSnapshotSchema.user)),
                selectinload(
                    jl_arg(PipelineSnapshotSchema.visualizations)
                ),
            ]
        )

    return options

get_step_configuration(step_name: str) -> StepConfigurationSchema

Get a step configuration of the snapshot.

Parameters:

Name	Type	Description	Default
step_name	str	The name of the step to get the configuration for.	required

Raises:

Type	Description
KeyError	If the step configuration is not found.

Returns:

Type	Description
StepConfigurationSchema	The step configuration.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

def get_step_configuration(
    self, step_name: str
) -> "StepConfigurationSchema":
    """Get a step configuration of the snapshot.

    Args:
        step_name: The name of the step to get the configuration for.

    Raises:
        KeyError: If the step configuration is not found.

    Returns:
        The step configuration.
    """
    step_configs = self.get_step_configurations(include=[step_name])
    if len(step_configs) == 0:
        raise KeyError(
            f"Step configuration for step `{step_name}` not found."
        )
    return step_configs[0]

get_step_configurations(include: Optional[List[str]] = None) -> List[StepConfigurationSchema]

Get step configurations for the snapshot.

Parameters:

Name	Type	Description	Default
include	Optional[List[str]]	List of step names to include. If not given, all step configurations will be included.	None

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
List[StepConfigurationSchema]	List of step configurations.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

def get_step_configurations(
    self, include: Optional[List[str]] = None
) -> List["StepConfigurationSchema"]:
    """Get step configurations for the snapshot.

    Args:
        include: List of step names to include. If not given, all step
            configurations will be included.

    Raises:
        RuntimeError: If no session for the schema exists.

    Returns:
        List of step configurations.
    """
    if session := object_session(self):
        query = (
            select(StepConfigurationSchema)
            .where(StepConfigurationSchema.snapshot_id == self.id)
            .order_by(asc(StepConfigurationSchema.index))
        )

        if include:
            query = query.where(
                col(StepConfigurationSchema.name).in_(include)
            )

        return list(session.execute(query).scalars().all())
    else:
        raise RuntimeError(
            "Missing DB session to fetch step configurations."
        )

to_model(include_metadata: bool = False, include_resources: bool = False, include_python_packages: bool = False, include_config_schema: Optional[bool] = None, step_configuration_filter: Optional[List[str]] = None, **kwargs: Any) -> PipelineSnapshotResponse

Convert schema to response.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
include_python_packages	bool	Whether the python packages will be filled.	False
include_config_schema	Optional[bool]	Whether the config schema will be filled.	None
step_configuration_filter	Optional[List[str]]	List of step configurations to include in the response. If not given, all step configurations will be included.	None
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
PipelineSnapshotResponse	The response.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    include_python_packages: bool = False,
    include_config_schema: Optional[bool] = None,
    step_configuration_filter: Optional[List[str]] = None,
    **kwargs: Any,
) -> PipelineSnapshotResponse:
    """Convert schema to response.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        include_python_packages: Whether the python packages will be filled.
        include_config_schema: Whether the config schema will be filled.
        step_configuration_filter: List of step configurations to include in
            the response. If not given, all step configurations will be
            included.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The response.
    """
    deployable = False
    if self.build and self.stack and self.stack.has_deployer:
        deployable = True

    body = PipelineSnapshotResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        runnable=self.is_runnable,
        deployable=deployable,
        is_dynamic=self.is_dynamic,
    )
    metadata = None
    if include_metadata:
        pipeline_configuration = PipelineConfiguration.model_validate_json(
            self.pipeline_configuration
        )
        step_configurations = {}
        for step_configuration in self.get_step_configurations(
            include=step_configuration_filter
        ):
            step_configurations[step_configuration.name] = Step.from_dict(
                json.loads(step_configuration.config),
                pipeline_configuration,
            )

        client_environment = json.loads(self.client_environment)
        if not include_python_packages:
            client_environment.pop("python_packages", None)

        config_template = None
        config_schema = None

        if include_config_schema and self.build and self.build.stack_id:
            from zenml.zen_stores import template_utils

            if step_configuration_filter:
                # If only a subset of step configurations is requested,
                # we still need to get all of them to generate the config
                # template and schema
                all_step_configurations = {
                    step_configuration.name: Step.from_dict(
                        json.loads(step_configuration.config),
                        pipeline_configuration,
                    )
                    for step_configuration in self.get_step_configurations()
                }
            else:
                all_step_configurations = step_configurations

            config_template = template_utils.generate_config_template(
                snapshot=self,
                pipeline_configuration=pipeline_configuration,
                step_configurations=all_step_configurations,
            )
            config_schema = template_utils.generate_config_schema(
                snapshot=self,
                pipeline_configuration=pipeline_configuration,
                step_configurations=all_step_configurations,
            )

        metadata = PipelineSnapshotResponseMetadata(
            description=self.description,
            source_code=self.source_code,
            run_name_template=self.run_name_template,
            pipeline_configuration=pipeline_configuration,
            step_configurations=step_configurations,
            client_environment=client_environment,
            client_version=self.client_version,
            server_version=self.server_version,
            pipeline_version_hash=self.pipeline_version_hash,
            pipeline_spec=PipelineSpec.model_validate_json(
                self.pipeline_spec
            )
            if self.pipeline_spec
            else None,
            code_path=self.code_path,
            template_id=self.template_id,
            source_snapshot_id=self.source_snapshot_id,
            config_schema=config_schema,
            config_template=config_template,
        )

    resources = None
    if include_resources:
        latest_run = self.latest_run
        latest_run_user = latest_run.user if latest_run else None

        resources = PipelineSnapshotResponseResources(
            user=self.user.to_model() if self.user else None,
            pipeline=self.pipeline.to_model(),
            stack=self.stack.to_model() if self.stack else None,
            build=self.build.to_model() if self.build else None,
            schedule=self.schedule.to_model() if self.schedule else None,
            code_reference=self.code_reference.to_model()
            if self.code_reference
            else None,
            deployment=self.deployment.to_model()
            if self.deployment
            else None,
            tags=[tag.to_model() for tag in self.tags],
            latest_run_id=latest_run.id if latest_run else None,
            latest_run_status=latest_run.status if latest_run else None,
            latest_run_user=latest_run_user.to_model()
            if latest_run_user
            else None,
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    return PipelineSnapshotResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: PipelineSnapshotUpdate) -> PipelineSnapshotSchema

Update the schema.

Parameters:

Name	Type	Description	Default
update	PipelineSnapshotUpdate	The update to apply.	required

Returns:

Type	Description
PipelineSnapshotSchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

def update(
    self, update: PipelineSnapshotUpdate
) -> "PipelineSnapshotSchema":
    """Update the schema.

    Args:
        update: The update to apply.

    Returns:
        The updated schema.
    """
    if isinstance(update.name, str):
        self.name = update.name
    elif update.name is False:
        self.name = None

    if update.description:
        self.description = update.description

    self.updated = utc_now()
    return self

ProjectSchema

Bases: NamedSchema

SQL Model for projects.

Functions

from_request(project: ProjectRequest) -> ProjectSchema classmethod

Create a ProjectSchema from a ProjectResponse.

Parameters:

Name	Type	Description	Default
project	ProjectRequest	The ProjectResponse from which to create the schema.	required

Returns:

Type	Description
ProjectSchema	The created ProjectSchema.

Source code in src/zenml/zen_stores/schemas/project_schemas.py

@classmethod
def from_request(cls, project: ProjectRequest) -> "ProjectSchema":
    """Create a `ProjectSchema` from a `ProjectResponse`.

    Args:
        project: The `ProjectResponse` from which to create the schema.

    Returns:
        The created `ProjectSchema`.
    """
    return cls(
        name=project.name,
        description=project.description,
        display_name=project.display_name,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ProjectResponse

Convert a ProjectSchema to a ProjectResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ProjectResponse	The converted ProjectResponseModel.

Source code in src/zenml/zen_stores/schemas/project_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ProjectResponse:
    """Convert a `ProjectSchema` to a `ProjectResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The converted `ProjectResponseModel`.
    """
    metadata = None
    if include_metadata:
        metadata = ProjectResponseMetadata(
            description=self.description,
        )
    return ProjectResponse(
        id=self.id,
        name=self.name,
        body=ProjectResponseBody(
            display_name=self.display_name,
            created=self.created,
            updated=self.updated,
        ),
        metadata=metadata,
    )

update(project_update: ProjectUpdate) -> ProjectSchema

Update a ProjectSchema from a ProjectUpdate.

Parameters:

Name	Type	Description	Default
project_update	ProjectUpdate	The ProjectUpdate from which to update the schema.	required

Returns:

Type	Description
ProjectSchema	The updated ProjectSchema.

Source code in src/zenml/zen_stores/schemas/project_schemas.py

def update(self, project_update: ProjectUpdate) -> "ProjectSchema":
    """Update a `ProjectSchema` from a `ProjectUpdate`.

    Args:
        project_update: The `ProjectUpdate` from which to update the
            schema.

    Returns:
        The updated `ProjectSchema`.
    """
    for field, value in project_update.model_dump(
        exclude_unset=True
    ).items():
        setattr(self, field, value)

    self.updated = utc_now()
    return self

ResourcePoolAllocationSchema

Bases: BaseSchema

Resource pool allocation schema.

Attributes

priority: Optional[int] property

Fetch the priority for this allocation.

Returns:

Type	Description
Optional[int]	The matching policy priority, if a policy exists.

ResourcePoolQueueSchema

Bases: BaseSchema

Resource pool queue schema.

ResourcePoolResourceSchema

Bases: BaseSchema

Resource pool resource schema.

ResourcePoolSchema

Bases: NamedSchema

Resource pool schema.

Functions

from_request(request: ResourcePoolRequest) -> ResourcePoolSchema classmethod

Create a resource pool schema from a request.

Parameters:

Name	Type	Description	Default
request	ResourcePoolRequest	The request from which to create the resource pool.	required

Returns:

Type	Description
ResourcePoolSchema	The resource pool schema.

Source code in src/zenml/zen_stores/schemas/resource_pool_schemas.py

@classmethod
def from_request(
    cls,
    request: "ResourcePoolRequest",
) -> "ResourcePoolSchema":
    """Create a resource pool schema from a request.

    Args:
        request: The request from which to create the resource pool.

    Returns:
        The resource pool schema.
    """
    return cls(
        name=request.name,
        user_id=request.user,
        description=request.description,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/resource_pool_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = [
        selectinload(jl_arg(ResourcePoolSchema.resources)),
        selectinload(jl_arg(ResourcePoolSchema.queue_items)).options(
            load_only(jl_arg(ResourcePoolQueueSchema.request_id))
        ),
    ]

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(ResourcePoolSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ResourcePoolResponse

Creates a ResourcePoolResponse from a ResourcePoolSchema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ResourcePoolResponse	A ResourcePoolResponse

Source code in src/zenml/zen_stores/schemas/resource_pool_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ResourcePoolResponse":
    """Creates a `ResourcePoolResponse` from a `ResourcePoolSchema`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A `ResourcePoolResponse`
    """
    body = ResourcePoolResponseBody(
        created=self.created,
        updated=self.updated,
        user_id=self.user_id,
        queue_length=len(self.queue_items),
        capacity={r.key: r.total for r in self.resources},
        occupied_resources={r.key: r.occupied for r in self.resources},
    )

    metadata = None
    if include_metadata:
        metadata = ResourcePoolResponseMetadata(
            description=self.description,
        )

    resources = None
    if include_resources:
        resources = ResourcePoolResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return ResourcePoolResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(resource_pool_update: ResourcePoolUpdate) -> ResourcePoolSchema

Updates a ResourcePoolSchema from a ResourcePoolUpdate.

Parameters:

Name	Type	Description	Default
resource_pool_update	ResourcePoolUpdate	The ResourcePoolUpdate to update from.	required

Returns:

Type	Description
ResourcePoolSchema	The updated ResourcePoolSchema.

Source code in src/zenml/zen_stores/schemas/resource_pool_schemas.py

def update(
    self, resource_pool_update: "ResourcePoolUpdate"
) -> "ResourcePoolSchema":
    """Updates a `ResourcePoolSchema` from a `ResourcePoolUpdate`.

    Args:
        resource_pool_update: The `ResourcePoolUpdate` to update from.

    Returns:
        The updated `ResourcePoolSchema`.
    """
    if resource_pool_update.description:
        self.description = resource_pool_update.description

    self.updated = utc_now()
    return self

ResourcePoolSubjectPolicyResourceSchema

Bases: BaseSchema

Resource pool subject policy resource schema.

ResourcePoolSubjectPolicySchema

Bases: BaseSchema

Resource pool subject policy schema.

Functions

from_request(request: ResourcePoolSubjectPolicyRequest) -> ResourcePoolSubjectPolicySchema classmethod

Creates a schema instance from a request model.

Parameters:

Name	Type	Description	Default
request	ResourcePoolSubjectPolicyRequest	The request model.	required

Returns:

Type	Description
ResourcePoolSubjectPolicySchema	The schema instance.

Source code in src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py

@classmethod
def from_request(
    cls, request: ResourcePoolSubjectPolicyRequest
) -> "ResourcePoolSubjectPolicySchema":
    """Creates a schema instance from a request model.

    Args:
        request: The request model.

    Returns:
        The schema instance.
    """
    return cls(
        user_id=request.user,
        component_id=request.component_id,
        pool_id=request.pool_id,
        priority=request.priority,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Gets query options for this schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	If metadata should be included.	False
include_resources	bool	If resources should be included.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
Sequence[ExecutableOption]	The query options.

Source code in src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Gets query options for this schema.

    Args:
        include_metadata: If metadata should be included.
        include_resources: If resources should be included.
        **kwargs: Additional keyword arguments.

    Returns:
        The query options.
    """
    options: List[ExecutableOption] = [
        selectinload(jl_arg(ResourcePoolSubjectPolicySchema.resources)),
    ]

    if include_resources:
        options.extend(
            [
                selectinload(
                    jl_arg(ResourcePoolSubjectPolicySchema.component)
                ),
                selectinload(jl_arg(ResourcePoolSubjectPolicySchema.pool)),
                selectinload(jl_arg(ResourcePoolSubjectPolicySchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ResourcePoolSubjectPolicyResponse

Converts this schema to a response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata.	False
include_resources	bool	Whether to include nested resources.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
ResourcePoolSubjectPolicyResponse	The response model.

Source code in src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ResourcePoolSubjectPolicyResponse":
    """Converts this schema to a response model.

    Args:
        include_metadata: Whether to include metadata.
        include_resources: Whether to include nested resources.
        **kwargs: Additional keyword arguments.

    Returns:
        The response model.
    """
    body = ResourcePoolSubjectPolicyResponseBody(
        created=self.created,
        updated=self.updated,
        user_id=self.user_id,
        priority=self.priority,
        reserved={
            resource.key: resource.reserved for resource in self.resources
        },
        limit={
            resource.key: resource.limit
            for resource in self.resources
            if resource.limit is not None
        },
    )

    metadata = None
    if include_metadata:
        metadata = ResourcePoolSubjectPolicyResponseMetadata()

    resources = None
    if include_resources:
        resources = ResourcePoolSubjectPolicyResponseResources(
            user=self.user.to_model() if self.user else None,
            component=self.component.to_model(),
            pool=self.pool.to_model(
                include_metadata=False, include_resources=False
            ),
        )

    return ResourcePoolSubjectPolicyResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: ResourcePoolSubjectPolicyUpdate) -> ResourcePoolSubjectPolicySchema

Updates this schema from an update model.

Parameters:

Name	Type	Description	Default
update	ResourcePoolSubjectPolicyUpdate	The update model.	required

Returns:

Type	Description
ResourcePoolSubjectPolicySchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py

def update(
    self, update: ResourcePoolSubjectPolicyUpdate
) -> "ResourcePoolSubjectPolicySchema":
    """Updates this schema from an update model.

    Args:
        update: The update model.

    Returns:
        The updated schema.
    """
    if update.priority is not None:
        self.priority = update.priority

    self.updated = utc_now()
    return self

ResourceRequestResourceSchema

Bases: BaseSchema

Resource request resource schema.

ResourceRequestSchema

Bases: BaseSchema

Resource request schema.

Functions

from_request(request: ResourceRequestRequest) -> ResourceRequestSchema classmethod

Create a resource request schema from a request.

Parameters:

Name	Type	Description	Default
request	ResourceRequestRequest	The ResourceRequestRequest to create from.	required

Returns:

Type	Description
ResourceRequestSchema	The created ResourceRequestSchema.

Source code in src/zenml/zen_stores/schemas/resource_request_schemas.py

@classmethod
def from_request(
    cls,
    request: "ResourceRequestRequest",
) -> "ResourceRequestSchema":
    """Create a resource request schema from a request.

    Args:
        request: The `ResourceRequestRequest` to create from.

    Returns:
        The created `ResourceRequestSchema`.
    """
    return cls(
        user_id=request.user,
        component_id=request.component_id,
        step_run_id=request.step_run_id,
        status=ResourceRequestStatus.PENDING.value,
        preemption_initiated_by_id=None,
        preemptible=request.preemptible,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/resource_request_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        **kwargs: Additional keyword arguments.

    Returns:
        A list of query options.
    """
    options = [
        selectinload(jl_arg(ResourceRequestSchema.requested_resources)),
    ]

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(ResourceRequestSchema.component)),
                selectinload(
                    jl_arg(ResourceRequestSchema.step_run)
                ).joinedload(jl_arg(StepRunSchema.pipeline_run)),
                selectinload(jl_arg(ResourceRequestSchema.pool)),
                selectinload(jl_arg(ResourceRequestSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ResourceRequestResponse

Creates a ResourceRequestResponse from a ResourceRequestSchema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
ResourceRequestResponse	A ResourceRequestResponse object.

Source code in src/zenml/zen_stores/schemas/resource_request_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ResourceRequestResponse":
    """Creates a `ResourceRequestResponse` from a `ResourceRequestSchema`.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        **kwargs: Additional keyword arguments.

    Returns:
        A `ResourceRequestResponse` object.
    """
    body = ResourceRequestResponseBody(
        created=self.created,
        updated=self.updated,
        user_id=self.user_id,
        requested_resources={
            r.key: r.amount for r in self.requested_resources
        },
        status=ResourceRequestStatus(self.status),
        status_reason=self.status_reason,
        preemptible=self.preemptible,
    )

    metadata = None
    if include_metadata:
        metadata = ResourceRequestResponseMetadata()

    resources = None
    if include_resources:
        resources = ResourceRequestResponseResources(
            user=self.user.to_model() if self.user else None,
            component=self.component.to_model()
            if self.component
            else None,
            step_run=self.step_run.to_model() if self.step_run else None,
            pipeline_run=self.step_run.pipeline_run.to_model()
            if self.step_run
            else None,
            pool=self.pool.to_model() if self.pool else None,
        )

    return ResourceRequestResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

RunMetadataResourceSchema

Bases: SQLModel

Table for linking resources to run metadata entries.

RunMetadataSchema

Bases: BaseSchema

SQL Model for run metadata.

RunTemplateSchema

Bases: NamedSchema

SQL Model for run templates.

Attributes

latest_run: Optional[PipelineRunSchema] property

Fetch the latest run for this template.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[PipelineRunSchema]	The latest run for this template.

Functions

from_request(request: RunTemplateRequest) -> RunTemplateSchema classmethod

Create a schema from a request.

Parameters:

Name	Type	Description	Default
request	RunTemplateRequest	The request to convert.	required

Returns:

Type	Description
RunTemplateSchema	The created schema.

Source code in src/zenml/zen_stores/schemas/run_template_schemas.py

@classmethod
def from_request(
    cls,
    request: RunTemplateRequest,
) -> "RunTemplateSchema":
    """Create a schema from a request.

    Args:
        request: The request to convert.


    Returns:
        The created schema.
    """
    return cls(
        user_id=request.user,
        project_id=request.project,
        name=request.name,
        description=request.description,
        hidden=request.hidden,
        source_snapshot_id=request.source_snapshot_id,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/run_template_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    from zenml.zen_stores.schemas import PipelineSnapshotSchema

    options = [
        joinedload(jl_arg(RunTemplateSchema.source_snapshot)).joinedload(
            jl_arg(PipelineSnapshotSchema.build)
        ),
    ]

    if include_metadata or include_resources:
        options.extend(
            [
                joinedload(
                    jl_arg(RunTemplateSchema.source_snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.pipeline)),
                joinedload(
                    jl_arg(RunTemplateSchema.source_snapshot)
                ).joinedload(
                    jl_arg(PipelineSnapshotSchema.code_reference)
                ),
            ]
        )
    if include_metadata:
        options.extend(
            [
                joinedload(
                    jl_arg(RunTemplateSchema.source_snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.stack)),
                joinedload(
                    jl_arg(RunTemplateSchema.source_snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.schedule)),
            ]
        )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(RunTemplateSchema.user)),
                # joinedload(jl_arg(RunTemplateSchema.tags)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> RunTemplateResponse

Convert the schema to a response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
RunTemplateResponse	Model representing this schema.

Source code in src/zenml/zen_stores/schemas/run_template_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> RunTemplateResponse:
    """Convert the schema to a response model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        Model representing this schema.
    """
    runnable = False
    if (
        self.source_snapshot
        and self.source_snapshot.build
        and not self.source_snapshot.build.is_local
        and self.source_snapshot.build.stack_id
    ):
        runnable = True

    body = RunTemplateResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        runnable=runnable,
        hidden=self.hidden,
    )

    metadata = None
    if include_metadata:
        pipeline_spec = None
        config_template = None
        config_schema = None

        if self.source_snapshot:
            from zenml.zen_stores import template_utils

            source_snapshot_model = self.source_snapshot.to_model(
                include_metadata=True
            )
            pipeline_spec = source_snapshot_model.pipeline_spec

            if (
                self.source_snapshot.build
                and self.source_snapshot.build.stack_id
            ):
                config_template = template_utils.generate_config_template(
                    snapshot=self.source_snapshot,
                    pipeline_configuration=source_snapshot_model.pipeline_configuration,
                    step_configurations=source_snapshot_model.step_configurations,
                )
                config_schema = template_utils.generate_config_schema(
                    snapshot=self.source_snapshot,
                    pipeline_configuration=source_snapshot_model.pipeline_configuration,
                    step_configurations=source_snapshot_model.step_configurations,
                )

        metadata = RunTemplateResponseMetadata(
            description=self.description,
            pipeline_spec=pipeline_spec,
            config_template=config_template,
            config_schema=config_schema,
        )

    resources = None
    if include_resources:
        if self.source_snapshot:
            pipeline = (
                self.source_snapshot.pipeline.to_model()
                if self.source_snapshot.pipeline
                else None
            )
            build = (
                self.source_snapshot.build.to_model()
                if self.source_snapshot.build
                else None
            )
            code_reference = (
                self.source_snapshot.code_reference.to_model()
                if self.source_snapshot.code_reference
                else None
            )
        else:
            pipeline = None
            build = None
            code_reference = None

        latest_run = self.latest_run

        resources = RunTemplateResponseResources(
            user=self.user.to_model() if self.user else None,
            source_snapshot=self.source_snapshot.to_model()
            if self.source_snapshot
            else None,
            pipeline=pipeline,
            build=build,
            code_reference=code_reference,
            tags=[tag.to_model() for tag in self.tags],
            latest_run_id=latest_run.id if latest_run else None,
            latest_run_status=latest_run.status if latest_run else None,
        )

    return RunTemplateResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: RunTemplateUpdate) -> RunTemplateSchema

Update the schema.

Parameters:

Name	Type	Description	Default
update	RunTemplateUpdate	The update model.	required

Returns:

Type	Description
RunTemplateSchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/run_template_schemas.py

def update(self, update: RunTemplateUpdate) -> "RunTemplateSchema":
    """Update the schema.

    Args:
        update: The update model.

    Returns:
        The updated schema.
    """
    for field, value in update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field in ["add_tags", "remove_tags"]:
            # Tags are handled separately
            continue
        setattr(self, field, value)

    self.updated = utc_now()
    return self

RunWaitConditionSchema

Bases: BaseSchema, RunMetadataInterface

SQLModel schema for persisted run wait conditions.

Functions

from_request(request: RunWaitConditionRequest) -> RunWaitConditionSchema classmethod

Create a schema object from a wait condition create request.

Parameters:

Name	Type	Description	Default
request	RunWaitConditionRequest	Wait condition creation request.	required

Returns:

Type	Description
RunWaitConditionSchema	The persisted wait condition schema object.

Source code in src/zenml/zen_stores/schemas/run_wait_condition_schemas.py

@classmethod
def from_request(
    cls, request: RunWaitConditionRequest
) -> "RunWaitConditionSchema":
    """Create a schema object from a wait condition create request.

    Args:
        request: Wait condition creation request.

    Returns:
        The persisted wait condition schema object.
    """
    return cls(
        run_id=request.run,
        project_id=request.project,
        user_id=request.user,
        name=request.name,
        type=request.type.value,
        status=RunWaitConditionStatus.PENDING.value,
        question=request.question,
        data_schema_json=(
            json.dumps(request.data_schema)
            if request.data_schema is not None
            else None
        ),
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get query options for converting schema rows to models.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included in responses.	False
include_resources	bool	Whether resources will be included in responses.	False
**kwargs	Any	Additional unused keyword arguments.	{}

Returns:

Type	Description
Sequence[ExecutableOption]	SQLAlchemy query options.

Source code in src/zenml/zen_stores/schemas/run_wait_condition_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get query options for converting schema rows to models.

    Args:
        include_metadata: Whether metadata will be included in responses.
        include_resources: Whether resources will be included in responses.
        **kwargs: Additional unused keyword arguments.

    Returns:
        SQLAlchemy query options.
    """
    options: List[ExecutableOption] = [joinedload(jl_arg(cls.run))]
    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(cls.user)),
                joinedload(jl_arg(cls.resolved_by_user)),
            ]
        )
    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> RunWaitConditionResponse

Convert the schema row to a response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata should be included.	False
include_resources	bool	Whether resources should be included.	False
**kwargs	Any	Additional unused keyword arguments.	{}

Returns:

Type	Description
RunWaitConditionResponse	The wait condition response model.

Source code in src/zenml/zen_stores/schemas/run_wait_condition_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> RunWaitConditionResponse:
    """Convert the schema row to a response model.

    Args:
        include_metadata: Whether metadata should be included.
        include_resources: Whether resources should be included.
        **kwargs: Additional unused keyword arguments.

    Returns:
        The wait condition response model.
    """
    data_schema: Optional[Dict[str, Any]] = None
    if self.data_schema_json:
        data_schema = json.loads(self.data_schema_json)

    result: Optional[Any] = None
    if self.result_json:
        result = json.loads(self.result_json)

    body = RunWaitConditionResponseBody(
        user_id=self.user_id,
        project_id=self.run.project_id,
        created=self.created,
        updated=self.updated,
        type=RunWaitConditionType(self.type),
        status=RunWaitConditionStatus(self.status),
        last_polled_at=self.last_polled_at,
        poller_instance_id=self.poller_instance_id,
        poller_lease_expires_at=self.poller_lease_expires_at,
        resolved_at=self.resolved_at,
        resolved_by_user_id=self.resolved_by_user_id,
    )

    metadata = None
    if include_metadata:
        metadata = RunWaitConditionResponseMetadata(
            question=self.question,
            run_metadata=self.fetch_metadata(),
            data_schema=data_schema,
            resolution=self.resolution,
            result=result,
        )

    resources = None
    if include_resources:
        resources = RunWaitConditionResponseResources(
            user=self.user.to_model() if self.user else None,
            run=self.run.to_model(
                include_metadata=False, include_resources=False
            ),
        )

    return RunWaitConditionResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

ScheduleSchema

Bases: NamedSchema, RunMetadataInterface

SQL Model for schedules.

Functions

from_request(schedule_request: ScheduleRequest) -> ScheduleSchema classmethod

Create a ScheduleSchema from a ScheduleRequest.

Parameters:

Name	Type	Description	Default
schedule_request	ScheduleRequest	The ScheduleRequest to create the schema from.	required

Returns:

Type	Description
ScheduleSchema	The created ScheduleSchema.

Source code in src/zenml/zen_stores/schemas/schedule_schema.py

@classmethod
def from_request(
    cls, schedule_request: ScheduleRequest
) -> "ScheduleSchema":
    """Create a `ScheduleSchema` from a `ScheduleRequest`.

    Args:
        schedule_request: The `ScheduleRequest` to create the schema from.

    Returns:
        The created `ScheduleSchema`.
    """
    if schedule_request.interval_second is not None:
        interval_second = schedule_request.interval_second.total_seconds()
    else:
        interval_second = None
    return cls(
        name=schedule_request.name,
        project_id=schedule_request.project,
        user_id=schedule_request.user,
        pipeline_id=schedule_request.pipeline_id,
        orchestrator_id=schedule_request.orchestrator_id,
        active=schedule_request.active,
        cron_expression=schedule_request.cron_expression,
        start_time=schedule_request.start_time,
        end_time=schedule_request.end_time,
        interval_second=interval_second,
        catchup=schedule_request.catchup,
        run_once_start_time=schedule_request.run_once_start_time,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/schedule_schema.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    # if include_metadata:
    #     options.extend(
    #         [
    #             joinedload(jl_arg(ScheduleSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend([joinedload(jl_arg(ScheduleSchema.user))])

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ScheduleResponse

Convert a ScheduleSchema to a ScheduleResponseModel.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ScheduleResponse	The created ScheduleResponseModel.

Source code in src/zenml/zen_stores/schemas/schedule_schema.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ScheduleResponse:
    """Convert a `ScheduleSchema` to a `ScheduleResponseModel`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ScheduleResponseModel`.
    """
    if self.interval_second is not None:
        interval_second = timedelta(seconds=self.interval_second)
    else:
        interval_second = None

    body = ScheduleResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        active=self.active,
        cron_expression=self.cron_expression,
        start_time=self.start_time,
        end_time=self.end_time,
        interval_second=interval_second,
        catchup=self.catchup,
        updated=self.updated,
        created=self.created,
        run_once_start_time=self.run_once_start_time,
        is_archived=self.is_archived,
    )
    metadata = None
    if include_metadata:
        metadata = ScheduleResponseMetadata(
            pipeline_id=self.pipeline_id,
            orchestrator_id=self.orchestrator_id,
            run_metadata=self.fetch_metadata(),
        )

    resources = None
    if include_resources:
        resources = ScheduleResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return ScheduleResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(schedule_update: ScheduleUpdate) -> ScheduleSchema

Update a ScheduleSchema from a ScheduleUpdateModel.

Parameters:

Name	Type	Description	Default
schedule_update	ScheduleUpdate	The ScheduleUpdateModel to update the schema from.	required

Returns:

Type	Description
ScheduleSchema	The updated ScheduleSchema.

Source code in src/zenml/zen_stores/schemas/schedule_schema.py

def update(self, schedule_update: ScheduleUpdate) -> "ScheduleSchema":
    """Update a `ScheduleSchema` from a `ScheduleUpdateModel`.

    Args:
        schedule_update: The `ScheduleUpdateModel` to update the schema from.

    Returns:
        The updated `ScheduleSchema`.
    """
    if schedule_update.name is not None:
        self.name = schedule_update.name

    if schedule_update.cron_expression:
        self.cron_expression = schedule_update.cron_expression

    if schedule_update.active is not None:
        self.active = schedule_update.active

    self.updated = utc_now()
    return self

SecretResourceSchema

Bases: BaseSchema

SQL Model for secret resource relationship.

SecretSchema

Bases: NamedSchema

SQL Model for secrets.

Attributes:

Name	Type	Description
name	str	The name of the secret.
values	Optional[bytes]	The values of the secret.

Functions

from_request(secret: SecretRequest, internal: bool = False) -> SecretSchema classmethod

Create a SecretSchema from a SecretRequest.

Parameters:

Name	Type	Description	Default
secret	SecretRequest	The SecretRequest from which to create the schema.	required
internal	bool	Whether the secret is internal.	False

Returns:

Type	Description
SecretSchema	The created SecretSchema.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

@classmethod
def from_request(
    cls,
    secret: SecretRequest,
    internal: bool = False,
) -> "SecretSchema":
    """Create a `SecretSchema` from a `SecretRequest`.

    Args:
        secret: The `SecretRequest` from which to create the schema.
        internal: Whether the secret is internal.

    Returns:
        The created `SecretSchema`.
    """
    assert secret.user is not None, "User must be set for secret creation."
    return cls(
        name=secret.name,
        private=secret.private,
        user_id=secret.user,
        # Don't store secret values implicitly in the secret. The
        # SQL secret store will call `store_secret_values` to store the
        # values separately if SQL is used as the secrets store.
        values=None,
        internal=internal,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend([joinedload(jl_arg(SecretSchema.user))])

    return options

get_secret_values(encryption_engine: Optional[AesGcmEngine] = None) -> Dict[str, str]

Get the secret values for this secret.

This method is used by the SQL secrets store to load the secret values from the database.

Parameters:

Name	Type	Description	Default
encryption_engine	Optional[AesGcmEngine]	The encryption engine to use to decrypt the secret values. If None, the values will be base64 decoded.	None

Returns:

Type	Description
Dict[str, str]	The secret values

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

def get_secret_values(
    self,
    encryption_engine: Optional[AesGcmEngine] = None,
) -> Dict[str, str]:
    """Get the secret values for this secret.

    This method is used by the SQL secrets store to load the secret values
    from the database.

    Args:
        encryption_engine: The encryption engine to use to decrypt the
            secret values. If None, the values will be base64 decoded.

    Returns:
        The secret values

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
    """
    if not self.values:
        raise KeyError(
            f"Secret values for secret {self.id} have not been stored in "
            f"the SQL secrets store."
        )
    return self._load_secret_values(self.values, encryption_engine)

set_secret_values(secret_values: Dict[str, str], encryption_engine: Optional[AesGcmEngine] = None) -> None

Create a SecretSchema from a SecretRequest.

This method is used by the SQL secrets store to store the secret values in the database.

Parameters:

Name	Type	Description	Default
secret_values	Dict[str, str]	The new secret values.	required
encryption_engine	Optional[AesGcmEngine]	The encryption engine to use to encrypt the secret values. If None, the values will be base64 encoded.	None

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

def set_secret_values(
    self,
    secret_values: Dict[str, str],
    encryption_engine: Optional[AesGcmEngine] = None,
) -> None:
    """Create a `SecretSchema` from a `SecretRequest`.

    This method is used by the SQL secrets store to store the secret values
    in the database.

    Args:
        secret_values: The new secret values.
        encryption_engine: The encryption engine to use to encrypt the
            secret values. If None, the values will be base64 encoded.
    """
    self.values = self._dump_secret_values(
        secret_values, encryption_engine
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> SecretResponse

Converts a secret schema to a secret model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
SecretResponse	The secret model.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> SecretResponse:
    """Converts a secret schema to a secret model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The secret model.
    """
    metadata = None
    if include_metadata:
        metadata = SecretResponseMetadata()

    resources = None
    if include_resources:
        resources = SecretResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    # Don't load the secret values implicitly in the secret. The
    # SQL secret store will call `get_secret_values` to load the
    # values separately if SQL is used as the secrets store.
    body = SecretResponseBody(
        user_id=self.user_id,
        created=self.created,
        updated=self.updated,
        private=self.private,
    )
    return SecretResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(secret_update: SecretUpdate) -> SecretSchema

Update a SecretSchema from a SecretUpdate.

Parameters:

Name	Type	Description	Default
secret_update	SecretUpdate	The SecretUpdate from which to update the schema.	required

Returns:

Type	Description
SecretSchema	The updated SecretSchema.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

def update(
    self,
    secret_update: SecretUpdate,
) -> "SecretSchema":
    """Update a `SecretSchema` from a `SecretUpdate`.

    Args:
        secret_update: The `SecretUpdate` from which to update the schema.

    Returns:
        The updated `SecretSchema`.
    """
    # Don't update the secret values implicitly in the secret. The
    # SQL secret store will call `set_secret_values` to update the
    # values separately if SQL is used as the secrets store.
    for field, value in secret_update.model_dump(
        exclude_unset=True, exclude={"user", "values"}
    ).items():
        setattr(self, field, value)

    self.updated = utc_now()
    return self

ServerSettingsSchema

Bases: SQLModel

SQL Model for the server settings.

Functions

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ServerSettingsResponse

Convert an ServerSettingsSchema to an ServerSettingsResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ServerSettingsResponse	The created SettingsResponse.

Source code in src/zenml/zen_stores/schemas/server_settings_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ServerSettingsResponse:
    """Convert an `ServerSettingsSchema` to an `ServerSettingsResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The created `SettingsResponse`.
    """
    body = ServerSettingsResponseBody(
        server_id=self.id,
        server_name=self.server_name,
        logo_url=self.logo_url,
        enable_analytics=self.enable_analytics,
        display_announcements=self.display_announcements,
        display_updates=self.display_updates,
        active=self.active,
        updated=self.updated,
        last_user_activity=self.last_user_activity,
    )

    metadata = None
    resources = None

    if include_metadata:
        metadata = ServerSettingsResponseMetadata()

    if include_resources:
        resources = ServerSettingsResponseResources()

    return ServerSettingsResponse(
        body=body, metadata=metadata, resources=resources
    )

update(settings_update: ServerSettingsUpdate) -> ServerSettingsSchema

Update a ServerSettingsSchema from a ServerSettingsUpdate.

Parameters:

Name	Type	Description	Default
settings_update	ServerSettingsUpdate	The ServerSettingsUpdate from which to update the schema.	required

Returns:

Type	Description
ServerSettingsSchema	The updated ServerSettingsSchema.

Source code in src/zenml/zen_stores/schemas/server_settings_schemas.py

def update(
    self, settings_update: ServerSettingsUpdate
) -> "ServerSettingsSchema":
    """Update a `ServerSettingsSchema` from a `ServerSettingsUpdate`.

    Args:
        settings_update: The `ServerSettingsUpdate` from which
            to update the schema.

    Returns:
        The updated `ServerSettingsSchema`.
    """
    for field, value in settings_update.model_dump(
        exclude_unset=True
    ).items():
        if hasattr(self, field):
            setattr(self, field, value)

    self.updated = utc_now()

    return self

update_onboarding_state(completed_steps: Set[str]) -> ServerSettingsSchema

Update the onboarding state.

Parameters:

Name	Type	Description	Default
completed_steps	Set[str]	Newly completed onboarding steps.	required

Returns:

Type	Description
ServerSettingsSchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/server_settings_schemas.py

def update_onboarding_state(
    self, completed_steps: Set[str]
) -> "ServerSettingsSchema":
    """Update the onboarding state.

    Args:
        completed_steps: Newly completed onboarding steps.

    Returns:
        The updated schema.
    """
    old_state = set(
        json.loads(self.onboarding_state) if self.onboarding_state else []
    )
    new_state = old_state.union(completed_steps)
    self.onboarding_state = json.dumps(list(new_state))
    self.updated = utc_now()

    return self

ServiceConnectorSchema

Bases: NamedSchema

SQL Model for service connectors.

Attributes

labels_dict: Dict[str, str] property

Returns the labels as a dictionary.

Returns:

Type	Description
Dict[str, str]	The labels as a dictionary.

resource_types_list: List[str] property

Returns the resource types as a list.

Returns:

Type	Description
List[str]	The resource types as a list.

Functions

from_request(connector_request: ServiceConnectorRequest, secret_id: Optional[UUID] = None) -> ServiceConnectorSchema classmethod

Create a ServiceConnectorSchema from a ServiceConnectorRequest.

Parameters:

Name	Type	Description	Default
connector_request	ServiceConnectorRequest	The ServiceConnectorRequest from which to create the schema.	required
secret_id	Optional[UUID]	The ID of the secret to use for this connector.	None

Returns:

Type	Description
ServiceConnectorSchema	The created ServiceConnectorSchema.

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

@classmethod
def from_request(
    cls,
    connector_request: ServiceConnectorRequest,
    secret_id: Optional[UUID] = None,
) -> "ServiceConnectorSchema":
    """Create a `ServiceConnectorSchema` from a `ServiceConnectorRequest`.

    Args:
        connector_request: The `ServiceConnectorRequest` from which to
            create the schema.
        secret_id: The ID of the secret to use for this connector.

    Returns:
        The created `ServiceConnectorSchema`.
    """
    assert connector_request.user is not None, "User must be set."
    configuration = connector_request.configuration.non_secrets
    return cls(
        user_id=connector_request.user,
        name=connector_request.name,
        description=connector_request.description,
        connector_type=connector_request.type,
        auth_method=connector_request.auth_method,
        resource_types=base64.b64encode(
            json.dumps(connector_request.resource_types).encode("utf-8")
        ),
        resource_id=connector_request.resource_id,
        supports_instances=connector_request.supports_instances,
        configuration=base64.b64encode(
            json.dumps(configuration).encode("utf-8")
        )
        if configuration
        else None,
        secret_id=secret_id,
        expires_at=connector_request.expires_at,
        expires_skew_tolerance=connector_request.expires_skew_tolerance,
        expiration_seconds=connector_request.expiration_seconds,
        labels=base64.b64encode(
            json.dumps(connector_request.labels).encode("utf-8")
        )
        if connector_request.labels
        else None,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend([joinedload(jl_arg(ServiceConnectorSchema.user))])

    return options

has_labels(labels: Dict[str, Optional[str]]) -> bool

Checks if the connector has the given labels.

Parameters:

Name	Type	Description	Default
labels	Dict[str, Optional[str]]	The labels to check for.	required

Returns:

Type	Description
bool	Whether the connector has the given labels.

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

def has_labels(self, labels: Dict[str, Optional[str]]) -> bool:
    """Checks if the connector has the given labels.

    Args:
        labels: The labels to check for.

    Returns:
        Whether the connector has the given labels.
    """
    return all(
        self.labels_dict.get(key, None) == value
        for key, value in labels.items()
        if value is not None
    ) and all(
        key in self.labels_dict
        for key, value in labels.items()
        if value is None
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ServiceConnectorResponse

Creates a ServiceConnector from a ServiceConnectorSchema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ServiceConnectorResponse	A ServiceConnectorModel

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ServiceConnectorResponse":
    """Creates a `ServiceConnector` from a `ServiceConnectorSchema`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        A `ServiceConnectorModel`
    """
    body = ServiceConnectorResponseBody(
        user_id=self.user_id,
        created=self.created,
        updated=self.updated,
        description=self.description,
        connector_type=self.connector_type,
        auth_method=self.auth_method,
        resource_types=self.resource_types_list,
        resource_id=self.resource_id,
        supports_instances=self.supports_instances,
        expires_at=self.expires_at,
        expires_skew_tolerance=self.expires_skew_tolerance,
    )
    metadata = None
    if include_metadata:
        metadata = ServiceConnectorResponseMetadata(
            configuration=ServiceConnectorConfiguration(
                **json.loads(base64.b64decode(self.configuration).decode())
            )
            if self.configuration
            else ServiceConnectorConfiguration(),
            expiration_seconds=self.expiration_seconds,
            labels=self.labels_dict,
        )
    resources = None
    if include_resources:
        resources = ServiceConnectorResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return ServiceConnectorResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(connector_update: ServiceConnectorUpdate, secret_id: Optional[UUID] = None) -> ServiceConnectorSchema

Updates a ServiceConnectorSchema from a ServiceConnectorUpdate.

Parameters:

Name	Type	Description	Default
connector_update	ServiceConnectorUpdate	The ServiceConnectorUpdate to update from.	required
secret_id	Optional[UUID]	The ID of the secret to use for this connector.	None

Returns:

Type	Description
ServiceConnectorSchema	The updated ServiceConnectorSchema.

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

def update(
    self,
    connector_update: ServiceConnectorUpdate,
    secret_id: Optional[UUID] = None,
) -> "ServiceConnectorSchema":
    """Updates a `ServiceConnectorSchema` from a `ServiceConnectorUpdate`.

    Args:
        connector_update: The `ServiceConnectorUpdate` to update from.
        secret_id: The ID of the secret to use for this connector.

    Returns:
        The updated `ServiceConnectorSchema`.
    """
    for field, value in connector_update.model_dump(
        exclude_unset=False,
        exclude={"user", "secrets"},
    ).items():
        if value is None:
            if field == "resource_id":
                # The resource ID field in the update is special: if set
                # to None in the update, it triggers the existing resource
                # ID to be cleared.
                self.resource_id = None
            if field == "expiration_seconds":
                # The expiration_seconds field in the update is special:
                # if set to None in the update, it triggers the existing
                # expiration_seconds to be cleared.
                self.expiration_seconds = None
            continue
        if field == "configuration":
            if connector_update.configuration is not None:
                configuration = connector_update.configuration.non_secrets
                if configuration is not None:
                    self.configuration = (
                        base64.b64encode(
                            json.dumps(configuration).encode("utf-8")
                        )
                        if configuration
                        else None
                    )
        elif field == "resource_types":
            self.resource_types = base64.b64encode(
                json.dumps(connector_update.resource_types).encode("utf-8")
            )
        elif field == "labels":
            self.labels = (
                base64.b64encode(
                    json.dumps(connector_update.labels).encode("utf-8")
                )
                if connector_update.labels
                else None
            )
        else:
            setattr(self, field, value)
    self.secret_id = secret_id
    self.updated = utc_now()
    return self

ServiceSchema

Bases: NamedSchema

SQL Model for service.

Functions

from_request(service_request: ServiceRequest) -> ServiceSchema classmethod

Convert a ServiceRequest to a ServiceSchema.

Parameters:

Name	Type	Description	Default
service_request	ServiceRequest	The request model to convert.	required

Returns:

Type	Description
ServiceSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/service_schemas.py

@classmethod
def from_request(
    cls, service_request: "ServiceRequest"
) -> "ServiceSchema":
    """Convert a `ServiceRequest` to a `ServiceSchema`.

    Args:
        service_request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=service_request.name,
        project_id=service_request.project,
        user_id=service_request.user,
        service_source=service_request.service_source,
        service_type=service_request.service_type.model_dump_json(),
        type=service_request.service_type.type,
        flavor=service_request.service_type.flavor,
        admin_state=service_request.admin_state,
        config=dict_to_bytes(service_request.config),
        labels=dict_to_bytes(service_request.labels)
        if service_request.labels
        else None,
        status=dict_to_bytes(service_request.status)
        if service_request.status
        else None,
        endpoint=dict_to_bytes(service_request.endpoint)
        if service_request.endpoint
        else None,
        state=service_request.status.get("state")
        if service_request.status
        else None,
        model_version_id=service_request.model_version_id,
        pipeline_run_id=service_request.pipeline_run_id,
        prediction_url=service_request.prediction_url,
        health_check_url=service_request.health_check_url,
        pipeline_name=service_request.config.get("pipeline_name"),
        pipeline_step_name=service_request.config.get(
            "pipeline_step_name"
        ),
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/service_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ServiceSchema.user)),
                joinedload(jl_arg(ServiceSchema.model_version)),
                joinedload(jl_arg(ServiceSchema.pipeline_run)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ServiceResponse

Convert an ServiceSchema to an ServiceResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
ServiceResponse	The created ServiceResponse.

Source code in src/zenml/zen_stores/schemas/service_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ServiceResponse:
    """Convert an `ServiceSchema` to an `ServiceResponse`.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        kwargs: Additional keyword arguments.

    Returns:
        The created `ServiceResponse`.
    """
    body = ServiceResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        service_type=json.loads(self.service_type),
        labels=json.loads(base64.b64decode(self.labels).decode())
        if self.labels
        else None,
        state=self.state,
    )
    metadata = None
    if include_metadata:
        metadata = ServiceResponseMetadata(
            service_source=self.service_source,
            config=json.loads(base64.b64decode(self.config).decode()),
            status=json.loads(base64.b64decode(self.status).decode())
            if self.status
            else None,
            endpoint=json.loads(base64.b64decode(self.endpoint).decode())
            if self.endpoint
            else None,
            admin_state=self.admin_state or None,
            prediction_url=self.prediction_url or None,
            health_check_url=self.health_check_url,
        )
    resources = None
    if include_resources:
        resources = ServiceResponseResources(
            user=self.user.to_model() if self.user else None,
            model_version=self.model_version.to_model()
            if self.model_version
            else None,
            pipeline_run=self.pipeline_run.to_model()
            if self.pipeline_run
            else None,
        )
    return ServiceResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: ServiceUpdate) -> ServiceSchema

Updates a ServiceSchema from a ServiceUpdate.

Parameters:

Name	Type	Description	Default
update	ServiceUpdate	The ServiceUpdate to update from.	required

Returns:

Type	Description
ServiceSchema	The updated ServiceSchema.

Source code in src/zenml/zen_stores/schemas/service_schemas.py

def update(
    self,
    update: ServiceUpdate,
) -> "ServiceSchema":
    """Updates a `ServiceSchema` from a `ServiceUpdate`.

    Args:
        update: The `ServiceUpdate` to update from.

    Returns:
        The updated `ServiceSchema`.
    """
    for field, value in update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field == "labels":
            self.labels = (
                dict_to_bytes(update.labels) if update.labels else None
            )
        elif field == "status":
            self.status = (
                dict_to_bytes(update.status) if update.status else None
            )
            self.state = (
                update.status.get("state") if update.status else None
            )
        elif field == "endpoint":
            self.endpoint = (
                dict_to_bytes(update.endpoint) if update.endpoint else None
            )
        else:
            setattr(self, field, value)
    self.updated = utc_now()
    return self

StackComponentSchema

Bases: NamedSchema

SQL Model for stack components.

Functions

from_request(request: ComponentRequest, service_connector: Optional[ServiceConnectorSchema] = None) -> StackComponentSchema classmethod

Create a component schema from a request.

Parameters:

Name	Type	Description	Default
request	ComponentRequest	The request from which to create the component.	required
service_connector	Optional[ServiceConnectorSchema]	Optional service connector to link to the component.	None

Returns:

Type	Description
StackComponentSchema	The component schema.

Source code in src/zenml/zen_stores/schemas/component_schemas.py

@classmethod
def from_request(
    cls,
    request: "ComponentRequest",
    service_connector: Optional[ServiceConnectorSchema] = None,
) -> "StackComponentSchema":
    """Create a component schema from a request.

    Args:
        request: The request from which to create the component.
        service_connector: Optional service connector to link to the
            component.

    Returns:
        The component schema.
    """
    return cls(
        name=request.name,
        user_id=request.user,
        type=request.type,
        flavor=request.flavor,
        configuration=base64.b64encode(
            json.dumps(request.configuration).encode("utf-8")
        ),
        labels=base64.b64encode(
            json.dumps(request.labels).encode("utf-8")
        ),
        environment=base64.b64encode(
            json.dumps(request.environment).encode("utf-8")
        ),
        connector=service_connector,
        connector_resource_id=request.connector_resource_id,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/component_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = [
        joinedload(jl_arg(StackComponentSchema.flavor_schema)),
    ]

    if include_metadata:
        options.extend(
            [joinedload(jl_arg(StackComponentSchema.connector))]
        )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(StackComponentSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ComponentResponse

Creates a ComponentModel from an instance of a StackComponentSchema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Raises:

Type	Description
RuntimeError	If the flavor for the component is missing in the DB.

Returns:

Type	Description
ComponentResponse	A ComponentModel

Source code in src/zenml/zen_stores/schemas/component_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ComponentResponse":
    """Creates a `ComponentModel` from an instance of a `StackComponentSchema`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Raises:
        RuntimeError: If the flavor for the component is missing in the DB.

    Returns:
        A `ComponentModel`
    """
    body = ComponentResponseBody(
        user_id=self.user_id,
        type=StackComponentType(self.type),
        flavor_name=self.flavor,
        created=self.created,
        updated=self.updated,
        logo_url=self.flavor_schema.logo_url
        if self.flavor_schema
        else None,
        integration=self.flavor_schema.integration
        if self.flavor_schema
        else None,
    )
    metadata = None
    if include_metadata:
        environment = None
        if self.environment:
            environment = json.loads(
                base64.b64decode(self.environment).decode()
            )
        metadata = ComponentResponseMetadata(
            configuration=json.loads(
                base64.b64decode(self.configuration).decode()
            ),
            labels=json.loads(base64.b64decode(self.labels).decode())
            if self.labels
            else None,
            environment=environment or {},
            connector_resource_id=self.connector_resource_id,
            connector=self.connector.to_model()
            if self.connector
            else None,
            secrets=[secret.id for secret in self.secrets],
        )
    resources = None
    if include_resources:
        if not self.flavor_schema:
            raise RuntimeError(
                f"Missing flavor {self.flavor} for component {self.name}."
            )

        resources = ComponentResponseResources(
            user=self.user.to_model() if self.user else None,
            flavor=self.flavor_schema.to_model(),
        )
    return ComponentResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(component_update: ComponentUpdate) -> StackComponentSchema

Updates a StackComponentSchema from a ComponentUpdate.

Parameters:

Name	Type	Description	Default
component_update	ComponentUpdate	The ComponentUpdate to update from.	required

Returns:

Type	Description
StackComponentSchema	The updated StackComponentSchema.

Source code in src/zenml/zen_stores/schemas/component_schemas.py

def update(
    self, component_update: "ComponentUpdate"
) -> "StackComponentSchema":
    """Updates a `StackComponentSchema` from a `ComponentUpdate`.

    Args:
        component_update: The `ComponentUpdate` to update from.

    Returns:
        The updated `StackComponentSchema`.
    """
    for field, value in component_update.model_dump(
        exclude_unset=True,
        exclude={
            "user",
            "connector",
            "add_secrets",
            "remove_secrets",
            "attach_resource_pools",
            "detach_resource_pools",
        },
    ).items():
        if field == "configuration":
            self.configuration = base64.b64encode(
                json.dumps(component_update.configuration).encode("utf-8")
            )
        elif field == "labels":
            self.labels = base64.b64encode(
                json.dumps(component_update.labels).encode("utf-8")
            )
        elif field == "environment":
            self.environment = base64.b64encode(
                json.dumps(component_update.environment).encode("utf-8")
            )
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

StackCompositionSchema

Bases: SQLModel

SQL Model for stack definitions.

Join table between Stacks and StackComponents.

StackSchema

Bases: NamedSchema

SQL Model for stacks.

Attributes

has_deployer: bool property

If the stack has a deployer component.

Returns:

Type	Description
bool	If the stack has a deployer component.

Raises:

Type	Description
RuntimeError	if the stack has no DB session.

Functions

from_request(request: StackRequest) -> StackSchema classmethod

Create a stack schema from a request.

Parameters:

Name	Type	Description	Default
request	StackRequest	The request from which to create the stack.	required

Returns:

Type	Description
StackSchema	The stack schema.

Source code in src/zenml/zen_stores/schemas/stack_schemas.py

@classmethod
def from_request(
    cls,
    request: "StackRequest",
) -> "StackSchema":
    """Create a stack schema from a request.

    Args:
        request: The request from which to create the stack.

    Returns:
        The stack schema.
    """
    return cls(
        user_id=request.user,
        stack_spec_path=request.stack_spec_path,
        name=request.name,
        description=request.description,
        labels=base64.b64encode(
            json.dumps(request.labels).encode("utf-8")
        ),
        environment=base64.b64encode(
            json.dumps(request.environment).encode("utf-8")
        ),
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/stack_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_metadata:
        options.extend(
            [
                joinedload(
                    jl_arg(StackSchema.stack_compositions)
                ).joinedload(jl_arg(StackCompositionSchema.component))
            ]
        )

    if include_resources:
        options.extend([joinedload(jl_arg(StackSchema.user))])

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> StackResponse

Converts the schema to a model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
StackResponse	The converted model.

Source code in src/zenml/zen_stores/schemas/stack_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "StackResponse":
    """Converts the schema to a model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The converted model.
    """
    body = StackResponseBody(
        user_id=self.user_id,
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        environment = None
        if self.environment:
            environment = json.loads(
                base64.b64decode(self.environment).decode()
            )

        components: Dict[
            StackComponentType, List["StackComponentSchema"]
        ] = defaultdict(list)

        sorted_compositions = sorted(
            self.stack_compositions,
            key=lambda composition: (
                composition.default_for_type is None,
            ),
        )
        for composition in sorted_compositions:
            component_type = StackComponentType(composition.component.type)
            components[component_type].append(composition.component)

        metadata = StackResponseMetadata(
            components={
                component_type: [
                    component.to_model() for component in component_list
                ]
                for component_type, component_list in components.items()
            },
            stack_spec_path=self.stack_spec_path,
            labels=json.loads(base64.b64decode(self.labels).decode())
            if self.labels
            else None,
            description=self.description,
            environment=environment or {},
            secrets=[secret.id for secret in self.secrets],
        )
    resources = None
    if include_resources:
        resources = StackResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return StackResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(stack_update: StackUpdate) -> StackSchema

Updates a stack schema with a stack update model.

Parameters:

Name	Type	Description	Default
stack_update	StackUpdate	StackUpdate to update the stack with.	required

Returns:

Type	Description
StackSchema	The updated StackSchema.

Source code in src/zenml/zen_stores/schemas/stack_schemas.py

def update(
    self,
    stack_update: "StackUpdate",
) -> "StackSchema":
    """Updates a stack schema with a stack update model.

    Args:
        stack_update: `StackUpdate` to update the stack with.

    Returns:
        The updated StackSchema.
    """
    for field, value in stack_update.model_dump(
        exclude_unset=True,
        exclude={"user", "components", "add_secrets", "remove_secrets"},
    ).items():
        if field == "labels":
            self.labels = base64.b64encode(
                json.dumps(stack_update.labels).encode("utf-8")
            )
        elif field == "environment":
            self.environment = base64.b64encode(
                json.dumps(stack_update.environment).encode("utf-8")
            )
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

StepConfigurationSchema

Bases: BaseSchema

SQL Model for step configurations.

StepRunInputArtifactSchema

Bases: SQLModel

SQL Model that defines which artifacts are inputs to which step.

StepRunOutputArtifactSchema

Bases: SQLModel

SQL Model that defines which artifacts are outputs of which step.

StepRunParentsSchema

Bases: SQLModel

SQL Model that defines the order of steps.

StepRunSchema

Bases: NamedSchema, RunMetadataInterface

SQL Model for steps of pipeline runs.

Functions

from_request(request: StepRunRequest, snapshot_id: Optional[UUID], version: int, is_retriable: bool) -> StepRunSchema classmethod

Create a step run schema from a step run request model.

Parameters:

Name	Type	Description	Default
request	StepRunRequest	The step run request model.	required
snapshot_id	Optional[UUID]	The snapshot ID.	required
version	int	The version of the step run.	required
is_retriable	bool	Whether the step run is retriable.	required

Returns:

Type	Description
StepRunSchema	The step run schema.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

@classmethod
def from_request(
    cls,
    request: StepRunRequest,
    snapshot_id: Optional[UUID],
    version: int,
    is_retriable: bool,
) -> "StepRunSchema":
    """Create a step run schema from a step run request model.

    Args:
        request: The step run request model.
        snapshot_id: The snapshot ID.
        version: The version of the step run.
        is_retriable: Whether the step run is retriable.

    Returns:
        The step run schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        start_time=request.start_time,
        end_time=request.end_time,
        status=request.status.value,
        snapshot_id=snapshot_id,
        original_step_run_id=request.original_step_run_id,
        pipeline_run_id=request.pipeline_run_id,
        docstring=request.docstring,
        cache_key=request.cache_key,
        cache_expires_at=request.cache_expires_at,
        code_hash=request.code_hash,
        source_code=request.source_code,
        version=version,
        is_retriable=is_retriable,
        exception_info=request.exception_info.model_dump_json()
        if request.exception_info
        else None,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    from zenml.zen_stores.schemas import (
        ArtifactVersionSchema,
        ModelVersionSchema,
    )

    options = [
        selectinload(jl_arg(StepRunSchema.snapshot)).load_only(
            jl_arg(PipelineSnapshotSchema.pipeline_configuration)
        ),
        selectinload(jl_arg(StepRunSchema.pipeline_run)).load_only(
            jl_arg(PipelineRunSchema.start_time)
        ),
        joinedload(jl_arg(StepRunSchema.static_config)),
        joinedload(jl_arg(StepRunSchema.dynamic_config)),
    ]

    # if include_metadata:
    #     options.extend(
    #         [
    #             joinedload(jl_arg(StepRunSchema.parents)),
    #             joinedload(jl_arg(StepRunSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend(
            [
                selectinload(
                    jl_arg(StepRunSchema.model_version)
                ).joinedload(
                    jl_arg(ModelVersionSchema.model), innerjoin=True
                ),
                selectinload(jl_arg(StepRunSchema.user)),
                selectinload(jl_arg(StepRunSchema.input_artifacts))
                .joinedload(
                    jl_arg(StepRunInputArtifactSchema.artifact_version),
                    innerjoin=True,
                )
                .joinedload(
                    jl_arg(ArtifactVersionSchema.artifact), innerjoin=True
                ),
                selectinload(jl_arg(StepRunSchema.output_artifacts))
                .joinedload(
                    jl_arg(StepRunOutputArtifactSchema.artifact_version),
                    innerjoin=True,
                )
                .joinedload(
                    jl_arg(ArtifactVersionSchema.artifact), innerjoin=True
                ),
                selectinload(jl_arg(StepRunSchema.logs)),
            ]
        )

    return options

get_step_configuration() -> Step

Get the step configuration for the step run.

Raises:

Type	Description
ValueError	If the step run has no step configuration.

Returns:

Type	Description
Step	The step configuration.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

def get_step_configuration(self) -> Step:
    """Get the step configuration for the step run.

    Raises:
        ValueError: If the step run has no step configuration.

    Returns:
        The step configuration.
    """
    step = None

    if self.snapshot is not None:
        if config_schema := (self.dynamic_config or self.static_config):
            pipeline_configuration = (
                PipelineConfiguration.model_validate_json(
                    self.snapshot.pipeline_configuration
                )
            )
            pipeline_configuration.finalize_substitutions(
                start_time=self.pipeline_run.start_time,
                inplace=True,
            )
            step = Step.from_dict(
                json.loads(config_schema.config),
                pipeline_configuration=pipeline_configuration,
            )
    if not step and self.step_configuration:
        # In this legacy case, we're guaranteed to have the merged
        # config stored in the DB, which means we can instantiate the
        # `Step` object directly without passing the pipeline
        # configuration.
        step = Step.model_validate_json(self.step_configuration)
    elif not step:
        raise ValueError(
            f"Unable to load the configuration for step `{self.name}` from "
            "the database. To solve this please delete the pipeline run "
            "that this step run belongs to. Pipeline Run ID: "
            f"`{self.pipeline_run_id}`."
        )

    return step

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> StepRunResponse

Convert a StepRunSchema to a StepRunResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
StepRunResponse	The created StepRunResponse.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> StepRunResponse:
    """Convert a `StepRunSchema` to a `StepRunResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created StepRunResponse.
    """
    step = self.get_step_configuration()

    body = StepRunResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        type=step.config.step_type,
        status=ExecutionStatus(self.status),
        version=self.version,
        is_retriable=self.is_retriable,
        start_time=self.start_time,
        end_time=self.end_time,
        latest_heartbeat=self.latest_heartbeat,
        created=self.created,
        updated=self.updated,
        model_version_id=self.model_version_id,
        substitutions=step.config.substitutions,
        heartbeat_threshold=self.heartbeat_threshold,
    )
    metadata = None
    if include_metadata:
        metadata = StepRunResponseMetadata(
            config=step.config,
            spec=step.spec,
            cache_key=self.cache_key,
            cache_expires_at=self.cache_expires_at,
            code_hash=self.code_hash,
            docstring=self.docstring,
            source_code=self.source_code,
            exception_info=ExceptionInfo.model_validate_json(
                self.exception_info
            )
            if self.exception_info
            else None,
            snapshot_id=self.snapshot_id,
            pipeline_run_id=self.pipeline_run_id,
            original_step_run_id=self.original_step_run_id,
            parent_step_ids=[p.parent_id for p in self.parents],
            run_metadata=self.fetch_metadata(),
        )

    resources = None
    if include_resources:
        model_version = None
        if self.model_version:
            model_version = self.model_version.to_model()

        input_artifacts: Dict[str, List[StepRunInputResponse]] = {}
        for input_artifact in self.input_artifacts:
            if input_artifact.name not in input_artifacts:
                input_artifacts[input_artifact.name] = []
            step_run_input = StepRunInputResponse(
                input_type=StepRunInputArtifactType(input_artifact.type),
                index=input_artifact.input_index,
                chunk_index=input_artifact.chunk_index,
                chunk_size=input_artifact.chunk_size,
                **input_artifact.artifact_version.to_model().model_dump(),
            )
            input_artifacts[input_artifact.name].append(step_run_input)

        for artifact_list in input_artifacts.values():
            artifact_list.sort(key=lambda a: a.index or 0)

        output_artifacts: Dict[str, List["ArtifactVersionResponse"]] = {}
        for output_artifact in self.output_artifacts:
            if output_artifact.name not in output_artifacts:
                output_artifacts[output_artifact.name] = []
            output_artifacts[output_artifact.name].append(
                output_artifact.artifact_version.to_model()
            )

        resources = StepRunResponseResources(
            user=self.user.to_model() if self.user else None,
            model_version=model_version,
            log_collection=[
                log.to_model()
                for log in sorted(self.logs, key=lambda log: log.created)
            ],
            inputs=input_artifacts,
            outputs=output_artifacts,
            resource_request=self.resource_request.to_model(
                include_metadata=True, include_resources=False
            )
            if self.resource_request
            else None,
        )

    return StepRunResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(step_update: StepRunUpdate) -> StepRunSchema

Update a step run schema with a step run update model.

Parameters:

Name	Type	Description	Default
step_update	StepRunUpdate	The step run update model.	required

Returns:

Type	Description
StepRunSchema	The updated step run schema.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

def update(self, step_update: "StepRunUpdate") -> "StepRunSchema":
    """Update a step run schema with a step run update model.

    Args:
        step_update: The step run update model.

    Returns:
        The updated step run schema.
    """
    for key, value in step_update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if key == "status":
            self.status = value.value
        if key == "end_time":
            self.end_time = value
        if key == "exception_info":
            self.exception_info = json.dumps(value)
        if key == "cache_expires_at":
            self.cache_expires_at = value

    self.updated = utc_now()

    return self

TagResourceSchema

Bases: BaseSchema

SQL Model for tag resource relationship.

Functions

from_request(request: TagResourceRequest) -> TagResourceSchema classmethod

Convert an TagResourceRequest to an TagResourceSchema.

Parameters:

Name	Type	Description	Default
request	TagResourceRequest	The request model version to convert.	required

Returns:

Type	Description
TagResourceSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

@classmethod
def from_request(cls, request: TagResourceRequest) -> "TagResourceSchema":
    """Convert an `TagResourceRequest` to an `TagResourceSchema`.

    Args:
        request: The request model version to convert.

    Returns:
        The converted schema.
    """
    return cls(
        tag_id=request.tag_id,
        resource_id=request.resource_id,
        resource_type=request.resource_type.value,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TagResourceResponse

Convert an TagResourceSchema to an TagResourceResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
TagResourceResponse	The created TagResourceResponse.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> TagResourceResponse:
    """Convert an `TagResourceSchema` to an `TagResourceResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `TagResourceResponse`.
    """
    return TagResourceResponse(
        id=self.id,
        body=TagResourceResponseBody(
            tag_id=self.tag_id,
            resource_id=self.resource_id,
            created=self.created,
            updated=self.updated,
            resource_type=TaggableResourceTypes(self.resource_type),
        ),
    )

TagSchema

Bases: NamedSchema

SQL Model for tag.

Attributes

tagged_count: int property

Fetch the number of resources tagged with this tag.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
int	The number of resources tagged with this tag.

Functions

from_request(request: TagRequest) -> TagSchema classmethod

Convert an TagRequest to an TagSchema.

Parameters:

Name	Type	Description	Default
request	TagRequest	The request model to convert.	required

Returns:

Type	Description
TagSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

@classmethod
def from_request(cls, request: TagRequest) -> "TagSchema":
    """Convert an `TagRequest` to an `TagSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        exclusive=request.exclusive,
        color=request.color.value,
        user_id=request.user,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend([joinedload(jl_arg(TagSchema.user))])

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TagResponse

Convert an TagSchema to an TagResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
TagResponse	The created TagResponse.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> TagResponse:
    """Convert an `TagSchema` to an `TagResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `TagResponse`.
    """
    metadata = None
    if include_metadata:
        metadata = TagResponseMetadata(
            tagged_count=self.tagged_count,
        )

    resources = None
    if include_resources:
        resources = TagResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return TagResponse(
        id=self.id,
        name=self.name,
        body=TagResponseBody(
            user_id=self.user_id,
            created=self.created,
            updated=self.updated,
            color=ColorVariants(self.color),
            exclusive=self.exclusive,
        ),
        metadata=metadata,
        resources=resources,
    )

update(update: TagUpdate) -> TagSchema

Updates a TagSchema from a TagUpdate.

Parameters:

Name	Type	Description	Default
update	TagUpdate	The TagUpdate to update from.	required

Returns:

Type	Description
TagSchema	The updated TagSchema.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

def update(self, update: TagUpdate) -> "TagSchema":
    """Updates a `TagSchema` from a `TagUpdate`.

    Args:
        update: The `TagUpdate` to update from.

    Returns:
        The updated `TagSchema`.
    """
    for field, value in update.model_dump(exclude_unset=True).items():
        if field == "color":
            setattr(self, field, value.value)
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

TriggerExecutionSchema

Bases: SQLModel

Association table linking triggers to pipeline snapshots.

• Enforces uniqueness per (trigger_id, snapshot_id)
• Cascades deletes from either parent row (DB-level ON DELETE CASCADE)

TriggerSchema

Bases: NamedSchema

SQL Model for schedules.

Attributes

latest_run: PipelineRunSchema | None property

Fetch the latest execution for this trigger.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
PipelineRunSchema | None	The latest run for this pipeline.

Functions

from_request(trigger_request: TriggerRequest) -> TriggerSchema classmethod

Creates a TriggerSchema object from a TriggerRequest.

Parameters:

Name	Type	Description	Default
trigger_request	TriggerRequest	A TriggerRequest object.	required

Returns:

Type	Description
TriggerSchema	A TriggerSchema object.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

@classmethod
def from_request(cls, trigger_request: TriggerRequest) -> "TriggerSchema":
    """Creates a TriggerSchema object from a TriggerRequest.

    Args:
        trigger_request: A TriggerRequest object.

    Returns:
        A TriggerSchema object.
    """
    extra_fields = trigger_request.get_extra_fields()

    schema = cls(
        name=trigger_request.name,
        project_id=trigger_request.project,
        user_id=trigger_request.user,
        active=trigger_request.active,
        configuration=trigger_request.get_config(),
        flavor=trigger_request.flavor,
        type=trigger_request.type,
        concurrency=trigger_request.concurrency,
    )

    for field_name, value in extra_fields.items():
        setattr(schema, field_name, value)

    return schema

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(TriggerSchema.snapshots)).selectinload(
                    jl_arg(PipelineSnapshotSchema.source_snapshot)
                ),
                selectinload(jl_arg(TriggerSchema.snapshot_links)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TRIGGER_RETURN_TYPE_UNION

Converts to Pydantic response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Flag - to include metadata.	False
include_resources	bool	Flag - include resources.	False
**kwargs	Any	Keyword arguments	{}

Returns:

Type	Description
TRIGGER_RETURN_TYPE_UNION	A TriggerResponse object.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> TRIGGER_RETURN_TYPE_UNION:
    """Converts to Pydantic response model.

    Args:
        include_metadata: Flag - to include metadata.
        include_resources: Flag - include resources.
        **kwargs: Keyword arguments

    Returns:
        A TriggerResponse object.
    """
    body_cls = TYPE_TO_RESPONSE_BODY_MAPPING[self.type]
    response_cls = TYPE_TO_RESPONSE_MAPPING[self.type]

    body = body_cls(
        user_id=self.user_id,
        project_id=self.project_id,
        active=self.active,
        updated=self.updated,
        created=self.created,
        is_archived=self.is_archived,
        type=TriggerType(self.type),
        flavor=TriggerFlavor(self.flavor),
        name=self.name,
        concurrency=self.concurrency,
        **json.loads(self.configuration),
    )

    for field in body.get_extra_fields():
        setattr(body, field, getattr(self, field))

    metadata = None
    if include_metadata:
        metadata = TriggerResponseMetadata()

    resources = None
    if include_resources:
        latest_run = self.latest_run
        display_snapshot_id_by_executable_id: dict[UUID, UUID] = {}
        snapshots = []
        executable_snapshots = []
        for snapshot in self.snapshots:
            snapshot_model = snapshot.to_model()
            executable_snapshots.append(snapshot_model)
            display_snapshot = (
                snapshot.source_snapshot.to_model()
                if snapshot.source_snapshot is not None
                else snapshot_model
            )
            snapshots.append(display_snapshot)
            display_snapshot_id_by_executable_id[snapshot.id] = (
                display_snapshot.id
            )

        snapshot_dispatch_states: dict[
            UUID, TriggerSnapshotDispatchState
        ] = {}
        for snapshot_link in self.snapshot_links:
            parsed_state = snapshot_link.parsed_dispatch_state
            display_snapshot_id = display_snapshot_id_by_executable_id.get(
                snapshot_link.snapshot_id
            )
            if (
                parsed_state is not None
                and display_snapshot_id is not None
            ):
                snapshot_dispatch_states[display_snapshot_id] = (
                    parsed_state
                )

        resources = TriggerResponseResources(
            user=self.user.to_model() if self.user else None,
            snapshots=snapshots,
            executable_snapshots=executable_snapshots,
            latest_run=latest_run.to_model()
            if latest_run is not None
            else None,
            snapshot_dispatch_states=snapshot_dispatch_states,
        )

    return response_cls(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(trigger_update: TriggerUpdate) -> TriggerSchema

Applies update operation (and validations).

Parameters:

Name	Type	Description	Default
trigger_update	TriggerUpdate	A TriggerUpdate object.	required

Returns:

Type	Description
TriggerSchema	The updated TriggerSchema.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

def update(self, trigger_update: TriggerUpdate) -> "TriggerSchema":
    """Applies update operation (and validations).

    Args:
        trigger_update: A TriggerUpdate object.

    Returns:
        The updated TriggerSchema.
    """
    for field, value in trigger_update.model_dump(
        exclude_unset=True,
        include=set(TriggerBase.model_fields.keys()),
    ).items():
        if field in ["type"]:
            continue
        setattr(self, field, value)

    self.configuration = trigger_update.get_config()

    for field_name, value in trigger_update.get_extra_fields().items():
        setattr(self, field_name, value)

    return self

TriggerSnapshotSchema

Bases: SQLModel

Association table linking triggers to pipeline snapshots.

• Enforces uniqueness per (trigger_id, snapshot_id)
• Cascades deletes from either parent row (DB-level ON DELETE CASCADE)

Attributes

parsed_dispatch_state: TriggerSnapshotDispatchState | None property

Parse persisted dispatch-state JSON into the typed model.

Returns:

Type	Description
TriggerSnapshotDispatchState | None	Parsed trigger dispatch state or None if missing/invalid.

UserSchema

Bases: NamedSchema

SQL Model for users.

Functions

from_service_account_request(model: Union[ServiceAccountRequest, ServiceAccountInternalRequest]) -> UserSchema classmethod

Create a UserSchema from a Service Account request.

Parameters:

Name	Type	Description	Default
model	Union[ServiceAccountRequest, ServiceAccountInternalRequest]	The ServiceAccountRequest or ServiceAccountInternalRequest from which to create the schema.	required

Returns:

Type	Description
UserSchema	The created UserSchema.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

@classmethod
def from_service_account_request(
    cls, model: Union[ServiceAccountRequest, ServiceAccountInternalRequest]
) -> "UserSchema":
    """Create a `UserSchema` from a Service Account request.

    Args:
        model: The `ServiceAccountRequest` or `ServiceAccountInternalRequest`
            from which to create the schema.

    Returns:
        The created `UserSchema`.
    """
    return cls(
        name=model.name,
        full_name=model.full_name,
        description=model.description or "",
        external_user_id=model.external_user_id
        if isinstance(model, ServiceAccountInternalRequest)
        else None,
        active=model.active,
        is_service_account=True,
        email_opted_in=False,
        is_admin=False,
        avatar_url=model.avatar_url,
    )

from_user_request(model: UserRequest) -> UserSchema classmethod

Create a UserSchema from a UserRequest.

Parameters:

Name	Type	Description	Default
model	UserRequest	The UserRequest from which to create the schema.	required

Returns:

Type	Description
UserSchema	The created UserSchema.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

@classmethod
def from_user_request(cls, model: UserRequest) -> "UserSchema":
    """Create a `UserSchema` from a `UserRequest`.

    Args:
        model: The `UserRequest` from which to create the schema.

    Returns:
        The created `UserSchema`.
    """
    return cls(
        name=model.name,
        full_name=model.full_name,
        avatar_url=model.avatar_url,
        active=model.active,
        password=model.create_hashed_password(),
        activation_token=model.create_hashed_activation_token(),
        external_user_id=model.external_user_id,
        email_opted_in=model.email_opted_in,
        email=model.email,
        is_service_account=False,
        is_admin=model.is_admin,
        user_metadata=json.dumps(model.user_metadata)
        if model.user_metadata
        else None,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, include_private: bool = False, **kwargs: Any) -> UserResponse

Convert a UserSchema to a UserResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}
include_private	bool	Whether to include the user private information this is to limit the amount of data one can get about other users.	False

Returns:

Type	Description
UserResponse	The converted UserResponse.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    include_private: bool = False,
    **kwargs: Any,
) -> UserResponse:
    """Convert a `UserSchema` to a `UserResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic
        include_private: Whether to include the user private information
            this is to limit the amount of data one can get about other
            users.

    Returns:
        The converted `UserResponse`.
    """
    metadata = None
    if include_metadata:
        metadata = UserResponseMetadata(
            email=self.email if include_private else None,
            external_user_id=self.external_user_id,
            user_metadata=json.loads(self.user_metadata)
            if self.user_metadata
            else {},
        )

    return UserResponse(
        id=self.id,
        name=self.name,
        body=UserResponseBody(
            active=self.active,
            full_name=self.full_name,
            email_opted_in=self.email_opted_in,
            is_service_account=self.is_service_account,
            created=self.created,
            updated=self.updated,
            is_admin=self.is_admin,
            default_project_id=self.default_project_id,
            avatar_url=self.avatar_url,
        ),
        metadata=metadata,
    )

to_service_account_model(include_metadata: bool = False, include_resources: bool = False) -> ServiceAccountResponse

Convert a UserSchema to a ServiceAccountResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False

Returns:

Type	Description
ServiceAccountResponse	The converted ServiceAccountResponse.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

def to_service_account_model(
    self, include_metadata: bool = False, include_resources: bool = False
) -> ServiceAccountResponse:
    """Convert a `UserSchema` to a `ServiceAccountResponse`.

    Args:
         include_metadata: Whether the metadata will be filled.
         include_resources: Whether the resources will be filled.

    Returns:
         The converted `ServiceAccountResponse`.
    """
    metadata = None
    if include_metadata:
        metadata = ServiceAccountResponseMetadata(
            description=self.description or "",
            external_user_id=self.external_user_id,
        )

    body = ServiceAccountResponseBody(
        full_name=self.full_name,
        created=self.created,
        updated=self.updated,
        active=self.active,
        avatar_url=self.avatar_url,
    )

    return ServiceAccountResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
    )

update_service_account(service_account_update: ServiceAccountUpdate) -> UserSchema

Update a UserSchema from a ServiceAccountUpdate.

Parameters:

Name	Type	Description	Default
service_account_update	ServiceAccountUpdate	The ServiceAccountUpdate from which to update the schema.	required

Returns:

Type	Description
UserSchema	The updated UserSchema.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

def update_service_account(
    self, service_account_update: ServiceAccountUpdate
) -> "UserSchema":
    """Update a `UserSchema` from a `ServiceAccountUpdate`.

    Args:
        service_account_update: The `ServiceAccountUpdate` from which
            to update the schema.

    Returns:
        The updated `UserSchema`.
    """
    for field, value in service_account_update.model_dump(
        exclude_none=True
    ).items():
        setattr(self, field, value)

    self.updated = utc_now()
    return self

update_user(user_update: UserUpdate) -> UserSchema

Update a UserSchema from a UserUpdate.

Parameters:

Name	Type	Description	Default
user_update	UserUpdate	The UserUpdate from which to update the schema.	required

Returns:

Type	Description
UserSchema	The updated UserSchema.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

def update_user(self, user_update: UserUpdate) -> "UserSchema":
    """Update a `UserSchema` from a `UserUpdate`.

    Args:
        user_update: The `UserUpdate` from which to update the schema.

    Returns:
        The updated `UserSchema`.
    """
    for field, value in user_update.model_dump(exclude_unset=True).items():
        if field == "old_password":
            continue

        if field == "password":
            setattr(self, field, user_update.create_hashed_password())
        elif field == "activation_token":
            setattr(
                self, field, user_update.create_hashed_activation_token()
            )
        elif field == "user_metadata":
            if value is not None:
                self.user_metadata = json.dumps(value)
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

Modules

api_key_schemas

SQLModel implementation of user tables.

Classes

APIKeySchema

Bases: NamedSchema

SQL Model for API keys.

Functions

from_request(service_account_id: UUID, request: APIKeyRequest) -> Tuple[APIKeySchema, str] classmethod

Convert a APIKeyRequest to a APIKeySchema.

Parameters:

Name	Type	Description	Default
service_account_id	UUID	The service account id to associate the key with.	required
request	APIKeyRequest	The request model to convert.	required

Returns:

Type	Description
Tuple[APIKeySchema, str]	The converted schema and the un-hashed API key.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

@classmethod
def from_request(
    cls,
    service_account_id: UUID,
    request: APIKeyRequest,
) -> Tuple["APIKeySchema", str]:
    """Convert a `APIKeyRequest` to a `APIKeySchema`.

    Args:
        service_account_id: The service account id to associate the key
            with.
        request: The request model to convert.

    Returns:
        The converted schema and the un-hashed API key.
    """
    key = cls._generate_jwt_secret_key()
    hashed_key = cls._get_hashed_key(key)
    now = utc_now()
    return (
        cls(
            name=request.name,
            description=request.description or "",
            key=hashed_key,
            service_account_id=service_account_id,
            created=now,
            updated=now,
        ),
        key,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = [
        joinedload(jl_arg(APIKeySchema.service_account), innerjoin=True),
    ]

    return options

internal_update(update: APIKeyInternalUpdate) -> APIKeySchema

Update an APIKeySchema with an APIKeyInternalUpdate.

The internal update can also update the last used timestamp.

Parameters:

Name	Type	Description	Default
update	APIKeyInternalUpdate	The update model.	required

Returns:

Type	Description
APIKeySchema	The updated APIKeySchema.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def internal_update(self, update: APIKeyInternalUpdate) -> "APIKeySchema":
    """Update an `APIKeySchema` with an `APIKeyInternalUpdate`.

    The internal update can also update the last used timestamp.

    Args:
        update: The update model.

    Returns:
        The updated `APIKeySchema`.
    """
    self.update(update)

    if update.update_last_login:
        self.last_login = self.updated

    return self

rotate(rotate_request: APIKeyRotateRequest) -> Tuple[APIKeySchema, str]

Rotate the key for an APIKeySchema.

Parameters:

Name	Type	Description	Default
rotate_request	APIKeyRotateRequest	The rotate request model.	required

Returns:

Type	Description
Tuple[APIKeySchema, str]	The updated APIKeySchema and the new un-hashed key.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def rotate(
    self,
    rotate_request: APIKeyRotateRequest,
) -> Tuple["APIKeySchema", str]:
    """Rotate the key for an `APIKeySchema`.

    Args:
        rotate_request: The rotate request model.

    Returns:
        The updated `APIKeySchema` and the new un-hashed key.
    """
    self.updated = utc_now()
    self.previous_key = self.key
    self.retain_period = rotate_request.retain_period_minutes
    new_key = self._generate_jwt_secret_key()
    self.key = self._get_hashed_key(new_key)
    self.last_rotated = self.updated

    return self, new_key

to_internal_model(include_metadata: bool = False, include_resources: bool = False) -> APIKeyInternalResponse

Convert a APIKeySchema to an APIKeyInternalResponse.

The internal response model includes the hashed key values.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False

Returns:

Type	Description
APIKeyInternalResponse	The created APIKeyInternalResponse.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def to_internal_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
) -> APIKeyInternalResponse:
    """Convert a `APIKeySchema` to an `APIKeyInternalResponse`.

    The internal response model includes the hashed key values.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.

    Returns:
        The created APIKeyInternalResponse.
    """
    model = self.to_model(
        include_metadata=include_metadata,
        include_resources=include_resources,
    )
    model.get_body().key = self.key

    return APIKeyInternalResponse(
        id=self.id,
        name=self.name,
        previous_key=self.previous_key,
        body=model.body,
        metadata=model.metadata,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> APIKeyResponse

Convert a APIKeySchema to an APIKeyResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic.	{}

Returns:

Type	Description
APIKeyResponse	The created APIKeyResponse.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> APIKeyResponse:
    """Convert a `APIKeySchema` to an `APIKeyResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic.

    Returns:
        The created APIKeyResponse.
    """
    metadata = None
    if include_metadata:
        metadata = APIKeyResponseMetadata(
            description=self.description,
            retain_period_minutes=self.retain_period,
            last_login=self.last_login,
            last_rotated=self.last_rotated,
        )

    body = APIKeyResponseBody(
        created=self.created,
        updated=self.updated,
        active=self.active,
        service_account=self.service_account.to_service_account_model(),
    )

    return APIKeyResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
    )

update(update: APIKeyUpdate) -> APIKeySchema

Update an APIKeySchema with an APIKeyUpdate.

Parameters:

Name	Type	Description	Default
update	APIKeyUpdate	The update model.	required

Returns:

Type	Description
APIKeySchema	The updated APIKeySchema.

Source code in src/zenml/zen_stores/schemas/api_key_schemas.py

def update(self, update: APIKeyUpdate) -> "APIKeySchema":
    """Update an `APIKeySchema` with an `APIKeyUpdate`.

    Args:
        update: The update model.

    Returns:
        The updated `APIKeySchema`.
    """
    for field, value in update.model_dump(exclude_none=True).items():
        if hasattr(self, field):
            setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

api_transaction_schemas

SQLModel implementation of idempotent API transaction tables.

Classes

ApiTransactionResultSchema

Bases: SQLModel

SQL Model for API transaction results.

ApiTransactionSchema

Bases: BaseSchema

SQL Model for API transactions.

The result payload is stored in a separate table to keep this table small and fast for cleanup operations. Deleting rows with large blobs is expensive because the entire row must be copied to the undo log.

Functions

from_request(request: ApiTransactionRequest) -> ApiTransactionSchema classmethod

Create a new API transaction from a request.

Parameters:

Name	Type	Description	Default
request	ApiTransactionRequest	The API transaction request.	required

Returns:

Type	Description
ApiTransactionSchema	The API transaction schema.

Source code in src/zenml/zen_stores/schemas/api_transaction_schemas.py

@classmethod
def from_request(
    cls, request: ApiTransactionRequest
) -> "ApiTransactionSchema":
    """Create a new API transaction from a request.

    Args:
        request: The API transaction request.

    Returns:
        The API transaction schema.
    """
    assert request.user is not None, "User must be set."
    return cls(
        id=request.transaction_id,
        user_id=request.user,
        method=request.method,
        url=request.url,
        completed=False,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ApiTransactionResponse

Convert the SQL model to a ZenML model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
ApiTransactionResponse	The API transaction response.

Source code in src/zenml/zen_stores/schemas/api_transaction_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ApiTransactionResponse:
    """Convert the SQL model to a ZenML model.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        **kwargs: Additional keyword arguments.

    Returns:
        The API transaction response.
    """
    response = ApiTransactionResponse(
        id=self.id,
        body=ApiTransactionResponseBody(
            method=self.method,
            url=self.url,
            created=self.created,
            updated=self.updated,
            user_id=self.user_id,
            completed=self.completed,
        ),
    )
    return response

update(update: ApiTransactionUpdate) -> ApiTransactionSchema

Update the API transaction.

Parameters:

Name	Type	Description	Default
update	ApiTransactionUpdate	The API transaction update.	required

Returns:

Type	Description
ApiTransactionSchema	The API transaction schema.

Source code in src/zenml/zen_stores/schemas/api_transaction_schemas.py

def update(self, update: ApiTransactionUpdate) -> "ApiTransactionSchema":
    """Update the API transaction.

    Args:
        update: The API transaction update.

    Returns:
        The API transaction schema.
    """
    self.updated = utc_now()
    self.expired = self.updated + timedelta(seconds=update.cache_time)
    return self

Functions

artifact_schemas

SQLModel implementation of artifact table.

Classes

ArtifactSchema

Bases: NamedSchema

SQL Model for artifacts.

Attributes

latest_version: Optional[ArtifactVersionSchema] property

Fetch the latest version for this artifact.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[ArtifactVersionSchema]	The latest version for this artifact.

Functions

from_request(artifact_request: ArtifactRequest) -> ArtifactSchema classmethod

Convert an ArtifactRequest to an ArtifactSchema.

Parameters:

Name	Type	Description	Default
artifact_request	ArtifactRequest	The request model to convert.	required

Returns:

Type	Description
ArtifactSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

@classmethod
def from_request(
    cls,
    artifact_request: ArtifactRequest,
) -> "ArtifactSchema":
    """Convert an `ArtifactRequest` to an `ArtifactSchema`.

    Args:
        artifact_request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=artifact_request.name,
        has_custom_name=artifact_request.has_custom_name,
        project_id=artifact_request.project,
        user_id=artifact_request.user,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ArtifactSchema.user)),
                # joinedload(jl_arg(ArtifactSchema.tags)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ArtifactResponse

Convert an ArtifactSchema to an ArtifactResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ArtifactResponse	The created ArtifactResponse.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ArtifactResponse:
    """Convert an `ArtifactSchema` to an `ArtifactResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic



    Returns:
        The created `ArtifactResponse`.
    """
    # Create the body of the model
    body = ArtifactResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
    )

    # Create the metadata of the model
    metadata = None
    if include_metadata:
        metadata = ArtifactResponseMetadata(
            has_custom_name=self.has_custom_name,
        )

    resources = None
    if include_resources:
        latest_id, latest_name = None, None
        if latest_version := self.latest_version:
            latest_id = latest_version.id
            latest_name = latest_version.version

        resources = ArtifactResponseResources(
            user=self.user.to_model() if self.user else None,
            tags=[tag.to_model() for tag in self.tags],
            latest_version_id=latest_id,
            latest_version_name=latest_name,
        )

    return ArtifactResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(artifact_update: ArtifactUpdate) -> ArtifactSchema

Update an ArtifactSchema with an ArtifactUpdate.

Parameters:

Name	Type	Description	Default
artifact_update	ArtifactUpdate	The update model to apply.	required

Returns:

Type	Description
ArtifactSchema	The updated ArtifactSchema.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

def update(self, artifact_update: ArtifactUpdate) -> "ArtifactSchema":
    """Update an `ArtifactSchema` with an `ArtifactUpdate`.

    Args:
        artifact_update: The update model to apply.

    Returns:
        The updated `ArtifactSchema`.
    """
    self.updated = utc_now()
    if artifact_update.name:
        self.name = artifact_update.name
        self.has_custom_name = True
    if artifact_update.has_custom_name is not None:
        self.has_custom_name = artifact_update.has_custom_name
    return self

ArtifactVersionSchema

Bases: BaseSchema, RunMetadataInterface

SQL Model for artifact versions.

Attributes

producer_run_ids: Optional[Tuple[UUID, UUID]] property

Fetch the producer run IDs for this artifact version.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[Tuple[UUID, UUID]]	The producer step run ID and pipeline run ID for this artifact
Optional[Tuple[UUID, UUID]]	version.

Functions

from_request(artifact_version_request: ArtifactVersionRequest) -> ArtifactVersionSchema classmethod

Convert an ArtifactVersionRequest to an ArtifactVersionSchema.

Parameters:

Name	Type	Description	Default
artifact_version_request	ArtifactVersionRequest	The request model to convert.	required

Raises:

Type	Description
ValueError	If the request does not specify a version number.

Returns:

Type	Description
ArtifactVersionSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

@classmethod
def from_request(
    cls,
    artifact_version_request: ArtifactVersionRequest,
) -> "ArtifactVersionSchema":
    """Convert an `ArtifactVersionRequest` to an `ArtifactVersionSchema`.

    Args:
        artifact_version_request: The request model to convert.

    Raises:
        ValueError: If the request does not specify a version number.

    Returns:
        The converted schema.
    """
    if not artifact_version_request.version:
        raise ValueError("Missing version for artifact version request.")

    try:
        version_number = int(artifact_version_request.version)
    except ValueError:
        version_number = None
    return cls(
        artifact_id=artifact_version_request.artifact_id,
        version=str(artifact_version_request.version),
        version_number=version_number,
        artifact_store_id=artifact_version_request.artifact_store_id,
        project_id=artifact_version_request.project,
        user_id=artifact_version_request.user,
        type=artifact_version_request.type.value,
        uri=artifact_version_request.uri,
        materializer=artifact_version_request.materializer.model_dump_json(),
        data_type=artifact_version_request.data_type.model_dump_json(),
        save_type=artifact_version_request.save_type.value,
        content_hash=artifact_version_request.content_hash,
        item_count=artifact_version_request.item_count,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    # if include_metadata:
    #     options.extend(
    #         [
    #             joinedload(jl_arg(ArtifactVersionSchema.visualizations)),
    #             joinedload(jl_arg(ArtifactVersionSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ArtifactVersionSchema.user)),
                # joinedload(jl_arg(ArtifactVersionSchema.tags)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ArtifactVersionResponse

Convert an ArtifactVersionSchema to an ArtifactVersionResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ArtifactVersionResponse	The created ArtifactVersionResponse.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ArtifactVersionResponse:
    """Convert an `ArtifactVersionSchema` to an `ArtifactVersionResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic



    Returns:
        The created `ArtifactVersionResponse`.
    """
    try:
        materializer = Source.model_validate_json(self.materializer)
    except ValidationError:
        # This is an old source which was an importable source path
        materializer = Source.from_import_path(self.materializer)

    try:
        data_type = Source.model_validate_json(self.data_type)
    except ValidationError:
        # This is an old source which was an importable source path
        data_type = Source.from_import_path(self.data_type)

    # Create the body of the model
    artifact = self.artifact.to_model()
    body = ArtifactVersionResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        artifact=artifact,
        version=self.version or str(self.version_number),
        uri=self.uri,
        type=ArtifactType(self.type),
        materializer=materializer,
        data_type=data_type,
        created=self.created,
        updated=self.updated,
        save_type=ArtifactSaveType(self.save_type),
        artifact_store_id=self.artifact_store_id,
        content_hash=self.content_hash,
        item_count=self.item_count,
    )

    # Create the metadata of the model
    metadata = None
    if include_metadata:
        metadata = ArtifactVersionResponseMetadata(
            visualizations=[v.to_model() for v in self.visualizations],
            run_metadata=self.fetch_metadata(),
        )

    resources = None
    if include_resources:
        producer_step_run_id, producer_pipeline_run_id = None, None
        if producer_run_ids := self.producer_run_ids:
            # TODO: Why was the producer_pipeline_run_id only set for one
            # of the cases before?
            producer_step_run_id, producer_pipeline_run_id = (
                producer_run_ids
            )

        resources = ArtifactVersionResponseResources(
            user=self.user.to_model() if self.user else None,
            tags=[tag.to_model() for tag in self.tags],
            producer_step_run_id=producer_step_run_id,
            producer_pipeline_run_id=producer_pipeline_run_id,
        )

    return ArtifactVersionResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(artifact_version_update: ArtifactVersionUpdate) -> ArtifactVersionSchema

Update an ArtifactVersionSchema with an ArtifactVersionUpdate.

Parameters:

Name	Type	Description	Default
artifact_version_update	ArtifactVersionUpdate	The update model to apply.	required

Returns:

Type	Description
ArtifactVersionSchema	The updated ArtifactVersionSchema.

Source code in src/zenml/zen_stores/schemas/artifact_schemas.py

def update(
    self, artifact_version_update: ArtifactVersionUpdate
) -> "ArtifactVersionSchema":
    """Update an `ArtifactVersionSchema` with an `ArtifactVersionUpdate`.

    Args:
        artifact_version_update: The update model to apply.

    Returns:
        The updated `ArtifactVersionSchema`.
    """
    self.updated = utc_now()
    return self

Functions

artifact_visualization_schemas

SQLModel implementation of artifact visualization table.

Classes

ArtifactVisualizationSchema

Bases: BaseSchema

SQL Model for visualizations of artifacts.

Functions

from_model(artifact_visualization_request: ArtifactVisualizationRequest, artifact_version_id: UUID) -> ArtifactVisualizationSchema classmethod

Convert a ArtifactVisualizationRequest to a ArtifactVisualizationSchema.

Parameters:

Name	Type	Description	Default
artifact_visualization_request	ArtifactVisualizationRequest	The visualization.	required
artifact_version_id	UUID	The UUID of the artifact version.	required

Returns:

Type	Description
ArtifactVisualizationSchema	The ArtifactVisualizationSchema.

Source code in src/zenml/zen_stores/schemas/artifact_visualization_schemas.py

@classmethod
def from_model(
    cls,
    artifact_visualization_request: ArtifactVisualizationRequest,
    artifact_version_id: UUID,
) -> "ArtifactVisualizationSchema":
    """Convert a `ArtifactVisualizationRequest` to a `ArtifactVisualizationSchema`.

    Args:
        artifact_visualization_request: The visualization.
        artifact_version_id: The UUID of the artifact version.

    Returns:
        The `ArtifactVisualizationSchema`.
    """
    return cls(
        type=artifact_visualization_request.type.value,
        uri=artifact_visualization_request.uri,
        artifact_version_id=artifact_version_id,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ArtifactVisualizationResponse

Convert an ArtifactVisualizationSchema to a Visualization.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ArtifactVisualizationResponse	The Visualization.

Source code in src/zenml/zen_stores/schemas/artifact_visualization_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ArtifactVisualizationResponse:
    """Convert an `ArtifactVisualizationSchema` to a `Visualization`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic



    Returns:
        The `Visualization`.
    """
    body = ArtifactVisualizationResponseBody(
        type=VisualizationType(self.type),
        uri=self.uri,
        created=self.created,
        updated=self.updated,
    )

    metadata = None
    if include_metadata:
        metadata = ArtifactVisualizationResponseMetadata(
            artifact_version_id=self.artifact_version_id,
        )

    resources = None
    if include_resources:
        if self.artifact_version is not None:
            artifact_version = self.artifact_version.to_model(
                include_metadata=False,
                include_resources=False,
            )
        else:
            artifact_version = None
        resources = ArtifactVisualizationResponseResources(
            artifact_version=artifact_version,
        )

    return ArtifactVisualizationResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

Functions

base_schemas

Base classes for SQLModel schemas.

Classes

BaseSchema

Bases: SQLModel

Base SQL Model for ZenML entities.

Functions

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

This method should return query options that improve the performance when trying to later on converting that schema to a model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/base_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    This method should return query options that improve the performance
    when trying to later on converting that schema to a model.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    return []

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Any

In case the Schema has a corresponding Model, this allows conversion to that model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Any	The model corresponding to this schema.

Raises:

Type	Description
NotImplementedError	When the base class fails to implement this.

Source code in src/zenml/zen_stores/schemas/base_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Any:
    """In case the Schema has a corresponding Model, this allows conversion to that model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The model corresponding to this schema.

    Raises:
        NotImplementedError: When the base class fails to implement this.
    """
    raise NotImplementedError(
        "No 'to_model()' method implemented for this"
        f"schema: '{self.__class__.__name__}'."
    )

NamedSchema

Bases: BaseSchema

Base Named SQL Model.

Functions

code_repository_schemas

SQL Model Implementations for code repositories.

Classes

CodeReferenceSchema

Bases: BaseSchema

SQL Model for code references.

Functions

from_request(request: CodeReferenceRequest, project_id: UUID) -> CodeReferenceSchema classmethod

Convert a CodeReferenceRequest to a CodeReferenceSchema.

Parameters:

Name	Type	Description	Default
request	CodeReferenceRequest	The request model to convert.	required
project_id	UUID	The project ID.	required

Returns:

Type	Description
CodeReferenceSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

@classmethod
def from_request(
    cls, request: "CodeReferenceRequest", project_id: UUID
) -> "CodeReferenceSchema":
    """Convert a `CodeReferenceRequest` to a `CodeReferenceSchema`.

    Args:
        request: The request model to convert.
        project_id: The project ID.

    Returns:
        The converted schema.
    """
    return cls(
        project_id=project_id,
        commit=request.commit,
        subdirectory=request.subdirectory,
        code_repository_id=request.code_repository,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> CodeReferenceResponse

Convert a CodeReferenceSchema to a CodeReferenceResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic.	{}

Returns:

Type	Description
CodeReferenceResponse	The converted model.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "CodeReferenceResponse":
    """Convert a `CodeReferenceSchema` to a `CodeReferenceResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic.

    Returns:
        The converted model.
    """
    body = CodeReferenceResponseBody(
        commit=self.commit,
        subdirectory=self.subdirectory,
        code_repository=self.code_repository.to_model(),
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = CodeReferenceResponseMetadata()

    return CodeReferenceResponse(
        id=self.id,
        body=body,
        metadata=metadata,
    )

CodeRepositorySchema

Bases: NamedSchema

SQL Model for code repositories.

Functions

from_request(request: CodeRepositoryRequest) -> CodeRepositorySchema classmethod

Convert a CodeRepositoryRequest to a CodeRepositorySchema.

Parameters:

Name	Type	Description	Default
request	CodeRepositoryRequest	The request model to convert.	required

Returns:

Type	Description
CodeRepositorySchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

@classmethod
def from_request(
    cls, request: "CodeRepositoryRequest"
) -> "CodeRepositorySchema":
    """Convert a `CodeRepositoryRequest` to a `CodeRepositorySchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        config=json.dumps(request.config),
        source=request.source.model_dump_json(),
        description=request.description,
        logo_url=request.logo_url,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(CodeRepositorySchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> CodeRepositoryResponse

Convert a CodeRepositorySchema to a CodeRepositoryResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic.	{}

Returns:

Type	Description
CodeRepositoryResponse	The created CodeRepositoryResponse.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "CodeRepositoryResponse":
    """Convert a `CodeRepositorySchema` to a `CodeRepositoryResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic.

    Returns:
        The created CodeRepositoryResponse.
    """
    body = CodeRepositoryResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        source=json.loads(self.source),
        logo_url=self.logo_url,
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = CodeRepositoryResponseMetadata(
            config=json.loads(self.config),
            description=self.description,
        )

    resources = None
    if include_resources:
        resources = CodeRepositoryResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return CodeRepositoryResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: CodeRepositoryUpdate) -> CodeRepositorySchema

Update a CodeRepositorySchema with a CodeRepositoryUpdate.

Parameters:

Name	Type	Description	Default
update	CodeRepositoryUpdate	The update model.	required

Returns:

Type	Description
CodeRepositorySchema	The updated CodeRepositorySchema.

Source code in src/zenml/zen_stores/schemas/code_repository_schemas.py

def update(self, update: "CodeRepositoryUpdate") -> "CodeRepositorySchema":
    """Update a `CodeRepositorySchema` with a `CodeRepositoryUpdate`.

    Args:
        update: The update model.

    Returns:
        The updated `CodeRepositorySchema`.
    """
    if update.name:
        self.name = update.name

    if update.description:
        self.description = update.description

    if update.logo_url:
        self.logo_url = update.logo_url

    if update.config:
        self.config = json.dumps(update.config)

    self.updated = utc_now()
    return self

Functions

component_schemas

SQL Model Implementations for Stack Components.

Classes

StackComponentSchema

Bases: NamedSchema

SQL Model for stack components.

Functions

from_request(request: ComponentRequest, service_connector: Optional[ServiceConnectorSchema] = None) -> StackComponentSchema classmethod

Create a component schema from a request.

Parameters:

Name	Type	Description	Default
request	ComponentRequest	The request from which to create the component.	required
service_connector	Optional[ServiceConnectorSchema]	Optional service connector to link to the component.	None

Returns:

Type	Description
StackComponentSchema	The component schema.

Source code in src/zenml/zen_stores/schemas/component_schemas.py

@classmethod
def from_request(
    cls,
    request: "ComponentRequest",
    service_connector: Optional[ServiceConnectorSchema] = None,
) -> "StackComponentSchema":
    """Create a component schema from a request.

    Args:
        request: The request from which to create the component.
        service_connector: Optional service connector to link to the
            component.

    Returns:
        The component schema.
    """
    return cls(
        name=request.name,
        user_id=request.user,
        type=request.type,
        flavor=request.flavor,
        configuration=base64.b64encode(
            json.dumps(request.configuration).encode("utf-8")
        ),
        labels=base64.b64encode(
            json.dumps(request.labels).encode("utf-8")
        ),
        environment=base64.b64encode(
            json.dumps(request.environment).encode("utf-8")
        ),
        connector=service_connector,
        connector_resource_id=request.connector_resource_id,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/component_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = [
        joinedload(jl_arg(StackComponentSchema.flavor_schema)),
    ]

    if include_metadata:
        options.extend(
            [joinedload(jl_arg(StackComponentSchema.connector))]
        )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(StackComponentSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ComponentResponse

Creates a ComponentModel from an instance of a StackComponentSchema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Raises:

Type	Description
RuntimeError	If the flavor for the component is missing in the DB.

Returns:

Type	Description
ComponentResponse	A ComponentModel

Source code in src/zenml/zen_stores/schemas/component_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ComponentResponse":
    """Creates a `ComponentModel` from an instance of a `StackComponentSchema`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Raises:
        RuntimeError: If the flavor for the component is missing in the DB.

    Returns:
        A `ComponentModel`
    """
    body = ComponentResponseBody(
        user_id=self.user_id,
        type=StackComponentType(self.type),
        flavor_name=self.flavor,
        created=self.created,
        updated=self.updated,
        logo_url=self.flavor_schema.logo_url
        if self.flavor_schema
        else None,
        integration=self.flavor_schema.integration
        if self.flavor_schema
        else None,
    )
    metadata = None
    if include_metadata:
        environment = None
        if self.environment:
            environment = json.loads(
                base64.b64decode(self.environment).decode()
            )
        metadata = ComponentResponseMetadata(
            configuration=json.loads(
                base64.b64decode(self.configuration).decode()
            ),
            labels=json.loads(base64.b64decode(self.labels).decode())
            if self.labels
            else None,
            environment=environment or {},
            connector_resource_id=self.connector_resource_id,
            connector=self.connector.to_model()
            if self.connector
            else None,
            secrets=[secret.id for secret in self.secrets],
        )
    resources = None
    if include_resources:
        if not self.flavor_schema:
            raise RuntimeError(
                f"Missing flavor {self.flavor} for component {self.name}."
            )

        resources = ComponentResponseResources(
            user=self.user.to_model() if self.user else None,
            flavor=self.flavor_schema.to_model(),
        )
    return ComponentResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(component_update: ComponentUpdate) -> StackComponentSchema

Updates a StackComponentSchema from a ComponentUpdate.

Parameters:

Name	Type	Description	Default
component_update	ComponentUpdate	The ComponentUpdate to update from.	required

Returns:

Type	Description
StackComponentSchema	The updated StackComponentSchema.

Source code in src/zenml/zen_stores/schemas/component_schemas.py

def update(
    self, component_update: "ComponentUpdate"
) -> "StackComponentSchema":
    """Updates a `StackComponentSchema` from a `ComponentUpdate`.

    Args:
        component_update: The `ComponentUpdate` to update from.

    Returns:
        The updated `StackComponentSchema`.
    """
    for field, value in component_update.model_dump(
        exclude_unset=True,
        exclude={
            "user",
            "connector",
            "add_secrets",
            "remove_secrets",
            "attach_resource_pools",
            "detach_resource_pools",
        },
    ).items():
        if field == "configuration":
            self.configuration = base64.b64encode(
                json.dumps(component_update.configuration).encode("utf-8")
            )
        elif field == "labels":
            self.labels = base64.b64encode(
                json.dumps(component_update.labels).encode("utf-8")
            )
        elif field == "environment":
            self.environment = base64.b64encode(
                json.dumps(component_update.environment).encode("utf-8")
            )
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

constants

Constant values needed by schema objects.

curated_visualization_schemas

SQLModel implementation of curated visualization tables.

Classes

CuratedVisualizationSchema

Bases: BaseSchema

SQL Model for curated visualizations.

Functions

from_request(request: CuratedVisualizationRequest) -> CuratedVisualizationSchema classmethod

Convert a request into a schema instance.

Parameters:

Name	Type	Description	Default
request	CuratedVisualizationRequest	The request to convert.	required

Returns:

Type	Description
CuratedVisualizationSchema	The created schema.

Source code in src/zenml/zen_stores/schemas/curated_visualization_schemas.py

@classmethod
def from_request(
    cls, request: CuratedVisualizationRequest
) -> "CuratedVisualizationSchema":
    """Convert a request into a schema instance.

    Args:
        request: The request to convert.

    Returns:
        The created schema.
    """
    return cls(
        project_id=request.project,
        artifact_visualization_id=request.artifact_visualization_id,
        display_name=request.display_name,
        display_order=request.display_order,
        layout_size=request.layout_size.value,
        resource_id=request.resource_id,
        resource_type=request.resource_type.value,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/curated_visualization_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options: List[ExecutableOption] = []

    if include_resources:
        options.append(selectinload(jl_arg(cls.artifact_visualization)))

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> CuratedVisualizationResponse

Convert schema into response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
CuratedVisualizationResponse	The created response model.

Source code in src/zenml/zen_stores/schemas/curated_visualization_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> CuratedVisualizationResponse:
    """Convert schema into response model.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        **kwargs: Additional keyword arguments.

    Returns:
        The created response model.
    """
    try:
        layout_size_enum = CuratedVisualizationSize(self.layout_size)
    except ValueError:
        layout_size_enum = CuratedVisualizationSize.FULL_WIDTH

    try:
        resource_type_enum = VisualizationResourceTypes(self.resource_type)
    except ValueError:
        resource_type_enum = VisualizationResourceTypes.PROJECT

    artifact_version_id = self.artifact_visualization.artifact_version_id

    body = CuratedVisualizationResponseBody(
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        artifact_visualization_id=self.artifact_visualization_id,
        artifact_version_id=artifact_version_id,
        display_name=self.display_name,
        display_order=self.display_order,
        layout_size=layout_size_enum,
        resource_id=self.resource_id,
        resource_type=resource_type_enum,
    )

    metadata = None
    if include_metadata:
        metadata = CuratedVisualizationResponseMetadata()

    resources = None
    if include_resources:
        artifact_visualization = self.artifact_visualization.to_model(
            include_metadata=False,
            include_resources=False,
        )
        resources = CuratedVisualizationResponseResources(
            artifact_visualization=artifact_visualization,
        )

    return CuratedVisualizationResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: CuratedVisualizationUpdate) -> CuratedVisualizationSchema

Update a schema instance from an update model.

Parameters:

Name	Type	Description	Default
update	CuratedVisualizationUpdate	The update definition.	required

Returns:

Type	Description
CuratedVisualizationSchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/curated_visualization_schemas.py

def update(
    self,
    update: CuratedVisualizationUpdate,
) -> "CuratedVisualizationSchema":
    """Update a schema instance from an update model.

    Args:
        update: The update definition.

    Returns:
        The updated schema.
    """
    changes = update.model_dump(exclude_unset=True)
    layout_size_update = changes.pop("layout_size", None)
    if layout_size_update is not None:
        self.layout_size = layout_size_update.value

    for field, value in changes.items():
        if hasattr(self, field):
            setattr(self, field, value)

    from zenml.utils.time_utils import utc_now

    self.updated = utc_now()
    return self

Functions

deployment_schemas

SQLModel implementation of pipeline deployments table.

Classes

DeploymentSchema

Bases: NamedSchema

SQL Model for pipeline deployment.

Functions

from_request(request: DeploymentRequest) -> DeploymentSchema classmethod

Convert a DeploymentRequest to a DeploymentSchema.

Parameters:

Name	Type	Description	Default
request	DeploymentRequest	The request model to convert.	required

Returns:

Type	Description
DeploymentSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/deployment_schemas.py

@classmethod
def from_request(cls, request: DeploymentRequest) -> "DeploymentSchema":
    """Convert a `DeploymentRequest` to a `DeploymentSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        status=DeploymentStatus.UNKNOWN.value,
        snapshot_id=request.snapshot_id,
        deployer_id=request.deployer_id,
        auth_key=request.auth_key,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/deployment_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(DeploymentSchema.user)),
                selectinload(jl_arg(DeploymentSchema.deployer)),
                selectinload(jl_arg(DeploymentSchema.snapshot)).joinedload(
                    jl_arg(PipelineSnapshotSchema.pipeline)
                ),
                selectinload(jl_arg(DeploymentSchema.snapshot)).joinedload(
                    jl_arg(PipelineSnapshotSchema.stack)
                ),
                selectinload(jl_arg(DeploymentSchema.visualizations)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> DeploymentResponse

Convert a DeploymentSchema to a DeploymentResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
DeploymentResponse	The created DeploymentResponse.

Source code in src/zenml/zen_stores/schemas/deployment_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> DeploymentResponse:
    """Convert a `DeploymentSchema` to a `DeploymentResponse`.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        kwargs: Additional keyword arguments.

    Returns:
        The created `DeploymentResponse`.
    """
    status: Optional[DeploymentStatus] = None
    if self.status in DeploymentStatus.values():
        status = DeploymentStatus(self.status)
    elif self.status is not None:
        status = DeploymentStatus.UNKNOWN
        logger.warning(
            f"Deployment status '{self.status}' used for deployment "
            f"{self.name} is not a valid DeploymentStatus value. "
            "Using UNKNOWN instead."
        )

    body = DeploymentResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        url=self.url,
        status=status,
    )

    metadata = None
    if include_metadata:
        metadata = DeploymentResponseMetadata(
            deployment_metadata=json.loads(self.deployment_metadata),
            auth_key=self.auth_key,
        )

    resources = None
    if include_resources:
        resources = DeploymentResponseResources(
            user=self.user.to_model() if self.user else None,
            tags=[tag.to_model() for tag in self.tags],
            snapshot=self.snapshot.to_model() if self.snapshot else None,
            deployer=self.deployer.to_model() if self.deployer else None,
            pipeline=self.snapshot.pipeline.to_model()
            if self.snapshot and self.snapshot.pipeline
            else None,
            stack=self.snapshot.stack.to_model()
            if self.snapshot and self.snapshot.stack
            else None,
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    return DeploymentResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: DeploymentUpdate) -> DeploymentSchema

Updates a DeploymentSchema from a DeploymentUpdate.

Parameters:

Name	Type	Description	Default
update	DeploymentUpdate	The DeploymentUpdate to update from.	required

Returns:

Type	Description
DeploymentSchema	The updated DeploymentSchema.

Source code in src/zenml/zen_stores/schemas/deployment_schemas.py

def update(
    self,
    update: DeploymentUpdate,
) -> "DeploymentSchema":
    """Updates a `DeploymentSchema` from a `DeploymentUpdate`.

    Args:
        update: The `DeploymentUpdate` to update from.

    Returns:
        The updated `DeploymentSchema`.
    """
    for field, value in update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field == "deployment_metadata":
            setattr(self, field, json.dumps(value))
        elif hasattr(self, field):
            setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

device_schemas

SQLModel implementation for authorized OAuth2 devices.

Classes

OAuthDeviceSchema

Bases: BaseSchema

SQL Model for authorized OAuth2 devices.

Functions

from_request(request: OAuthDeviceInternalRequest) -> Tuple[OAuthDeviceSchema, str, str] classmethod

Create an authorized device DB entry from a device authorization request.

Parameters:

Name	Type	Description	Default
request	OAuthDeviceInternalRequest	The device authorization request.	required

Returns:

Type	Description
Tuple[OAuthDeviceSchema, str, str]	The created OAuthDeviceSchema, the user code and the device code.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

@classmethod
def from_request(
    cls, request: OAuthDeviceInternalRequest
) -> Tuple["OAuthDeviceSchema", str, str]:
    """Create an authorized device DB entry from a device authorization request.

    Args:
        request: The device authorization request.

    Returns:
        The created `OAuthDeviceSchema`, the user code and the device code.
    """
    user_code = cls._generate_user_code()
    device_code = cls._generate_device_code()
    hashed_user_code = cls._get_hashed_code(user_code)
    hashed_device_code = cls._get_hashed_code(device_code)
    now = utc_now()
    return (
        cls(
            client_id=request.client_id,
            user_code=hashed_user_code,
            device_code=hashed_device_code,
            status=OAuthDeviceStatus.PENDING.value,
            failed_auth_attempts=0,
            expires=now + timedelta(seconds=request.expires_in),
            os=request.os,
            ip_address=request.ip_address,
            hostname=request.hostname,
            python_version=request.python_version,
            zenml_version=request.zenml_version,
            city=request.city,
            region=request.region,
            country=request.country,
            created=now,
            updated=now,
        ),
        user_code,
        device_code,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(OAuthDeviceSchema.user)),
            ]
        )

    return options

internal_update(device_update: OAuthDeviceInternalUpdate) -> Tuple[OAuthDeviceSchema, Optional[str], Optional[str]]

Update an authorized device from an internal device update model.

Parameters:

Name	Type	Description	Default
device_update	OAuthDeviceInternalUpdate	The internal device update model.	required

Returns:

Type	Description
OAuthDeviceSchema	The updated OAuthDeviceSchema and the new user code and device
Optional[str]	code, if they were generated.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

def internal_update(
    self, device_update: OAuthDeviceInternalUpdate
) -> Tuple["OAuthDeviceSchema", Optional[str], Optional[str]]:
    """Update an authorized device from an internal device update model.

    Args:
        device_update: The internal device update model.

    Returns:
        The updated `OAuthDeviceSchema` and the new user code and device
        code, if they were generated.
    """
    now = utc_now()
    user_code: Optional[str] = None
    device_code: Optional[str] = None

    # This call also takes care of setting fields that have the same
    # name in the internal model and the schema.
    self.update(device_update)

    if device_update.expires_in is not None:
        if device_update.expires_in <= 0:
            self.expires = None
        else:
            self.expires = now + timedelta(
                seconds=device_update.expires_in
            )
    if device_update.update_last_login:
        self.last_login = now
    if device_update.generate_new_codes:
        user_code = self._generate_user_code()
        device_code = self._generate_device_code()
        self.user_code = self._get_hashed_code(user_code)
        self.device_code = self._get_hashed_code(device_code)
    self.updated = now
    return self, user_code, device_code

to_internal_model(include_metadata: bool = False, include_resources: bool = False) -> OAuthDeviceInternalResponse

Convert a device schema to an internal device response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False

Returns:

Type	Description
OAuthDeviceInternalResponse	The converted internal device response model.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

def to_internal_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
) -> OAuthDeviceInternalResponse:
    """Convert a device schema to an internal device response model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.

    Returns:
        The converted internal device response model.
    """
    device_model = self.to_model(
        include_metadata=include_metadata,
        include_resources=include_resources,
    )
    return OAuthDeviceInternalResponse(
        id=device_model.id,
        body=device_model.body,
        metadata=device_model.metadata,
        resources=device_model.resources,
        user_code=self.user_code,
        device_code=self.device_code,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> OAuthDeviceResponse

Convert a device schema to a device response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
OAuthDeviceResponse	The converted device response model.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> OAuthDeviceResponse:
    """Convert a device schema to a device response model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The converted device response model.
    """
    metadata = None
    if include_metadata:
        metadata = OAuthDeviceResponseMetadata(
            python_version=self.python_version,
            zenml_version=self.zenml_version,
            city=self.city,
            region=self.region,
            country=self.country,
            failed_auth_attempts=self.failed_auth_attempts,
            last_login=self.last_login,
        )

    body = OAuthDeviceResponseBody(
        user_id=self.user_id,
        created=self.created,
        updated=self.updated,
        client_id=self.client_id,
        expires=self.expires,
        trusted_device=self.trusted_device,
        status=OAuthDeviceStatus(self.status),
        os=self.os,
        ip_address=self.ip_address,
        hostname=self.hostname,
    )
    resources = None
    if include_resources:
        resources = OAuthDeviceResponseResources(
            user=self.user.to_model() if self.user else None,
        )
    return OAuthDeviceResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(device_update: OAuthDeviceUpdate) -> OAuthDeviceSchema

Update an authorized device from a device update model.

Parameters:

Name	Type	Description	Default
device_update	OAuthDeviceUpdate	The device update model.	required

Returns:

Type	Description
OAuthDeviceSchema	The updated OAuthDeviceSchema.

Source code in src/zenml/zen_stores/schemas/device_schemas.py

def update(self, device_update: OAuthDeviceUpdate) -> "OAuthDeviceSchema":
    """Update an authorized device from a device update model.

    Args:
        device_update: The device update model.

    Returns:
        The updated `OAuthDeviceSchema`.
    """
    for field, value in device_update.model_dump(
        exclude_none=True
    ).items():
        if hasattr(self, field):
            setattr(self, field, value)

    if device_update.locked is True:
        self.status = OAuthDeviceStatus.LOCKED.value
    elif device_update.locked is False:
        self.status = OAuthDeviceStatus.ACTIVE.value

    self.updated = utc_now()
    return self

Functions

flavor_schemas

SQL Model Implementations for Flavors.

Classes

FlavorSchema

Bases: NamedSchema

SQL Model for flavors.

Attributes:

Name	Type	Description
type	str	The type of the flavor.
source	str	The source of the flavor.
config_schema	str	The config schema of the flavor.
integration	Optional[str]	The integration associated with the flavor.

Functions

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/flavor_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(FlavorSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> FlavorResponse

Converts a flavor schema to a flavor model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
FlavorResponse	The flavor model.

Source code in src/zenml/zen_stores/schemas/flavor_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "FlavorResponse":
    """Converts a flavor schema to a flavor model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The flavor model.
    """
    body = FlavorResponseBody(
        user_id=self.user_id,
        type=StackComponentType(self.type),
        display_name=self.display_name
        or self.name.replace("_", " ").title(),
        integration=self.integration,
        source=self.source,
        logo_url=self.logo_url,
        is_custom=self.is_custom,
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = FlavorResponseMetadata(
            config_schema=json.loads(self.config_schema),
            connector_type=self.connector_type,
            connector_resource_type=self.connector_resource_type,
            connector_resource_id_attr=self.connector_resource_id_attr,
            docs_url=self.docs_url,
            sdk_docs_url=self.sdk_docs_url,
        )
    resources = None
    if include_resources:
        resources = FlavorResponseResources(
            user=self.user.to_model() if self.user else None,
        )
    return FlavorResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(flavor_update: FlavorUpdate) -> FlavorSchema

Update a FlavorSchema from a FlavorUpdate.

Parameters:

Name	Type	Description	Default
flavor_update	FlavorUpdate	The FlavorUpdate from which to update the schema.	required

Returns:

Type	Description
FlavorSchema	The updated FlavorSchema.

Source code in src/zenml/zen_stores/schemas/flavor_schemas.py

def update(
    self,
    flavor_update: "FlavorUpdate",
) -> "FlavorSchema":
    """Update a `FlavorSchema` from a `FlavorUpdate`.

    Args:
        flavor_update: The `FlavorUpdate` from which to update the schema.

    Returns:
        The updated `FlavorSchema`.
    """
    for field, value in flavor_update.model_dump(
        exclude_unset=True, exclude={"user"}
    ).items():
        if field == "config_schema":
            setattr(self, field, json.dumps(value))
        elif field == "type":
            setattr(self, field, value.value)
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

logs_schemas

SQLModel implementation of pipeline logs tables.

Classes

LogsSchema

Bases: BaseSchema

SQL Model for logs.

Functions

from_request(request: LogsRequest) -> LogsSchema classmethod

Create a LogsSchema from a LogsRequest.

Parameters:

Name	Type	Description	Default
request	LogsRequest	The LogsRequest to create the LogsSchema from.	required

Returns:

Type	Description
LogsSchema	The created LogsSchema.

Source code in src/zenml/zen_stores/schemas/logs_schemas.py

@classmethod
def from_request(cls, request: LogsRequest) -> "LogsSchema":
    """Create a `LogsSchema` from a `LogsRequest`.

    Args:
        request: The `LogsRequest` to create the `LogsSchema` from.

    Returns:
        The created `LogsSchema`.
    """
    return LogsSchema(
        id=request.id,
        uri=request.uri,
        source=request.source,
        project_id=request.project,
        user_id=request.user,
        pipeline_run_id=request.pipeline_run_id,
        step_run_id=request.step_run_id,
        artifact_store_id=request.artifact_store_id,
        log_store_id=request.log_store_id,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> LogsResponse

Convert a LogsSchema to a LogsResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
LogsResponse	The created LogsResponse.

Source code in src/zenml/zen_stores/schemas/logs_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "LogsResponse":
    """Convert a `LogsSchema` to a `LogsResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The created `LogsResponse`.
    """
    body = LogsResponseBody(
        uri=self.uri,
        source=self.source,
        created=self.created,
        updated=self.updated,
        project_id=self.project_id,
        user_id=self.user_id,
    )

    metadata = None
    if include_metadata:
        metadata = LogsResponseMetadata(
            step_run_id=self.step_run_id,
            pipeline_run_id=self.pipeline_run_id,
            artifact_store_id=self.artifact_store_id,
            log_store_id=self.log_store_id,
        )

    resources = None
    if include_resources:
        resources = LogsResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return LogsResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

Functions

model_schemas

SQLModel implementation of model tables.

Classes

ModelSchema

Bases: NamedSchema

SQL Model for model.

Attributes

latest_version: Optional[ModelVersionSchema] property

Fetch the latest version for this model.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[ModelVersionSchema]	The latest version for this model.

Functions

from_request(model_request: ModelRequest) -> ModelSchema classmethod

Convert an ModelRequest to an ModelSchema.

Parameters:

Name	Type	Description	Default
model_request	ModelRequest	The request model to convert.	required

Returns:

Type	Description
ModelSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def from_request(cls, model_request: ModelRequest) -> "ModelSchema":
    """Convert an `ModelRequest` to an `ModelSchema`.

    Args:
        model_request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=model_request.name,
        project_id=model_request.project,
        user_id=model_request.user,
        license=model_request.license,
        description=model_request.description,
        audience=model_request.audience,
        use_cases=model_request.use_cases,
        limitations=model_request.limitations,
        trade_offs=model_request.trade_offs,
        ethics=model_request.ethics,
        save_models_to_registry=model_request.save_models_to_registry,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ModelSchema.user)),
                # joinedload(jl_arg(ModelSchema.tags)),
                selectinload(jl_arg(ModelSchema.visualizations)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ModelResponse

Convert an ModelSchema to an ModelResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ModelResponse	The created ModelResponse.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ModelResponse:
    """Convert an `ModelSchema` to an `ModelResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ModelResponse`.
    """
    metadata = None
    if include_metadata:
        metadata = ModelResponseMetadata(
            license=self.license,
            description=self.description,
            audience=self.audience,
            use_cases=self.use_cases,
            limitations=self.limitations,
            trade_offs=self.trade_offs,
            ethics=self.ethics,
            save_models_to_registry=self.save_models_to_registry,
        )

    resources = None
    if include_resources:
        if latest_version := self.latest_version:
            latest_version_name = latest_version.name
            latest_version_id = latest_version.id
        else:
            latest_version_name = None
            latest_version_id = None

        resources = ModelResponseResources(
            user=self.user.to_model() if self.user else None,
            tags=[tag.to_model() for tag in self.tags],
            latest_version_name=latest_version_name,
            latest_version_id=latest_version_id,
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    body = ModelResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
    )

    return ModelResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(model_update: ModelUpdate) -> ModelSchema

Updates a ModelSchema from a ModelUpdate.

Parameters:

Name	Type	Description	Default
model_update	ModelUpdate	The ModelUpdate to update from.	required

Returns:

Type	Description
ModelSchema	The updated ModelSchema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def update(
    self,
    model_update: ModelUpdate,
) -> "ModelSchema":
    """Updates a `ModelSchema` from a `ModelUpdate`.

    Args:
        model_update: The `ModelUpdate` to update from.

    Returns:
        The updated `ModelSchema`.
    """
    for field, value in model_update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field in ["add_tags", "remove_tags"]:
            # Tags are handled separately
            continue
        setattr(self, field, value)
    self.updated = utc_now()
    return self

ModelVersionArtifactSchema

Bases: BaseSchema

SQL Model for linking of Model Versions and Artifacts M:M.

Functions

from_request(model_version_artifact_request: ModelVersionArtifactRequest) -> ModelVersionArtifactSchema classmethod

Convert an ModelVersionArtifactRequest to a ModelVersionArtifactSchema.

Parameters:

Name	Type	Description	Default
model_version_artifact_request	ModelVersionArtifactRequest	The request link to convert.	required

Returns:

Type	Description
ModelVersionArtifactSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def from_request(
    cls,
    model_version_artifact_request: ModelVersionArtifactRequest,
) -> "ModelVersionArtifactSchema":
    """Convert an `ModelVersionArtifactRequest` to a `ModelVersionArtifactSchema`.

    Args:
        model_version_artifact_request: The request link to convert.

    Returns:
        The converted schema.
    """
    return cls(
        model_version_id=model_version_artifact_request.model_version,
        artifact_version_id=model_version_artifact_request.artifact_version,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ModelVersionArtifactResponse

Convert an ModelVersionArtifactSchema to an ModelVersionArtifactResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ModelVersionArtifactResponse	The created ModelVersionArtifactResponseModel.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ModelVersionArtifactResponse:
    """Convert an `ModelVersionArtifactSchema` to an `ModelVersionArtifactResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ModelVersionArtifactResponseModel`.
    """
    return ModelVersionArtifactResponse(
        id=self.id,
        body=ModelVersionArtifactResponseBody(
            created=self.created,
            updated=self.updated,
            model_version=self.model_version_id,
            artifact_version=self.artifact_version.to_model(),
        ),
        metadata=BaseResponseMetadata() if include_metadata else None,
    )

ModelVersionPipelineRunSchema

Bases: BaseSchema

SQL Model for linking of Model Versions and Pipeline Runs M:M.

Functions

from_request(model_version_pipeline_run_request: ModelVersionPipelineRunRequest) -> ModelVersionPipelineRunSchema classmethod

Convert an ModelVersionPipelineRunRequest to an ModelVersionPipelineRunSchema.

Parameters:

Name	Type	Description	Default
model_version_pipeline_run_request	ModelVersionPipelineRunRequest	The request link to convert.	required

Returns:

Type	Description
ModelVersionPipelineRunSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def from_request(
    cls,
    model_version_pipeline_run_request: ModelVersionPipelineRunRequest,
) -> "ModelVersionPipelineRunSchema":
    """Convert an `ModelVersionPipelineRunRequest` to an `ModelVersionPipelineRunSchema`.

    Args:
        model_version_pipeline_run_request: The request link to convert.

    Returns:
        The converted schema.
    """
    return cls(
        model_version_id=model_version_pipeline_run_request.model_version,
        pipeline_run_id=model_version_pipeline_run_request.pipeline_run,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ModelVersionPipelineRunResponse

Convert an ModelVersionPipelineRunSchema to an ModelVersionPipelineRunResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ModelVersionPipelineRunResponse	The created ModelVersionPipelineRunResponse.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ModelVersionPipelineRunResponse:
    """Convert an `ModelVersionPipelineRunSchema` to an `ModelVersionPipelineRunResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ModelVersionPipelineRunResponse`.
    """
    return ModelVersionPipelineRunResponse(
        id=self.id,
        body=ModelVersionPipelineRunResponseBody(
            created=self.created,
            updated=self.updated,
            model_version=self.model_version_id,
            pipeline_run=self.pipeline_run.to_model(),
        ),
        metadata=BaseResponseMetadata() if include_metadata else None,
    )

ModelVersionSchema

Bases: NamedSchema, RunMetadataInterface

SQL Model for model version.

Functions

from_request(model_version_request: ModelVersionRequest, model_version_number: int, producer_run_id: Optional[UUID] = None) -> ModelVersionSchema classmethod

Convert an ModelVersionRequest to an ModelVersionSchema.

Parameters:

Name	Type	Description	Default
model_version_request	ModelVersionRequest	The request model version to convert.	required
model_version_number	int	The model version number.	required
producer_run_id	Optional[UUID]	The ID of the producer run.	None

Returns:

Type	Description
ModelVersionSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def from_request(
    cls,
    model_version_request: ModelVersionRequest,
    model_version_number: int,
    producer_run_id: Optional[UUID] = None,
) -> "ModelVersionSchema":
    """Convert an `ModelVersionRequest` to an `ModelVersionSchema`.

    Args:
        model_version_request: The request model version to convert.
        model_version_number: The model version number.
        producer_run_id: The ID of the producer run.

    Returns:
        The converted schema.
    """
    id_ = uuid4()
    is_numeric = str(model_version_number) == model_version_request.name

    return cls(
        id=id_,
        project_id=model_version_request.project,
        user_id=model_version_request.user,
        model_id=model_version_request.model,
        name=model_version_request.name,
        number=model_version_number,
        description=model_version_request.description,
        stage=model_version_request.stage,
        producer_run_id_if_numeric=producer_run_id
        if (producer_run_id and is_numeric)
        else id_,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = [
        joinedload(jl_arg(ModelVersionSchema.model), innerjoin=True),
    ]

    # if include_metadata:
    #     options.extend(
    #         [
    #             joinedload(jl_arg(ModelVersionSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ModelVersionSchema.user)),
                # joinedload(jl_arg(ModelVersionSchema.services)),
                # joinedload(jl_arg(ModelVersionSchema.tags)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ModelVersionResponse

Convert an ModelVersionSchema to an ModelVersionResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ModelVersionResponse	The created ModelVersionResponse.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ModelVersionResponse:
    """Convert an `ModelVersionSchema` to an `ModelVersionResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ModelVersionResponse`.
    """
    from zenml.models import ServiceResponse

    metadata = None
    if include_metadata:
        metadata = ModelVersionResponseMetadata(
            description=self.description,
            run_metadata=self.fetch_metadata(),
        )

    resources = None
    if include_resources:
        services = cast(
            Page[ServiceResponse],
            get_page_from_list(
                items_list=self.services,
                response_model=ServiceResponse,
                include_resources=include_resources,
                include_metadata=include_metadata,
            ),
        )
        resources = ModelVersionResponseResources(
            user=self.user.to_model() if self.user else None,
            services=services,
            tags=[tag.to_model() for tag in self.tags],
        )

    body = ModelVersionResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        stage=self.stage,
        number=self.number,
        model=self.model.to_model(),
    )

    return ModelVersionResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(target_stage: Optional[str] = None, target_name: Optional[str] = None, target_description: Optional[str] = None) -> ModelVersionSchema

Updates a ModelVersionSchema to a target stage.

Parameters:

Name	Type	Description	Default
target_stage	Optional[str]	The stage to be updated.	None
target_name	Optional[str]	The version name to be updated.	None
target_description	Optional[str]	The version description to be updated.	None

Returns:

Type	Description
ModelVersionSchema	The updated ModelVersionSchema.

Source code in src/zenml/zen_stores/schemas/model_schemas.py

def update(
    self,
    target_stage: Optional[str] = None,
    target_name: Optional[str] = None,
    target_description: Optional[str] = None,
) -> "ModelVersionSchema":
    """Updates a `ModelVersionSchema` to a target stage.

    Args:
        target_stage: The stage to be updated.
        target_name: The version name to be updated.
        target_description: The version description to be updated.

    Returns:
        The updated `ModelVersionSchema`.
    """
    if target_stage is not None:
        self.stage = target_stage
    if target_name is not None:
        self.name = target_name
    if target_description is not None:
        self.description = target_description
    self.updated = utc_now()
    return self

Functions

pipeline_build_schemas

SQLModel implementation of pipeline build tables.

Classes

PipelineBuildSchema

Bases: BaseSchema

SQL Model for pipeline builds.

Functions

from_request(request: PipelineBuildRequest) -> PipelineBuildSchema classmethod

Convert a PipelineBuildRequest to a PipelineBuildSchema.

Parameters:

Name	Type	Description	Default
request	PipelineBuildRequest	The request to convert.	required

Returns:

Type	Description
PipelineBuildSchema	The created PipelineBuildSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_build_schemas.py

@classmethod
def from_request(
    cls, request: PipelineBuildRequest
) -> "PipelineBuildSchema":
    """Convert a `PipelineBuildRequest` to a `PipelineBuildSchema`.

    Args:
        request: The request to convert.

    Returns:
        The created `PipelineBuildSchema`.
    """
    return cls(
        stack_id=request.stack,
        project_id=request.project,
        user_id=request.user,
        pipeline_id=request.pipeline,
        images=json.dumps(request.images, default=pydantic_encoder),
        is_local=request.is_local,
        contains_code=request.contains_code,
        zenml_version=request.zenml_version,
        python_version=request.python_version,
        checksum=request.checksum,
        stack_checksum=request.stack_checksum,
        duration=request.duration,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/pipeline_build_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_metadata:
        options.extend(
            [
                joinedload(jl_arg(PipelineBuildSchema.pipeline)),
                joinedload(jl_arg(PipelineBuildSchema.stack)),
            ]
        )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(PipelineBuildSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> PipelineBuildResponse

Convert a PipelineBuildSchema to a PipelineBuildResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
PipelineBuildResponse	The created PipelineBuildResponse.

Source code in src/zenml/zen_stores/schemas/pipeline_build_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> PipelineBuildResponse:
    """Convert a `PipelineBuildSchema` to a `PipelineBuildResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `PipelineBuildResponse`.
    """
    body = PipelineBuildResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = PipelineBuildResponseMetadata(
            pipeline=self.pipeline.to_model() if self.pipeline else None,
            stack=self.stack.to_model() if self.stack else None,
            images=json.loads(self.images),
            zenml_version=self.zenml_version,
            python_version=self.python_version,
            checksum=self.checksum,
            stack_checksum=self.stack_checksum,
            is_local=self.is_local,
            contains_code=self.contains_code,
            duration=self.duration,
        )

    resources = None
    if include_resources:
        resources = PipelineBuildResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return PipelineBuildResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

Functions

pipeline_run_schemas

SQLModel implementation of pipeline run tables.

Classes

PipelineRunOutputSchema

Bases: BaseSchema

SQL model defining pipeline run outputs.

PipelineRunSchema

Bases: NamedSchema, RunMetadataInterface

SQL Model for pipeline runs.

Functions

fetch_metadata_collection(include_full_metadata: bool = False, **kwargs: Any) -> Dict[str, List[RunMetadataEntry]]

Fetches all the metadata entries related to the pipeline run.

Parameters:

Name	Type	Description	Default
include_full_metadata	bool	Whether the full metadata will be included.	False
**kwargs	Any	Keyword arguments.	{}

Returns:

Type	Description
Dict[str, List[RunMetadataEntry]]	a dictionary, where the key is the key of the metadata entry and the values represent the list of entries with this key.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def fetch_metadata_collection(
    self, include_full_metadata: bool = False, **kwargs: Any
) -> Dict[str, List[RunMetadataEntry]]:
    """Fetches all the metadata entries related to the pipeline run.

    Args:
        include_full_metadata: Whether the full metadata will be included.
        **kwargs: Keyword arguments.

    Returns:
        a dictionary, where the key is the key of the metadata entry
            and the values represent the list of entries with this key.
    """
    # Fetch the metadata related to this run
    metadata_collection = super().fetch_metadata_collection(**kwargs)

    if include_full_metadata:
        # Fetch the metadata related to the steps of this run
        for s in self.step_runs:
            step_metadata = s.fetch_metadata_collection()
            for k, v in step_metadata.items():
                metadata_collection[f"{s.name}::{k}"] = v

        # Fetch the metadata related to the schedule of this run
        if self.snapshot is not None:
            if schedule := self.snapshot.schedule:
                schedule_metadata = schedule.fetch_metadata_collection()
                for k, v in schedule_metadata.items():
                    metadata_collection[f"schedule:{k}"] = v

    return metadata_collection

from_request(request: PipelineRunRequest, pipeline_id: UUID, index: int, enable_heartbeat: bool, root_run_id: Optional[UUID] = None) -> PipelineRunSchema classmethod

Convert a PipelineRunRequest to a PipelineRunSchema.

Parameters:

Name	Type	Description	Default
request	PipelineRunRequest	The request to convert.	required
pipeline_id	UUID	The ID of the pipeline.	required
index	int	The index of the pipeline run.	required
enable_heartbeat	bool	Whether the heartbeat should be enabled.	required
root_run_id	Optional[UUID]	The root_run_id of the parent run, if this run is a child run.	None

Returns:

Type	Description
PipelineRunSchema	The created PipelineRunSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

@classmethod
def from_request(
    cls,
    request: "PipelineRunRequest",
    pipeline_id: UUID,
    index: int,
    enable_heartbeat: bool,
    root_run_id: Optional[UUID] = None,
) -> "PipelineRunSchema":
    """Convert a `PipelineRunRequest` to a `PipelineRunSchema`.

    Args:
        request: The request to convert.
        pipeline_id: The ID of the pipeline.
        index: The index of the pipeline run.
        enable_heartbeat: Whether the heartbeat should be enabled.
        root_run_id: The `root_run_id` of the parent run, if this
            run is a child run.

    Returns:
        The created `PipelineRunSchema`.
    """
    orchestrator_environment = json.dumps(request.orchestrator_environment)
    if len(orchestrator_environment) > TEXT_FIELD_MAX_LENGTH:
        logger.warning(
            "Orchestrator environment is too large to be stored in the "
            "database. Skipping."
        )
        orchestrator_environment = "{}"

    triggered_by = None
    triggered_by_type = None
    if request.trigger_info:
        if request.trigger_info.step_run_id:
            triggered_by = request.trigger_info.step_run_id
            triggered_by_type = PipelineRunTriggeredByType.STEP_RUN.value
        elif request.trigger_info.deployment_id:
            triggered_by = request.trigger_info.deployment_id
            triggered_by_type = PipelineRunTriggeredByType.DEPLOYMENT.value

    return cls(
        project_id=request.project,
        user_id=request.user,
        name=request.name,
        orchestrator_run_id=request.orchestrator_run_id,
        orchestrator_environment=orchestrator_environment,
        start_time=request.start_time,
        end_time=request.end_time,
        status=request.status.value,
        index=index,
        in_progress=not request.status.is_finished,
        status_reason=request.status_reason,
        pipeline_id=pipeline_id,
        snapshot_id=request.snapshot,
        triggered_by=triggered_by,
        triggered_by_type=triggered_by_type,
        enable_heartbeat=enable_heartbeat,
        exception_info=request.exception_info.model_dump_json()
        if request.exception_info
        else None,
        original_run_id=request.original_run_id,
        parent_run_id=request.parent_run_id,
        child_key=request.child_key,
        root_run_id=root_run_id,
    )

get_pipeline_configuration() -> PipelineConfiguration

Get the pipeline configuration for the pipeline run.

Raises:

Type	Description
RuntimeError	if the pipeline run has no snapshot and no pipeline configuration.

Returns:

Type	Description
PipelineConfiguration	The pipeline configuration.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def get_pipeline_configuration(self) -> PipelineConfiguration:
    """Get the pipeline configuration for the pipeline run.

    Raises:
        RuntimeError: if the pipeline run has no snapshot and no pipeline
            configuration.

    Returns:
        The pipeline configuration.
    """
    if self.snapshot:
        pipeline_config = PipelineConfiguration.model_validate_json(
            self.snapshot.pipeline_configuration
        )
    elif self.pipeline_configuration:
        pipeline_config = PipelineConfiguration.model_validate_json(
            self.pipeline_configuration
        )
    else:
        raise RuntimeError(
            "Pipeline run has no snapshot and no pipeline configuration."
        )

    pipeline_config.finalize_substitutions(
        start_time=self.start_time, inplace=True
    )
    return pipeline_config

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    from zenml.zen_stores.schemas import ModelVersionSchema

    options = []

    if include_metadata:
        options.extend(
            [
                selectinload(jl_arg(PipelineRunSchema.trigger_execution)),
            ]
        )

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(PipelineRunSchema.outputs)),
                selectinload(jl_arg(PipelineRunSchema.parent_run)),
                selectinload(
                    jl_arg(PipelineRunSchema.model_version)
                ).joinedload(
                    jl_arg(ModelVersionSchema.model), innerjoin=True
                ),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(
                    jl_arg(PipelineSnapshotSchema.source_snapshot)
                ),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.pipeline)),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.stack)),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.build)),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.schedule)),
                selectinload(
                    jl_arg(PipelineRunSchema.snapshot)
                ).joinedload(
                    jl_arg(PipelineSnapshotSchema.code_reference)
                ),
                selectinload(jl_arg(PipelineRunSchema.logs)),
                selectinload(jl_arg(PipelineRunSchema.wait_conditions)),
                selectinload(jl_arg(PipelineRunSchema.user)),
                selectinload(jl_arg(PipelineRunSchema.tags)),
                selectinload(jl_arg(PipelineRunSchema.visualizations)),
                joinedload(jl_arg(PipelineRunSchema.trigger)),
            ]
        )

    return options

get_step_configuration(step_name: str) -> Step

Get the step configuration for the pipeline run.

Parameters:

Name	Type	Description	Default
step_name	str	The name of the step to get the configuration for.	required

Raises:

Type	Description
RuntimeError	If the pipeline run has no snapshot.

Returns:

Type	Description
Step	The step configuration.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def get_step_configuration(self, step_name: str) -> Step:
    """Get the step configuration for the pipeline run.

    Args:
        step_name: The name of the step to get the configuration for.

    Raises:
        RuntimeError: If the pipeline run has no snapshot.

    Returns:
        The step configuration.
    """
    if self.snapshot:
        pipeline_configuration = self.get_pipeline_configuration()
        return Step.from_dict(
            data=json.loads(
                self.snapshot.get_step_configuration(step_name).config
            ),
            pipeline_configuration=pipeline_configuration,
        )
    else:
        raise RuntimeError("Pipeline run has no snapshot.")

get_upstream_steps() -> Dict[str, List[str]]

Get the list of all the upstream steps for each step.

Returns:

Type	Description
Dict[str, List[str]]	The list of upstream steps for each step.

Raises:

Type	Description
RuntimeError	If the pipeline run has no snapshot or the snapshot has no pipeline spec.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def get_upstream_steps(self) -> Dict[str, List[str]]:
    """Get the list of all the upstream steps for each step.

    Returns:
        The list of upstream steps for each step.

    Raises:
        RuntimeError: If the pipeline run has no snapshot or
            the snapshot has no pipeline spec.
    """
    if self.snapshot and self.snapshot.pipeline_spec:
        pipeline_spec = PipelineSpec.model_validate_json(
            self.snapshot.pipeline_spec
        )
        steps = {}
        for step_spec in pipeline_spec.steps:
            steps[step_spec.invocation_id] = step_spec.upstream_steps
        return steps
    else:
        raise RuntimeError("Pipeline run has no snapshot.")

is_placeholder_run() -> bool

Whether the pipeline run is a placeholder run.

Returns:

Type	Description
bool	Whether the pipeline run is a placeholder run.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def is_placeholder_run(self) -> bool:
    """Whether the pipeline run is a placeholder run.

    Returns:
        Whether the pipeline run is a placeholder run.
    """
    return self.status in {
        ExecutionStatus.INITIALIZING.value,
        ExecutionStatus.PROVISIONING.value,
    }

to_model(include_metadata: bool = False, include_resources: bool = False, include_python_packages: bool = False, include_full_metadata: bool = False, **kwargs: Any) -> PipelineRunResponse

Convert a PipelineRunSchema to a PipelineRunResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
include_python_packages	bool	Whether the python packages will be filled.	False
include_full_metadata	bool	Whether the full metadata will be included.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
PipelineRunResponse	The created PipelineRunResponse.

Raises:

Type	Description
RuntimeError	if the model creation fails.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    include_python_packages: bool = False,
    include_full_metadata: bool = False,
    **kwargs: Any,
) -> "PipelineRunResponse":
    """Convert a `PipelineRunSchema` to a `PipelineRunResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        include_python_packages: Whether the python packages will be filled.
        include_full_metadata: Whether the full metadata will be included.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `PipelineRunResponse`.

    Raises:
        RuntimeError: if the model creation fails.
    """
    if self.snapshot is not None:
        config = PipelineConfiguration.model_validate_json(
            self.snapshot.pipeline_configuration
        )
        client_environment = json.loads(self.snapshot.client_environment)
    elif self.pipeline_configuration is not None:
        config = PipelineConfiguration.model_validate_json(
            self.pipeline_configuration
        )
        client_environment = (
            json.loads(self.client_environment)
            if self.client_environment
            else {}
        )
    else:
        raise RuntimeError(
            "Pipeline run model creation has failed. Each pipeline run "
            "entry should either have a snapshot_id or "
            "pipeline_configuration."
        )

    config.finalize_substitutions(start_time=self.start_time, inplace=True)

    body = PipelineRunResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        status=ExecutionStatus(self.status),
        status_reason=self.status_reason,
        created=self.created,
        updated=self.updated,
        in_progress=self.in_progress,
        index=self.index,
        pipeline_id=self.pipeline_id,
        child_key=self.child_key,
        root_run_id=self.root_run_id,
    )
    metadata = None
    if include_metadata:
        is_templatable = False
        if (
            self.snapshot
            and self.snapshot.build
            and not self.snapshot.build.is_local
            and self.snapshot.build.stack_id
        ):
            is_templatable = True

        orchestrator_environment = (
            json.loads(self.orchestrator_environment)
            if self.orchestrator_environment
            else {}
        )

        if not include_python_packages:
            client_environment.pop("python_packages", None)
            orchestrator_environment.pop("python_packages", None)

        trigger_info: Optional[PipelineRunTriggerInfo] = None
        if self.triggered_by and self.triggered_by_type:
            if (
                self.triggered_by_type
                == PipelineRunTriggeredByType.STEP_RUN.value
            ):
                trigger_info = PipelineRunTriggerInfo(
                    step_run_id=self.triggered_by,
                )
            elif (
                self.triggered_by_type
                == PipelineRunTriggeredByType.DEPLOYMENT.value
            ):
                trigger_info = PipelineRunTriggerInfo(
                    deployment_id=self.triggered_by,
                )

        metadata = PipelineRunResponseMetadata(
            run_metadata=self.fetch_metadata(
                include_full_metadata=include_full_metadata
            ),
            config=config,
            start_time=self.start_time,
            end_time=self.end_time,
            client_environment=client_environment,
            orchestrator_environment=orchestrator_environment,
            orchestrator_run_id=self.orchestrator_run_id,
            code_path=self.snapshot.code_path if self.snapshot else None,
            template_id=self.snapshot.template_id
            if self.snapshot
            else None,
            is_templatable=is_templatable,
            trigger_info=trigger_info,
            enable_heartbeat=self.enable_heartbeat,
            exception_info=json.loads(self.exception_info)
            if self.exception_info
            else None,
            trigger_execution_info=json.loads(self.trigger_execution.info)
            if self.trigger_execution and self.trigger_execution.info
            else None,
        )

    resources = None
    if include_resources:
        if self.snapshot:
            source_snapshot = (
                self.snapshot.source_snapshot.to_model()
                if self.snapshot.source_snapshot
                else None
            )
            stack = (
                self.snapshot.stack.to_model()
                if self.snapshot.stack
                else None
            )
            pipeline: Optional["PipelineResponse"] = (
                self.snapshot.pipeline.to_model()
            )
            build = (
                self.snapshot.build.to_model()
                if self.snapshot.build
                else None
            )
            schedule = (
                self.snapshot.schedule.to_model()
                if self.snapshot.schedule
                else None
            )
            code_reference = (
                self.snapshot.code_reference.to_model()
                if self.snapshot.code_reference
                else None
            )
        else:
            source_snapshot = None
            stack = self.stack.to_model() if self.stack else None
            pipeline = self.pipeline.to_model() if self.pipeline else None
            build = self.build.to_model() if self.build else None
            schedule = self.schedule.to_model() if self.schedule else None
            code_reference = None

        resources = PipelineRunResponseResources(
            user=self.user.to_model() if self.user else None,
            snapshot=self.snapshot.to_model() if self.snapshot else None,
            source_snapshot=source_snapshot,
            stack=stack,
            pipeline=pipeline,
            build=build,
            schedule=schedule,
            code_reference=code_reference,
            model_version=self.model_version.to_model()
            if self.model_version
            else None,
            tags=[tag.to_model() for tag in self.tags],
            log_collection=[
                log.to_model()
                for log in sorted(self.logs, key=lambda log: log.created)
            ],
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
            trigger=self.trigger.to_model() if self.trigger else None,
            original_run=self.original_run.to_model()
            if self.original_run
            else None,
            parent_run=self.parent_run.to_model()
            if self.parent_run
            else None,
            active_wait_condition=next(
                (
                    condition.to_model()
                    for condition in self.wait_conditions
                    if condition.status
                    == RunWaitConditionStatus.PENDING.value
                ),
                None,
            ),
            outputs={
                output.name: output.artifact_version.to_model()
                for output in sorted(
                    self.outputs, key=lambda output: output.output_index
                )
            },
        )

    return PipelineRunResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(run_update: PipelineRunUpdate) -> PipelineRunSchema

Update a PipelineRunSchema with a PipelineRunUpdate.

Parameters:

Name	Type	Description	Default
run_update	PipelineRunUpdate	The PipelineRunUpdate to update with.	required

Raises:

Type	Description
ValueError	When trying to update the orchestrator run ID of a run that already has a different one.

Returns:

Type	Description
PipelineRunSchema	The updated PipelineRunSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def update(self, run_update: "PipelineRunUpdate") -> "PipelineRunSchema":
    """Update a `PipelineRunSchema` with a `PipelineRunUpdate`.

    Args:
        run_update: The `PipelineRunUpdate` to update with.

    Raises:
        ValueError: When trying to update the orchestrator run ID of a
            run that already has a different one.

    Returns:
        The updated `PipelineRunSchema`.
    """
    if run_update.orchestrator_run_id:
        if (
            self.orchestrator_run_id
            and self.orchestrator_run_id != run_update.orchestrator_run_id
        ):
            raise ValueError(
                "Updating the orchestrator run ID of a run with an "
                "existing orchestrator run ID "
                f"({self.orchestrator_run_id}) is not allowed."
            )
        self.orchestrator_run_id = run_update.orchestrator_run_id

    if run_update.exception_info:
        self.exception_info = run_update.exception_info.model_dump_json()

    self.updated = utc_now()
    return self

update_placeholder(request: PipelineRunRequest) -> PipelineRunSchema

Update a placeholder run.

Parameters:

Name	Type	Description	Default
request	PipelineRunRequest	The pipeline run request which should replace the placeholder.	required

Raises:

Type	Description
RuntimeError	If the DB entry does not represent a placeholder run.
ValueError	If the run request is not a valid request to replace the placeholder run.

Returns:

Type	Description
PipelineRunSchema	The updated PipelineRunSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def update_placeholder(
    self, request: "PipelineRunRequest"
) -> "PipelineRunSchema":
    """Update a placeholder run.

    Args:
        request: The pipeline run request which should replace the
            placeholder.

    Raises:
        RuntimeError: If the DB entry does not represent a placeholder run.
        ValueError: If the run request is not a valid request to replace the
            placeholder run.

    Returns:
        The updated `PipelineRunSchema`.
    """
    if not self.is_placeholder_run():
        raise RuntimeError(
            f"Unable to replace pipeline run {self.id} which is not a "
            "placeholder run."
        )

    if request.is_placeholder_request:
        raise ValueError(
            "Cannot replace a placeholder run with another placeholder run."
        )

    if (
        self.snapshot_id != request.snapshot
        or self.project_id != request.project
    ):
        raise ValueError(
            "Snapshot or project ID of placeholder run "
            "do not match the IDs of the run request."
        )

    if not request.orchestrator_run_id:
        raise ValueError(
            "Orchestrator run ID is required to replace a placeholder run."
        )

    if (
        self.orchestrator_run_id
        and self.orchestrator_run_id != request.orchestrator_run_id
    ):
        raise ValueError(
            "Orchestrator run ID of placeholder run does not match the "
            "ID of the run request."
        )

    orchestrator_environment = json.dumps(request.orchestrator_environment)

    self.orchestrator_run_id = request.orchestrator_run_id
    self.orchestrator_environment = orchestrator_environment
    self.status = request.status.value
    self.in_progress = not request.status.is_finished

    self.updated = utc_now()

    return self

update_status(requested_status: Optional[ExecutionStatus], status_reason: Optional[str] = None) -> bool

Update the status of the pipeline run.

Parameters:

Name	Type	Description	Default
requested_status	Optional[ExecutionStatus]	The requested status of the pipeline run.	required
status_reason	Optional[str]	The reason for the status of the pipeline run.	None

Raises:

Type	Description
IllegalOperationError	If the requested status transition is invalid.

Returns:

Type	Description
bool	Whether the status was updated.

Source code in src/zenml/zen_stores/schemas/pipeline_run_schemas.py

def update_status(
    self,
    requested_status: Optional[ExecutionStatus],
    status_reason: Optional[str] = None,
) -> bool:
    """Update the status of the pipeline run.

    Args:
        requested_status: The requested status of the pipeline run.
        status_reason: The reason for the status of the pipeline run.

    Raises:
        IllegalOperationError: If the requested status transition is
            invalid.

    Returns:
        Whether the status was updated.
    """
    if requested_status in {
        ExecutionStatus.CACHED,
        ExecutionStatus.SKIPPED,
        ExecutionStatus.RETRYING,
        ExecutionStatus.RETRIED,
    }:
        raise IllegalOperationError(
            f"Execution status `{requested_status}` is not valid for "
            "pipeline runs."
        )

    current_status = ExecutionStatus(self.status)

    if (
        requested_status == ExecutionStatus.RESUMING
        and current_status
        not in {
            ExecutionStatus.PAUSED,
            ExecutionStatus.FAILED,
        }
    ):
        raise IllegalOperationError(
            "Only failed or paused runs can be resumed."
        )

    if (
        requested_status == ExecutionStatus.RESUMING
        and self.parent_run_id is not None
    ):
        raise IllegalOperationError(
            "Cannot resume a child run. Resume the parent run instead."
        )

    if (
        requested_status == ExecutionStatus.PROVISIONING
        and current_status != ExecutionStatus.INITIALIZING
    ):
        # Ignore transitions to provisioning from non-initializing states.
        # This could happen if the orchestrator starts running the pipeline
        # before the client environment can update the status to provisioning.
        return False

    # Snapshot always exists for pipeline runs of newer versions
    assert self.snapshot
    is_dynamic_pipeline = self.snapshot.is_dynamic

    if is_dynamic_pipeline:
        # In dynamic pipelines, the run status is only updated on manual
        # status updates and does not depend on step statuses.
        if requested_status is None and status_reason is None:
            return False

        new_status = requested_status or current_status
    else:
        # For static pipelines we compute the run status based on the step
        # statuses.
        new_status = _compute_static_pipeline_run_status(
            run_status=requested_status or current_status,
            step_statuses=self._get_step_run_statuses(),
            num_steps=self.snapshot.step_count,
        )

    if current_status.is_finished:
        if (
            current_status == ExecutionStatus.FAILED
            and new_status == ExecutionStatus.RESUMING
        ):
            # Allow failed -> resuming transition for retries.
            pass
        elif current_status != new_status:
            raise IllegalOperationError(
                "Cannot update the status of a finished run."
            )

    self.status = new_status.value

    if is_dynamic_pipeline:
        self.in_progress = not new_status.is_finished
    else:
        self.in_progress = self._check_if_run_in_progress()

    now = utc_now()
    if not self.in_progress and self.end_time is None:
        self.end_time = now
    elif new_status == ExecutionStatus.RESUMING:
        self.end_time = None

    if status_reason:
        self.status_reason = status_reason
    elif (
        current_status == ExecutionStatus.STOPPING
        and new_status == ExecutionStatus.STOPPED
    ):
        # Don't clear status reason for stopping -> stopped transition
        pass
    elif current_status != new_status:
        # Clear status reason when the status changes.
        self.status_reason = None

    self.updated = now
    return True

Functions

pipeline_schemas

SQL Model Implementations for Pipelines and Pipeline Runs.

Classes

PipelineSchema

Bases: NamedSchema

SQL Model for pipelines.

Attributes

latest_run: Optional[PipelineRunSchema] property

Fetch the latest run for this pipeline.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[PipelineRunSchema]	The latest run for this pipeline.

Functions

from_request(pipeline_request: PipelineRequest) -> PipelineSchema classmethod

Convert a PipelineRequest to a PipelineSchema.

Parameters:

Name	Type	Description	Default
pipeline_request	PipelineRequest	The request model to convert.	required

Returns:

Type	Description
PipelineSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/pipeline_schemas.py

@classmethod
def from_request(
    cls,
    pipeline_request: "PipelineRequest",
) -> "PipelineSchema":
    """Convert a `PipelineRequest` to a `PipelineSchema`.

    Args:
        pipeline_request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=pipeline_request.name,
        description=pipeline_request.description,
        project_id=pipeline_request.project,
        user_id=pipeline_request.user,
        run_count=0,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/pipeline_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(PipelineSchema.user)),
                # joinedload(jl_arg(PipelineSchema.tags)),
                selectinload(jl_arg(PipelineSchema.visualizations)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> PipelineResponse

Convert a PipelineSchema to a PipelineResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
PipelineResponse	The created PipelineResponse.

Source code in src/zenml/zen_stores/schemas/pipeline_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "PipelineResponse":
    """Convert a `PipelineSchema` to a `PipelineResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The created PipelineResponse.
    """
    body = PipelineResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
    )

    metadata = None
    if include_metadata:
        metadata = PipelineResponseMetadata(
            description=self.description,
        )

    resources = None
    if include_resources:
        latest_run = self.latest_run
        latest_run_user = latest_run.user if latest_run else None

        resources = PipelineResponseResources(
            user=self.user.to_model() if self.user else None,
            latest_run_user=latest_run_user.to_model()
            if latest_run_user
            else None,
            latest_run_id=latest_run.id if latest_run else None,
            latest_run_status=latest_run.status if latest_run else None,
            tags=[tag.to_model() for tag in self.tags],
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    return PipelineResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(pipeline_update: PipelineUpdate) -> PipelineSchema

Update a PipelineSchema with a PipelineUpdate.

Parameters:

Name	Type	Description	Default
pipeline_update	PipelineUpdate	The update model.	required

Returns:

Type	Description
PipelineSchema	The updated PipelineSchema.

Source code in src/zenml/zen_stores/schemas/pipeline_schemas.py

def update(self, pipeline_update: "PipelineUpdate") -> "PipelineSchema":
    """Update a `PipelineSchema` with a `PipelineUpdate`.

    Args:
        pipeline_update: The update model.

    Returns:
        The updated `PipelineSchema`.
    """
    self.description = pipeline_update.description
    self.updated = utc_now()
    return self

Functions

pipeline_snapshot_schemas

Pipeline snapshot schemas.

Classes

PipelineSnapshotSchema

Bases: BaseSchema

SQL Model for pipeline snapshots.

Attributes

is_runnable: bool property

Implements the is_runnable property.

Returns:

Type	Description
bool	True if the snapshot is runnable from server.

latest_run: Optional[PipelineRunSchema] property

Fetch the latest run for this snapshot.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[PipelineRunSchema]	The latest run for this snapshot.

Functions

from_request(request: PipelineSnapshotRequest, code_reference_id: Optional[UUID]) -> PipelineSnapshotSchema classmethod

Create schema from request.

Parameters:

Name	Type	Description	Default
request	PipelineSnapshotRequest	The request to convert.	required
code_reference_id	Optional[UUID]	Optional ID of the code reference for the snapshot.	required

Returns:

Type	Description
PipelineSnapshotSchema	The created schema.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

@classmethod
def from_request(
    cls,
    request: PipelineSnapshotRequest,
    code_reference_id: Optional[UUID],
) -> "PipelineSnapshotSchema":
    """Create schema from request.

    Args:
        request: The request to convert.
        code_reference_id: Optional ID of the code reference for the
            snapshot.

    Returns:
        The created schema.
    """
    client_env = json.dumps(request.client_environment)
    if len(client_env) > TEXT_FIELD_MAX_LENGTH:
        logger.warning(
            "Client environment is too large to be stored in the database. "
            "Skipping."
        )
        client_env = "{}"

    name = None
    if isinstance(request.name, str):
        name = request.name

    return cls(
        name=name,
        description=request.description,
        source_code=request.source_code,
        is_dynamic=request.is_dynamic,
        stack_id=request.stack,
        project_id=request.project,
        pipeline_id=request.pipeline,
        build_id=request.build,
        user_id=request.user,
        schedule_id=request.schedule,
        template_id=request.template,
        source_snapshot_id=request.source_snapshot,
        code_reference_id=code_reference_id,
        run_name_template=request.run_name_template,
        pipeline_configuration=request.pipeline_configuration.model_dump_json(),
        step_count=len(request.step_configurations),
        client_environment=client_env,
        client_version=request.client_version,
        server_version=request.server_version,
        pipeline_version_hash=request.pipeline_version_hash,
        pipeline_spec=json.dumps(
            request.pipeline_spec.model_dump(mode="json"), sort_keys=True
        )
        if request.pipeline_spec
        else None,
        code_path=request.code_path,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_metadata:
        options.extend(
            [
                joinedload(jl_arg(PipelineSnapshotSchema.stack)),
                joinedload(jl_arg(PipelineSnapshotSchema.build)),
                joinedload(jl_arg(PipelineSnapshotSchema.pipeline)),
                joinedload(jl_arg(PipelineSnapshotSchema.schedule)),
                joinedload(jl_arg(PipelineSnapshotSchema.code_reference)),
            ]
        )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(PipelineSnapshotSchema.user)),
                selectinload(
                    jl_arg(PipelineSnapshotSchema.visualizations)
                ),
            ]
        )

    return options

get_step_configuration(step_name: str) -> StepConfigurationSchema

Get a step configuration of the snapshot.

Parameters:

Name	Type	Description	Default
step_name	str	The name of the step to get the configuration for.	required

Raises:

Type	Description
KeyError	If the step configuration is not found.

Returns:

Type	Description
StepConfigurationSchema	The step configuration.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

def get_step_configuration(
    self, step_name: str
) -> "StepConfigurationSchema":
    """Get a step configuration of the snapshot.

    Args:
        step_name: The name of the step to get the configuration for.

    Raises:
        KeyError: If the step configuration is not found.

    Returns:
        The step configuration.
    """
    step_configs = self.get_step_configurations(include=[step_name])
    if len(step_configs) == 0:
        raise KeyError(
            f"Step configuration for step `{step_name}` not found."
        )
    return step_configs[0]

get_step_configurations(include: Optional[List[str]] = None) -> List[StepConfigurationSchema]

Get step configurations for the snapshot.

Parameters:

Name	Type	Description	Default
include	Optional[List[str]]	List of step names to include. If not given, all step configurations will be included.	None

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
List[StepConfigurationSchema]	List of step configurations.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

def get_step_configurations(
    self, include: Optional[List[str]] = None
) -> List["StepConfigurationSchema"]:
    """Get step configurations for the snapshot.

    Args:
        include: List of step names to include. If not given, all step
            configurations will be included.

    Raises:
        RuntimeError: If no session for the schema exists.

    Returns:
        List of step configurations.
    """
    if session := object_session(self):
        query = (
            select(StepConfigurationSchema)
            .where(StepConfigurationSchema.snapshot_id == self.id)
            .order_by(asc(StepConfigurationSchema.index))
        )

        if include:
            query = query.where(
                col(StepConfigurationSchema.name).in_(include)
            )

        return list(session.execute(query).scalars().all())
    else:
        raise RuntimeError(
            "Missing DB session to fetch step configurations."
        )

to_model(include_metadata: bool = False, include_resources: bool = False, include_python_packages: bool = False, include_config_schema: Optional[bool] = None, step_configuration_filter: Optional[List[str]] = None, **kwargs: Any) -> PipelineSnapshotResponse

Convert schema to response.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
include_python_packages	bool	Whether the python packages will be filled.	False
include_config_schema	Optional[bool]	Whether the config schema will be filled.	None
step_configuration_filter	Optional[List[str]]	List of step configurations to include in the response. If not given, all step configurations will be included.	None
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
PipelineSnapshotResponse	The response.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    include_python_packages: bool = False,
    include_config_schema: Optional[bool] = None,
    step_configuration_filter: Optional[List[str]] = None,
    **kwargs: Any,
) -> PipelineSnapshotResponse:
    """Convert schema to response.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        include_python_packages: Whether the python packages will be filled.
        include_config_schema: Whether the config schema will be filled.
        step_configuration_filter: List of step configurations to include in
            the response. If not given, all step configurations will be
            included.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The response.
    """
    deployable = False
    if self.build and self.stack and self.stack.has_deployer:
        deployable = True

    body = PipelineSnapshotResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        runnable=self.is_runnable,
        deployable=deployable,
        is_dynamic=self.is_dynamic,
    )
    metadata = None
    if include_metadata:
        pipeline_configuration = PipelineConfiguration.model_validate_json(
            self.pipeline_configuration
        )
        step_configurations = {}
        for step_configuration in self.get_step_configurations(
            include=step_configuration_filter
        ):
            step_configurations[step_configuration.name] = Step.from_dict(
                json.loads(step_configuration.config),
                pipeline_configuration,
            )

        client_environment = json.loads(self.client_environment)
        if not include_python_packages:
            client_environment.pop("python_packages", None)

        config_template = None
        config_schema = None

        if include_config_schema and self.build and self.build.stack_id:
            from zenml.zen_stores import template_utils

            if step_configuration_filter:
                # If only a subset of step configurations is requested,
                # we still need to get all of them to generate the config
                # template and schema
                all_step_configurations = {
                    step_configuration.name: Step.from_dict(
                        json.loads(step_configuration.config),
                        pipeline_configuration,
                    )
                    for step_configuration in self.get_step_configurations()
                }
            else:
                all_step_configurations = step_configurations

            config_template = template_utils.generate_config_template(
                snapshot=self,
                pipeline_configuration=pipeline_configuration,
                step_configurations=all_step_configurations,
            )
            config_schema = template_utils.generate_config_schema(
                snapshot=self,
                pipeline_configuration=pipeline_configuration,
                step_configurations=all_step_configurations,
            )

        metadata = PipelineSnapshotResponseMetadata(
            description=self.description,
            source_code=self.source_code,
            run_name_template=self.run_name_template,
            pipeline_configuration=pipeline_configuration,
            step_configurations=step_configurations,
            client_environment=client_environment,
            client_version=self.client_version,
            server_version=self.server_version,
            pipeline_version_hash=self.pipeline_version_hash,
            pipeline_spec=PipelineSpec.model_validate_json(
                self.pipeline_spec
            )
            if self.pipeline_spec
            else None,
            code_path=self.code_path,
            template_id=self.template_id,
            source_snapshot_id=self.source_snapshot_id,
            config_schema=config_schema,
            config_template=config_template,
        )

    resources = None
    if include_resources:
        latest_run = self.latest_run
        latest_run_user = latest_run.user if latest_run else None

        resources = PipelineSnapshotResponseResources(
            user=self.user.to_model() if self.user else None,
            pipeline=self.pipeline.to_model(),
            stack=self.stack.to_model() if self.stack else None,
            build=self.build.to_model() if self.build else None,
            schedule=self.schedule.to_model() if self.schedule else None,
            code_reference=self.code_reference.to_model()
            if self.code_reference
            else None,
            deployment=self.deployment.to_model()
            if self.deployment
            else None,
            tags=[tag.to_model() for tag in self.tags],
            latest_run_id=latest_run.id if latest_run else None,
            latest_run_status=latest_run.status if latest_run else None,
            latest_run_user=latest_run_user.to_model()
            if latest_run_user
            else None,
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    return PipelineSnapshotResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: PipelineSnapshotUpdate) -> PipelineSnapshotSchema

Update the schema.

Parameters:

Name	Type	Description	Default
update	PipelineSnapshotUpdate	The update to apply.	required

Returns:

Type	Description
PipelineSnapshotSchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/pipeline_snapshot_schemas.py

def update(
    self, update: PipelineSnapshotUpdate
) -> "PipelineSnapshotSchema":
    """Update the schema.

    Args:
        update: The update to apply.

    Returns:
        The updated schema.
    """
    if isinstance(update.name, str):
        self.name = update.name
    elif update.name is False:
        self.name = None

    if update.description:
        self.description = update.description

    self.updated = utc_now()
    return self

StepConfigurationSchema

Bases: BaseSchema

SQL Model for step configurations.

Functions

project_schemas

SQL Model Implementations for projects.

Classes

ProjectSchema

Bases: NamedSchema

SQL Model for projects.

Functions

from_request(project: ProjectRequest) -> ProjectSchema classmethod

Create a ProjectSchema from a ProjectResponse.

Parameters:

Name	Type	Description	Default
project	ProjectRequest	The ProjectResponse from which to create the schema.	required

Returns:

Type	Description
ProjectSchema	The created ProjectSchema.

Source code in src/zenml/zen_stores/schemas/project_schemas.py

@classmethod
def from_request(cls, project: ProjectRequest) -> "ProjectSchema":
    """Create a `ProjectSchema` from a `ProjectResponse`.

    Args:
        project: The `ProjectResponse` from which to create the schema.

    Returns:
        The created `ProjectSchema`.
    """
    return cls(
        name=project.name,
        description=project.description,
        display_name=project.display_name,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ProjectResponse

Convert a ProjectSchema to a ProjectResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ProjectResponse	The converted ProjectResponseModel.

Source code in src/zenml/zen_stores/schemas/project_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ProjectResponse:
    """Convert a `ProjectSchema` to a `ProjectResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The converted `ProjectResponseModel`.
    """
    metadata = None
    if include_metadata:
        metadata = ProjectResponseMetadata(
            description=self.description,
        )
    return ProjectResponse(
        id=self.id,
        name=self.name,
        body=ProjectResponseBody(
            display_name=self.display_name,
            created=self.created,
            updated=self.updated,
        ),
        metadata=metadata,
    )

update(project_update: ProjectUpdate) -> ProjectSchema

Update a ProjectSchema from a ProjectUpdate.

Parameters:

Name	Type	Description	Default
project_update	ProjectUpdate	The ProjectUpdate from which to update the schema.	required

Returns:

Type	Description
ProjectSchema	The updated ProjectSchema.

Source code in src/zenml/zen_stores/schemas/project_schemas.py

def update(self, project_update: ProjectUpdate) -> "ProjectSchema":
    """Update a `ProjectSchema` from a `ProjectUpdate`.

    Args:
        project_update: The `ProjectUpdate` from which to update the
            schema.

    Returns:
        The updated `ProjectSchema`.
    """
    for field, value in project_update.model_dump(
        exclude_unset=True
    ).items():
        setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

resource_pool_policy_schemas

Resource pool subject policy schemas.

Classes

ResourcePoolSubjectPolicyResourceSchema

Bases: BaseSchema

Resource pool subject policy resource schema.

ResourcePoolSubjectPolicySchema

Bases: BaseSchema

Resource pool subject policy schema.

Functions

from_request(request: ResourcePoolSubjectPolicyRequest) -> ResourcePoolSubjectPolicySchema classmethod

Creates a schema instance from a request model.

Parameters:

Name	Type	Description	Default
request	ResourcePoolSubjectPolicyRequest	The request model.	required

Returns:

Type	Description
ResourcePoolSubjectPolicySchema	The schema instance.

Source code in src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py

@classmethod
def from_request(
    cls, request: ResourcePoolSubjectPolicyRequest
) -> "ResourcePoolSubjectPolicySchema":
    """Creates a schema instance from a request model.

    Args:
        request: The request model.

    Returns:
        The schema instance.
    """
    return cls(
        user_id=request.user,
        component_id=request.component_id,
        pool_id=request.pool_id,
        priority=request.priority,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Gets query options for this schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	If metadata should be included.	False
include_resources	bool	If resources should be included.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
Sequence[ExecutableOption]	The query options.

Source code in src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Gets query options for this schema.

    Args:
        include_metadata: If metadata should be included.
        include_resources: If resources should be included.
        **kwargs: Additional keyword arguments.

    Returns:
        The query options.
    """
    options: List[ExecutableOption] = [
        selectinload(jl_arg(ResourcePoolSubjectPolicySchema.resources)),
    ]

    if include_resources:
        options.extend(
            [
                selectinload(
                    jl_arg(ResourcePoolSubjectPolicySchema.component)
                ),
                selectinload(jl_arg(ResourcePoolSubjectPolicySchema.pool)),
                selectinload(jl_arg(ResourcePoolSubjectPolicySchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ResourcePoolSubjectPolicyResponse

Converts this schema to a response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata.	False
include_resources	bool	Whether to include nested resources.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
ResourcePoolSubjectPolicyResponse	The response model.

Source code in src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ResourcePoolSubjectPolicyResponse":
    """Converts this schema to a response model.

    Args:
        include_metadata: Whether to include metadata.
        include_resources: Whether to include nested resources.
        **kwargs: Additional keyword arguments.

    Returns:
        The response model.
    """
    body = ResourcePoolSubjectPolicyResponseBody(
        created=self.created,
        updated=self.updated,
        user_id=self.user_id,
        priority=self.priority,
        reserved={
            resource.key: resource.reserved for resource in self.resources
        },
        limit={
            resource.key: resource.limit
            for resource in self.resources
            if resource.limit is not None
        },
    )

    metadata = None
    if include_metadata:
        metadata = ResourcePoolSubjectPolicyResponseMetadata()

    resources = None
    if include_resources:
        resources = ResourcePoolSubjectPolicyResponseResources(
            user=self.user.to_model() if self.user else None,
            component=self.component.to_model(),
            pool=self.pool.to_model(
                include_metadata=False, include_resources=False
            ),
        )

    return ResourcePoolSubjectPolicyResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: ResourcePoolSubjectPolicyUpdate) -> ResourcePoolSubjectPolicySchema

Updates this schema from an update model.

Parameters:

Name	Type	Description	Default
update	ResourcePoolSubjectPolicyUpdate	The update model.	required

Returns:

Type	Description
ResourcePoolSubjectPolicySchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/resource_pool_policy_schemas.py

def update(
    self, update: ResourcePoolSubjectPolicyUpdate
) -> "ResourcePoolSubjectPolicySchema":
    """Updates this schema from an update model.

    Args:
        update: The update model.

    Returns:
        The updated schema.
    """
    if update.priority is not None:
        self.priority = update.priority

    self.updated = utc_now()
    return self

Functions

resource_pool_schemas

Resource Pool schemas.

Classes

ResourcePoolAllocationSchema

Bases: BaseSchema

Resource pool allocation schema.

Attributes

priority: Optional[int] property

Fetch the priority for this allocation.

Returns:

Type	Description
Optional[int]	The matching policy priority, if a policy exists.

ResourcePoolQueueSchema

Bases: BaseSchema

Resource pool queue schema.

ResourcePoolResourceSchema

Bases: BaseSchema

Resource pool resource schema.

ResourcePoolSchema

Bases: NamedSchema

Resource pool schema.

Functions

from_request(request: ResourcePoolRequest) -> ResourcePoolSchema classmethod

Create a resource pool schema from a request.

Parameters:

Name	Type	Description	Default
request	ResourcePoolRequest	The request from which to create the resource pool.	required

Returns:

Type	Description
ResourcePoolSchema	The resource pool schema.

Source code in src/zenml/zen_stores/schemas/resource_pool_schemas.py

@classmethod
def from_request(
    cls,
    request: "ResourcePoolRequest",
) -> "ResourcePoolSchema":
    """Create a resource pool schema from a request.

    Args:
        request: The request from which to create the resource pool.

    Returns:
        The resource pool schema.
    """
    return cls(
        name=request.name,
        user_id=request.user,
        description=request.description,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/resource_pool_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = [
        selectinload(jl_arg(ResourcePoolSchema.resources)),
        selectinload(jl_arg(ResourcePoolSchema.queue_items)).options(
            load_only(jl_arg(ResourcePoolQueueSchema.request_id))
        ),
    ]

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(ResourcePoolSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ResourcePoolResponse

Creates a ResourcePoolResponse from a ResourcePoolSchema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ResourcePoolResponse	A ResourcePoolResponse

Source code in src/zenml/zen_stores/schemas/resource_pool_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ResourcePoolResponse":
    """Creates a `ResourcePoolResponse` from a `ResourcePoolSchema`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A `ResourcePoolResponse`
    """
    body = ResourcePoolResponseBody(
        created=self.created,
        updated=self.updated,
        user_id=self.user_id,
        queue_length=len(self.queue_items),
        capacity={r.key: r.total for r in self.resources},
        occupied_resources={r.key: r.occupied for r in self.resources},
    )

    metadata = None
    if include_metadata:
        metadata = ResourcePoolResponseMetadata(
            description=self.description,
        )

    resources = None
    if include_resources:
        resources = ResourcePoolResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return ResourcePoolResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(resource_pool_update: ResourcePoolUpdate) -> ResourcePoolSchema

Updates a ResourcePoolSchema from a ResourcePoolUpdate.

Parameters:

Name	Type	Description	Default
resource_pool_update	ResourcePoolUpdate	The ResourcePoolUpdate to update from.	required

Returns:

Type	Description
ResourcePoolSchema	The updated ResourcePoolSchema.

Source code in src/zenml/zen_stores/schemas/resource_pool_schemas.py

def update(
    self, resource_pool_update: "ResourcePoolUpdate"
) -> "ResourcePoolSchema":
    """Updates a `ResourcePoolSchema` from a `ResourcePoolUpdate`.

    Args:
        resource_pool_update: The `ResourcePoolUpdate` to update from.

    Returns:
        The updated `ResourcePoolSchema`.
    """
    if resource_pool_update.description:
        self.description = resource_pool_update.description

    self.updated = utc_now()
    return self

Functions

resource_request_schemas

Resource request schemas.

Classes

ResourceRequestResourceSchema

Bases: BaseSchema

Resource request resource schema.

ResourceRequestSchema

Bases: BaseSchema

Resource request schema.

Functions

from_request(request: ResourceRequestRequest) -> ResourceRequestSchema classmethod

Create a resource request schema from a request.

Parameters:

Name	Type	Description	Default
request	ResourceRequestRequest	The ResourceRequestRequest to create from.	required

Returns:

Type	Description
ResourceRequestSchema	The created ResourceRequestSchema.

Source code in src/zenml/zen_stores/schemas/resource_request_schemas.py

@classmethod
def from_request(
    cls,
    request: "ResourceRequestRequest",
) -> "ResourceRequestSchema":
    """Create a resource request schema from a request.

    Args:
        request: The `ResourceRequestRequest` to create from.

    Returns:
        The created `ResourceRequestSchema`.
    """
    return cls(
        user_id=request.user,
        component_id=request.component_id,
        step_run_id=request.step_run_id,
        status=ResourceRequestStatus.PENDING.value,
        preemption_initiated_by_id=None,
        preemptible=request.preemptible,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/resource_request_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        **kwargs: Additional keyword arguments.

    Returns:
        A list of query options.
    """
    options = [
        selectinload(jl_arg(ResourceRequestSchema.requested_resources)),
    ]

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(ResourceRequestSchema.component)),
                selectinload(
                    jl_arg(ResourceRequestSchema.step_run)
                ).joinedload(jl_arg(StepRunSchema.pipeline_run)),
                selectinload(jl_arg(ResourceRequestSchema.pool)),
                selectinload(jl_arg(ResourceRequestSchema.user)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ResourceRequestResponse

Creates a ResourceRequestResponse from a ResourceRequestSchema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
ResourceRequestResponse	A ResourceRequestResponse object.

Source code in src/zenml/zen_stores/schemas/resource_request_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ResourceRequestResponse":
    """Creates a `ResourceRequestResponse` from a `ResourceRequestSchema`.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        **kwargs: Additional keyword arguments.

    Returns:
        A `ResourceRequestResponse` object.
    """
    body = ResourceRequestResponseBody(
        created=self.created,
        updated=self.updated,
        user_id=self.user_id,
        requested_resources={
            r.key: r.amount for r in self.requested_resources
        },
        status=ResourceRequestStatus(self.status),
        status_reason=self.status_reason,
        preemptible=self.preemptible,
    )

    metadata = None
    if include_metadata:
        metadata = ResourceRequestResponseMetadata()

    resources = None
    if include_resources:
        resources = ResourceRequestResponseResources(
            user=self.user.to_model() if self.user else None,
            component=self.component.to_model()
            if self.component
            else None,
            step_run=self.step_run.to_model() if self.step_run else None,
            pipeline_run=self.step_run.pipeline_run.to_model()
            if self.step_run
            else None,
            pool=self.pool.to_model() if self.pool else None,
        )

    return ResourceRequestResponse(
        id=self.id,
        body=body,
        metadata=metadata,
        resources=resources,
    )

Functions

run_metadata_schemas

SQLModel implementation of run metadata tables.

Classes

RunMetadataResourceSchema

Bases: SQLModel

Table for linking resources to run metadata entries.

RunMetadataSchema

Bases: BaseSchema

SQL Model for run metadata.

Functions

run_template_schemas

SQLModel implementation of run template tables.

Classes

RunTemplateSchema

Bases: NamedSchema

SQL Model for run templates.

Attributes

latest_run: Optional[PipelineRunSchema] property

Fetch the latest run for this template.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
Optional[PipelineRunSchema]	The latest run for this template.

Functions

from_request(request: RunTemplateRequest) -> RunTemplateSchema classmethod

Create a schema from a request.

Parameters:

Name	Type	Description	Default
request	RunTemplateRequest	The request to convert.	required

Returns:

Type	Description
RunTemplateSchema	The created schema.

Source code in src/zenml/zen_stores/schemas/run_template_schemas.py

@classmethod
def from_request(
    cls,
    request: RunTemplateRequest,
) -> "RunTemplateSchema":
    """Create a schema from a request.

    Args:
        request: The request to convert.


    Returns:
        The created schema.
    """
    return cls(
        user_id=request.user,
        project_id=request.project,
        name=request.name,
        description=request.description,
        hidden=request.hidden,
        source_snapshot_id=request.source_snapshot_id,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/run_template_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    from zenml.zen_stores.schemas import PipelineSnapshotSchema

    options = [
        joinedload(jl_arg(RunTemplateSchema.source_snapshot)).joinedload(
            jl_arg(PipelineSnapshotSchema.build)
        ),
    ]

    if include_metadata or include_resources:
        options.extend(
            [
                joinedload(
                    jl_arg(RunTemplateSchema.source_snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.pipeline)),
                joinedload(
                    jl_arg(RunTemplateSchema.source_snapshot)
                ).joinedload(
                    jl_arg(PipelineSnapshotSchema.code_reference)
                ),
            ]
        )
    if include_metadata:
        options.extend(
            [
                joinedload(
                    jl_arg(RunTemplateSchema.source_snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.stack)),
                joinedload(
                    jl_arg(RunTemplateSchema.source_snapshot)
                ).joinedload(jl_arg(PipelineSnapshotSchema.schedule)),
            ]
        )

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(RunTemplateSchema.user)),
                # joinedload(jl_arg(RunTemplateSchema.tags)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> RunTemplateResponse

Convert the schema to a response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
RunTemplateResponse	Model representing this schema.

Source code in src/zenml/zen_stores/schemas/run_template_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> RunTemplateResponse:
    """Convert the schema to a response model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        Model representing this schema.
    """
    runnable = False
    if (
        self.source_snapshot
        and self.source_snapshot.build
        and not self.source_snapshot.build.is_local
        and self.source_snapshot.build.stack_id
    ):
        runnable = True

    body = RunTemplateResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        runnable=runnable,
        hidden=self.hidden,
    )

    metadata = None
    if include_metadata:
        pipeline_spec = None
        config_template = None
        config_schema = None

        if self.source_snapshot:
            from zenml.zen_stores import template_utils

            source_snapshot_model = self.source_snapshot.to_model(
                include_metadata=True
            )
            pipeline_spec = source_snapshot_model.pipeline_spec

            if (
                self.source_snapshot.build
                and self.source_snapshot.build.stack_id
            ):
                config_template = template_utils.generate_config_template(
                    snapshot=self.source_snapshot,
                    pipeline_configuration=source_snapshot_model.pipeline_configuration,
                    step_configurations=source_snapshot_model.step_configurations,
                )
                config_schema = template_utils.generate_config_schema(
                    snapshot=self.source_snapshot,
                    pipeline_configuration=source_snapshot_model.pipeline_configuration,
                    step_configurations=source_snapshot_model.step_configurations,
                )

        metadata = RunTemplateResponseMetadata(
            description=self.description,
            pipeline_spec=pipeline_spec,
            config_template=config_template,
            config_schema=config_schema,
        )

    resources = None
    if include_resources:
        if self.source_snapshot:
            pipeline = (
                self.source_snapshot.pipeline.to_model()
                if self.source_snapshot.pipeline
                else None
            )
            build = (
                self.source_snapshot.build.to_model()
                if self.source_snapshot.build
                else None
            )
            code_reference = (
                self.source_snapshot.code_reference.to_model()
                if self.source_snapshot.code_reference
                else None
            )
        else:
            pipeline = None
            build = None
            code_reference = None

        latest_run = self.latest_run

        resources = RunTemplateResponseResources(
            user=self.user.to_model() if self.user else None,
            source_snapshot=self.source_snapshot.to_model()
            if self.source_snapshot
            else None,
            pipeline=pipeline,
            build=build,
            code_reference=code_reference,
            tags=[tag.to_model() for tag in self.tags],
            latest_run_id=latest_run.id if latest_run else None,
            latest_run_status=latest_run.status if latest_run else None,
        )

    return RunTemplateResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: RunTemplateUpdate) -> RunTemplateSchema

Update the schema.

Parameters:

Name	Type	Description	Default
update	RunTemplateUpdate	The update model.	required

Returns:

Type	Description
RunTemplateSchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/run_template_schemas.py

def update(self, update: RunTemplateUpdate) -> "RunTemplateSchema":
    """Update the schema.

    Args:
        update: The update model.

    Returns:
        The updated schema.
    """
    for field, value in update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field in ["add_tags", "remove_tags"]:
            # Tags are handled separately
            continue
        setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

run_wait_condition_schemas

SQLModel implementation of run wait condition tables.

Classes

RunWaitConditionSchema

Bases: BaseSchema, RunMetadataInterface

SQLModel schema for persisted run wait conditions.

Functions

from_request(request: RunWaitConditionRequest) -> RunWaitConditionSchema classmethod

Create a schema object from a wait condition create request.

Parameters:

Name	Type	Description	Default
request	RunWaitConditionRequest	Wait condition creation request.	required

Returns:

Type	Description
RunWaitConditionSchema	The persisted wait condition schema object.

Source code in src/zenml/zen_stores/schemas/run_wait_condition_schemas.py

@classmethod
def from_request(
    cls, request: RunWaitConditionRequest
) -> "RunWaitConditionSchema":
    """Create a schema object from a wait condition create request.

    Args:
        request: Wait condition creation request.

    Returns:
        The persisted wait condition schema object.
    """
    return cls(
        run_id=request.run,
        project_id=request.project,
        user_id=request.user,
        name=request.name,
        type=request.type.value,
        status=RunWaitConditionStatus.PENDING.value,
        question=request.question,
        data_schema_json=(
            json.dumps(request.data_schema)
            if request.data_schema is not None
            else None
        ),
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get query options for converting schema rows to models.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included in responses.	False
include_resources	bool	Whether resources will be included in responses.	False
**kwargs	Any	Additional unused keyword arguments.	{}

Returns:

Type	Description
Sequence[ExecutableOption]	SQLAlchemy query options.

Source code in src/zenml/zen_stores/schemas/run_wait_condition_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get query options for converting schema rows to models.

    Args:
        include_metadata: Whether metadata will be included in responses.
        include_resources: Whether resources will be included in responses.
        **kwargs: Additional unused keyword arguments.

    Returns:
        SQLAlchemy query options.
    """
    options: List[ExecutableOption] = [joinedload(jl_arg(cls.run))]
    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(cls.user)),
                joinedload(jl_arg(cls.resolved_by_user)),
            ]
        )
    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> RunWaitConditionResponse

Convert the schema row to a response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata should be included.	False
include_resources	bool	Whether resources should be included.	False
**kwargs	Any	Additional unused keyword arguments.	{}

Returns:

Type	Description
RunWaitConditionResponse	The wait condition response model.

Source code in src/zenml/zen_stores/schemas/run_wait_condition_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> RunWaitConditionResponse:
    """Convert the schema row to a response model.

    Args:
        include_metadata: Whether metadata should be included.
        include_resources: Whether resources should be included.
        **kwargs: Additional unused keyword arguments.

    Returns:
        The wait condition response model.
    """
    data_schema: Optional[Dict[str, Any]] = None
    if self.data_schema_json:
        data_schema = json.loads(self.data_schema_json)

    result: Optional[Any] = None
    if self.result_json:
        result = json.loads(self.result_json)

    body = RunWaitConditionResponseBody(
        user_id=self.user_id,
        project_id=self.run.project_id,
        created=self.created,
        updated=self.updated,
        type=RunWaitConditionType(self.type),
        status=RunWaitConditionStatus(self.status),
        last_polled_at=self.last_polled_at,
        poller_instance_id=self.poller_instance_id,
        poller_lease_expires_at=self.poller_lease_expires_at,
        resolved_at=self.resolved_at,
        resolved_by_user_id=self.resolved_by_user_id,
    )

    metadata = None
    if include_metadata:
        metadata = RunWaitConditionResponseMetadata(
            question=self.question,
            run_metadata=self.fetch_metadata(),
            data_schema=data_schema,
            resolution=self.resolution,
            result=result,
        )

    resources = None
    if include_resources:
        resources = RunWaitConditionResponseResources(
            user=self.user.to_model() if self.user else None,
            run=self.run.to_model(
                include_metadata=False, include_resources=False
            ),
        )

    return RunWaitConditionResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

Functions

schedule_schema

SQL Model Implementations for Pipeline Schedules.

Classes

ScheduleSchema

Bases: NamedSchema, RunMetadataInterface

SQL Model for schedules.

Functions

from_request(schedule_request: ScheduleRequest) -> ScheduleSchema classmethod

Create a ScheduleSchema from a ScheduleRequest.

Parameters:

Name	Type	Description	Default
schedule_request	ScheduleRequest	The ScheduleRequest to create the schema from.	required

Returns:

Type	Description
ScheduleSchema	The created ScheduleSchema.

Source code in src/zenml/zen_stores/schemas/schedule_schema.py

@classmethod
def from_request(
    cls, schedule_request: ScheduleRequest
) -> "ScheduleSchema":
    """Create a `ScheduleSchema` from a `ScheduleRequest`.

    Args:
        schedule_request: The `ScheduleRequest` to create the schema from.

    Returns:
        The created `ScheduleSchema`.
    """
    if schedule_request.interval_second is not None:
        interval_second = schedule_request.interval_second.total_seconds()
    else:
        interval_second = None
    return cls(
        name=schedule_request.name,
        project_id=schedule_request.project,
        user_id=schedule_request.user,
        pipeline_id=schedule_request.pipeline_id,
        orchestrator_id=schedule_request.orchestrator_id,
        active=schedule_request.active,
        cron_expression=schedule_request.cron_expression,
        start_time=schedule_request.start_time,
        end_time=schedule_request.end_time,
        interval_second=interval_second,
        catchup=schedule_request.catchup,
        run_once_start_time=schedule_request.run_once_start_time,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/schedule_schema.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    # if include_metadata:
    #     options.extend(
    #         [
    #             joinedload(jl_arg(ScheduleSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend([joinedload(jl_arg(ScheduleSchema.user))])

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ScheduleResponse

Convert a ScheduleSchema to a ScheduleResponseModel.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ScheduleResponse	The created ScheduleResponseModel.

Source code in src/zenml/zen_stores/schemas/schedule_schema.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ScheduleResponse:
    """Convert a `ScheduleSchema` to a `ScheduleResponseModel`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `ScheduleResponseModel`.
    """
    if self.interval_second is not None:
        interval_second = timedelta(seconds=self.interval_second)
    else:
        interval_second = None

    body = ScheduleResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        active=self.active,
        cron_expression=self.cron_expression,
        start_time=self.start_time,
        end_time=self.end_time,
        interval_second=interval_second,
        catchup=self.catchup,
        updated=self.updated,
        created=self.created,
        run_once_start_time=self.run_once_start_time,
        is_archived=self.is_archived,
    )
    metadata = None
    if include_metadata:
        metadata = ScheduleResponseMetadata(
            pipeline_id=self.pipeline_id,
            orchestrator_id=self.orchestrator_id,
            run_metadata=self.fetch_metadata(),
        )

    resources = None
    if include_resources:
        resources = ScheduleResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return ScheduleResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(schedule_update: ScheduleUpdate) -> ScheduleSchema

Update a ScheduleSchema from a ScheduleUpdateModel.

Parameters:

Name	Type	Description	Default
schedule_update	ScheduleUpdate	The ScheduleUpdateModel to update the schema from.	required

Returns:

Type	Description
ScheduleSchema	The updated ScheduleSchema.

Source code in src/zenml/zen_stores/schemas/schedule_schema.py

def update(self, schedule_update: ScheduleUpdate) -> "ScheduleSchema":
    """Update a `ScheduleSchema` from a `ScheduleUpdateModel`.

    Args:
        schedule_update: The `ScheduleUpdateModel` to update the schema from.

    Returns:
        The updated `ScheduleSchema`.
    """
    if schedule_update.name is not None:
        self.name = schedule_update.name

    if schedule_update.cron_expression:
        self.cron_expression = schedule_update.cron_expression

    if schedule_update.active is not None:
        self.active = schedule_update.active

    self.updated = utc_now()
    return self

Functions

schema_utils

Utility functions for SQLModel schemas.

Functions

build_foreign_key_field(source: str, target: str, source_column: str, target_column: str, ondelete: str, nullable: bool, custom_constraint_name: Optional[str] = None, **sa_column_kwargs: Any) -> Any

Build a SQLModel foreign key field.

Parameters:

Name	Type	Description	Default
source	str	Source table name.	required
target	str	Target table name.	required
source_column	str	Source column name.	required
target_column	str	Target column name.	required
ondelete	str	On delete behavior.	required
nullable	bool	Whether the field is nullable.	required
custom_constraint_name	Optional[str]	Custom name for the foreign key constraint.	None
**sa_column_kwargs	Any	Keyword arguments for the SQLAlchemy column.	{}

Returns:

Type	Description
Any	SQLModel foreign key field.

Raises:

Type	Description
ValueError	If the ondelete and nullable arguments are not compatible, or if the foreign key constraint name is too long.

Source code in src/zenml/zen_stores/schemas/schema_utils.py

def build_foreign_key_field(
    source: str,
    target: str,
    source_column: str,
    target_column: str,
    ondelete: str,
    nullable: bool,
    custom_constraint_name: Optional[str] = None,
    **sa_column_kwargs: Any,
) -> Any:
    """Build a SQLModel foreign key field.

    Args:
        source: Source table name.
        target: Target table name.
        source_column: Source column name.
        target_column: Target column name.
        ondelete: On delete behavior.
        nullable: Whether the field is nullable.
        custom_constraint_name: Custom name for the foreign key constraint.
        **sa_column_kwargs: Keyword arguments for the SQLAlchemy column.

    Returns:
        SQLModel foreign key field.

    Raises:
        ValueError: If the ondelete and nullable arguments are not compatible,
            or if the foreign key constraint name is too long.
    """
    if not nullable and ondelete == "SET NULL":
        raise ValueError(
            "Cannot set ondelete to SET NULL if the field is not nullable."
        )
    constraint_name = custom_constraint_name or foreign_key_constraint_name(
        source=source,
        target=target,
        source_column=source_column,
    )
    if len(constraint_name) > 64:
        raise ValueError(
            f"Foreign key constraint name {constraint_name} is too long. "
            "The maximum length is 64 characters."
        )
    return Field(
        sa_column=Column(
            ForeignKey(
                f"{target}.{target_column}",
                name=constraint_name,
                ondelete=ondelete,
            ),
            nullable=nullable,
            **sa_column_kwargs,
        ),
    )

build_index(table_name: str, column_names: List[str], **kwargs: Any) -> Index

Build an index object.

Parameters:

Name	Type	Description	Default
table_name	str	The name of the table for which the index will be created.	required
column_names	List[str]	Names of the columns on which the index will be created.	required
**kwargs	Any	Additional keyword arguments to pass to the Index.	{}

Returns:

Type	Description
Index	The index.

Source code in src/zenml/zen_stores/schemas/schema_utils.py

def build_index(
    table_name: str, column_names: List[str], **kwargs: Any
) -> Index:
    """Build an index object.

    Args:
        table_name: The name of the table for which the index will be created.
        column_names: Names of the columns on which the index will be created.
        **kwargs: Additional keyword arguments to pass to the Index.

    Returns:
        The index.
    """
    name = get_index_name(table_name=table_name, column_names=column_names)
    return Index(name, *column_names, **kwargs)

foreign_key_constraint_name(source: str, target: str, source_column: str) -> str

Defines the name of a foreign key constraint.

For simplicity, we use the naming convention used by alembic here: https://alembic.sqlalchemy.org/en/latest/batch.html#dropping-unnamed-or-named-foreign-key-constraints.

Parameters:

Name	Type	Description	Default
source	str	Source table name.	required
target	str	Target table name.	required
source_column	str	Source column name.	required

Returns:

Type	Description
str	Name of the foreign key constraint.

Source code in src/zenml/zen_stores/schemas/schema_utils.py

def foreign_key_constraint_name(
    source: str, target: str, source_column: str
) -> str:
    """Defines the name of a foreign key constraint.

    For simplicity, we use the naming convention used by alembic here:
    https://alembic.sqlalchemy.org/en/latest/batch.html#dropping-unnamed-or-named-foreign-key-constraints.

    Args:
        source: Source table name.
        target: Target table name.
        source_column: Source column name.

    Returns:
        Name of the foreign key constraint.
    """
    return f"fk_{source}_{source_column}_{target}"

get_index_name(table_name: str, column_names: List[str]) -> str

Get the name for an index.

Parameters:

Name	Type	Description	Default
table_name	str	The name of the table for which the index will be created.	required
column_names	List[str]	Names of the columns on which the index will be created.	required

Returns:

Type	Description
str	The index name.

Source code in src/zenml/zen_stores/schemas/schema_utils.py

def get_index_name(table_name: str, column_names: List[str]) -> str:
    """Get the name for an index.

    Args:
        table_name: The name of the table for which the index will be created.
        column_names: Names of the columns on which the index will be created.

    Returns:
        The index name.
    """
    columns = "_".join(column_names)
    # MySQL allows a maximum of 64 characters in identifiers
    return f"ix_{table_name}_{columns}"[:64]

secret_schemas

SQL Model Implementations for Secrets.

Classes

SecretDecodeError

Bases: Exception

Raised when a secret cannot be decoded or decrypted.

SecretResourceSchema

Bases: BaseSchema

SQL Model for secret resource relationship.

SecretSchema

Bases: NamedSchema

SQL Model for secrets.

Attributes:

Name	Type	Description
name	str	The name of the secret.
values	Optional[bytes]	The values of the secret.

Functions

from_request(secret: SecretRequest, internal: bool = False) -> SecretSchema classmethod

Create a SecretSchema from a SecretRequest.

Parameters:

Name	Type	Description	Default
secret	SecretRequest	The SecretRequest from which to create the schema.	required
internal	bool	Whether the secret is internal.	False

Returns:

Type	Description
SecretSchema	The created SecretSchema.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

@classmethod
def from_request(
    cls,
    secret: SecretRequest,
    internal: bool = False,
) -> "SecretSchema":
    """Create a `SecretSchema` from a `SecretRequest`.

    Args:
        secret: The `SecretRequest` from which to create the schema.
        internal: Whether the secret is internal.

    Returns:
        The created `SecretSchema`.
    """
    assert secret.user is not None, "User must be set for secret creation."
    return cls(
        name=secret.name,
        private=secret.private,
        user_id=secret.user,
        # Don't store secret values implicitly in the secret. The
        # SQL secret store will call `store_secret_values` to store the
        # values separately if SQL is used as the secrets store.
        values=None,
        internal=internal,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend([joinedload(jl_arg(SecretSchema.user))])

    return options

get_secret_values(encryption_engine: Optional[AesGcmEngine] = None) -> Dict[str, str]

Get the secret values for this secret.

This method is used by the SQL secrets store to load the secret values from the database.

Parameters:

Name	Type	Description	Default
encryption_engine	Optional[AesGcmEngine]	The encryption engine to use to decrypt the secret values. If None, the values will be base64 decoded.	None

Returns:

Type	Description
Dict[str, str]	The secret values

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

def get_secret_values(
    self,
    encryption_engine: Optional[AesGcmEngine] = None,
) -> Dict[str, str]:
    """Get the secret values for this secret.

    This method is used by the SQL secrets store to load the secret values
    from the database.

    Args:
        encryption_engine: The encryption engine to use to decrypt the
            secret values. If None, the values will be base64 decoded.

    Returns:
        The secret values

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
    """
    if not self.values:
        raise KeyError(
            f"Secret values for secret {self.id} have not been stored in "
            f"the SQL secrets store."
        )
    return self._load_secret_values(self.values, encryption_engine)

set_secret_values(secret_values: Dict[str, str], encryption_engine: Optional[AesGcmEngine] = None) -> None

Create a SecretSchema from a SecretRequest.

This method is used by the SQL secrets store to store the secret values in the database.

Parameters:

Name	Type	Description	Default
secret_values	Dict[str, str]	The new secret values.	required
encryption_engine	Optional[AesGcmEngine]	The encryption engine to use to encrypt the secret values. If None, the values will be base64 encoded.	None

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

def set_secret_values(
    self,
    secret_values: Dict[str, str],
    encryption_engine: Optional[AesGcmEngine] = None,
) -> None:
    """Create a `SecretSchema` from a `SecretRequest`.

    This method is used by the SQL secrets store to store the secret values
    in the database.

    Args:
        secret_values: The new secret values.
        encryption_engine: The encryption engine to use to encrypt the
            secret values. If None, the values will be base64 encoded.
    """
    self.values = self._dump_secret_values(
        secret_values, encryption_engine
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> SecretResponse

Converts a secret schema to a secret model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
SecretResponse	The secret model.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> SecretResponse:
    """Converts a secret schema to a secret model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The secret model.
    """
    metadata = None
    if include_metadata:
        metadata = SecretResponseMetadata()

    resources = None
    if include_resources:
        resources = SecretResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    # Don't load the secret values implicitly in the secret. The
    # SQL secret store will call `get_secret_values` to load the
    # values separately if SQL is used as the secrets store.
    body = SecretResponseBody(
        user_id=self.user_id,
        created=self.created,
        updated=self.updated,
        private=self.private,
    )
    return SecretResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(secret_update: SecretUpdate) -> SecretSchema

Update a SecretSchema from a SecretUpdate.

Parameters:

Name	Type	Description	Default
secret_update	SecretUpdate	The SecretUpdate from which to update the schema.	required

Returns:

Type	Description
SecretSchema	The updated SecretSchema.

Source code in src/zenml/zen_stores/schemas/secret_schemas.py

def update(
    self,
    secret_update: SecretUpdate,
) -> "SecretSchema":
    """Update a `SecretSchema` from a `SecretUpdate`.

    Args:
        secret_update: The `SecretUpdate` from which to update the schema.

    Returns:
        The updated `SecretSchema`.
    """
    # Don't update the secret values implicitly in the secret. The
    # SQL secret store will call `set_secret_values` to update the
    # values separately if SQL is used as the secrets store.
    for field, value in secret_update.model_dump(
        exclude_unset=True, exclude={"user", "values"}
    ).items():
        setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

server_settings_schemas

SQLModel implementation for the server settings table.

Classes

ServerSettingsSchema

Bases: SQLModel

SQL Model for the server settings.

Functions

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ServerSettingsResponse

Convert an ServerSettingsSchema to an ServerSettingsResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ServerSettingsResponse	The created SettingsResponse.

Source code in src/zenml/zen_stores/schemas/server_settings_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ServerSettingsResponse:
    """Convert an `ServerSettingsSchema` to an `ServerSettingsResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The created `SettingsResponse`.
    """
    body = ServerSettingsResponseBody(
        server_id=self.id,
        server_name=self.server_name,
        logo_url=self.logo_url,
        enable_analytics=self.enable_analytics,
        display_announcements=self.display_announcements,
        display_updates=self.display_updates,
        active=self.active,
        updated=self.updated,
        last_user_activity=self.last_user_activity,
    )

    metadata = None
    resources = None

    if include_metadata:
        metadata = ServerSettingsResponseMetadata()

    if include_resources:
        resources = ServerSettingsResponseResources()

    return ServerSettingsResponse(
        body=body, metadata=metadata, resources=resources
    )

update(settings_update: ServerSettingsUpdate) -> ServerSettingsSchema

Update a ServerSettingsSchema from a ServerSettingsUpdate.

Parameters:

Name	Type	Description	Default
settings_update	ServerSettingsUpdate	The ServerSettingsUpdate from which to update the schema.	required

Returns:

Type	Description
ServerSettingsSchema	The updated ServerSettingsSchema.

Source code in src/zenml/zen_stores/schemas/server_settings_schemas.py

def update(
    self, settings_update: ServerSettingsUpdate
) -> "ServerSettingsSchema":
    """Update a `ServerSettingsSchema` from a `ServerSettingsUpdate`.

    Args:
        settings_update: The `ServerSettingsUpdate` from which
            to update the schema.

    Returns:
        The updated `ServerSettingsSchema`.
    """
    for field, value in settings_update.model_dump(
        exclude_unset=True
    ).items():
        if hasattr(self, field):
            setattr(self, field, value)

    self.updated = utc_now()

    return self

update_onboarding_state(completed_steps: Set[str]) -> ServerSettingsSchema

Update the onboarding state.

Parameters:

Name	Type	Description	Default
completed_steps	Set[str]	Newly completed onboarding steps.	required

Returns:

Type	Description
ServerSettingsSchema	The updated schema.

Source code in src/zenml/zen_stores/schemas/server_settings_schemas.py

def update_onboarding_state(
    self, completed_steps: Set[str]
) -> "ServerSettingsSchema":
    """Update the onboarding state.

    Args:
        completed_steps: Newly completed onboarding steps.

    Returns:
        The updated schema.
    """
    old_state = set(
        json.loads(self.onboarding_state) if self.onboarding_state else []
    )
    new_state = old_state.union(completed_steps)
    self.onboarding_state = json.dumps(list(new_state))
    self.updated = utc_now()

    return self

Functions

service_connector_schemas

SQL Model Implementations for Service Connectors.

Classes

ServiceConnectorSchema

Bases: NamedSchema

SQL Model for service connectors.

Attributes

labels_dict: Dict[str, str] property

Returns the labels as a dictionary.

Returns:

Type	Description
Dict[str, str]	The labels as a dictionary.

resource_types_list: List[str] property

Returns the resource types as a list.

Returns:

Type	Description
List[str]	The resource types as a list.

Functions

from_request(connector_request: ServiceConnectorRequest, secret_id: Optional[UUID] = None) -> ServiceConnectorSchema classmethod

Create a ServiceConnectorSchema from a ServiceConnectorRequest.

Parameters:

Name	Type	Description	Default
connector_request	ServiceConnectorRequest	The ServiceConnectorRequest from which to create the schema.	required
secret_id	Optional[UUID]	The ID of the secret to use for this connector.	None

Returns:

Type	Description
ServiceConnectorSchema	The created ServiceConnectorSchema.

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

@classmethod
def from_request(
    cls,
    connector_request: ServiceConnectorRequest,
    secret_id: Optional[UUID] = None,
) -> "ServiceConnectorSchema":
    """Create a `ServiceConnectorSchema` from a `ServiceConnectorRequest`.

    Args:
        connector_request: The `ServiceConnectorRequest` from which to
            create the schema.
        secret_id: The ID of the secret to use for this connector.

    Returns:
        The created `ServiceConnectorSchema`.
    """
    assert connector_request.user is not None, "User must be set."
    configuration = connector_request.configuration.non_secrets
    return cls(
        user_id=connector_request.user,
        name=connector_request.name,
        description=connector_request.description,
        connector_type=connector_request.type,
        auth_method=connector_request.auth_method,
        resource_types=base64.b64encode(
            json.dumps(connector_request.resource_types).encode("utf-8")
        ),
        resource_id=connector_request.resource_id,
        supports_instances=connector_request.supports_instances,
        configuration=base64.b64encode(
            json.dumps(configuration).encode("utf-8")
        )
        if configuration
        else None,
        secret_id=secret_id,
        expires_at=connector_request.expires_at,
        expires_skew_tolerance=connector_request.expires_skew_tolerance,
        expiration_seconds=connector_request.expiration_seconds,
        labels=base64.b64encode(
            json.dumps(connector_request.labels).encode("utf-8")
        )
        if connector_request.labels
        else None,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend([joinedload(jl_arg(ServiceConnectorSchema.user))])

    return options

has_labels(labels: Dict[str, Optional[str]]) -> bool

Checks if the connector has the given labels.

Parameters:

Name	Type	Description	Default
labels	Dict[str, Optional[str]]	The labels to check for.	required

Returns:

Type	Description
bool	Whether the connector has the given labels.

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

def has_labels(self, labels: Dict[str, Optional[str]]) -> bool:
    """Checks if the connector has the given labels.

    Args:
        labels: The labels to check for.

    Returns:
        Whether the connector has the given labels.
    """
    return all(
        self.labels_dict.get(key, None) == value
        for key, value in labels.items()
        if value is not None
    ) and all(
        key in self.labels_dict
        for key, value in labels.items()
        if value is None
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ServiceConnectorResponse

Creates a ServiceConnector from a ServiceConnectorSchema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ServiceConnectorResponse	A ServiceConnectorModel

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ServiceConnectorResponse":
    """Creates a `ServiceConnector` from a `ServiceConnectorSchema`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        A `ServiceConnectorModel`
    """
    body = ServiceConnectorResponseBody(
        user_id=self.user_id,
        created=self.created,
        updated=self.updated,
        description=self.description,
        connector_type=self.connector_type,
        auth_method=self.auth_method,
        resource_types=self.resource_types_list,
        resource_id=self.resource_id,
        supports_instances=self.supports_instances,
        expires_at=self.expires_at,
        expires_skew_tolerance=self.expires_skew_tolerance,
    )
    metadata = None
    if include_metadata:
        metadata = ServiceConnectorResponseMetadata(
            configuration=ServiceConnectorConfiguration(
                **json.loads(base64.b64decode(self.configuration).decode())
            )
            if self.configuration
            else ServiceConnectorConfiguration(),
            expiration_seconds=self.expiration_seconds,
            labels=self.labels_dict,
        )
    resources = None
    if include_resources:
        resources = ServiceConnectorResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return ServiceConnectorResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(connector_update: ServiceConnectorUpdate, secret_id: Optional[UUID] = None) -> ServiceConnectorSchema

Updates a ServiceConnectorSchema from a ServiceConnectorUpdate.

Parameters:

Name	Type	Description	Default
connector_update	ServiceConnectorUpdate	The ServiceConnectorUpdate to update from.	required
secret_id	Optional[UUID]	The ID of the secret to use for this connector.	None

Returns:

Type	Description
ServiceConnectorSchema	The updated ServiceConnectorSchema.

Source code in src/zenml/zen_stores/schemas/service_connector_schemas.py

def update(
    self,
    connector_update: ServiceConnectorUpdate,
    secret_id: Optional[UUID] = None,
) -> "ServiceConnectorSchema":
    """Updates a `ServiceConnectorSchema` from a `ServiceConnectorUpdate`.

    Args:
        connector_update: The `ServiceConnectorUpdate` to update from.
        secret_id: The ID of the secret to use for this connector.

    Returns:
        The updated `ServiceConnectorSchema`.
    """
    for field, value in connector_update.model_dump(
        exclude_unset=False,
        exclude={"user", "secrets"},
    ).items():
        if value is None:
            if field == "resource_id":
                # The resource ID field in the update is special: if set
                # to None in the update, it triggers the existing resource
                # ID to be cleared.
                self.resource_id = None
            if field == "expiration_seconds":
                # The expiration_seconds field in the update is special:
                # if set to None in the update, it triggers the existing
                # expiration_seconds to be cleared.
                self.expiration_seconds = None
            continue
        if field == "configuration":
            if connector_update.configuration is not None:
                configuration = connector_update.configuration.non_secrets
                if configuration is not None:
                    self.configuration = (
                        base64.b64encode(
                            json.dumps(configuration).encode("utf-8")
                        )
                        if configuration
                        else None
                    )
        elif field == "resource_types":
            self.resource_types = base64.b64encode(
                json.dumps(connector_update.resource_types).encode("utf-8")
            )
        elif field == "labels":
            self.labels = (
                base64.b64encode(
                    json.dumps(connector_update.labels).encode("utf-8")
                )
                if connector_update.labels
                else None
            )
        else:
            setattr(self, field, value)
    self.secret_id = secret_id
    self.updated = utc_now()
    return self

Functions

service_schemas

SQLModel implementation of service table.

Classes

ServiceSchema

Bases: NamedSchema

SQL Model for service.

Functions

from_request(service_request: ServiceRequest) -> ServiceSchema classmethod

Convert a ServiceRequest to a ServiceSchema.

Parameters:

Name	Type	Description	Default
service_request	ServiceRequest	The request model to convert.	required

Returns:

Type	Description
ServiceSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/service_schemas.py

@classmethod
def from_request(
    cls, service_request: "ServiceRequest"
) -> "ServiceSchema":
    """Convert a `ServiceRequest` to a `ServiceSchema`.

    Args:
        service_request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=service_request.name,
        project_id=service_request.project,
        user_id=service_request.user,
        service_source=service_request.service_source,
        service_type=service_request.service_type.model_dump_json(),
        type=service_request.service_type.type,
        flavor=service_request.service_type.flavor,
        admin_state=service_request.admin_state,
        config=dict_to_bytes(service_request.config),
        labels=dict_to_bytes(service_request.labels)
        if service_request.labels
        else None,
        status=dict_to_bytes(service_request.status)
        if service_request.status
        else None,
        endpoint=dict_to_bytes(service_request.endpoint)
        if service_request.endpoint
        else None,
        state=service_request.status.get("state")
        if service_request.status
        else None,
        model_version_id=service_request.model_version_id,
        pipeline_run_id=service_request.pipeline_run_id,
        prediction_url=service_request.prediction_url,
        health_check_url=service_request.health_check_url,
        pipeline_name=service_request.config.get("pipeline_name"),
        pipeline_step_name=service_request.config.get(
            "pipeline_step_name"
        ),
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/service_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(ServiceSchema.user)),
                joinedload(jl_arg(ServiceSchema.model_version)),
                joinedload(jl_arg(ServiceSchema.pipeline_run)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ServiceResponse

Convert an ServiceSchema to an ServiceResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether to include metadata in the response.	False
include_resources	bool	Whether to include resources in the response.	False
kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
ServiceResponse	The created ServiceResponse.

Source code in src/zenml/zen_stores/schemas/service_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> ServiceResponse:
    """Convert an `ServiceSchema` to an `ServiceResponse`.

    Args:
        include_metadata: Whether to include metadata in the response.
        include_resources: Whether to include resources in the response.
        kwargs: Additional keyword arguments.

    Returns:
        The created `ServiceResponse`.
    """
    body = ServiceResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        service_type=json.loads(self.service_type),
        labels=json.loads(base64.b64decode(self.labels).decode())
        if self.labels
        else None,
        state=self.state,
    )
    metadata = None
    if include_metadata:
        metadata = ServiceResponseMetadata(
            service_source=self.service_source,
            config=json.loads(base64.b64decode(self.config).decode()),
            status=json.loads(base64.b64decode(self.status).decode())
            if self.status
            else None,
            endpoint=json.loads(base64.b64decode(self.endpoint).decode())
            if self.endpoint
            else None,
            admin_state=self.admin_state or None,
            prediction_url=self.prediction_url or None,
            health_check_url=self.health_check_url,
        )
    resources = None
    if include_resources:
        resources = ServiceResponseResources(
            user=self.user.to_model() if self.user else None,
            model_version=self.model_version.to_model()
            if self.model_version
            else None,
            pipeline_run=self.pipeline_run.to_model()
            if self.pipeline_run
            else None,
        )
    return ServiceResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: ServiceUpdate) -> ServiceSchema

Updates a ServiceSchema from a ServiceUpdate.

Parameters:

Name	Type	Description	Default
update	ServiceUpdate	The ServiceUpdate to update from.	required

Returns:

Type	Description
ServiceSchema	The updated ServiceSchema.

Source code in src/zenml/zen_stores/schemas/service_schemas.py

def update(
    self,
    update: ServiceUpdate,
) -> "ServiceSchema":
    """Updates a `ServiceSchema` from a `ServiceUpdate`.

    Args:
        update: The `ServiceUpdate` to update from.

    Returns:
        The updated `ServiceSchema`.
    """
    for field, value in update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field == "labels":
            self.labels = (
                dict_to_bytes(update.labels) if update.labels else None
            )
        elif field == "status":
            self.status = (
                dict_to_bytes(update.status) if update.status else None
            )
            self.state = (
                update.status.get("state") if update.status else None
            )
        elif field == "endpoint":
            self.endpoint = (
                dict_to_bytes(update.endpoint) if update.endpoint else None
            )
        else:
            setattr(self, field, value)
    self.updated = utc_now()
    return self

Functions

stack_schemas

SQL Model Implementations for Stacks.

Classes

StackCompositionSchema

Bases: SQLModel

SQL Model for stack definitions.

Join table between Stacks and StackComponents.

StackSchema

Bases: NamedSchema

SQL Model for stacks.

Attributes

has_deployer: bool property

If the stack has a deployer component.

Returns:

Type	Description
bool	If the stack has a deployer component.

Raises:

Type	Description
RuntimeError	if the stack has no DB session.

Functions

from_request(request: StackRequest) -> StackSchema classmethod

Create a stack schema from a request.

Parameters:

Name	Type	Description	Default
request	StackRequest	The request from which to create the stack.	required

Returns:

Type	Description
StackSchema	The stack schema.

Source code in src/zenml/zen_stores/schemas/stack_schemas.py

@classmethod
def from_request(
    cls,
    request: "StackRequest",
) -> "StackSchema":
    """Create a stack schema from a request.

    Args:
        request: The request from which to create the stack.

    Returns:
        The stack schema.
    """
    return cls(
        user_id=request.user,
        stack_spec_path=request.stack_spec_path,
        name=request.name,
        description=request.description,
        labels=base64.b64encode(
            json.dumps(request.labels).encode("utf-8")
        ),
        environment=base64.b64encode(
            json.dumps(request.environment).encode("utf-8")
        ),
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/stack_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_metadata:
        options.extend(
            [
                joinedload(
                    jl_arg(StackSchema.stack_compositions)
                ).joinedload(jl_arg(StackCompositionSchema.component))
            ]
        )

    if include_resources:
        options.extend([joinedload(jl_arg(StackSchema.user))])

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> StackResponse

Converts the schema to a model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
StackResponse	The converted model.

Source code in src/zenml/zen_stores/schemas/stack_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "StackResponse":
    """Converts the schema to a model.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The converted model.
    """
    body = StackResponseBody(
        user_id=self.user_id,
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        environment = None
        if self.environment:
            environment = json.loads(
                base64.b64decode(self.environment).decode()
            )

        components: Dict[
            StackComponentType, List["StackComponentSchema"]
        ] = defaultdict(list)

        sorted_compositions = sorted(
            self.stack_compositions,
            key=lambda composition: (
                composition.default_for_type is None,
            ),
        )
        for composition in sorted_compositions:
            component_type = StackComponentType(composition.component.type)
            components[component_type].append(composition.component)

        metadata = StackResponseMetadata(
            components={
                component_type: [
                    component.to_model() for component in component_list
                ]
                for component_type, component_list in components.items()
            },
            stack_spec_path=self.stack_spec_path,
            labels=json.loads(base64.b64decode(self.labels).decode())
            if self.labels
            else None,
            description=self.description,
            environment=environment or {},
            secrets=[secret.id for secret in self.secrets],
        )
    resources = None
    if include_resources:
        resources = StackResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return StackResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(stack_update: StackUpdate) -> StackSchema

Updates a stack schema with a stack update model.

Parameters:

Name	Type	Description	Default
stack_update	StackUpdate	StackUpdate to update the stack with.	required

Returns:

Type	Description
StackSchema	The updated StackSchema.

Source code in src/zenml/zen_stores/schemas/stack_schemas.py

def update(
    self,
    stack_update: "StackUpdate",
) -> "StackSchema":
    """Updates a stack schema with a stack update model.

    Args:
        stack_update: `StackUpdate` to update the stack with.

    Returns:
        The updated StackSchema.
    """
    for field, value in stack_update.model_dump(
        exclude_unset=True,
        exclude={"user", "components", "add_secrets", "remove_secrets"},
    ).items():
        if field == "labels":
            self.labels = base64.b64encode(
                json.dumps(stack_update.labels).encode("utf-8")
            )
        elif field == "environment":
            self.environment = base64.b64encode(
                json.dumps(stack_update.environment).encode("utf-8")
            )
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

step_run_schemas

SQLModel implementation of step run tables.

Classes

StepRunInputArtifactSchema

Bases: SQLModel

SQL Model that defines which artifacts are inputs to which step.

StepRunOutputArtifactSchema

Bases: SQLModel

SQL Model that defines which artifacts are outputs of which step.

StepRunParentsSchema

Bases: SQLModel

SQL Model that defines the order of steps.

StepRunSchema

Bases: NamedSchema, RunMetadataInterface

SQL Model for steps of pipeline runs.

Functions

from_request(request: StepRunRequest, snapshot_id: Optional[UUID], version: int, is_retriable: bool) -> StepRunSchema classmethod

Create a step run schema from a step run request model.

Parameters:

Name	Type	Description	Default
request	StepRunRequest	The step run request model.	required
snapshot_id	Optional[UUID]	The snapshot ID.	required
version	int	The version of the step run.	required
is_retriable	bool	Whether the step run is retriable.	required

Returns:

Type	Description
StepRunSchema	The step run schema.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

@classmethod
def from_request(
    cls,
    request: StepRunRequest,
    snapshot_id: Optional[UUID],
    version: int,
    is_retriable: bool,
) -> "StepRunSchema":
    """Create a step run schema from a step run request model.

    Args:
        request: The step run request model.
        snapshot_id: The snapshot ID.
        version: The version of the step run.
        is_retriable: Whether the step run is retriable.

    Returns:
        The step run schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        start_time=request.start_time,
        end_time=request.end_time,
        status=request.status.value,
        snapshot_id=snapshot_id,
        original_step_run_id=request.original_step_run_id,
        pipeline_run_id=request.pipeline_run_id,
        docstring=request.docstring,
        cache_key=request.cache_key,
        cache_expires_at=request.cache_expires_at,
        code_hash=request.code_hash,
        source_code=request.source_code,
        version=version,
        is_retriable=is_retriable,
        exception_info=request.exception_info.model_dump_json()
        if request.exception_info
        else None,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    from zenml.zen_stores.schemas import (
        ArtifactVersionSchema,
        ModelVersionSchema,
    )

    options = [
        selectinload(jl_arg(StepRunSchema.snapshot)).load_only(
            jl_arg(PipelineSnapshotSchema.pipeline_configuration)
        ),
        selectinload(jl_arg(StepRunSchema.pipeline_run)).load_only(
            jl_arg(PipelineRunSchema.start_time)
        ),
        joinedload(jl_arg(StepRunSchema.static_config)),
        joinedload(jl_arg(StepRunSchema.dynamic_config)),
    ]

    # if include_metadata:
    #     options.extend(
    #         [
    #             joinedload(jl_arg(StepRunSchema.parents)),
    #             joinedload(jl_arg(StepRunSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend(
            [
                selectinload(
                    jl_arg(StepRunSchema.model_version)
                ).joinedload(
                    jl_arg(ModelVersionSchema.model), innerjoin=True
                ),
                selectinload(jl_arg(StepRunSchema.user)),
                selectinload(jl_arg(StepRunSchema.input_artifacts))
                .joinedload(
                    jl_arg(StepRunInputArtifactSchema.artifact_version),
                    innerjoin=True,
                )
                .joinedload(
                    jl_arg(ArtifactVersionSchema.artifact), innerjoin=True
                ),
                selectinload(jl_arg(StepRunSchema.output_artifacts))
                .joinedload(
                    jl_arg(StepRunOutputArtifactSchema.artifact_version),
                    innerjoin=True,
                )
                .joinedload(
                    jl_arg(ArtifactVersionSchema.artifact), innerjoin=True
                ),
                selectinload(jl_arg(StepRunSchema.logs)),
            ]
        )

    return options

get_step_configuration() -> Step

Get the step configuration for the step run.

Raises:

Type	Description
ValueError	If the step run has no step configuration.

Returns:

Type	Description
Step	The step configuration.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

def get_step_configuration(self) -> Step:
    """Get the step configuration for the step run.

    Raises:
        ValueError: If the step run has no step configuration.

    Returns:
        The step configuration.
    """
    step = None

    if self.snapshot is not None:
        if config_schema := (self.dynamic_config or self.static_config):
            pipeline_configuration = (
                PipelineConfiguration.model_validate_json(
                    self.snapshot.pipeline_configuration
                )
            )
            pipeline_configuration.finalize_substitutions(
                start_time=self.pipeline_run.start_time,
                inplace=True,
            )
            step = Step.from_dict(
                json.loads(config_schema.config),
                pipeline_configuration=pipeline_configuration,
            )
    if not step and self.step_configuration:
        # In this legacy case, we're guaranteed to have the merged
        # config stored in the DB, which means we can instantiate the
        # `Step` object directly without passing the pipeline
        # configuration.
        step = Step.model_validate_json(self.step_configuration)
    elif not step:
        raise ValueError(
            f"Unable to load the configuration for step `{self.name}` from "
            "the database. To solve this please delete the pipeline run "
            "that this step run belongs to. Pipeline Run ID: "
            f"`{self.pipeline_run_id}`."
        )

    return step

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> StepRunResponse

Convert a StepRunSchema to a StepRunResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
StepRunResponse	The created StepRunResponse.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> StepRunResponse:
    """Convert a `StepRunSchema` to a `StepRunResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created StepRunResponse.
    """
    step = self.get_step_configuration()

    body = StepRunResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        type=step.config.step_type,
        status=ExecutionStatus(self.status),
        version=self.version,
        is_retriable=self.is_retriable,
        start_time=self.start_time,
        end_time=self.end_time,
        latest_heartbeat=self.latest_heartbeat,
        created=self.created,
        updated=self.updated,
        model_version_id=self.model_version_id,
        substitutions=step.config.substitutions,
        heartbeat_threshold=self.heartbeat_threshold,
    )
    metadata = None
    if include_metadata:
        metadata = StepRunResponseMetadata(
            config=step.config,
            spec=step.spec,
            cache_key=self.cache_key,
            cache_expires_at=self.cache_expires_at,
            code_hash=self.code_hash,
            docstring=self.docstring,
            source_code=self.source_code,
            exception_info=ExceptionInfo.model_validate_json(
                self.exception_info
            )
            if self.exception_info
            else None,
            snapshot_id=self.snapshot_id,
            pipeline_run_id=self.pipeline_run_id,
            original_step_run_id=self.original_step_run_id,
            parent_step_ids=[p.parent_id for p in self.parents],
            run_metadata=self.fetch_metadata(),
        )

    resources = None
    if include_resources:
        model_version = None
        if self.model_version:
            model_version = self.model_version.to_model()

        input_artifacts: Dict[str, List[StepRunInputResponse]] = {}
        for input_artifact in self.input_artifacts:
            if input_artifact.name not in input_artifacts:
                input_artifacts[input_artifact.name] = []
            step_run_input = StepRunInputResponse(
                input_type=StepRunInputArtifactType(input_artifact.type),
                index=input_artifact.input_index,
                chunk_index=input_artifact.chunk_index,
                chunk_size=input_artifact.chunk_size,
                **input_artifact.artifact_version.to_model().model_dump(),
            )
            input_artifacts[input_artifact.name].append(step_run_input)

        for artifact_list in input_artifacts.values():
            artifact_list.sort(key=lambda a: a.index or 0)

        output_artifacts: Dict[str, List["ArtifactVersionResponse"]] = {}
        for output_artifact in self.output_artifacts:
            if output_artifact.name not in output_artifacts:
                output_artifacts[output_artifact.name] = []
            output_artifacts[output_artifact.name].append(
                output_artifact.artifact_version.to_model()
            )

        resources = StepRunResponseResources(
            user=self.user.to_model() if self.user else None,
            model_version=model_version,
            log_collection=[
                log.to_model()
                for log in sorted(self.logs, key=lambda log: log.created)
            ],
            inputs=input_artifacts,
            outputs=output_artifacts,
            resource_request=self.resource_request.to_model(
                include_metadata=True, include_resources=False
            )
            if self.resource_request
            else None,
        )

    return StepRunResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(step_update: StepRunUpdate) -> StepRunSchema

Update a step run schema with a step run update model.

Parameters:

Name	Type	Description	Default
step_update	StepRunUpdate	The step run update model.	required

Returns:

Type	Description
StepRunSchema	The updated step run schema.

Source code in src/zenml/zen_stores/schemas/step_run_schemas.py

def update(self, step_update: "StepRunUpdate") -> "StepRunSchema":
    """Update a step run schema with a step run update model.

    Args:
        step_update: The step run update model.

    Returns:
        The updated step run schema.
    """
    for key, value in step_update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if key == "status":
            self.status = value.value
        if key == "end_time":
            self.end_time = value
        if key == "exception_info":
            self.exception_info = json.dumps(value)
        if key == "cache_expires_at":
            self.cache_expires_at = value

    self.updated = utc_now()

    return self

Functions

tag_schemas

SQLModel implementation of tag tables.

Classes

TagResourceSchema

Bases: BaseSchema

SQL Model for tag resource relationship.

Functions

from_request(request: TagResourceRequest) -> TagResourceSchema classmethod

Convert an TagResourceRequest to an TagResourceSchema.

Parameters:

Name	Type	Description	Default
request	TagResourceRequest	The request model version to convert.	required

Returns:

Type	Description
TagResourceSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

@classmethod
def from_request(cls, request: TagResourceRequest) -> "TagResourceSchema":
    """Convert an `TagResourceRequest` to an `TagResourceSchema`.

    Args:
        request: The request model version to convert.

    Returns:
        The converted schema.
    """
    return cls(
        tag_id=request.tag_id,
        resource_id=request.resource_id,
        resource_type=request.resource_type.value,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TagResourceResponse

Convert an TagResourceSchema to an TagResourceResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
TagResourceResponse	The created TagResourceResponse.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> TagResourceResponse:
    """Convert an `TagResourceSchema` to an `TagResourceResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `TagResourceResponse`.
    """
    return TagResourceResponse(
        id=self.id,
        body=TagResourceResponseBody(
            tag_id=self.tag_id,
            resource_id=self.resource_id,
            created=self.created,
            updated=self.updated,
            resource_type=TaggableResourceTypes(self.resource_type),
        ),
    )

TagSchema

Bases: NamedSchema

SQL Model for tag.

Attributes

tagged_count: int property

Fetch the number of resources tagged with this tag.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
int	The number of resources tagged with this tag.

Functions

from_request(request: TagRequest) -> TagSchema classmethod

Convert an TagRequest to an TagSchema.

Parameters:

Name	Type	Description	Default
request	TagRequest	The request model to convert.	required

Returns:

Type	Description
TagSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

@classmethod
def from_request(cls, request: TagRequest) -> "TagSchema":
    """Convert an `TagRequest` to an `TagSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        exclusive=request.exclusive,
        color=request.color.value,
        user_id=request.user,
    )

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend([joinedload(jl_arg(TagSchema.user))])

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TagResponse

Convert an TagSchema to an TagResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
TagResponse	The created TagResponse.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> TagResponse:
    """Convert an `TagSchema` to an `TagResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic


    Returns:
        The created `TagResponse`.
    """
    metadata = None
    if include_metadata:
        metadata = TagResponseMetadata(
            tagged_count=self.tagged_count,
        )

    resources = None
    if include_resources:
        resources = TagResponseResources(
            user=self.user.to_model() if self.user else None,
        )

    return TagResponse(
        id=self.id,
        name=self.name,
        body=TagResponseBody(
            user_id=self.user_id,
            created=self.created,
            updated=self.updated,
            color=ColorVariants(self.color),
            exclusive=self.exclusive,
        ),
        metadata=metadata,
        resources=resources,
    )

update(update: TagUpdate) -> TagSchema

Updates a TagSchema from a TagUpdate.

Parameters:

Name	Type	Description	Default
update	TagUpdate	The TagUpdate to update from.	required

Returns:

Type	Description
TagSchema	The updated TagSchema.

Source code in src/zenml/zen_stores/schemas/tag_schemas.py

def update(self, update: TagUpdate) -> "TagSchema":
    """Updates a `TagSchema` from a `TagUpdate`.

    Args:
        update: The `TagUpdate` to update from.

    Returns:
        The updated `TagSchema`.
    """
    for field, value in update.model_dump(exclude_unset=True).items():
        if field == "color":
            setattr(self, field, value.value)
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

trigger_assoc

SQL Model Implementations for Triggers Associations.

Classes

TriggerExecutionSchema

Bases: SQLModel

Association table linking triggers to pipeline snapshots.

• Enforces uniqueness per (trigger_id, snapshot_id)
• Cascades deletes from either parent row (DB-level ON DELETE CASCADE)

TriggerSnapshotSchema

Bases: SQLModel

Association table linking triggers to pipeline snapshots.

• Enforces uniqueness per (trigger_id, snapshot_id)
• Cascades deletes from either parent row (DB-level ON DELETE CASCADE)

Attributes

parsed_dispatch_state: TriggerSnapshotDispatchState | None property

Parse persisted dispatch-state JSON into the typed model.

Returns:

Type	Description
TriggerSnapshotDispatchState | None	Parsed trigger dispatch state or None if missing/invalid.

Functions

trigger_schemas

SQL Model Implementations for Triggers.

Classes

TriggerSchema

Bases: NamedSchema

SQL Model for schedules.

Attributes

latest_run: PipelineRunSchema | None property

Fetch the latest execution for this trigger.

Raises:

Type	Description
RuntimeError	If no session for the schema exists.

Returns:

Type	Description
PipelineRunSchema | None	The latest run for this pipeline.

Functions

from_request(trigger_request: TriggerRequest) -> TriggerSchema classmethod

Creates a TriggerSchema object from a TriggerRequest.

Parameters:

Name	Type	Description	Default
trigger_request	TriggerRequest	A TriggerRequest object.	required

Returns:

Type	Description
TriggerSchema	A TriggerSchema object.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

@classmethod
def from_request(cls, trigger_request: TriggerRequest) -> "TriggerSchema":
    """Creates a TriggerSchema object from a TriggerRequest.

    Args:
        trigger_request: A TriggerRequest object.

    Returns:
        A TriggerSchema object.
    """
    extra_fields = trigger_request.get_extra_fields()

    schema = cls(
        name=trigger_request.name,
        project_id=trigger_request.project,
        user_id=trigger_request.user,
        active=trigger_request.active,
        configuration=trigger_request.get_config(),
        flavor=trigger_request.flavor,
        type=trigger_request.type,
        concurrency=trigger_request.concurrency,
    )

    for field_name, value in extra_fields.items():
        setattr(schema, field_name, value)

    return schema

get_query_options(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> Sequence[ExecutableOption] classmethod

Get the query options for the schema.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether metadata will be included when converting the schema to a model.	False
include_resources	bool	Whether resources will be included when converting the schema to a model.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
Sequence[ExecutableOption]	A list of query options.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

@classmethod
def get_query_options(
    cls,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> Sequence[ExecutableOption]:
    """Get the query options for the schema.

    Args:
        include_metadata: Whether metadata will be included when converting
            the schema to a model.
        include_resources: Whether resources will be included when
            converting the schema to a model.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        A list of query options.
    """
    options = []

    if include_resources:
        options.extend(
            [
                selectinload(jl_arg(TriggerSchema.snapshots)).selectinload(
                    jl_arg(PipelineSnapshotSchema.source_snapshot)
                ),
                selectinload(jl_arg(TriggerSchema.snapshot_links)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TRIGGER_RETURN_TYPE_UNION

Converts to Pydantic response model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Flag - to include metadata.	False
include_resources	bool	Flag - include resources.	False
**kwargs	Any	Keyword arguments	{}

Returns:

Type	Description
TRIGGER_RETURN_TYPE_UNION	A TriggerResponse object.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> TRIGGER_RETURN_TYPE_UNION:
    """Converts to Pydantic response model.

    Args:
        include_metadata: Flag - to include metadata.
        include_resources: Flag - include resources.
        **kwargs: Keyword arguments

    Returns:
        A TriggerResponse object.
    """
    body_cls = TYPE_TO_RESPONSE_BODY_MAPPING[self.type]
    response_cls = TYPE_TO_RESPONSE_MAPPING[self.type]

    body = body_cls(
        user_id=self.user_id,
        project_id=self.project_id,
        active=self.active,
        updated=self.updated,
        created=self.created,
        is_archived=self.is_archived,
        type=TriggerType(self.type),
        flavor=TriggerFlavor(self.flavor),
        name=self.name,
        concurrency=self.concurrency,
        **json.loads(self.configuration),
    )

    for field in body.get_extra_fields():
        setattr(body, field, getattr(self, field))

    metadata = None
    if include_metadata:
        metadata = TriggerResponseMetadata()

    resources = None
    if include_resources:
        latest_run = self.latest_run
        display_snapshot_id_by_executable_id: dict[UUID, UUID] = {}
        snapshots = []
        executable_snapshots = []
        for snapshot in self.snapshots:
            snapshot_model = snapshot.to_model()
            executable_snapshots.append(snapshot_model)
            display_snapshot = (
                snapshot.source_snapshot.to_model()
                if snapshot.source_snapshot is not None
                else snapshot_model
            )
            snapshots.append(display_snapshot)
            display_snapshot_id_by_executable_id[snapshot.id] = (
                display_snapshot.id
            )

        snapshot_dispatch_states: dict[
            UUID, TriggerSnapshotDispatchState
        ] = {}
        for snapshot_link in self.snapshot_links:
            parsed_state = snapshot_link.parsed_dispatch_state
            display_snapshot_id = display_snapshot_id_by_executable_id.get(
                snapshot_link.snapshot_id
            )
            if (
                parsed_state is not None
                and display_snapshot_id is not None
            ):
                snapshot_dispatch_states[display_snapshot_id] = (
                    parsed_state
                )

        resources = TriggerResponseResources(
            user=self.user.to_model() if self.user else None,
            snapshots=snapshots,
            executable_snapshots=executable_snapshots,
            latest_run=latest_run.to_model()
            if latest_run is not None
            else None,
            snapshot_dispatch_states=snapshot_dispatch_states,
        )

    return response_cls(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(trigger_update: TriggerUpdate) -> TriggerSchema

Applies update operation (and validations).

Parameters:

Name	Type	Description	Default
trigger_update	TriggerUpdate	A TriggerUpdate object.	required

Returns:

Type	Description
TriggerSchema	The updated TriggerSchema.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

def update(self, trigger_update: TriggerUpdate) -> "TriggerSchema":
    """Applies update operation (and validations).

    Args:
        trigger_update: A TriggerUpdate object.

    Returns:
        The updated TriggerSchema.
    """
    for field, value in trigger_update.model_dump(
        exclude_unset=True,
        include=set(TriggerBase.model_fields.keys()),
    ).items():
        if field in ["type"]:
            continue
        setattr(self, field, value)

    self.configuration = trigger_update.get_config()

    for field_name, value in trigger_update.get_extra_fields().items():
        setattr(self, field_name, value)

    return self

Functions

user_schemas

SQLModel implementation of user tables.

Classes

UserSchema

Bases: NamedSchema

SQL Model for users.

Functions

from_service_account_request(model: Union[ServiceAccountRequest, ServiceAccountInternalRequest]) -> UserSchema classmethod

Create a UserSchema from a Service Account request.

Parameters:

Name	Type	Description	Default
model	Union[ServiceAccountRequest, ServiceAccountInternalRequest]	The ServiceAccountRequest or ServiceAccountInternalRequest from which to create the schema.	required

Returns:

Type	Description
UserSchema	The created UserSchema.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

@classmethod
def from_service_account_request(
    cls, model: Union[ServiceAccountRequest, ServiceAccountInternalRequest]
) -> "UserSchema":
    """Create a `UserSchema` from a Service Account request.

    Args:
        model: The `ServiceAccountRequest` or `ServiceAccountInternalRequest`
            from which to create the schema.

    Returns:
        The created `UserSchema`.
    """
    return cls(
        name=model.name,
        full_name=model.full_name,
        description=model.description or "",
        external_user_id=model.external_user_id
        if isinstance(model, ServiceAccountInternalRequest)
        else None,
        active=model.active,
        is_service_account=True,
        email_opted_in=False,
        is_admin=False,
        avatar_url=model.avatar_url,
    )

from_user_request(model: UserRequest) -> UserSchema classmethod

Create a UserSchema from a UserRequest.

Parameters:

Name	Type	Description	Default
model	UserRequest	The UserRequest from which to create the schema.	required

Returns:

Type	Description
UserSchema	The created UserSchema.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

@classmethod
def from_user_request(cls, model: UserRequest) -> "UserSchema":
    """Create a `UserSchema` from a `UserRequest`.

    Args:
        model: The `UserRequest` from which to create the schema.

    Returns:
        The created `UserSchema`.
    """
    return cls(
        name=model.name,
        full_name=model.full_name,
        avatar_url=model.avatar_url,
        active=model.active,
        password=model.create_hashed_password(),
        activation_token=model.create_hashed_activation_token(),
        external_user_id=model.external_user_id,
        email_opted_in=model.email_opted_in,
        email=model.email,
        is_service_account=False,
        is_admin=model.is_admin,
        user_metadata=json.dumps(model.user_metadata)
        if model.user_metadata
        else None,
    )

to_model(include_metadata: bool = False, include_resources: bool = False, include_private: bool = False, **kwargs: Any) -> UserResponse

Convert a UserSchema to a UserResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}
include_private	bool	Whether to include the user private information this is to limit the amount of data one can get about other users.	False

Returns:

Type	Description
UserResponse	The converted UserResponse.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    include_private: bool = False,
    **kwargs: Any,
) -> UserResponse:
    """Convert a `UserSchema` to a `UserResponse`.

    Args:
        include_metadata: Whether the metadata will be filled.
        include_resources: Whether the resources will be filled.
        **kwargs: Keyword arguments to allow schema specific logic
        include_private: Whether to include the user private information
            this is to limit the amount of data one can get about other
            users.

    Returns:
        The converted `UserResponse`.
    """
    metadata = None
    if include_metadata:
        metadata = UserResponseMetadata(
            email=self.email if include_private else None,
            external_user_id=self.external_user_id,
            user_metadata=json.loads(self.user_metadata)
            if self.user_metadata
            else {},
        )

    return UserResponse(
        id=self.id,
        name=self.name,
        body=UserResponseBody(
            active=self.active,
            full_name=self.full_name,
            email_opted_in=self.email_opted_in,
            is_service_account=self.is_service_account,
            created=self.created,
            updated=self.updated,
            is_admin=self.is_admin,
            default_project_id=self.default_project_id,
            avatar_url=self.avatar_url,
        ),
        metadata=metadata,
    )

to_service_account_model(include_metadata: bool = False, include_resources: bool = False) -> ServiceAccountResponse

Convert a UserSchema to a ServiceAccountResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Whether the metadata will be filled.	False
include_resources	bool	Whether the resources will be filled.	False

Returns:

Type	Description
ServiceAccountResponse	The converted ServiceAccountResponse.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

def to_service_account_model(
    self, include_metadata: bool = False, include_resources: bool = False
) -> ServiceAccountResponse:
    """Convert a `UserSchema` to a `ServiceAccountResponse`.

    Args:
         include_metadata: Whether the metadata will be filled.
         include_resources: Whether the resources will be filled.

    Returns:
         The converted `ServiceAccountResponse`.
    """
    metadata = None
    if include_metadata:
        metadata = ServiceAccountResponseMetadata(
            description=self.description or "",
            external_user_id=self.external_user_id,
        )

    body = ServiceAccountResponseBody(
        full_name=self.full_name,
        created=self.created,
        updated=self.updated,
        active=self.active,
        avatar_url=self.avatar_url,
    )

    return ServiceAccountResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
    )

update_service_account(service_account_update: ServiceAccountUpdate) -> UserSchema

Update a UserSchema from a ServiceAccountUpdate.

Parameters:

Name	Type	Description	Default
service_account_update	ServiceAccountUpdate	The ServiceAccountUpdate from which to update the schema.	required

Returns:

Type	Description
UserSchema	The updated UserSchema.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

def update_service_account(
    self, service_account_update: ServiceAccountUpdate
) -> "UserSchema":
    """Update a `UserSchema` from a `ServiceAccountUpdate`.

    Args:
        service_account_update: The `ServiceAccountUpdate` from which
            to update the schema.

    Returns:
        The updated `UserSchema`.
    """
    for field, value in service_account_update.model_dump(
        exclude_none=True
    ).items():
        setattr(self, field, value)

    self.updated = utc_now()
    return self

update_user(user_update: UserUpdate) -> UserSchema

Update a UserSchema from a UserUpdate.

Parameters:

Name	Type	Description	Default
user_update	UserUpdate	The UserUpdate from which to update the schema.	required

Returns:

Type	Description
UserSchema	The updated UserSchema.

Source code in src/zenml/zen_stores/schemas/user_schemas.py

def update_user(self, user_update: UserUpdate) -> "UserSchema":
    """Update a `UserSchema` from a `UserUpdate`.

    Args:
        user_update: The `UserUpdate` from which to update the schema.

    Returns:
        The updated `UserSchema`.
    """
    for field, value in user_update.model_dump(exclude_unset=True).items():
        if field == "old_password":
            continue

        if field == "password":
            setattr(self, field, user_update.create_hashed_password())
        elif field == "activation_token":
            setattr(
                self, field, user_update.create_hashed_activation_token()
            )
        elif field == "user_metadata":
            if value is not None:
                self.user_metadata = json.dumps(value)
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

utils

Utils for schemas.

Classes

RunMetadataInterface

The interface for entities with run metadata.

Functions

fetch_metadata(**kwargs: Any) -> Dict[str, MetadataType]

Fetches the latest metadata entry related to the entity.

Parameters:

Name	Type	Description	Default
**kwargs	Any	Keyword arguments to pass to the metadata collection.	{}

Returns:

Type	Description
Dict[str, MetadataType]	A dictionary, where the key is the key of the metadata entry and the values represent the latest entry with this key.

Source code in src/zenml/zen_stores/schemas/utils.py

def fetch_metadata(self, **kwargs: Any) -> Dict[str, MetadataType]:
    """Fetches the latest metadata entry related to the entity.

    Args:
        **kwargs: Keyword arguments to pass to the metadata collection.

    Returns:
        A dictionary, where the key is the key of the metadata entry
            and the values represent the latest entry with this key.
    """
    metadata_collection = self.fetch_metadata_collection(**kwargs)
    metadata: Dict[str, MetadataType] = {}

    for key, values in metadata_collection.items():
        values = sorted(values, key=lambda x: x.created, reverse=False)

        if all(isinstance(item.value, dict) for item in values):
            # All metadata values for this key are dictionaries, so we can
            # merge them into a single dictionary
            metadata[key] = {
                k: v
                for item in values
                for k, v in item.value.items()  # type: ignore[union-attr]
            }
        else:
            metadata[key] = values[-1].value

    return metadata

fetch_metadata_collection(**kwargs: Any) -> Dict[str, List[RunMetadataEntry]]

Fetches all the metadata entries related to the entity.

Parameters:

Name	Type	Description	Default
**kwargs	Any	Keyword arguments.	{}

Returns:

Type	Description
Dict[str, List[RunMetadataEntry]]	A dictionary, where the key is the key of the metadata entry and the values represent the list of entries with this key.

Source code in src/zenml/zen_stores/schemas/utils.py

def fetch_metadata_collection(
    self, **kwargs: Any
) -> Dict[str, List[RunMetadataEntry]]:
    """Fetches all the metadata entries related to the entity.

    Args:
        **kwargs: Keyword arguments.

    Returns:
        A dictionary, where the key is the key of the metadata entry
            and the values represent the list of entries with this key.
    """
    metadata_collection: Dict[str, List[RunMetadataEntry]] = {}

    for rm in self.run_metadata:
        if rm.key not in metadata_collection:
            metadata_collection[rm.key] = []
        metadata_collection[rm.key].append(
            RunMetadataEntry(
                value=json.loads(rm.value),
                created=rm.created,
            )
        )

    return metadata_collection

Functions

get_page_from_list(items_list: List[S], response_model: Type[BaseResponse], size: int = 5, page: int = 1, include_resources: bool = False, include_metadata: bool = False) -> Page[BaseResponse]

Converts list of schemas into page of response models.

Parameters:

Name	Type	Description	Default
items_list	List[S]	List of schemas	required
response_model	Type[BaseResponse]	Response model	required
size	int	Page size	5
page	int	Page number	1
include_metadata	bool	Whether metadata should be included in response models	False
include_resources	bool	Whether resources should be included in response models	False

Returns:

Type	Description
Page[BaseResponse]	A page of list items.

Source code in src/zenml/zen_stores/schemas/utils.py

def get_page_from_list(
    items_list: List[S],
    response_model: Type[BaseResponse],  # type: ignore[type-arg]
    size: int = 5,
    page: int = 1,
    include_resources: bool = False,
    include_metadata: bool = False,
) -> Page[BaseResponse]:  # type: ignore[type-arg]
    """Converts list of schemas into page of response models.

    Args:
        items_list: List of schemas
        response_model: Response model
        size: Page size
        page: Page number
        include_metadata: Whether metadata should be included in response models
        include_resources: Whether resources should be included in response models

    Returns:
        A page of list items.
    """
    total = len(items_list)
    if total == 0:
        total_pages = 1
    else:
        total_pages = math.ceil(total / size)

    start = (page - 1) * size
    end = start + size

    page_items = [
        item.to_model(
            include_metadata=include_metadata,
            include_resources=include_resources,
        )
        for item in items_list[start:end]
    ]
    return Page[response_model](  # type: ignore[valid-type]
        index=page,
        max_size=size,
        total_pages=total_pages,
        total=total,
        items=page_items,
    )

get_resource_type_name(schema_class: Type[BaseSchema]) -> str

Get the name of a resource from a schema class.

Parameters:

Name	Type	Description	Default
schema_class	Type[BaseSchema]	The schema class to get the name of.	required

Returns:

Type	Description
str	The name of the resource.

Source code in src/zenml/zen_stores/schemas/utils.py

def get_resource_type_name(schema_class: Type[BaseSchema]) -> str:
    """Get the name of a resource from a schema class.

    Args:
        schema_class: The schema class to get the name of.

    Returns:
        The name of the resource.
    """
    entity_name = schema_class.__tablename__
    assert isinstance(entity_name, str)
    # Some entities are plural, some are singular, some have multiple words
    # in their table name connected by underscores (e.g. pipeline_run)
    return entity_name.replace("_", " ").rstrip("s")

jl_arg(column: Any) -> InstrumentedAttribute[Any]

Cast a SQLModel column to a joinedload argument.

Parameters:

Name	Type	Description	Default
column	Any	The column.	required

Returns:

Type	Description
InstrumentedAttribute[Any]	The column cast to a joinedload argument.

Source code in src/zenml/zen_stores/schemas/utils.py

def jl_arg(column: Any) -> InstrumentedAttribute[Any]:
    """Cast a SQLModel column to a joinedload argument.

    Args:
        column: The column.

    Returns:
        The column cast to a joinedload argument.
    """
    return cast(InstrumentedAttribute[Any], column)

secrets_stores

Centralized secrets management.

Modules

aws_secrets_store

AWS Secrets Store implementation.

Classes

AWSSecretsStore(zen_store: BaseZenStore, **kwargs: Any)

Bases: ServiceConnectorSecretsStore

Secrets store implementation that uses the AWS Secrets Manager API.

This secrets store implementation uses the AWS Secrets Manager API to store secrets. It allows a single AWS Secrets Manager region "instance" to be shared with other ZenML deployments as well as other third party users and applications.

Here are some implementation highlights:

• the name/ID of an AWS secret is derived from the ZenML secret UUID and a zenml prefix in the form zenml/{zenml_secret_uuid}. This clearly identifies a secret as being managed by ZenML in the AWS console.

• the Secrets Store also uses AWS secret tags to store additional metadata associated with a ZenML secret. The zenml tag in particular is used to identify and group all secrets that belong to the same ZenML deployment.

• all secret key-values configured in a ZenML secret are stored as a single JSON string value in the AWS secret value.

Source code in src/zenml/zen_stores/secrets_stores/base_secrets_store.py

def __init__(
    self,
    zen_store: "BaseZenStore",
    **kwargs: Any,
) -> None:
    """Create and initialize a secrets store.

    Args:
        zen_store: The ZenML store that owns this secrets store.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.

    Raises:
        RuntimeError: If the store cannot be initialized.
    """
    super().__init__(**kwargs)
    self._zen_store = zen_store

    try:
        self._initialize()
    except Exception as e:
        raise RuntimeError(
            f"Error initializing {self.type.value} secrets store: {str(e)}"
        ) from e

Functions

delete_secret_values(secret_id: UUID) -> None

Deletes secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret.	required

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	If the AWS Secrets Manager API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/aws_secrets_store.py

def delete_secret_values(self, secret_id: UUID) -> None:
    """Deletes secret values for an existing secret.

    Args:
        secret_id: The ID of the secret.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: If the AWS Secrets Manager API returns an unexpected
            error.
    """
    aws_secret_id = self._get_aws_secret_id(secret_id)

    try:
        self.client.delete_secret(
            SecretId=aws_secret_id,
            # We set this to force immediate deletion of the AWS secret
            # instead of waiting for the recovery window to expire.
            ForceDeleteWithoutRecovery=True,
        )
    except ClientError as e:
        if e.response["Error"]["Code"] == "ResourceNotFoundException":
            raise KeyError(f"Secret with ID {secret_id} not found")

        if (
            e.response["Error"]["Code"] == "InvalidRequestException"
            and "marked for deletion" in e.response["Error"]["Message"]
        ):
            raise KeyError(f"Secret with ID {secret_id} not found")

        raise RuntimeError(
            f"Error deleting secret with ID {secret_id}: {e}"
        )

    logger.debug(f"Deleted AWS secret: {aws_secret_id}")

get_secret_values(secret_id: UUID) -> Dict[str, str]

Get the secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required

Returns:

Type	Description
Dict[str, str]	The secret values.

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	If the AWS Secrets Manager API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/aws_secrets_store.py

def get_secret_values(self, secret_id: UUID) -> Dict[str, str]:
    """Get the secret values for an existing secret.

    Args:
        secret_id: ID of the secret.

    Returns:
        The secret values.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: If the AWS Secrets Manager API returns an unexpected
            error.
    """
    aws_secret_id = self._get_aws_secret_id(secret_id)

    try:
        get_secret_value_response = self.client.get_secret_value(
            SecretId=aws_secret_id
        )
        # We need a separate AWS API call to get the AWS secret tags which
        # contain the ZenML secret metadata, since the get_secret_ value API
        # does not return them.
        describe_secret_response = self.client.describe_secret(
            SecretId=aws_secret_id
        )
    except ClientError as e:
        if e.response["Error"]["Code"] == "ResourceNotFoundException" or (
            e.response["Error"]["Code"] == "InvalidRequestException"
            and "marked for deletion" in e.response["Error"]["Message"]
        ):
            raise KeyError(
                f"Can't find the secret values for secret ID '{secret_id}' "
                f"in the secrets store back-end: {str(e)}"
            ) from e

        raise RuntimeError(
            f"Error fetching secret with ID {secret_id} {e}"
        )

    # Convert the AWS secret tags to a metadata dictionary.
    metadata: Dict[str, str] = {
        tag["Key"]: tag["Value"]
        for tag in describe_secret_response["Tags"]
    }

    # The _verify_secret_metadata method raises a KeyError if the
    # secret is not valid or does not belong to this server. Here we
    # simply pass the exception up the stack, as if the secret was not found
    # in the first place.
    self._verify_secret_metadata(
        secret_id=secret_id,
        metadata=metadata,
    )

    values = get_secret_value_response["SecretString"]

    logger.debug(f"Fetched AWS secret: {aws_secret_id}")

    secret_values = json.loads(values)

    if not isinstance(secret_values, dict):
        raise RuntimeError(
            f"AWS secret values for secret ID {aws_secret_id} could not be "
            "decoded: expected a dictionary."
        )

    return secret_values

store_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Store secret values for a new secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required
secret_values	Dict[str, str]	Values for the secret.	required

Raises:

Type	Description
RuntimeError	If the AWS Secrets Manager API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/aws_secrets_store.py

def store_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Store secret values for a new secret.

    Args:
        secret_id: ID of the secret.
        secret_values: Values for the secret.

    Raises:
        RuntimeError: If the AWS Secrets Manager API returns an unexpected
            error.
    """
    aws_secret_id = self._get_aws_secret_id(secret_id)
    secret_value = json.dumps(secret_values)

    # Convert the ZenML secret metadata to AWS tags
    metadata = self._get_secret_metadata(secret_id=secret_id)
    tags = self._get_aws_secret_tags(metadata)

    try:
        self.client.create_secret(
            Name=aws_secret_id,
            SecretString=secret_value,
            Tags=tags,
        )
    except ClientError as e:
        raise RuntimeError(f"Error creating secret: {e}")

    logger.debug(f"Created AWS secret: {aws_secret_id}")

update_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Updates secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret to be updated.	required
secret_values	Dict[str, str]	The new secret values.	required

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	If the AWS Secrets Manager API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/aws_secrets_store.py

def update_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Updates secret values for an existing secret.

    Args:
        secret_id: The ID of the secret to be updated.
        secret_values: The new secret values.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: If the AWS Secrets Manager API returns an unexpected
            error.
    """
    aws_secret_id = self._get_aws_secret_id(secret_id)
    secret_value = json.dumps(secret_values)

    # Convert the ZenML secret metadata to AWS tags
    metadata = self._get_secret_metadata(secret_id)
    tags = self._get_aws_secret_tags(metadata)

    try:
        # One call to update the secret values
        self.client.put_secret_value(
            SecretId=aws_secret_id,
            SecretString=secret_value,
        )
        # Another call to update the tags
        self.client.tag_resource(
            SecretId=aws_secret_id,
            Tags=tags,
        )
    except ClientError as e:
        if e.response["Error"]["Code"] == "ResourceNotFoundException":
            raise KeyError(f"Secret with ID {secret_id} not found")
        raise RuntimeError(f"Error updating secret: {e}")

    logger.debug(f"Updated AWS secret: {aws_secret_id}")

AWSSecretsStoreConfiguration

Bases: ServiceConnectorSecretsStoreConfiguration

AWS secrets store configuration.

Attributes:

Name	Type	Description
type	SecretsStoreType	The type of the store.

Attributes

region: str property

The AWS region to use.

Returns:

Type	Description
str	The AWS region to use.

Raises:

Type	Description
ValueError	If the region is not configured.

Functions

populate_config(data: Dict[str, Any]) -> Dict[str, Any] classmethod

Populate the connector configuration from legacy attributes.

Parameters:

Name	Type	Description	Default
data	Dict[str, Any]	Dict representing user-specified runtime settings.	required

Returns:

Type	Description
Dict[str, Any]	Validated settings.

Source code in src/zenml/zen_stores/secrets_stores/aws_secrets_store.py

@model_validator(mode="before")
@classmethod
@before_validator_handler
def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
    """Populate the connector configuration from legacy attributes.

    Args:
        data: Dict representing user-specified runtime settings.

    Returns:
        Validated settings.
    """
    # Search for legacy attributes and populate the connector configuration
    # from them, if they exist.
    if data.get("region_name"):
        if not data.get("aws_access_key_id") or not data.get(
            "aws_secret_access_key"
        ):
            logger.warning(
                "The `region_name` AWS secrets store attribute is deprecated "
                "and will be removed in a future version of ZenML. Please use "
                "the `auth_method` and `auth_config` attributes instead. "
                "Using an implicit authentication method for AWS Secrets."
            )
            data["auth_method"] = AWSAuthenticationMethods.IMPLICIT
            data["auth_config"] = dict(
                region=data.get("region_name"),
            )
        else:
            logger.warning(
                "The `aws_access_key_id`, `aws_secret_access_key` and "
                "`region_name` AWS secrets store attributes are deprecated and "
                "will be removed in a future version of ZenML. Please use the "
                "`auth_method` and `auth_config` attributes instead."
            )
            data["auth_method"] = AWSAuthenticationMethods.SECRET_KEY
            data["auth_config"] = dict(
                aws_access_key_id=data.get("aws_access_key_id"),
                aws_secret_access_key=data.get("aws_secret_access_key"),
                region=data.get("region_name"),
            )

    return data

Functions

azure_secrets_store

Azure Secrets Store implementation.

Classes

AzureSecretsStore(zen_store: BaseZenStore, **kwargs: Any)

Bases: ServiceConnectorSecretsStore

Secrets store implementation that uses the Azure Key Vault API.

This secrets store implementation uses the Azure Key Vault API to store secrets. It allows a single Azure Key Vault to be shared with other ZenML deployments as well as other third party users and applications.

Here are some implementation highlights:

• the name/ID of an Azure secret is derived from the ZenML secret UUID and a zenml prefix in the form zenml-{zenml_secret_uuid}. This clearly identifies a secret as being managed by ZenML in the Azure console.

• the Secrets Store also uses Azure Key Vault secret tags to store metadata associated with a ZenML secret. The zenml tag in particular is used to identify and group all secrets that belong to the same ZenML deployment.

• all secret key-values configured in a ZenML secret are stored as a single JSON string value in the Azure Key Vault secret value.

Source code in src/zenml/zen_stores/secrets_stores/base_secrets_store.py

def __init__(
    self,
    zen_store: "BaseZenStore",
    **kwargs: Any,
) -> None:
    """Create and initialize a secrets store.

    Args:
        zen_store: The ZenML store that owns this secrets store.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.

    Raises:
        RuntimeError: If the store cannot be initialized.
    """
    super().__init__(**kwargs)
    self._zen_store = zen_store

    try:
        self._initialize()
    except Exception as e:
        raise RuntimeError(
            f"Error initializing {self.type.value} secrets store: {str(e)}"
        ) from e

Attributes

client: SecretClient property

Initialize and return the Azure Key Vault client.

Returns:

Type	Description
SecretClient	The Azure Key Vault client.

Functions

delete_secret_values(secret_id: UUID) -> None

Deletes secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret.	required

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	if the Azure Key Vault API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/azure_secrets_store.py

def delete_secret_values(self, secret_id: UUID) -> None:
    """Deletes secret values for an existing secret.

    Args:
        secret_id: The ID of the secret.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: if the Azure Key Vault API returns an unexpected
            error.
    """
    azure_secret_id = self._get_azure_secret_id(secret_id)

    try:
        self.client.begin_delete_secret(
            azure_secret_id,
        ).wait()
    except ResourceNotFoundError:
        raise KeyError(f"Secret with ID {secret_id} not found")
    except HttpResponseError as e:
        raise RuntimeError(
            f"Error deleting secret with ID {secret_id}: {e}"
        )

    logger.debug(f"Deleted Azure secret: {azure_secret_id}")

get_secret_values(secret_id: UUID) -> Dict[str, str]

Get the secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required

Returns:

Type	Description
Dict[str, str]	The secret values.

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	if the Azure Key Vault API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/azure_secrets_store.py

def get_secret_values(self, secret_id: UUID) -> Dict[str, str]:
    """Get the secret values for an existing secret.

    Args:
        secret_id: ID of the secret.

    Returns:
        The secret values.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: if the Azure Key Vault API returns an unexpected
            error.
    """
    azure_secret_id = self._get_azure_secret_id(secret_id)

    try:
        azure_secret = self.client.get_secret(
            azure_secret_id,
        )
    except ResourceNotFoundError as e:
        raise KeyError(
            f"Can't find the secret values for secret ID '{secret_id}' "
            f"in the secrets store back-end: {str(e)}"
        ) from e
    except HttpResponseError as e:
        raise RuntimeError(
            f"Error fetching secret with ID {secret_id} {e}"
        )

    # The _verify_secret_metadata method raises a KeyError if the
    # secret is not valid or does not belong to this server. Here we
    # simply pass the exception up the stack, as if the secret was not found
    # in the first place.
    assert azure_secret.properties.tags is not None
    self._verify_secret_metadata(
        secret_id=secret_id,
        metadata=azure_secret.properties.tags,
    )

    values = json.loads(azure_secret.value) if azure_secret.value else {}

    if not isinstance(values, dict):
        raise RuntimeError(
            f"Azure Key Vault secret values for secret {azure_secret_id} "
            "could not be retrieved: invalid type for values"
        )

    logger.debug(f"Retrieved Azure secret: {azure_secret_id}")

    return values

store_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Store secret values for a new secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required
secret_values	Dict[str, str]	Values for the secret.	required

Raises:

Type	Description
RuntimeError	if the Azure Key Vault API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/azure_secrets_store.py

def store_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Store secret values for a new secret.

    Args:
        secret_id: ID of the secret.
        secret_values: Values for the secret.

    Raises:
        RuntimeError: if the Azure Key Vault API returns an unexpected
            error.
    """
    azure_secret_id = self._get_azure_secret_id(secret_id)
    secret_value = json.dumps(secret_values)

    # Use the ZenML secret metadata as Azure tags
    metadata = self._get_secret_metadata(secret_id=secret_id)

    try:
        self.client.set_secret(
            azure_secret_id,
            secret_value,
            tags=metadata,
            content_type="application/json",
        )
    except HttpResponseError as e:
        raise RuntimeError(f"Error creating secret: {e}")

    logger.debug(f"Created Azure secret: {azure_secret_id}")

update_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Updates secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret to be updated.	required
secret_values	Dict[str, str]	The new secret values.	required

Raises:

Type	Description
RuntimeError	if the Azure Key Vault API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/azure_secrets_store.py

def update_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Updates secret values for an existing secret.

    Args:
        secret_id: The ID of the secret to be updated.
        secret_values: The new secret values.

    Raises:
        RuntimeError: if the Azure Key Vault API returns an unexpected
            error.
    """
    azure_secret_id = self._get_azure_secret_id(secret_id)
    secret_value = json.dumps(secret_values)

    # Convert the ZenML secret metadata to Azure tags
    metadata = self._get_secret_metadata(secret_id=secret_id)

    try:
        self.client.set_secret(
            azure_secret_id,
            secret_value,
            tags=metadata,
            content_type="application/json",
        )
    except HttpResponseError as e:
        raise RuntimeError(f"Error updating secret {secret_id}: {e}")

    logger.debug(f"Updated Azure secret: {azure_secret_id}")

AzureSecretsStoreConfiguration

Bases: ServiceConnectorSecretsStoreConfiguration

Azure secrets store configuration.

Attributes:

Name	Type	Description
type	SecretsStoreType	The type of the store.
key_vault_name	str	Name of the Azure Key Vault that this secrets store will use to store secrets.

Functions

populate_config(data: Dict[str, Any]) -> Dict[str, Any] classmethod

Populate the connector configuration from legacy attributes.

Parameters:

Name	Type	Description	Default
data	Dict[str, Any]	Dict representing user-specified runtime settings.	required

Returns:

Type	Description
Dict[str, Any]	Validated settings.

Source code in src/zenml/zen_stores/secrets_stores/azure_secrets_store.py

@model_validator(mode="before")
@classmethod
@before_validator_handler
def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
    """Populate the connector configuration from legacy attributes.

    Args:
        data: Dict representing user-specified runtime settings.

    Returns:
        Validated settings.
    """
    # Search for legacy attributes and populate the connector configuration
    # from them, if they exist.
    if (
        data.get("azure_client_id")
        and data.get("azure_client_secret")
        and data.get("azure_tenant_id")
    ):
        logger.warning(
            "The `azure_client_id`, `azure_client_secret` and "
            "`azure_tenant_id` attributes are deprecated and will be "
            "removed in a future version or ZenML. Please use the "
            "`auth_method` and `auth_config` attributes instead."
        )
        data["auth_method"] = AzureAuthenticationMethods.SERVICE_PRINCIPAL
        data["auth_config"] = dict(
            client_id=data.get("azure_client_id"),
            client_secret=data.get("azure_client_secret"),
            tenant_id=data.get("azure_tenant_id"),
        )

    return data

Functions

base_secrets_store

Base Secrets Store implementation.

Classes

BaseSecretsStore(zen_store: BaseZenStore, **kwargs: Any)

Bases: BaseModel, SecretsStoreInterface, ABC

Base class for accessing and persisting ZenML secret values.

Attributes:

Name	Type	Description
config	SecretsStoreConfiguration	The configuration of the secret store.
_zen_store	Optional[BaseZenStore]	The ZenML store that owns this secrets store.

Create and initialize a secrets store.

Parameters:

Name	Type	Description	Default
zen_store	BaseZenStore	The ZenML store that owns this secrets store.	required
**kwargs	Any	Additional keyword arguments to pass to the Pydantic constructor.	{}

Raises:

Type	Description
RuntimeError	If the store cannot be initialized.

Source code in src/zenml/zen_stores/secrets_stores/base_secrets_store.py

def __init__(
    self,
    zen_store: "BaseZenStore",
    **kwargs: Any,
) -> None:
    """Create and initialize a secrets store.

    Args:
        zen_store: The ZenML store that owns this secrets store.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.

    Raises:
        RuntimeError: If the store cannot be initialized.
    """
    super().__init__(**kwargs)
    self._zen_store = zen_store

    try:
        self._initialize()
    except Exception as e:
        raise RuntimeError(
            f"Error initializing {self.type.value} secrets store: {str(e)}"
        ) from e

Attributes

type: SecretsStoreType property

The type of the secrets store.

Returns:

Type	Description
SecretsStoreType	The type of the secrets store.

zen_store: BaseZenStore property

The ZenML store that owns this secrets store.

Returns:

Type	Description
BaseZenStore	The ZenML store that owns this secrets store.

Raises:

Type	Description
ValueError	If the store is not initialized.

Functions

convert_config(data: Dict[str, Any]) -> Dict[str, Any] classmethod

Method to infer the correct type of the config and convert.

Parameters:

Name	Type	Description	Default
data	Dict[str, Any]	The provided configuration object, can potentially be a generic object	required

Raises:

Type	Description
ValueError	If the provided config object's type does not match any of the current implementations.

Returns:

Type	Description
Dict[str, Any]	The converted configuration object.

Source code in src/zenml/zen_stores/secrets_stores/base_secrets_store.py

@model_validator(mode="before")
@classmethod
@before_validator_handler
def convert_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
    """Method to infer the correct type of the config and convert.

    Args:
        data: The provided configuration object, can potentially be a
            generic object

    Raises:
        ValueError: If the provided config object's type does not match
            any of the current implementations.

    Returns:
        The converted configuration object.
    """
    if data["config"].type == SecretsStoreType.SQL:
        from zenml.zen_stores.secrets_stores.sql_secrets_store import (
            SqlSecretsStoreConfiguration,
        )

        data["config"] = SqlSecretsStoreConfiguration(
            **data["config"].model_dump()
        )

    elif data["config"].type == SecretsStoreType.GCP:
        from zenml.zen_stores.secrets_stores.gcp_secrets_store import (
            GCPSecretsStoreConfiguration,
        )

        data["config"] = GCPSecretsStoreConfiguration(
            **data["config"].model_dump()
        )

    elif data["config"].type == SecretsStoreType.AWS:
        from zenml.zen_stores.secrets_stores.aws_secrets_store import (
            AWSSecretsStoreConfiguration,
        )

        data["config"] = AWSSecretsStoreConfiguration(
            **data["config"].model_dump()
        )

    elif data["config"].type == SecretsStoreType.AZURE:
        from zenml.zen_stores.secrets_stores.azure_secrets_store import (
            AzureSecretsStoreConfiguration,
        )

        data["config"] = AzureSecretsStoreConfiguration(
            **data["config"].model_dump()
        )

    elif data["config"].type == SecretsStoreType.HASHICORP:
        from zenml.zen_stores.secrets_stores.hashicorp_secrets_store import (
            HashiCorpVaultSecretsStoreConfiguration,
        )

        data["config"] = HashiCorpVaultSecretsStoreConfiguration(
            **data["config"].model_dump()
        )
    elif (
        data["config"].type == SecretsStoreType.CUSTOM
        or data["config"].type == SecretsStoreType.NONE
    ):
        pass
    else:
        raise ValueError(
            f"Unknown type '{data['config'].type}' for the configuration."
        )

    return data

create_store(config: SecretsStoreConfiguration, **kwargs: Any) -> BaseSecretsStore staticmethod

Create and initialize a secrets store from a secrets store configuration.

Parameters:

Name	Type	Description	Default
config	SecretsStoreConfiguration	The secrets store configuration to use.	required
**kwargs	Any	Additional keyword arguments to pass to the store class	{}

Returns:

Type	Description
BaseSecretsStore	The initialized secrets store.

Source code in src/zenml/zen_stores/secrets_stores/base_secrets_store.py

@staticmethod
def create_store(
    config: SecretsStoreConfiguration,
    **kwargs: Any,
) -> "BaseSecretsStore":
    """Create and initialize a secrets store from a secrets store configuration.

    Args:
        config: The secrets store configuration to use.
        **kwargs: Additional keyword arguments to pass to the store class

    Returns:
        The initialized secrets store.
    """
    logger.debug(
        f"Creating secrets store with type '{config.type.value}'..."
    )
    store_class = BaseSecretsStore.get_store_class(config)
    store = store_class(
        config=config,
        **kwargs,
    )
    return store

get_store_class(store_config: SecretsStoreConfiguration) -> Type[BaseSecretsStore] staticmethod

Returns the class of the given secrets store type.

Parameters:

Name	Type	Description	Default
store_config	SecretsStoreConfiguration	The configuration of the secrets store.	required

Returns:

Type	Description
Type[BaseSecretsStore]	The class corresponding to the configured secrets store or None if
Type[BaseSecretsStore]	the type is unknown.

Raises:

Type	Description
TypeError	If the secrets store type is unsupported.

Source code in src/zenml/zen_stores/secrets_stores/base_secrets_store.py

@staticmethod
def get_store_class(
    store_config: SecretsStoreConfiguration,
) -> Type["BaseSecretsStore"]:
    """Returns the class of the given secrets store type.

    Args:
        store_config: The configuration of the secrets store.

    Returns:
        The class corresponding to the configured secrets store or None if
        the type is unknown.

    Raises:
        TypeError: If the secrets store type is unsupported.
    """
    if store_config.type == SecretsStoreType.SQL:
        from zenml.zen_stores.secrets_stores.sql_secrets_store import (
            SqlSecretsStore,
        )

        return SqlSecretsStore

    if store_config.type == SecretsStoreType.AWS:
        from zenml.zen_stores.secrets_stores.aws_secrets_store import (
            AWSSecretsStore,
        )

        return AWSSecretsStore
    elif store_config.type == SecretsStoreType.GCP:
        from zenml.zen_stores.secrets_stores.gcp_secrets_store import (
            GCPSecretsStore,
        )

        return GCPSecretsStore
    elif store_config.type == SecretsStoreType.AZURE:
        from zenml.zen_stores.secrets_stores.azure_secrets_store import (
            AzureSecretsStore,
        )

        return AzureSecretsStore
    elif store_config.type == SecretsStoreType.HASHICORP:
        from zenml.zen_stores.secrets_stores.hashicorp_secrets_store import (
            HashiCorpVaultSecretsStore,
        )

        return HashiCorpVaultSecretsStore
    elif store_config.type != SecretsStoreType.CUSTOM:
        raise TypeError(
            f"No store implementation found for secrets store type "
            f"`{store_config.type.value}`."
        )

    return BaseSecretsStore._load_custom_store_class(store_config)

Functions Modules

gcp_secrets_store

Implementation of the GCP Secrets Store.

Classes

GCPSecretsStore(zen_store: BaseZenStore, **kwargs: Any)

Bases: ServiceConnectorSecretsStore

Secrets store implementation that uses the GCP Secrets Manager API.

Source code in src/zenml/zen_stores/secrets_stores/base_secrets_store.py

def __init__(
    self,
    zen_store: "BaseZenStore",
    **kwargs: Any,
) -> None:
    """Create and initialize a secrets store.

    Args:
        zen_store: The ZenML store that owns this secrets store.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.

    Raises:
        RuntimeError: If the store cannot be initialized.
    """
    super().__init__(**kwargs)
    self._zen_store = zen_store

    try:
        self._initialize()
    except Exception as e:
        raise RuntimeError(
            f"Error initializing {self.type.value} secrets store: {str(e)}"
        ) from e

Attributes

client: SecretManagerServiceClient property

Initialize and return the GCP Secrets Manager client.

Returns:

Type	Description
SecretManagerServiceClient	The GCP Secrets Manager client instance.

parent_name: str property

Construct the GCP parent path to the secret manager.

Returns:

Type	Description
str	The parent path to the secret manager

Functions

delete_secret_values(secret_id: UUID) -> None

Deletes secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret.	required

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	if the GCP Secrets Manager API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py

def delete_secret_values(self, secret_id: UUID) -> None:
    """Deletes secret values for an existing secret.

    Args:
        secret_id: The ID of the secret.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: if the GCP Secrets Manager API returns an unexpected
            error.
    """
    gcp_secret_name = self.client.secret_path(
        self.config.project_id,
        self._get_gcp_secret_name(secret_id=secret_id),
    )

    try:
        self.client.delete_secret(request={"name": gcp_secret_name})
    except google_exceptions.NotFound:
        raise KeyError(f"Secret with ID {secret_id} not found")
    except Exception as e:
        raise RuntimeError(f"Failed to delete secret: {str(e)}") from e

    logger.debug(f"Deleted GCP secret: {gcp_secret_name}")

get_secret_values(secret_id: UUID) -> Dict[str, str]

Get the secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required

Returns:

Type	Description
Dict[str, str]	The secret values.

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	if the GCP Secrets Manager API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py

def get_secret_values(self, secret_id: UUID) -> Dict[str, str]:
    """Get the secret values for an existing secret.

    Args:
        secret_id: ID of the secret.

    Returns:
        The secret values.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: if the GCP Secrets Manager API returns an unexpected
            error.
    """
    gcp_secret_name = self.client.secret_path(
        self.config.project_id,
        self._get_gcp_secret_name(secret_id=secret_id),
    )

    try:
        secret = self.client.get_secret(name=gcp_secret_name)
        secret_version_values = self.client.access_secret_version(
            name=f"{gcp_secret_name}/versions/latest"
        )
    except google_exceptions.NotFound as e:
        raise KeyError(
            f"Can't find the secret values for secret ID '{secret_id}' "
            f"in the secrets store back-end: {str(e)}"
        ) from e
    except Exception as e:
        raise RuntimeError(
            f"Error fetching secret with ID {secret_id} {e}"
        )

    # The GCP secret labels do not really behave like a dictionary: when
    # a key is not found, it does not raise a KeyError, but instead
    # returns an empty string. That's why we make this conversion.
    metadata = dict(secret.labels)

    # The _verify_secret_metadata method raises a KeyError if the
    # secret is not valid or does not belong to this server. Here we
    # simply pass the exception up the stack, as if the secret was not found
    # in the first place.
    self._verify_secret_metadata(
        secret_id=secret_id,
        metadata=metadata,
    )

    secret_values = json.loads(
        secret_version_values.payload.data.decode("UTF-8")
    )

    if not isinstance(secret_values, dict):
        raise RuntimeError(
            f"Google secret values for secret ID {gcp_secret_name} could "
            "not be decoded: expected a dictionary."
        )

    logger.debug(f"Fetched GCP secret: {gcp_secret_name}")

    return secret_values

store_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Store secret values for a new secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required
secret_values	Dict[str, str]	Values for the secret.	required

Raises:

Type	Description
RuntimeError	if the GCP Secrets Manager API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py

def store_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Store secret values for a new secret.

    Args:
        secret_id: ID of the secret.
        secret_values: Values for the secret.

    Raises:
        RuntimeError: if the GCP Secrets Manager API returns an unexpected
            error.
    """
    secret_value = json.dumps(secret_values)

    labels = self._get_secret_metadata(secret_id=secret_id)

    try:
        gcp_secret = self.client.create_secret(
            request={
                "parent": self.parent_name,
                "secret_id": self._get_gcp_secret_name(secret_id),
                "secret": {
                    "replication": {"automatic": {}},
                    "labels": labels,
                },
            }
        )

        logger.debug(f"Created empty GCP parent secret: {gcp_secret.name}")

        self.client.add_secret_version(
            request={
                "parent": gcp_secret.name,
                "payload": {"data": secret_value.encode()},
            }
        )

        logger.debug(f"Added value to GCP secret {gcp_secret.name}")
    except Exception as e:
        raise RuntimeError(f"Failed to create secret.: {str(e)}") from e

    logger.debug(f"Created GCP secret {gcp_secret.name}")

update_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Updates secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret to be updated.	required
secret_values	Dict[str, str]	The new secret values.	required

Raises:

Type	Description
RuntimeError	if the GCP Secrets Manager API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py

def update_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Updates secret values for an existing secret.

    Args:
        secret_id: The ID of the secret to be updated.
        secret_values: The new secret values.

    Raises:
        RuntimeError: if the GCP Secrets Manager API returns an unexpected
            error.
    """
    gcp_secret_name = self.client.secret_path(
        self.config.project_id,
        self._get_gcp_secret_name(secret_id=secret_id),
    )

    # Convert the ZenML secret metadata to GCP labels
    metadata = self._get_secret_metadata(secret_id)

    try:
        # Update the secret metadata
        update_secret = {
            "name": gcp_secret_name,
            "labels": metadata,
        }
        update_mask = {"paths": ["labels"]}
        gcp_updated_secret = self.client.update_secret(
            request={
                "secret": update_secret,
                "update_mask": update_mask,
            }
        )
        # Add a new secret version
        secret_value = json.dumps(secret_values)
        self.client.add_secret_version(
            request={
                "parent": gcp_updated_secret.name,
                "payload": {"data": secret_value.encode()},
            }
        )
    except Exception as e:
        raise RuntimeError(f"Error updating secret: {e}") from e

    logger.debug(f"Updated GCP secret: {gcp_secret_name}")

GCPSecretsStoreConfiguration

Bases: ServiceConnectorSecretsStoreConfiguration

GCP secrets store configuration.

Attributes:

Name	Type	Description
type	SecretsStoreType	The type of the store.

Attributes

project_id: str property

Get the GCP project ID.

Returns:

Type	Description
str	The GCP project ID.

Raises:

Type	Description
ValueError	If the project ID is not set.

Functions

populate_config(data: Dict[str, Any]) -> Dict[str, Any] classmethod

Populate the connector configuration from legacy attributes.

Parameters:

Name	Type	Description	Default
data	Dict[str, Any]	Dict representing user-specified runtime settings.	required

Returns:

Type	Description
Dict[str, Any]	Validated settings.

Source code in src/zenml/zen_stores/secrets_stores/gcp_secrets_store.py

@model_validator(mode="before")
@classmethod
@before_validator_handler
def populate_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
    """Populate the connector configuration from legacy attributes.

    Args:
        data: Dict representing user-specified runtime settings.

    Returns:
        Validated settings.
    """
    # Search for legacy attributes and populate the connector configuration
    # from them, if they exist.
    if data.get("project_id"):
        if not os.environ.get("GOOGLE_APPLICATION_CREDENTIALS"):
            logger.warning(
                "The `project_id` GCP secrets store attribute is "
                "deprecated and will be removed in a future version of ZenML. "
                "Please use the `auth_method` and `auth_config` attributes "
                "instead. Using an implicit GCP authentication to access "
                "the GCP Secrets Manager API."
            )
            data["auth_method"] = GCPAuthenticationMethods.IMPLICIT
            data["auth_config"] = dict(
                project_id=data.get("project_id"),
            )
        else:
            logger.warning(
                "The `project_id` GCP secrets store attribute and the "
                "`GOOGLE_APPLICATION_CREDENTIALS` environment variable are "
                "deprecated and will be removed in a future version of ZenML. "
                "Please use the `auth_method` and `auth_config` attributes "
                "instead."
            )
            data["auth_method"] = GCPAuthenticationMethods.SERVICE_ACCOUNT
            data["auth_config"] = dict(
                project_id=data.get("project_id"),
            )
            # Load the service account credentials from the file
            with open(os.environ["GOOGLE_APPLICATION_CREDENTIALS"]) as f:
                data["auth_config"]["service_account_json"] = f.read()

    return data

Functions

hashicorp_secrets_store

HashiCorp Vault Secrets Store implementation.

Classes

HashiCorpVaultAuthMethod

Bases: StrEnum

HashiCorp Vault authentication methods.

HashiCorpVaultSecretsStore(zen_store: BaseZenStore, **kwargs: Any)

Bases: BaseSecretsStore

Secrets store implementation that uses the HashiCorp Vault API.

This secrets store implementation uses the HashiCorp Vault API to store secrets. It allows a single HashiCorp Vault server to be shared with other ZenML deployments as well as other third party users and applications.

Here are some implementation highlights:

• the name/ID of an HashiCorp Vault secret is derived from the ZenML secret UUID and a zenml prefix in the form zenml/{zenml_secret_uuid}. This clearly identifies a secret as being managed by ZenML in the HashiCorp Vault server.

• given that HashiCorp Vault secrets do not support attaching arbitrary metadata in the form of label or tags, we store the entire ZenML secret metadata alongside the secret values in the HashiCorp Vault secret value.

Attributes:

Name	Type	Description
config	HashiCorpVaultSecretsStoreConfiguration	The configuration of the HashiCorp Vault secrets store.
TYPE	SecretsStoreType	The type of the store.
CONFIG_TYPE	Type[SecretsStoreConfiguration]	The type of the store configuration.

Source code in src/zenml/zen_stores/secrets_stores/base_secrets_store.py

def __init__(
    self,
    zen_store: "BaseZenStore",
    **kwargs: Any,
) -> None:
    """Create and initialize a secrets store.

    Args:
        zen_store: The ZenML store that owns this secrets store.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.

    Raises:
        RuntimeError: If the store cannot be initialized.
    """
    super().__init__(**kwargs)
    self._zen_store = zen_store

    try:
        self._initialize()
    except Exception as e:
        raise RuntimeError(
            f"Error initializing {self.type.value} secrets store: {str(e)}"
        ) from e

Attributes

client: hvac.Client property

Initialize and return the HashiCorp Vault client.

Returns:

Type	Description
Client	The HashiCorp Vault client.

Raises:

Type	Description
ValueError	If the configuration is invalid.

Functions

delete_secret_values(secret_id: UUID) -> None

Deletes secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret.	required

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	If the HashiCorp Vault API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py

def delete_secret_values(self, secret_id: UUID) -> None:
    """Deletes secret values for an existing secret.

    Args:
        secret_id: The ID of the secret.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: If the HashiCorp Vault API returns an unexpected
            error.
    """
    vault_secret_id = self._get_vault_secret_id(secret_id)

    try:
        self.client.secrets.kv.v2.delete_metadata_and_all_versions(
            path=vault_secret_id,
            mount_point=self.config.mount_point or DEFAULT_MOUNT_POINT,
        )
    except InvalidPath:
        raise KeyError(f"Secret with ID {secret_id} does not exist.")
    except VaultError as e:
        raise RuntimeError(
            f"Error deleting secret with ID {secret_id}: {e}"
        )

    logger.debug(f"Deleted HashiCorp Vault secret: {vault_secret_id}")

get_secret_values(secret_id: UUID) -> Dict[str, str]

Get the secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required

Returns:

Type	Description
Dict[str, str]	The secret values.

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	If the HashiCorp Vault API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py

def get_secret_values(self, secret_id: UUID) -> Dict[str, str]:
    """Get the secret values for an existing secret.

    Args:
        secret_id: ID of the secret.

    Returns:
        The secret values.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: If the HashiCorp Vault API returns an unexpected
            error.
    """
    vault_secret_id = self._get_vault_secret_id(secret_id)

    try:
        vault_secret = (
            self.client.secrets.kv.v2.read_secret(
                path=vault_secret_id,
                mount_point=self.config.mount_point or DEFAULT_MOUNT_POINT,
            )
            .get("data", {})
            .get("data", {})
        )
    except InvalidPath as e:
        raise KeyError(
            f"Can't find the secret values for secret ID '{secret_id}' "
            f"in the secrets store back-end: {str(e)}"
        ) from e
    except VaultError as e:
        raise RuntimeError(
            f"Error fetching secret with ID {secret_id} {e}"
        )

    try:
        metadata = vault_secret[ZENML_VAULT_SECRET_METADATA_KEY]
        values = vault_secret[ZENML_VAULT_SECRET_VALUES_KEY]
    except (KeyError, ValueError) as e:
        raise KeyError(
            f"Secret could not be retrieved: missing required metadata: {e}"
        )

    if not isinstance(values, dict) or not isinstance(metadata, dict):
        raise RuntimeError(
            f"HashiCorp Vault secret values for secret {vault_secret_id} "
            "could not be retrieved: invalid type for metadata or values"
        )

    # The _verify_secret_metadata method raises a KeyError if the
    # secret is not valid or does not belong to this server. Here we
    # simply pass the exception up the stack, as if the secret was not found
    # in the first place.
    self._verify_secret_metadata(
        secret_id=secret_id,
        metadata=metadata,
    )

    logger.debug(f"Fetched HashiCorp Vault secret: {vault_secret_id}")

    return values

store_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Store secret values for a new secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required
secret_values	Dict[str, str]	Values for the secret.	required

Raises:

Type	Description
RuntimeError	If the HashiCorp Vault API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py

def store_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Store secret values for a new secret.

    Args:
        secret_id: ID of the secret.
        secret_values: Values for the secret.

    Raises:
        RuntimeError: If the HashiCorp Vault API returns an unexpected
            error.
    """
    vault_secret_id = self._get_vault_secret_id(secret_id)

    metadata = self._get_secret_metadata(secret_id=secret_id)

    try:
        self.client.secrets.kv.v2.create_or_update_secret(
            path=vault_secret_id,
            # Store the ZenML secret metadata alongside the secret values
            secret={
                ZENML_VAULT_SECRET_VALUES_KEY: secret_values,
                ZENML_VAULT_SECRET_METADATA_KEY: metadata,
            },
            # Do not allow overwriting an existing secret
            cas=0,
            mount_point=self.config.mount_point or DEFAULT_MOUNT_POINT,
        )
    except VaultError as e:
        raise RuntimeError(f"Error creating secret: {e}")

    logger.debug(f"Created HashiCorp Vault secret: {vault_secret_id}")

update_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Updates secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret to be updated.	required
secret_values	Dict[str, str]	The new secret values.	required

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.
RuntimeError	If the HashiCorp Vault API returns an unexpected error.

Source code in src/zenml/zen_stores/secrets_stores/hashicorp_secrets_store.py

def update_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Updates secret values for an existing secret.

    Args:
        secret_id: The ID of the secret to be updated.
        secret_values: The new secret values.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
        RuntimeError: If the HashiCorp Vault API returns an unexpected
            error.
    """
    vault_secret_id = self._get_vault_secret_id(secret_id)

    # Convert the ZenML secret metadata to HashiCorp Vault tags
    metadata = self._get_secret_metadata(secret_id=secret_id)

    try:
        self.client.secrets.kv.v2.create_or_update_secret(
            path=vault_secret_id,
            # Store the ZenML secret metadata alongside the secret values
            secret={
                ZENML_VAULT_SECRET_VALUES_KEY: secret_values,
                ZENML_VAULT_SECRET_METADATA_KEY: metadata,
            },
            mount_point=self.config.mount_point or DEFAULT_MOUNT_POINT,
        )
    except InvalidPath:
        raise KeyError(f"Secret with ID {secret_id} does not exist.")
    except VaultError as e:
        raise RuntimeError(f"Error updating secret {secret_id}: {e}")

    logger.debug(f"Updated HashiCorp Vault secret: {vault_secret_id}")

HashiCorpVaultSecretsStoreConfiguration

Bases: SecretsStoreConfiguration

HashiCorp Vault secrets store configuration.

Attributes:

Name	Type	Description
type	SecretsStoreType	The type of the store.
vault_addr	str	The url of the Vault server. If not set, the value will be loaded from the VAULT_ADDR environment variable, if configured.
vault_namespace	Optional[str]	The Vault Enterprise namespace.
mount_point	Optional[str]	The mount point to use for all secrets.
auth_method	HashiCorpVaultAuthMethod	The authentication method to use to authenticate with the Vault server.
auth_mount_point	Optional[str]	Custom mount point to use for the authentication method.
vault_token	Optional[PlainSerializedSecretStr]	The token used to authenticate with the Vault server. If not set, the token will be loaded from the VAULT_TOKEN environment variable or from the ~/.vault-token file, if configured.
app_role_id	Optional[str]	The Vault role ID to use. Only used if the authentication method is APP_ROLE.
app_secret_id	Optional[str]	The Vault secret ID to use. Only used if the authentication method is APP_ROLE.
aws_role	Optional[str]	The AWS role to use. Only used if the authentication method is AWS.
aws_header_value	Optional[str]	The AWS header value to use. Only used if the authentication method is AWS and the mount point enforces it.
max_versions	int	The maximum number of secret versions to keep.

Functions

secrets_store_interface

ZenML secrets store interface.

Classes

SecretsStoreInterface

Bases: ABC

ZenML secrets store interface.

All ZenML secrets stores must implement the methods in this interface.

Functions

delete_secret_values(secret_id: UUID) -> None abstractmethod

Deletes secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret.	required

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.

Source code in src/zenml/zen_stores/secrets_stores/secrets_store_interface.py

@abstractmethod
def delete_secret_values(self, secret_id: UUID) -> None:
    """Deletes secret values for an existing secret.

    Args:
        secret_id: The ID of the secret.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
    """

get_secret_values(secret_id: UUID) -> Dict[str, str] abstractmethod

Get the secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required

Returns:

Type	Description
Dict[str, str]	The secret values.

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.

Source code in src/zenml/zen_stores/secrets_stores/secrets_store_interface.py

@abstractmethod
def get_secret_values(self, secret_id: UUID) -> Dict[str, str]:
    """Get the secret values for an existing secret.

    Args:
        secret_id: ID of the secret.

    Returns:
        The secret values.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
    """

store_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None abstractmethod

Store secret values for a new secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required
secret_values	Dict[str, str]	Values for the secret.	required

Source code in src/zenml/zen_stores/secrets_stores/secrets_store_interface.py

@abstractmethod
def store_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Store secret values for a new secret.

    Args:
        secret_id: ID of the secret.
        secret_values: Values for the secret.
    """

update_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None abstractmethod

Updates secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret to be updated.	required
secret_values	Dict[str, str]	The new secret values.	required

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.

Source code in src/zenml/zen_stores/secrets_stores/secrets_store_interface.py

@abstractmethod
def update_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Updates secret values for an existing secret.

    Args:
        secret_id: The ID of the secret to be updated.
        secret_values: The new secret values.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
    """

service_connector_secrets_store

Base secrets store class used for all secrets stores that use a service connector.

Classes

ServiceConnectorSecretsStore(zen_store: BaseZenStore, **kwargs: Any)

Bases: BaseSecretsStore

Base secrets store class for service connector-based secrets stores.

All secrets store implementations that use a Service Connector to authenticate and connect to the secrets store back-end should inherit from this class and:

• implement the _initialize_client_from_connector method
• use a configuration class that inherits from ServiceConnectorSecretsStoreConfiguration
• set the SERVICE_CONNECTOR_TYPE to the service connector type used to connect to the secrets store back-end
• set the SERVICE_CONNECTOR_RESOURCE_TYPE to the resource type used to connect to the secrets store back-end

Source code in src/zenml/zen_stores/secrets_stores/base_secrets_store.py

def __init__(
    self,
    zen_store: "BaseZenStore",
    **kwargs: Any,
) -> None:
    """Create and initialize a secrets store.

    Args:
        zen_store: The ZenML store that owns this secrets store.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.

    Raises:
        RuntimeError: If the store cannot be initialized.
    """
    super().__init__(**kwargs)
    self._zen_store = zen_store

    try:
        self._initialize()
    except Exception as e:
        raise RuntimeError(
            f"Error initializing {self.type.value} secrets store: {str(e)}"
        ) from e

Attributes

client: Any property

Get the secrets store API client.

Returns:

Type	Description
Any	The secrets store API client instance.

lock: Lock property

Get the lock used to treat the client initialization as a critical section.

Returns:

Type	Description
Lock	The lock instance.

ServiceConnectorSecretsStoreConfiguration

Bases: SecretsStoreConfiguration

Base configuration for secrets stores that use a service connector.

Attributes:

Name	Type	Description
auth_method	str	The service connector authentication method to use.
auth_config	Dict[str, Any]	The service connector authentication configuration.

Functions

validate_auth_config(data: Dict[str, Any]) -> Dict[str, Any] classmethod

Convert the authentication configuration if given in JSON format.

Parameters:

Name	Type	Description	Default
data	Dict[str, Any]	The configuration values.	required

Returns:

Type	Description
Dict[str, Any]	The validated configuration values.

Raises:

Type	Description
ValueError	If the authentication configuration is not a valid JSON object.

Source code in src/zenml/zen_stores/secrets_stores/service_connector_secrets_store.py

@model_validator(mode="before")
@classmethod
@before_validator_handler
def validate_auth_config(cls, data: Dict[str, Any]) -> Dict[str, Any]:
    """Convert the authentication configuration if given in JSON format.

    Args:
        data: The configuration values.

    Returns:
        The validated configuration values.

    Raises:
        ValueError: If the authentication configuration is not a valid
            JSON object.
    """
    if isinstance(data.get("auth_config"), str):
        try:
            data["auth_config"] = json.loads(data["auth_config"])
        except json.JSONDecodeError as e:
            raise ValueError(
                f"The authentication configuration is not a valid JSON "
                f"object: {e}"
            )
    return data

Functions

sql_secrets_store

SQL Secrets Store implementation.

Classes

SqlSecretsStore(zen_store: BaseZenStore, **kwargs: Any)

Bases: BaseSecretsStore

Secrets store implementation that uses the SQL ZenML store as a backend.

This secrets store piggybacks on the SQL ZenML store. It uses the same database and configuration as the SQL ZenML store.

Attributes:

Name	Type	Description
config	SqlSecretsStoreConfiguration	The configuration of the SQL secrets store.
TYPE	SecretsStoreType	The type of the store.
CONFIG_TYPE	Type[SecretsStoreConfiguration]	The type of the store configuration.

Create and initialize the SQL secrets store.

Parameters:

Name	Type	Description	Default
zen_store	BaseZenStore	The ZenML store that owns this SQL secrets store.	required
**kwargs	Any	Additional keyword arguments to pass to the Pydantic constructor.	{}

Raises:

Type	Description
IllegalOperationError	If the ZenML store to which this secrets store belongs is not a SQL ZenML store.

Source code in src/zenml/zen_stores/secrets_stores/sql_secrets_store.py

def __init__(
    self,
    zen_store: "BaseZenStore",
    **kwargs: Any,
) -> None:
    """Create and initialize the SQL secrets store.

    Args:
        zen_store: The ZenML store that owns this SQL secrets store.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.

    Raises:
        IllegalOperationError: If the ZenML store to which this secrets
            store belongs is not a SQL ZenML store.
    """
    from zenml.zen_stores.sql_zen_store import SqlZenStore

    if not isinstance(zen_store, SqlZenStore):
        raise IllegalOperationError(
            "The SQL secrets store can only be used with the SQL ZenML "
            "store."
        )
    super().__init__(zen_store, **kwargs)

Attributes

engine: Engine property

The SQLAlchemy engine.

Returns:

Type	Description
Engine	The SQLAlchemy engine.

zen_store: SqlZenStore property

The ZenML store that this SQL secrets store is using as a back-end.

Returns:

Type	Description
SqlZenStore	The ZenML store that this SQL secrets store is using as a back-end.

Raises:

Type	Description
ValueError	If the store is not initialized.

Functions

delete_secret_values(secret_id: UUID) -> None

Deletes secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret.	required

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.

Source code in src/zenml/zen_stores/secrets_stores/sql_secrets_store.py

def delete_secret_values(self, secret_id: UUID) -> None:
    """Deletes secret values for an existing secret.

    Args:
        secret_id: The ID of the secret.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
    """
    with Session(self.engine) as session:
        try:
            secret_in_db = session.exec(
                select(SecretSchema).where(SecretSchema.id == secret_id)
            ).one()
            secret_in_db.values = None
            session.commit()
        except NoResultFound:
            raise KeyError(f"Secret with ID {secret_id} not found.")

get_secret_values(secret_id: UUID) -> Dict[str, str]

Get the secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required

Returns:

Type	Description
Dict[str, str]	The secret values.

Raises:

Type	Description
KeyError	if no secret values for the given ID are stored in the secrets store.

Source code in src/zenml/zen_stores/secrets_stores/sql_secrets_store.py

def get_secret_values(self, secret_id: UUID) -> Dict[str, str]:
    """Get the secret values for an existing secret.

    Args:
        secret_id: ID of the secret.

    Returns:
        The secret values.

    Raises:
        KeyError: if no secret values for the given ID are stored in the
            secrets store.
    """
    with Session(self.engine) as session:
        secret_in_db = session.exec(
            select(SecretSchema).where(SecretSchema.id == secret_id)
        ).first()
        if secret_in_db is None:
            raise KeyError(f"Secret with ID {secret_id} not found.")
        try:
            return secret_in_db.get_secret_values(
                encryption_engine=self._encryption_engine,
            )
        except SecretDecodeError:
            raise KeyError(
                f"Secret values for secret {secret_id} could not be "
                f"decoded. This can happen if encryption has "
                f"been enabled/disabled or if the encryption key has been "
                "reconfigured without proper secrets migration."
            )

store_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Store secret values for a new secret.

The secret is already created in the database by the SQL Zen store, this method only stores the secret values.

Parameters:

Name	Type	Description	Default
secret_id	UUID	ID of the secret.	required
secret_values	Dict[str, str]	Values for the secret.	required

Raises:

Type	Description
KeyError	if a secret for the given ID is not found.

Source code in src/zenml/zen_stores/secrets_stores/sql_secrets_store.py

def store_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Store secret values for a new secret.

    The secret is already created in the database by the SQL Zen store, this
    method only stores the secret values.

    Args:
        secret_id: ID of the secret.
        secret_values: Values for the secret.

    Raises:
        KeyError: if a secret for the given ID is not found.
    """
    with Session(self.engine) as session:
        secret_in_db = session.exec(
            select(SecretSchema).where(SecretSchema.id == secret_id)
        ).first()
        if secret_in_db is None:
            raise KeyError(f"Secret with ID {secret_id} not found.")
        secret_in_db.set_secret_values(
            secret_values=secret_values,
            encryption_engine=self._encryption_engine,
        )
        session.add(secret_in_db)
        session.commit()

update_secret_values(secret_id: UUID, secret_values: Dict[str, str]) -> None

Updates secret values for an existing secret.

Parameters:

Name	Type	Description	Default
secret_id	UUID	The ID of the secret to be updated.	required
secret_values	Dict[str, str]	The new secret values.	required

Source code in src/zenml/zen_stores/secrets_stores/sql_secrets_store.py

def update_secret_values(
    self,
    secret_id: UUID,
    secret_values: Dict[str, str],
) -> None:
    """Updates secret values for an existing secret.

    Args:
        secret_id: The ID of the secret to be updated.
        secret_values: The new secret values.
    """
    self.store_secret_values(secret_id, secret_values)

SqlSecretsStoreConfiguration

Bases: SecretsStoreConfiguration

SQL secrets store configuration.

Attributes:

Name	Type	Description
type	SecretsStoreType	The type of the store.
encryption_key	Optional[PlainSerializedSecretStr]	The encryption key to use for the SQL secrets store. If not set, the passwords will not be encrypted in the database.

Functions

sql_run_statistics

SQL aggregation query for pipeline run statistics.

Classes

Functions

compute_run_statistics(session: Session, request: RunStatisticsRequest, driver: Optional[SQLDatabaseDriver]) -> RunStatisticsResponse

Aggregate pipeline run statistics for the request.

Parameters:

Name	Type	Description	Default
session	Session	Open SQL session, scoped to the request's project.	required
request	RunStatisticsRequest	Statistics request.	required
driver	Optional[SQLDatabaseDriver]	SQL dialect.	required

Returns:

Type	Description
RunStatisticsResponse	Grouped statistics.

Source code in src/zenml/zen_stores/sql_run_statistics.py

def compute_run_statistics(
    session: Session,
    request: RunStatisticsRequest,
    driver: Optional[SQLDatabaseDriver],
) -> RunStatisticsResponse:
    """Aggregate pipeline run statistics for the request.

    Args:
        session: Open SQL session, scoped to the request's project.
        request: Statistics request.
        driver: SQL dialect.

    Returns:
        Grouped statistics.
    """
    runs_subquery = _filtered_runs_subquery(request=request, driver=driver)
    query, decoders, reverse_for_display = _aggregation_query(
        request=request, runs_subquery=runs_subquery, driver=driver
    )

    rows = session.execute(query).all()
    truncated = len(rows) > request.max_groups
    if truncated:
        rows = rows[: request.max_groups]

    if reverse_for_display:
        rows = list(reversed(rows))

    return RunStatisticsResponse(
        groups=[
            _row_to_group(row=row, request=request, decoders=decoders)
            for row in rows
        ],
        truncated=truncated,
    )

sql_zen_store

SQL Zen Store implementation.

Classes

SQLDatabaseDriver

Bases: StrEnum

SQL database drivers supported by the SQL ZenML store.

Session

Bases: Session

Session subclass that automatically tracks duration and calling context.

SqlZenStore(skip_default_registrations: bool = False, **kwargs: Any)

Bases: BaseZenStore

Store Implementation that uses SQL database backend.

Attributes:

Name	Type	Description
config	SqlZenStoreConfiguration	The configuration of the SQL ZenML store.
skip_migrations	bool	Whether to skip migrations when initializing the store.
TYPE	StoreType	The type of the store.
CONFIG_TYPE	Type[StoreConfiguration]	The type of the store configuration.
_engine	Optional[Engine]	The SQLAlchemy engine.

Source code in src/zenml/zen_stores/base_zen_store.py

def __init__(
    self,
    skip_default_registrations: bool = False,
    **kwargs: Any,
) -> None:
    """Create and initialize a store.

    Args:
        skip_default_registrations: If `True`, the creation of the default
            stack and user in the store will be skipped.
        **kwargs: Additional keyword arguments to pass to the Pydantic
            constructor.
    """
    super().__init__(**kwargs)

    self._initialize()

    if not skip_default_registrations:
        logger.debug("Initializing database")
        self._initialize_database()
    else:
        logger.debug("Skipping database initialization")

Attributes

alembic: Alembic property

The Alembic wrapper.

Returns:

Type	Description
Alembic	The Alembic wrapper.

Raises:

Type	Description
ValueError	If the store is not initialized.

backup_secrets_store: Optional[BaseSecretsStore] property

The backup secrets store associated with this store.

Returns:

Type	Description
Optional[BaseSecretsStore]	The backup secrets store associated with this store.

db_backup_engine: BaseDatabaseBackupEngine property

The database backup engine.

Returns:

Type	Description
BaseDatabaseBackupEngine	The database backup engine.

Raises:

Type	Description
ValueError	If the database backup engine is not initialized.

engine: Engine property

The SQLAlchemy engine.

Returns:

Type	Description
Engine	The SQLAlchemy engine.

Raises:

Type	Description
ValueError	If the store is not initialized.

resource_pools: ResourcePoolsSQLStoreInterface property writable

The resource pools store associated with this store.

Returns:

Type	Description
ResourcePoolsSQLStoreInterface	The resource pools store associated with this store.

Raises:

Type	Description
NotImplementedError	If the resource pool functionality is not enabled

resource_pools_enabled: bool property

Whether the resource pools functionality is enabled.

Returns:

Type	Description
bool	True if the resource pools functionality is enabled, False otherwise.

secrets_store: BaseSecretsStore property

The secrets store associated with this store.

Returns:

Type	Description
BaseSecretsStore	The secrets store associated with this store.

Raises:

Type	Description
SecretsStoreNotConfiguredError	If no secrets store is configured.

Functions

activate_server(request: ServerActivationRequest) -> Optional[UserResponse]

Activate the server and optionally create the default admin user.

Parameters:

Name	Type	Description	Default
request	ServerActivationRequest	The server activation request.	required

Returns:

Type	Description
Optional[UserResponse]	The default admin user that was created, if any.

Raises:

Type	Description
IllegalOperationError	If the server is already active.

Source code in src/zenml/zen_stores/sql_zen_store.py

def activate_server(
    self, request: ServerActivationRequest
) -> Optional[UserResponse]:
    """Activate the server and optionally create the default admin user.

    Args:
        request: The server activation request.

    Returns:
        The default admin user that was created, if any.

    Raises:
        IllegalOperationError: If the server is already active.
    """
    with Session(self.engine) as session:
        settings = self._get_server_settings(session=session)

        if settings.active:
            # The server can only be activated once
            raise IllegalOperationError("The server is already active.")

        settings.update(request)
        settings.active = True
        session.add(settings)
        session.commit()

    # Update the server settings to reflect the activation
    self.update_server_settings(request)

    if request.admin_username and request.admin_password is not None:
        # Create the default admin user
        return self.create_user(
            UserRequest(
                name=request.admin_username,
                active=True,
                password=request.admin_password,
                is_admin=True,
            )
        )

    return None

attach_trigger_to_snapshot(trigger_id: UUID, snapshot_id: UUID, run_configuration: PipelineRunConfiguration | None = None, allow_replace: bool = False) -> None

Attaches (links) a trigger to a snapshot.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	The ID of the trigger.	required
snapshot_id	UUID	The ID of the snapshot.	required
run_configuration	PipelineRunConfiguration | None	The configuration applied to subsequent runs.	None
allow_replace	bool	Allow replacement if attachment already exists.	False

Raises:

Type	Description
IllegalOperationError	if the trigger is archived.
KeyError	If associated entities do not exist.

Source code in src/zenml/zen_stores/sql_zen_store.py

def attach_trigger_to_snapshot(
    self,
    trigger_id: UUID,
    snapshot_id: UUID,
    run_configuration: PipelineRunConfiguration | None = None,
    allow_replace: bool = False,
) -> None:
    """Attaches (links) a trigger to a snapshot.

    Args:
        trigger_id: The ID of the trigger.
        snapshot_id: The ID of the snapshot.
        run_configuration: The configuration applied to subsequent runs.
        allow_replace: Allow replacement if attachment already exists.

    Raises:
        IllegalOperationError: if the trigger is archived.
        KeyError: If associated entities do not exist.
    """
    with Session(self.engine) as session:
        snapshot = session.get(PipelineSnapshotSchema, snapshot_id)

        if not snapshot:
            raise KeyError(f"Snapshot {snapshot_id} doesn't exist.")

        if not snapshot.is_runnable:
            raise IllegalOperationError(
                f"Can not attach trigger {trigger_id} to non-runnable snapshot {snapshot_id}"
            )

        trigger = session.get(TriggerSchema, trigger_id)

        if not trigger:
            raise KeyError(f"Trigger {trigger_id} doesn't exist.")

        if trigger.is_archived:
            raise IllegalOperationError(
                f"Can not attach snapshot {snapshot_id} to archived trigger {trigger_id}."
            )

        new_assoc = TriggerSnapshotSchema(
            trigger_id=trigger_id,
            snapshot_id=snapshot_id,
        )

        session.add(new_assoc)
        session.commit()

backup_secrets(ignore_errors: bool = True, delete_secrets: bool = False) -> None

Backs up all secrets to the configured backup secrets store.

Parameters:

Name	Type	Description	Default
ignore_errors	bool	Whether to ignore individual errors during the backup process and attempt to backup all secrets.	True
delete_secrets	bool	Whether to delete the secrets that have been successfully backed up from the primary secrets store. Setting this flag effectively moves all secrets from the primary secrets store to the backup secrets store.	False

Raises:

Type	Description
BackupSecretsStoreNotConfiguredError	if no backup secrets store is configured.
Exception	If a secret backup operation fails and ignore_errors is False.

Source code in src/zenml/zen_stores/sql_zen_store.py

def backup_secrets(
    self, ignore_errors: bool = True, delete_secrets: bool = False
) -> None:
    """Backs up all secrets to the configured backup secrets store.

    Args:
        ignore_errors: Whether to ignore individual errors during the backup
            process and attempt to backup all secrets.
        delete_secrets: Whether to delete the secrets that have been
            successfully backed up from the primary secrets store. Setting
            this flag effectively moves all secrets from the primary secrets
            store to the backup secrets store.

    Raises:
        BackupSecretsStoreNotConfiguredError: if no backup secrets store is
            configured.
        Exception: If a secret backup operation fails and
            ignore_errors is False.
    """  # noqa: DOC502
    if not self.backup_secrets_store:
        raise BackupSecretsStoreNotConfiguredError(
            "Unable to backup secrets: No backup secrets store is "
            "configured."
        )

    with Session(self.engine) as session:
        secrets_in_db = session.exec(select(SecretSchema)).all()

    for secret in secrets_in_db:
        try:
            values = self._get_secret_values(
                secret_id=secret.id, use_backup=False
            )
        except Exception:
            logger.exception(
                f"Failed to get secret values for secret with ID "
                f"{secret.id}."
            )
            if ignore_errors:
                continue
            raise

        try:
            self._backup_secret_values(secret_id=secret.id, values=values)
        except Exception:
            logger.exception(
                f"Failed to backup secret with ID {secret.id}. "
            )
            if ignore_errors:
                continue
            raise

        if delete_secrets:
            try:
                self._delete_secret_values(
                    secret_id=secret.id, delete_backup=False
                )
            except Exception:
                logger.exception(
                    f"Failed to delete secret with ID {secret.id} from the "
                    f"primary secrets store after backing it up to the "
                    f"backup secrets store."
                )
                if ignore_errors:
                    continue
                raise

batch_create_artifact_versions(artifact_versions: List[ArtifactVersionRequest]) -> List[ArtifactVersionResponse]

Creates a batch of artifact versions.

Parameters:

Name	Type	Description	Default
artifact_versions	List[ArtifactVersionRequest]	The artifact versions to create.	required

Returns:

Type	Description
List[ArtifactVersionResponse]	The created artifact versions.

Source code in src/zenml/zen_stores/sql_zen_store.py

def batch_create_artifact_versions(
    self, artifact_versions: List[ArtifactVersionRequest]
) -> List[ArtifactVersionResponse]:
    """Creates a batch of artifact versions.

    Args:
        artifact_versions: The artifact versions to create.

    Returns:
        The created artifact versions.
    """
    return [
        self.create_artifact_version(artifact_version)
        for artifact_version in artifact_versions
    ]

batch_create_tag_resource(tag_resources: List[TagResourceRequest]) -> List[TagResourceResponse]

Create a batch of tag resource relationships.

Parameters:

Name	Type	Description	Default
tag_resources	List[TagResourceRequest]	The tag resource relationships to be created.	required

Returns:

Type	Description
List[TagResourceResponse]	The newly created tag resource relationships.

Source code in src/zenml/zen_stores/sql_zen_store.py

def batch_create_tag_resource(
    self, tag_resources: List[TagResourceRequest]
) -> List[TagResourceResponse]:
    """Create a batch of tag resource relationships.

    Args:
        tag_resources: The tag resource relationships to be created.

    Returns:
        The newly created tag resource relationships.
    """
    with Session(self.engine) as session:
        resources: List[
            Tuple[TagSchema, TaggableResourceTypes, BaseSchema]
        ] = []
        for tag_resource in tag_resources:
            resource_schema = self._get_schema_from_resource_type(
                tag_resource.resource_type
            )
            resource = self._get_schema_by_id(
                resource_id=tag_resource.resource_id,
                schema_class=resource_schema,
                session=session,
            )
            tag_schema = self._get_tag_schema(
                tag_name_or_id=tag_resource.tag_id,
                session=session,
            )
            resources.append(
                (
                    tag_schema,
                    tag_resource.resource_type,
                    resource,
                )
            )
        return [
            r.to_model()
            for r in self._create_tag_resource_schemas(
                tag_resources=resources, session=session
            )
        ]

batch_delete_tag_resource(tag_resources: List[TagResourceRequest]) -> None

Delete a batch of tag resource relationships.

Parameters:

Name	Type	Description	Default
tag_resources	List[TagResourceRequest]	The tag resource relationships to be deleted.	required

Source code in src/zenml/zen_stores/sql_zen_store.py

def batch_delete_tag_resource(
    self, tag_resources: List[TagResourceRequest]
) -> None:
    """Delete a batch of tag resource relationships.

    Args:
        tag_resources: The tag resource relationships to be deleted.
    """
    with Session(self.engine) as session:
        self._delete_tag_resource_schemas(
            tag_resources=tag_resources,
            session=session,
        )

cleanup_expired_api_transactions() -> None

Delete completed API transactions that have expired.

Source code in src/zenml/zen_stores/sql_zen_store.py

def cleanup_expired_api_transactions(self) -> None:
    """Delete completed API transactions that have expired."""
    with Session(self.engine) as session:
        session.execute(
            delete(ApiTransactionSchema).where(
                col(ApiTransactionSchema.completed),
                col(ApiTransactionSchema.expired) < utc_now(),
            )
        )

        session.commit()

clear_trigger_dispatch_error(trigger_id: UUID, snapshot_id: UUID | None = None) -> None

Clear dispatch errors for one or all trigger snapshot associations.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	Trigger ID.	required
snapshot_id	UUID | None	Optional display/source snapshot ID.	None

Raises:

Type	Description
KeyError	If the snapshot is not attached to the trigger.

Source code in src/zenml/zen_stores/sql_zen_store.py

def clear_trigger_dispatch_error(
    self,
    trigger_id: UUID,
    snapshot_id: UUID | None = None,
) -> None:
    """Clear dispatch errors for one or all trigger snapshot associations.

    Args:
        trigger_id: Trigger ID.
        snapshot_id: Optional display/source snapshot ID.

    Raises:
        KeyError: If the snapshot is not attached to the trigger.
    """
    with Session(self.engine) as session:
        trigger = self._get_schema_by_id(
            resource_id=trigger_id,
            schema_class=TriggerSchema,
            session=session,
        )

        if snapshot_id is None:
            executable_snapshot_ids = [
                snapshot.id for snapshot in trigger.snapshots
            ]
        else:
            executable_snapshot_ids = [
                snapshot.id
                for snapshot in trigger.snapshots
                if snapshot.source_snapshot_id == snapshot_id
            ]
            if not executable_snapshot_ids:
                raise KeyError(
                    f"Snapshot {snapshot_id} is not attached to trigger "
                    f"{trigger_id}"
                )

        for executable_snapshot_id in executable_snapshot_ids:
            row = self._get_trigger_snapshot_association(
                trigger_id=trigger_id,
                snapshot_id=executable_snapshot_id,
                session=session,
            )
            state = row.parsed_dispatch_state
            if state is None:
                continue

            if state.last_status == TriggerDispatchStatusCode.ERROR:
                row.dispatch_state = None
            else:
                state.clear_error_details()
                row.dispatch_state = state.model_dump_json()
            session.add(row)

        session.commit()

count_pipelines(filter_model: PipelineFilter) -> int

Count all pipelines.

Parameters:

Name	Type	Description	Default
filter_model	PipelineFilter	The filter model to use for counting pipelines.	required

Returns:

Type	Description
int	The number of pipelines.

Source code in src/zenml/zen_stores/sql_zen_store.py

def count_pipelines(self, filter_model: PipelineFilter) -> int:
    """Count all pipelines.

    Args:
        filter_model: The filter model to use for counting pipelines.

    Returns:
        The number of pipelines.
    """
    return self._count_entity(
        schema=PipelineSchema, filter_model=filter_model
    )

count_projects(filter_model: Optional[ProjectFilter] = None) -> int

Count all projects.

Parameters:

Name	Type	Description	Default
filter_model	Optional[ProjectFilter]	The filter model to use for counting projects.	None

Returns:

Type	Description
int	The number of projects.

Source code in src/zenml/zen_stores/sql_zen_store.py

def count_projects(
    self, filter_model: Optional[ProjectFilter] = None
) -> int:
    """Count all projects.

    Args:
        filter_model: The filter model to use for counting projects.

    Returns:
        The number of projects.
    """
    return self._count_entity(
        schema=ProjectSchema, filter_model=filter_model
    )

count_runs(filter_model: PipelineRunFilter) -> int

Count all pipeline runs.

Parameters:

Name	Type	Description	Default
filter_model	PipelineRunFilter	The filter model to filter the runs.	required

Returns:

Type	Description
int	The number of pipeline runs.

Source code in src/zenml/zen_stores/sql_zen_store.py

def count_runs(self, filter_model: PipelineRunFilter) -> int:
    """Count all pipeline runs.

    Args:
        filter_model: The filter model to filter the runs.

    Returns:
        The number of pipeline runs.
    """
    return self._count_entity(
        schema=PipelineRunSchema, filter_model=filter_model
    )

count_stack_components(filter_model: Optional[ComponentFilter] = None) -> int

Count all components.

Parameters:

Name	Type	Description	Default
filter_model	Optional[ComponentFilter]	The filter model to use for counting components.	None

Returns:

Type	Description
int	The number of components.

Source code in src/zenml/zen_stores/sql_zen_store.py

def count_stack_components(
    self, filter_model: Optional[ComponentFilter] = None
) -> int:
    """Count all components.

    Args:
        filter_model: The filter model to use for counting components.

    Returns:
        The number of components.
    """
    return self._count_entity(
        schema=StackComponentSchema, filter_model=filter_model
    )

count_stacks(filter_model: Optional[StackFilter]) -> int

Count all stacks.

Parameters:

Name	Type	Description	Default
filter_model	Optional[StackFilter]	The filter model to filter the stacks.	required

Returns:

Type	Description
int	The number of stacks.

Source code in src/zenml/zen_stores/sql_zen_store.py

def count_stacks(self, filter_model: Optional[StackFilter]) -> int:
    """Count all stacks.

    Args:
        filter_model: The filter model to filter the stacks.

    Returns:
        The number of stacks.
    """
    return self._count_entity(
        schema=StackSchema, filter_model=filter_model
    )

create_api_key(service_account_id: UUID, api_key: APIKeyRequest) -> APIKeyResponse

Create a new API key for a service account.

Parameters:

Name	Type	Description	Default
service_account_id	UUID	The ID of the service account for which to create the API key.	required
api_key	APIKeyRequest	The API key to create.	required

Returns:

Type	Description
APIKeyResponse	The created API key.

Raises:

Type	Description
EntityExistsError	If an API key with the same name is already configured for the same service account.

Source code in src/zenml/zen_stores/sql_zen_store.py

def create_api_key(
    self, service_account_id: UUID, api_key: APIKeyRequest
) -> APIKeyResponse:
    """Create a new API key for a service account.

    Args:
        service_account_id: The ID of the service account for which to
            create the API key.
        api_key: The API key to create.

    Returns:
        The created API key.

    Raises:
        EntityExistsError: If an API key with the same name is already
            configured for the same service account.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=api_key, session=session)

        # Fetch the service account
        service_account = self._get_account_schema(
            service_account_id, session=session, service_account=True
        )

        # Check if a key with the same name already exists for the same
        # service account
        try:
            self._get_api_key(
                service_account_id=service_account.id,
                api_key_name_or_id=api_key.name,
                session=session,
            )
            raise EntityExistsError(
                f"Unable to register API key with name '{api_key.name}': "
                "Found an existing API key with the same name configured "
                f"for the same '{service_account.name}' service account."
            )
        except KeyError:
            pass

        new_api_key, key_value = APIKeySchema.from_request(
            service_account_id=service_account.id,
            request=api_key,
        )
        session.add(new_api_key)
        session.commit()

        api_key_model = new_api_key.to_model(
            include_metadata=True, include_resources=True
        )
        api_key_model.set_key(key_value)
        return api_key_model

create_artifact(artifact: ArtifactRequest) -> ArtifactResponse

Creates a new artifact.

Parameters:

Name	Type	Description	Default
artifact	ArtifactRequest	The artifact to create.	required

Returns:

Type	Description
ArtifactResponse	The newly created artifact.

Source code in src/zenml/zen_stores/sql_zen_store.py

def create_artifact(self, artifact: ArtifactRequest) -> ArtifactResponse:
    """Creates a new artifact.

    Args:
        artifact: The artifact to create.

    Returns:
        The newly created artifact.
    """
    validate_name(artifact)
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=artifact, session=session)

        # Check if an artifact with the given name already exists
        self._verify_name_uniqueness(
            resource=artifact,
            schema=ArtifactSchema,
            session=session,
        )

        # Create the artifact.
        artifact_schema = ArtifactSchema.from_request(artifact)

        session.add(artifact_schema)
        session.commit()

        # Save tags of the artifact.
        self._attach_tags_to_resources(
            tags=artifact.tags,
            resources=artifact_schema,
            session=session,
        )
        session.refresh(artifact_schema)

        return artifact_schema.to_model(
            include_metadata=True, include_resources=True
        )

create_artifact_version(artifact_version: ArtifactVersionRequest) -> ArtifactVersionResponse

Create an artifact version.

Parameters:

Name	Type	Description	Default
artifact_version	ArtifactVersionRequest	The artifact version to create.	required

Raises:

Type	Description
EntityExistsError	If an artifact version with the same name already exists.
EntityCreationError	If the artifact version creation failed.

Returns:

Type	Description
ArtifactVersionResponse	The created artifact version.

Source code in src/zenml/zen_stores/sql_zen_store.py

def create_artifact_version(
    self, artifact_version: ArtifactVersionRequest
) -> ArtifactVersionResponse:
    """Create an artifact version.

    Args:
        artifact_version: The artifact version to create.

    Raises:
        EntityExistsError: If an artifact version with the same name
            already exists.
        EntityCreationError: If the artifact version creation failed.

    Returns:
        The created artifact version.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=artifact_version, session=session
        )

        self._get_reference_schema_by_id(
            resource=artifact_version,
            reference_schema=StackComponentSchema,
            reference_id=artifact_version.artifact_store_id,
            session=session,
            reference_type="artifact store",
        )

        if artifact_name := artifact_version.artifact_name:
            artifact_schema = self._get_or_create_artifact_for_name(
                name=artifact_name,
                project_id=artifact_version.project,
                has_custom_name=artifact_version.has_custom_name,
                session=session,
            )
            artifact_version.artifact_id = artifact_schema.id

        assert artifact_version.artifact_id

        artifact_version_schema: Optional[ArtifactVersionSchema] = None

        if artifact_version.version is None:
            # No explicit version in the request -> We will try to
            # auto-increment the numeric version of the artifact version
            remaining_tries = MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION
            while remaining_tries > 0:
                remaining_tries -= 1
                try:
                    artifact_version.version = str(
                        self._get_next_numeric_version_for_artifact(
                            session=session,
                            artifact_id=artifact_version.artifact_id,
                        )
                    )

                    artifact_version_schema = (
                        ArtifactVersionSchema.from_request(
                            artifact_version
                        )
                    )
                    session.add(artifact_version_schema)
                    session.commit()
                except IntegrityError:
                    # We have to rollback the failed session first in order
                    # to continue using it
                    session.rollback()
                    if remaining_tries == 0:
                        raise EntityCreationError(
                            f"Failed to create version for artifact "
                            f"{artifact_schema.name}. This is most likely "
                            "caused by multiple parallel requests that try "
                            "to create versions for this artifact in the "
                            "database."
                        )
                    else:
                        attempt = (
                            MAX_RETRIES_FOR_VERSIONED_ENTITY_CREATION
                            - remaining_tries
                        )
                        sleep_duration = exponential_backoff_with_jitter(
                            attempt=attempt
                        )

                        logger.debug(
                            "Failed to create artifact version %s "
                            "(version %s) due to an integrity error. "
                            "Retrying in %f seconds.",
                            artifact_schema.name,
                            artifact_version.version,
                            sleep_duration,
                        )
                        time.sleep(sleep_duration)
                else:
                    break
        else:
            # An explicit version was specified for the artifact version.
            # We don't do any incrementing and fail immediately if the
            # version already exists.
            try:
                artifact_version_schema = (
                    ArtifactVersionSchema.from_request(artifact_version)
                )
                session.add(artifact_version_schema)
                session.commit()
            except IntegrityError:
                # We have to rollback the failed session first in order
                # to continue using it
                session.rollback()
                raise EntityExistsError(
                    f"Unable to create artifact version "
                    f"{artifact_schema.name} (version "
                    f"{artifact_version.version}): An artifact with the "
                    "same name and version already exists."
                )

        assert artifact_version_schema is not None

        # Save visualizations of the artifact
        if artifact_version.visualizations:
            for vis in artifact_version.visualizations:
                vis_schema = ArtifactVisualizationSchema.from_model(
                    artifact_visualization_request=vis,
                    artifact_version_id=artifact_version_schema.id,
                )
                session.add(vis_schema)
            session.commit()

        # Save tags of the artifact
        self._attach_tags_to_resources(
            tags=artifact_version.tags,
            resources=artifact_version_schema,
            session=session,
        )

        # Save metadata of the artifact
        if artifact_version.metadata:
            values: Dict[str, "MetadataType"] = {}
            types: Dict[str, "MetadataTypeEnum"] = {}
            for key, value in artifact_version.metadata.items():
                # Skip metadata that is too large to be stored in the DB.
                if len(json.dumps(value)) > TEXT_FIELD_MAX_LENGTH:
                    logger.warning(
                        f"Metadata value for key '{key}' is too large to be "
                        "stored in the database. Skipping."
                    )
                    continue
                # Skip metadata that is not of a supported type.
                try:
                    metadata_type = get_metadata_type(value)
                except ValueError as e:
                    logger.warning(
                        f"Metadata value for key '{key}' is not of a "
                        f"supported type. Skipping. Full error: {e}"
                    )
                    continue
                values[key] = value
                types[key] = metadata_type
            self.create_run_metadata(
                RunMetadataRequest(
                    project=artifact_version.project,
                    resources=[
                        RunMetadataResource(
                            id=artifact_version_schema.id,
                            type=MetadataResourceTypes.ARTIFACT_VERSION,
                        )
                    ],
                    values=values,
                    types=types,
                )
            )

        session.commit()
        session.refresh(artifact_version_schema)

        return artifact_version_schema.to_model(
            include_metadata=True, include_resources=True
        )

create_authorized_device(device: OAuthDeviceInternalRequest) -> OAuthDeviceInternalResponse

Creates a new OAuth 2.0 authorized device.

Parameters:

Name	Type	Description	Default
device	OAuthDeviceInternalRequest	The device to be created.	required

Returns:

Type	Description
OAuthDeviceInternalResponse	The newly created device.

Raises:

Type	Description
EntityExistsError	If a device for the same client ID already exists.

Source code in src/zenml/zen_stores/sql_zen_store.py

def create_authorized_device(
    self, device: OAuthDeviceInternalRequest
) -> OAuthDeviceInternalResponse:
    """Creates a new OAuth 2.0 authorized device.

    Args:
        device: The device to be created.

    Returns:
        The newly created device.

    Raises:
        EntityExistsError: If a device for the same client ID already
            exists.
    """
    with Session(self.engine) as session:
        existing_device = session.exec(
            select(OAuthDeviceSchema).where(
                # We search for a device with the same client ID
                # because the client ID is the one that is used to
                # identify the device
                OAuthDeviceSchema.client_id == device.client_id
            )
        ).first()
        if existing_device is not None:
            raise EntityExistsError(
                f"Unable to create device with client ID "
                f"'{device.client_id}': A device with this client ID "
                "already exists."
            )

        (
            new_device,
            user_code,
            device_code,
        ) = OAuthDeviceSchema.from_request(device)
        session.add(new_device)
        session.commit()
        session.refresh(new_device)

        device_model = new_device.to_internal_model(
            include_metadata=True, include_resources=True
        )
        # Replace the hashed user code with the original user code
        device_model.user_code = user_code
        # Replace the hashed device code with the original device code
        device_model.device_code = device_code

        return device_model

create_build(build: PipelineBuildRequest) -> PipelineBuildResponse

Creates a new build.

Parameters:

Name	Type	Description	Default
build	PipelineBuildRequest	The build to create.	required

Returns:

Type	Description
PipelineBuildResponse	The newly created build.

Source code in src/zenml/zen_stores/sql_zen_store.py

def create_build(
    self,
    build: PipelineBuildRequest,
) -> PipelineBuildResponse:
    """Creates a new build.

    Args:
        build: The build to create.

    Returns:
        The newly created build.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=build, session=session)
        self._get_reference_schema_by_id(
            resource=build,
            reference_schema=StackSchema,
            reference_id=build.stack,
            session=session,
        )

        self._get_reference_schema_by_id(
            resource=build,
            reference_schema=PipelineSchema,
            reference_id=build.pipeline,
            session=session,
        )

        new_build = PipelineBuildSchema.from_request(build)
        session.add(new_build)
        session.commit()
        session.refresh(new_build)

        return new_build.to_model(
            include_metadata=True, include_resources=True
        )

create_code_repository(code_repository: CodeRepositoryRequest) -> CodeRepositoryResponse

Creates a new code repository.

Parameters:

Name	Type	Description	Default
code_repository	CodeRepositoryRequest	Code repository to be created.	required

Returns:

Type	Description
CodeRepositoryResponse	The newly created code repository.

Source code in src/zenml/zen_stores/sql_zen_store.py

@track_decorator(AnalyticsEvent.REGISTERED_CODE_REPOSITORY)
def create_code_repository(
    self, code_repository: CodeRepositoryRequest
) -> CodeRepositoryResponse:
    """Creates a new code repository.

    Args:
        code_repository: Code repository to be created.

    Returns:
        The newly created code repository.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=code_repository, session=session
        )

        self._verify_name_uniqueness(
            resource=code_repository,
            schema=CodeRepositorySchema,
            session=session,
        )

        new_repo = CodeRepositorySchema.from_request(code_repository)
        session.add(new_repo)
        session.commit()
        session.refresh(new_repo)

        return new_repo.to_model(
            include_metadata=True, include_resources=True
        )

create_curated_visualization(visualization: CuratedVisualizationRequest) -> CuratedVisualizationResponse

Persist a curated visualization link.

Parameters:

Name	Type	Description	Default
visualization	CuratedVisualizationRequest	The curated visualization to create.	required

Returns:

Type	Description
CuratedVisualizationResponse	The created curated visualization.

Raises:

Type	Description
IllegalOperationError	If the curated visualization does not target the same project as the artifact visualization.
ValueError	If the resource type is invalid.
KeyError	If the resource is not found.

Source code in src/zenml/zen_stores/sql_zen_store.py

def create_curated_visualization(
    self, visualization: CuratedVisualizationRequest
) -> CuratedVisualizationResponse:
    """Persist a curated visualization link.

    Args:
        visualization: The curated visualization to create.

    Returns:
        The created curated visualization.

    Raises:
        IllegalOperationError: If the curated visualization does not target the same project as the artifact visualization.
        ValueError: If the resource type is invalid.
        KeyError: If the resource is not found.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=visualization, session=session
        )

        artifact_visualization: ArtifactVisualizationSchema = (
            self._get_reference_schema_by_id(
                resource=visualization,
                reference_schema=ArtifactVisualizationSchema,
                reference_id=visualization.artifact_visualization_id,
                session=session,
            )
        )

        artifact_version = artifact_visualization.artifact_version
        project_id = artifact_version.project_id

        if visualization.project != project_id:
            raise IllegalOperationError(
                "Curated visualizations must target the same project as "
                "the artifact visualization."
            )
        project_id = visualization.project

        resource_schema_map: Dict[
            VisualizationResourceTypes, Type[BaseSchema]
        ] = {
            VisualizationResourceTypes.DEPLOYMENT: DeploymentSchema,
            VisualizationResourceTypes.MODEL: ModelSchema,
            VisualizationResourceTypes.PIPELINE: PipelineSchema,
            VisualizationResourceTypes.PIPELINE_RUN: PipelineRunSchema,
            VisualizationResourceTypes.PIPELINE_SNAPSHOT: PipelineSnapshotSchema,
            VisualizationResourceTypes.PROJECT: ProjectSchema,
        }

        if visualization.resource_type not in resource_schema_map:
            raise ValueError(
                f"Invalid resource type: {visualization.resource_type}"
            )

        schema_class = resource_schema_map[visualization.resource_type]
        resource_schema