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.triggered_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_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

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.triggered_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_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}"

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

utils

ZenML database migration, backup and recovery utilities.

Classes

MigrationUtils

Bases: BaseModel

Utilities for database migration, backup and recovery.

Attributes

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_to_db(backup_db_name: str) -> None

Backup the database to a backup database.

Parameters:

Name	Type	Description	Default
backup_db_name	str	Backup database name to backup to.	required

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

def backup_database_to_db(self, backup_db_name: str) -> None:
    """Backup the database to a backup database.

    Args:
        backup_db_name: Backup database name to backup to.
    """
    # Re-create the backup database
    self.create_database(
        database=backup_db_name,
        drop=True,
    )

    backup_engine = self.create_engine(database=backup_db_name)

    self._copy_database(self.engine, backup_engine)

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

backup_database_to_file(dump_file: str) -> None

Backup the database to a file.

This method dumps the entire database into a JSON file. Instead of using a SQL dump, we use a proprietary JSON dump because:

* it is (mostly) not dependent on the SQL dialect or database version
* it is safer with respect to SQL injection attacks
* it is easier to read and debug

The JSON file contains a list of JSON objects instead of a single JSON object, because it allows for buffered reading and writing of the file and thus reduces the memory footprint. Each JSON object can contain either schema or data information about a single table. For tables with a large amount of data, the data is split into multiple JSON objects with the first object always containing the schema.

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
dump_file	str	The path to the dump file.	required

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

def backup_database_to_file(self, dump_file: str) -> None:
    """Backup the database to a file.

    This method dumps the entire database into a JSON file. Instead of
    using a SQL dump, we use a proprietary JSON dump because:

        * it is (mostly) not dependent on the SQL dialect or database version
        * it is safer with respect to SQL injection attacks
        * it is easier to read and debug

    The JSON file contains a list of JSON objects instead of a single JSON
    object, because it allows for buffered reading and writing of the file
    and thus reduces the memory footprint. Each JSON object can contain
    either schema or data information about a single table. For tables with
    a large amount of data, the data is split into multiple JSON objects
    with the first object always containing the schema.

    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:
        dump_file: The path to the dump file.
    """
    # create the directory if it does not exist
    dump_path = os.path.dirname(os.path.abspath(dump_file))
    if not os.path.exists(dump_path):
        os.makedirs(dump_path)

    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,
            dump_file,
        )
        return

    with open(dump_file, "w") as f:

        def json_dump(obj: Dict[str, Any]) -> None:
            """Dump a JSON object to the dump file.

            Args:
                obj: The JSON object to dump.
            """
            # Write the data to the JSON file. Use an encoder that
            # can handle datetime, Decimal and other types.
            json.dump(
                obj,
                f,
                indent=4,
                default=pydantic_encoder,
            )
            f.write("\n")

        # Call the generic backup method with a function that dumps the
        # JSON objects to the dump file
        self.backup_database_to_storage(json_dump)

    logger.debug(f"Database backed up to {dump_file}")

backup_database_to_memory() -> List[Dict[str, Any]]

Backup the database in memory.

Returns:

Type	Description
List[Dict[str, Any]]	The in-memory representation of the database backup.

Raises:

Type	Description
RuntimeError	If the database cannot be backed up successfully.

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

def backup_database_to_memory(self) -> List[Dict[str, Any]]:
    """Backup the database in memory.

    Returns:
        The in-memory representation of the database backup.

    Raises:
        RuntimeError: If the database cannot be backed up successfully.
    """
    if self.url.drivername == "sqlite":
        # For a sqlite database, this is not supported.
        raise RuntimeError(
            "In-memory backup is not supported for sqlite databases."
        )

    db_dump: List[Dict[str, Any]] = []

    def store_in_mem(obj: Dict[str, Any]) -> None:
        """Store a JSON object in the in-memory database backup.

        Args:
            obj: The JSON object to store.
        """
        db_dump.append(obj)

    # Call the generic backup method with a function that stores the
    # JSON objects in the in-memory database backup
    self.backup_database_to_storage(store_in_mem)

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

    return db_dump

backup_database_to_storage(store_db_info: Callable[[Dict[str, Any]], None]) -> None

Backup the database to a storage location.

Backup the database to an abstract storage location. The storage location is specified by a function that is called repeatedly to store the database information. The function 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.

Parameters:

Name	Type	Description	Default
store_db_info	Callable[[Dict[str, Any]], None]	The function to call to store the database information.	required

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

def backup_database_to_storage(
    self, store_db_info: Callable[[Dict[str, Any]], None]
) -> None:
    """Backup the database to a storage location.

    Backup the database to an abstract storage location. The storage
    location is specified by a function that is called repeatedly to
    store the database information. The function 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.

    Args:
        store_db_info: The function to call to store the database
            information.
    """
    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
            store_db_info(
                dict(
                    table=table.name,
                    create_stmt=create_table_stmt,
                    self_references=has_self_referential_foreign_keys,
                )
            )

            for stmt in index_create_statements:
                store_db_info(
                    dict(
                        table=table.name,
                        index_create_stmt=stmt,
                    )
                )

            # 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()

                    store_db_info(
                        dict(
                            table=table.name,
                            data=[row._asdict() for row in rows],
                        ),
                    )

create_database(database: Optional[str] = None, drop: bool = False) -> None

Creates a mysql database.

Parameters:

Name	Type	Description	Default
database	Optional[str]	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/utils.py

def create_database(
    self,
    database: Optional[str] = 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: Optional[str] = None) -> Engine

Get the SQLAlchemy engine for a database.

Parameters:

Name	Type	Description	Default
database	Optional[str]	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/utils.py

def create_engine(self, database: Optional[str] = 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: Optional[str] = None) -> bool

Check if a database exists.

Parameters:

Name	Type	Description	Default
database	Optional[str]	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/utils.py

def database_exists(
    self,
    database: Optional[str] = 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: Optional[str] = None) -> None

Drops a mysql database.

Parameters:

Name	Type	Description	Default
database	Optional[str]	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/utils.py

def drop_database(
    self,
    database: Optional[str] = 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/utils.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_from_db(backup_db_name: str) -> None

Restore the database from the backup database.

Parameters:

Name	Type	Description	Default
backup_db_name	str	Backup database name to restore from.	required

Raises:

Type	Description
RuntimeError	If the backup database does not exist.

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

def restore_database_from_db(self, backup_db_name: str) -> None:
    """Restore the database from the backup database.

    Args:
        backup_db_name: Backup database name to restore from.

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

    backup_engine = self.create_engine(database=backup_db_name)

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

    self._copy_database(backup_engine, self.engine)

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

restore_database_from_file(dump_file: str) -> None

Restore the database from a backup dump file.

See the documentation of the backup_database_to_file method for details on the format of the dump file.

Parameters:

Name	Type	Description	Default
dump_file	str	The path to the dump file.	required

Raises:

Type	Description
RuntimeError	If the database cannot be restored successfully.

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

def restore_database_from_file(self, dump_file: str) -> None:
    """Restore the database from a backup dump file.

    See the documentation of the `backup_database_to_file` method for
    details on the format of the dump file.

    Args:
        dump_file: The path to the dump file.

    Raises:
        RuntimeError: If the database cannot be restored successfully.
    """
    if not os.path.exists(dump_file):
        raise RuntimeError(
            f"Database backup file '{dump_file}' does not "
            f"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(
            dump_file,
            self.url.database,
        )
        return

    # read the DB dump file one JSON object at a time
    with open(dump_file, "r") as f:

        def json_load() -> Generator[Dict[str, Any], None, None]:
            """Generator that loads the JSON objects in the dump file.

            Yields:
                The loaded JSON objects.
            """
            buffer = ""
            while True:
                chunk = f.readline()
                if not chunk:
                    break
                buffer += chunk
                if chunk.rstrip() == "}":
                    yield json.loads(buffer)
                    buffer = ""

        # Call the generic restore method with a function that loads the
        # JSON objects from the dump file
        self.restore_database_from_storage(json_load)

    logger.info(f"Database successfully restored from '{dump_file}'")

restore_database_from_memory(db_dump: List[Dict[str, Any]]) -> None

Restore the database from an in-memory backup.

Parameters:

Name	Type	Description	Default
db_dump	List[Dict[str, Any]]	The in-memory database backup to restore from generated by the backup_database_to_memory method.	required

Raises:

Type	Description
RuntimeError	If the database cannot be restored successfully.

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

def restore_database_from_memory(
    self, db_dump: List[Dict[str, Any]]
) -> None:
    """Restore the database from an in-memory backup.

    Args:
        db_dump: The in-memory database backup to restore from generated
            by the `backup_database_to_memory` method.

    Raises:
        RuntimeError: If the database cannot be restored successfully.
    """
    if self.url.drivername == "sqlite":
        # For a sqlite database, this is not supported.
        raise RuntimeError(
            "In-memory backup is not supported for sqlite databases."
        )

    def load_from_mem() -> Generator[Dict[str, Any], None, None]:
        """Generator that loads the JSON objects from the in-memory backup.

        Yields:
            The loaded JSON objects.
        """
        for obj in db_dump:
            yield obj

    # Call the generic restore method with a function that loads the
    # JSON objects from the in-memory database backup
    self.restore_database_from_storage(load_from_mem)

    logger.info("Database successfully restored from memory")

restore_database_from_storage(load_db_info: Callable[[], Generator[Dict[str, Any], None, None]]) -> None

Restore the database from a backup storage location.

Restores the database from an abstract storage location. The storage location is specified by a function that is called repeatedly to load the database information from the external storage chunk by chunk. The function must yield 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 function must return None when there is no more data to load.

Parameters:

Name	Type	Description	Default
load_db_info	Callable[[], Generator[Dict[str, Any], None, None]]	The function to call to load the database information.	required

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

def restore_database_from_storage(
    self, load_db_info: Callable[[], Generator[Dict[str, Any], None, None]]
) -> None:
    """Restore the database from a backup storage location.

    Restores the database from an abstract storage location. The storage
    location is specified by a function that is called repeatedly to
    load the database information from the external storage chunk by chunk.
    The function must yield 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 function must return `None` when there is no more data to load.

    Args:
        load_db_info: The function to call to load the database
            information.
    """
    # 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 load_db_info():
            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"))

Functions

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

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,
    )

create_action(action: ActionRequest) -> ActionResponse

Create an action.

Parameters:

Name	Type	Description	Default
action	ActionRequest	The action to create.	required

Returns:

Type	Description
ActionResponse	The created action.

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

def create_action(self, action: ActionRequest) -> ActionResponse:
    """Create an action.

    Args:
        action: The action to create.

    Returns:
        The created action.
    """
    return self._create_resource(
        resource=action,
        route=ACTIONS,
        response_model=ActionResponse,
    )

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_event_source(event_source: EventSourceRequest) -> EventSourceResponse

Create an event_source.

Parameters:

Name	Type	Description	Default
event_source	EventSourceRequest	The event_source to create.	required

Returns:

Type	Description
EventSourceResponse	The created event_source.

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

def create_event_source(
    self, event_source: EventSourceRequest
) -> EventSourceResponse:
    """Create an event_source.

    Args:
        event_source: The event_source to create.

    Returns:
        The created event_source.
    """
    return self._create_resource(
        resource=event_source,
        route=EVENT_SOURCES,
        response_model=EventSourceResponse,
    )

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_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_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_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) -> TriggerResponse

Create an trigger.

Parameters:

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

Returns:

Type	Description
TriggerResponse	The created trigger.

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

def create_trigger(self, trigger: TriggerRequest) -> TriggerResponse:
    """Create an trigger.

    Args:
        trigger: The trigger to create.

    Returns:
        The created trigger.
    """
    return self._create_resource(
        resource=trigger,
        route=TRIGGERS,
        response_model=TriggerResponse,
    )

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_action(action_id: UUID) -> None

Delete an action.

Parameters:

Name	Type	Description	Default
action_id	UUID	The ID of the action to delete.	required

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

def delete_action(self, action_id: UUID) -> None:
    """Delete an action.

    Args:
        action_id: The ID of the action to delete.
    """
    self._delete_resource(
        resource_id=action_id,
        route=ACTIONS,
    )

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_event_source(event_source_id: UUID) -> None

Delete an event_source.

Parameters:

Name	Type	Description	Default
event_source_id	UUID	The ID of the event_source to delete.	required

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

def delete_event_source(self, event_source_id: UUID) -> None:
    """Delete an event_source.

    Args:
        event_source_id: The ID of the event_source to delete.
    """
    self._delete_resource(
        resource_id=event_source_id,
        route=EVENT_SOURCES,
    )

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_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) -> None

Deletes a schedule.

Parameters:

Name	Type	Description	Default
schedule_id	UUID	The ID of the schedule to delete.	required

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

def delete_schedule(self, schedule_id: UUID) -> None:
    """Deletes a schedule.

    Args:
        schedule_id: The ID of the schedule to delete.
    """
    self._delete_resource(
        resource_id=schedule_id,
        route=SCHEDULES,
    )

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) -> None

Delete an trigger.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	The ID of the trigger to delete.	required

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

def delete_trigger(self, trigger_id: UUID) -> None:
    """Delete an trigger.

    Args:
        trigger_id: The ID of the trigger to delete.
    """
    self._delete_resource(
        resource_id=trigger_id,
        route=TRIGGERS,
    )

delete_trigger_execution(trigger_execution_id: UUID) -> None

Delete a trigger execution.

Parameters:

Name	Type	Description	Default
trigger_execution_id	UUID	The ID of the trigger execution to delete.	required

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

def delete_trigger_execution(self, trigger_execution_id: UUID) -> None:
    """Delete a trigger execution.

    Args:
        trigger_execution_id: The ID of the trigger execution to delete.
    """
    self._delete_resource(
        resource_id=trigger_execution_id,
        route=TRIGGER_EXECUTIONS,
    )

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,
    )

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_action(action_id: UUID, hydrate: bool = True) -> ActionResponse

Get an action by ID.

Parameters:

Name	Type	Description	Default
action_id	UUID	The ID of the action 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
ActionResponse	The action.

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

def get_action(
    self,
    action_id: UUID,
    hydrate: bool = True,
) -> ActionResponse:
    """Get an action by ID.

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

    Returns:
        The action.
    """
    return self._get_resource(
        resource_id=action_id,
        route=ACTIONS,
        response_model=ActionResponse,
        params={"hydrate": hydrate},
    )

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_event_source(event_source_id: UUID, hydrate: bool = True) -> EventSourceResponse

Get an event_source by ID.

Parameters:

Name	Type	Description	Default
event_source_id	UUID	The ID of the event_source 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
EventSourceResponse	The event_source.

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

def get_event_source(
    self,
    event_source_id: UUID,
    hydrate: bool = True,
) -> EventSourceResponse:
    """Get an event_source by ID.

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

    Returns:
        The event_source.
    """
    return self._get_resource(
        resource_id=event_source_id,
        route=EVENT_SOURCES,
        response_model=EventSourceResponse,
        params={"hydrate": hydrate},
    )

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"--url {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_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_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_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) -> TriggerResponse

Get a trigger by ID.

Parameters:

Name	Type	Description	Default
trigger_id	UUID	The ID of the trigger 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
TriggerResponse	The trigger.

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

def get_trigger(
    self,
    trigger_id: UUID,
    hydrate: bool = True,
) -> TriggerResponse:
    """Get a trigger by ID.

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

    Returns:
        The trigger.
    """
    return self._get_resource(
        resource_id=trigger_id,
        route=TRIGGERS,
        response_model=TriggerResponse,
        params={"hydrate": hydrate},
    )

get_trigger_execution(trigger_execution_id: UUID, hydrate: bool = True) -> TriggerExecutionResponse

Get an trigger execution by ID.

Parameters:

Name	Type	Description	Default
trigger_execution_id	UUID	The ID of the trigger execution 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
TriggerExecutionResponse	The trigger execution.

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

def get_trigger_execution(
    self,
    trigger_execution_id: UUID,
    hydrate: bool = True,
) -> TriggerExecutionResponse:
    """Get an trigger execution by ID.

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

    Returns:
        The trigger execution.
    """
    return self._get_resource(
        resource_id=trigger_execution_id,
        route=TRIGGER_EXECUTIONS,
        response_model=TriggerExecutionResponse,
        params={"hydrate": hydrate},
    )

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_actions(action_filter_model: ActionFilter, hydrate: bool = False) -> Page[ActionResponse]

List all actions matching the given filter criteria.

Parameters:

Name	Type	Description	Default
action_filter_model	ActionFilter	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[ActionResponse]	A list of all actions matching the filter criteria.

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

def list_actions(
    self,
    action_filter_model: ActionFilter,
    hydrate: bool = False,
) -> Page[ActionResponse]:
    """List all actions matching the given filter criteria.

    Args:
        action_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 actions matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=ACTIONS,
        response_model=ActionResponse,
        filter_model=action_filter_model,
        params={"hydrate": hydrate},
    )

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_event_sources(event_source_filter_model: EventSourceFilter, hydrate: bool = False) -> Page[EventSourceResponse]

List all event_sources matching the given filter criteria.

Parameters:

Name	Type	Description	Default
event_source_filter_model	EventSourceFilter	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[EventSourceResponse]	A list of all event_sources matching the filter criteria.

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

def list_event_sources(
    self,
    event_source_filter_model: EventSourceFilter,
    hydrate: bool = False,
) -> Page[EventSourceResponse]:
    """List all event_sources matching the given filter criteria.

    Args:
        event_source_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 event_sources matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=EVENT_SOURCES,
        response_model=EventSourceResponse,
        filter_model=event_source_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_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_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_trigger_executions(trigger_execution_filter_model: TriggerExecutionFilter, hydrate: bool = False) -> Page[TriggerExecutionResponse]

List all trigger executions matching the given filter criteria.

Parameters:

Name	Type	Description	Default
trigger_execution_filter_model	TriggerExecutionFilter	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[TriggerExecutionResponse]	A list of all trigger executions matching the filter criteria.

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

def list_trigger_executions(
    self,
    trigger_execution_filter_model: TriggerExecutionFilter,
    hydrate: bool = False,
) -> Page[TriggerExecutionResponse]:
    """List all trigger executions matching the given filter criteria.

    Args:
        trigger_execution_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 trigger executions matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=TRIGGER_EXECUTIONS,
        response_model=TriggerExecutionResponse,
        filter_model=trigger_execution_filter_model,
        params={"hydrate": hydrate},
    )

list_triggers(trigger_filter_model: TriggerFilter, hydrate: bool = False) -> Page[TriggerResponse]

List all triggers matching the given filter criteria.

Parameters:

Name	Type	Description	Default
trigger_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[TriggerResponse]	A list of all triggers matching the filter criteria.

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

def list_triggers(
    self,
    trigger_filter_model: TriggerFilter,
    hydrate: bool = False,
) -> Page[TriggerResponse]:
    """List all triggers matching the given filter criteria.

    Args:
        trigger_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 triggers matching the filter criteria.
    """
    return self._list_paginated_resources(
        route=TRIGGERS,
        response_model=TriggerResponse,
        filter_model=trigger_filter_model,
        params={"hydrate": hydrate},
    )

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,
        },
    )

put(path: str, body: Optional[BaseModel] = None, params: Optional[Dict[str, Any]] = None, timeout: Optional[int] = None, **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
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,
    **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.
        kwargs: Additional keyword arguments to pass to the request.

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

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_action(action_id: UUID, action_update: ActionUpdate) -> ActionResponse

Update an existing action.

Parameters:

Name	Type	Description	Default
action_id	UUID	The ID of the action to update.	required
action_update	ActionUpdate	The update to be applied to the action.	required

Returns:

Type	Description
ActionResponse	The updated action.

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

def update_action(
    self,
    action_id: UUID,
    action_update: ActionUpdate,
) -> ActionResponse:
    """Update an existing action.

    Args:
        action_id: The ID of the action to update.
        action_update: The update to be applied to the action.

    Returns:
        The updated action.
    """
    return self._update_resource(
        resource_id=action_id,
        resource_update=action_update,
        route=ACTIONS,
        response_model=ActionResponse,
    )

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_event_source(event_source_id: UUID, event_source_update: EventSourceUpdate) -> EventSourceResponse

Update an existing event_source.

Parameters:

Name	Type	Description	Default
event_source_id	UUID	The ID of the event_source to update.	required
event_source_update	EventSourceUpdate	The update to be applied to the event_source.	required

Returns:

Type	Description
EventSourceResponse	The updated event_source.

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

def update_event_source(
    self,
    event_source_id: UUID,
    event_source_update: EventSourceUpdate,
) -> EventSourceResponse:
    """Update an existing event_source.

    Args:
        event_source_id: The ID of the event_source to update.
        event_source_update: The update to be applied to the event_source.

    Returns:
        The updated event_source.
    """
    return self._update_resource(
        resource_id=event_source_id,
        resource_update=event_source_update,
        route=EVENT_SOURCES,
        response_model=EventSourceResponse,
    )

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_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_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_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_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) -> TriggerResponse

Update an existing 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
TriggerResponse	The updated trigger.

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

def update_trigger(
    self,
    trigger_id: UUID,
    trigger_update: TriggerUpdate,
) -> TriggerResponse:
    """Update an existing 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.
    """
    return self._update_resource(
        resource_id=trigger_id,
        resource_update=trigger_update,
        route=TRIGGERS,
        response_model=TriggerResponse,
    )

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.

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()

    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	{}
**kwargs	Any	Keyword arguments to filter models.	{}

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

        **kwargs: Keyword arguments to filter models.

    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

ActionSchema

Bases: NamedSchema

SQL Model for actions.

Functions

from_request(request: ActionRequest) -> ActionSchema classmethod

Convert a ActionRequest to a ActionSchema.

Parameters:

Name	Type	Description	Default
request	ActionRequest	The request model to convert.	required

Returns:

Type	Description
ActionSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/action_schemas.py

@classmethod
def from_request(cls, request: "ActionRequest") -> "ActionSchema":
    """Convert a `ActionRequest` to a `ActionSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        configuration=base64.b64encode(
            json.dumps(
                request.configuration, default=pydantic_encoder
            ).encode("utf-8"),
        ),
        flavor=request.flavor,
        plugin_subtype=request.plugin_subtype,
        description=request.description,
        service_account_id=request.service_account_id,
        auth_window=request.auth_window,
    )

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/action_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(ActionSchema.user)),
                joinedload(
                    jl_arg(ActionSchema.service_account), innerjoin=True
                ),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ActionResponse

Converts the action schema to a model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
include_resources	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ActionResponse	The converted model.

Source code in src/zenml/zen_stores/schemas/action_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ActionResponse":
    """Converts the action schema to a model.

    Args:
        include_metadata: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        include_resources: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The converted model.
    """
    body = ActionResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        flavor=self.flavor,
        plugin_subtype=self.plugin_subtype,
    )
    metadata = None
    if include_metadata:
        metadata = ActionResponseMetadata(
            configuration=json.loads(
                base64.b64decode(self.configuration).decode()
            ),
            description=self.description,
            auth_window=self.auth_window,
        )
    resources = None
    if include_resources:
        resources = ActionResponseResources(
            user=self.user.to_model() if self.user else None,
            service_account=self.service_account.to_model(),
        )
    return ActionResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(action_update: ActionUpdate) -> ActionSchema

Updates a action schema with a action update model.

Parameters:

Name	Type	Description	Default
action_update	ActionUpdate	ActionUpdate to update the action with.	required

Returns:

Type	Description
ActionSchema	The updated ActionSchema.

Source code in src/zenml/zen_stores/schemas/action_schemas.py

def update(self, action_update: "ActionUpdate") -> "ActionSchema":
    """Updates a action schema with a action update model.

    Args:
        action_update: `ActionUpdate` to update the action with.

    Returns:
        The updated ActionSchema.
    """
    for field, value in action_update.dict(
        exclude_unset=True,
        exclude_none=True,
    ).items():
        if field == "configuration":
            self.configuration = base64.b64encode(
                json.dumps(
                    action_update.configuration, default=pydantic_encoder
                ).encode("utf-8")
            )
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

ApiTransactionSchema

Bases: BaseSchema

SQL Model for API transactions.

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,
        ),
    )
    if self.result is not None:
        response.set_result(self.result)
    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.
    """
    if update.result is not None:
        self.result = update.get_result()
    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,
    )

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,
    )

    # 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	{}

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

    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	{}
kwargs	Any	Additional keyword arguments.	{}

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

        kwargs: Additional keyword arguments.

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

EventSourceSchema

Bases: NamedSchema

SQL Model for tag.

Functions

from_request(request: EventSourceRequest) -> EventSourceSchema classmethod

Convert an EventSourceRequest to an EventSourceSchema.

Parameters:

Name	Type	Description	Default
request	EventSourceRequest	The request model to convert.	required

Returns:

Type	Description
EventSourceSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/event_source_schemas.py

@classmethod
def from_request(cls, request: EventSourceRequest) -> "EventSourceSchema":
    """Convert an `EventSourceRequest` to an `EventSourceSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        project_id=request.project,
        user_id=request.user,
        flavor=request.flavor,
        plugin_subtype=request.plugin_subtype,
        name=request.name,
        description=request.description,
        configuration=base64.b64encode(
            json.dumps(
                request.configuration,
                sort_keys=False,
                default=pydantic_encoder,
            ).encode("utf-8")
        ),
        is_active=True,  # Makes no sense to create an inactive event source
    )

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/event_source_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(EventSourceSchema.user)),
                # joinedload(jl_arg(EventSourceSchema.triggers)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> EventSourceResponse

Convert an EventSourceSchema to an EventSourceResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
include_resources	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
EventSourceResponse	The created EventSourceResponse.

Source code in src/zenml/zen_stores/schemas/event_source_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> EventSourceResponse:
    """Convert an `EventSourceSchema` to an `EventSourceResponse`.

    Args:
        include_metadata: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        include_resources: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The created `EventSourceResponse`.
    """
    from zenml.models import TriggerResponse

    body = EventSourceResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        flavor=self.flavor,
        plugin_subtype=self.plugin_subtype,
        is_active=self.is_active,
    )
    resources = None
    if include_resources:
        triggers = cast(
            Page[TriggerResponse],
            get_page_from_list(
                items_list=self.triggers,
                response_model=TriggerResponse,
                include_resources=include_resources,
                include_metadata=include_metadata,
            ),
        )
        resources = EventSourceResponseResources(
            user=self.user.to_model() if self.user else None,
            triggers=triggers,
        )
    metadata = None
    if include_metadata:
        metadata = EventSourceResponseMetadata(
            description=self.description,
            configuration=json.loads(
                base64.b64decode(self.configuration).decode()
            ),
        )
    return EventSourceResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: EventSourceUpdate) -> EventSourceSchema

Updates a EventSourceSchema from a EventSourceUpdate.

Parameters:

Name	Type	Description	Default
update	EventSourceUpdate	The EventSourceUpdate to update from.	required

Returns:

Type	Description
EventSourceSchema	The updated EventSourceSchema.

Source code in src/zenml/zen_stores/schemas/event_source_schemas.py

def update(self, update: EventSourceUpdate) -> "EventSourceSchema":
    """Updates a `EventSourceSchema` from a `EventSourceUpdate`.

    Args:
        update: The `EventSourceUpdate` to update from.

    Returns:
        The updated `EventSourceSchema`.
    """
    for field, value in update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field == "configuration":
            self.configuration = base64.b64encode(
                json.dumps(
                    update.configuration, default=pydantic_encoder
                ).encode("utf-8")
            )
        else:
            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),
        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

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,
    )
    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,
        )
    return LogsResponse(
        id=self.id,
        body=body,
        metadata=metadata,
    )

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,
    )

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) -> PipelineRunSchema classmethod

Convert a PipelineRunRequest to a PipelineRunSchema.

Parameters:

Name	Type	Description	Default
request	PipelineRunRequest	The request to convert.	required

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"
) -> "PipelineRunSchema":
    """Convert a `PipelineRunRequest` to a `PipelineRunSchema`.

    Args:
        request: The request to convert.

    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,
        status=request.status.value,
        in_progress=not request.status.is_finished,
        status_reason=request.status_reason,
        pipeline_id=request.pipeline,
        snapshot_id=request.snapshot,
        trigger_execution_id=request.trigger_execution_id,
        triggered_by=triggered_by,
        triggered_by_type=triggered_by_type,
    )

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(
    #         [
    #             joinedload(jl_arg(PipelineRunSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend(
            [
                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.user)),
                selectinload(jl_arg(PipelineRunSchema.tags)),
                selectinload(jl_arg(PipelineRunSchema.visualizations)),
            ]
        )

    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,
    )
    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)

        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,
        )

    resources = None
    if include_resources:
        # Add the client logs as "logs" if they exist, for backwards compatibility
        # TODO: This will be safe to remove in future releases (>0.84.0).
        client_logs = [
            log_entry
            for log_entry in self.logs
            if log_entry.source == "client"
        ]

        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,
            trigger_execution=(
                self.trigger_execution.to_model()
                if self.trigger_execution
                else None
            ),
            model_version=self.model_version.to_model()
            if self.model_version
            else None,
            tags=[tag.to_model() for tag in self.tags],
            logs=client_logs[0].to_model() if client_logs else None,
            log_collection=[log.to_model() for log in self.logs],
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    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

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.

    Returns:
        The updated `PipelineRunSchema`.
    """
    if run_update.status:
        if (
            run_update.status == ExecutionStatus.PROVISIONING
            and self.status != ExecutionStatus.INITIALIZING.value
        ):
            # This run is already past the provisioning status, so we ignore
            # the update.
            pass
        else:
            self.status = run_update.status.value
            self.end_time = run_update.end_time

            if run_update.status_reason:
                self.status_reason = run_update.status_reason

        self.in_progress = self._check_if_run_in_progress()

    if run_update.orchestrator_run_id:
        self.orchestrator_run_id = run_update.orchestrator_run_id

    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.pipeline_id != request.pipeline
        or self.project_id != request.project
    ):
        raise ValueError(
            "Snapshot, project or pipeline 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

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,
    )

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

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,
        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.
    """
    runnable = False
    if self.build and not self.build.is_local and self.build.stack_id:
        runnable = True

    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=runnable,
        deployable=deployable,
    )
    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,
                step_configurations=all_step_configurations,
            )

        metadata = PipelineSnapshotResponseMetadata(
            description=self.description,
            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

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

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,
    )
    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

    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"},
    ).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, components: Sequence[StackComponentSchema]) -> 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
components	Sequence[StackComponentSchema]	List of components to link to 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",
    components: Sequence["StackComponentSchema"],
) -> "StackSchema":
    """Create a stack schema from a request.

    Args:
        request: The request from which to create the stack.
        components: List of components to link to 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,
        components=components,
        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.components)).joinedload(
    #                 jl_arg(StackComponentSchema.flavor_schema)
    #             ),
    #         ]
    #     )

    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()
            )
        metadata = StackResponseMetadata(
            components={c.type: [c.to_model()] for c in self.components},
            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, components: List[StackComponentSchema]) -> 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
components	List[StackComponentSchema]	List of StackComponentSchema 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",
    components: List["StackComponentSchema"],
) -> "StackSchema":
    """Updates a stack schema with a stack update model.

    Args:
        stack_update: `StackUpdate` to update the stack with.
        components: List of `StackComponentSchema` to update the stack with.

    Returns:
        The updated StackSchema.
    """
    for field, value in stack_update.model_dump(
        exclude_unset=True,
        exclude={"user", "add_secrets", "remove_secrets"},
    ).items():
        if field == "components":
            self.components = components
        elif 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.step_configuration_schema)),
    ]

    if include_metadata:
        options.extend(
            [
                selectinload(jl_arg(StepRunSchema.logs)),
                # 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
                ),
            ]
        )

    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 self.step_configuration_schema:
            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(self.step_configuration_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,
        status=ExecutionStatus(self.status),
        version=self.version,
        is_retriable=self.is_retriable,
        start_time=self.start_time,
        end_time=self.end_time,
        created=self.created,
        updated=self.updated,
        model_version_id=self.model_version_id,
        substitutions=step.config.substitutions,
    )
    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,
            logs=self.logs.to_model() if self.logs 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),
                **input_artifact.artifact_version.to_model().model_dump(),
            )
            input_artifacts[input_artifact.name].append(step_run_input)

        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,
            inputs=input_artifacts,
            outputs=output_artifacts,
        )

    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: BaseSchema

SQL Model for trigger executions.

Functions

from_request(request: TriggerExecutionRequest) -> TriggerExecutionSchema classmethod

Convert a TriggerExecutionRequest to a TriggerExecutionSchema.

Parameters:

Name	Type	Description	Default
request	TriggerExecutionRequest	The request model to convert.	required

Returns:

Type	Description
TriggerExecutionSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

@classmethod
def from_request(
    cls, request: "TriggerExecutionRequest"
) -> "TriggerExecutionSchema":
    """Convert a `TriggerExecutionRequest` to a `TriggerExecutionSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        trigger_id=request.trigger,
        event_metadata=base64.b64encode(
            json.dumps(request.event_metadata).encode("utf-8")
        ),
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TriggerExecutionResponse

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
TriggerExecutionResponse	The converted model.

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,
) -> "TriggerExecutionResponse":
    """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 = TriggerExecutionResponseBody(
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = TriggerExecutionResponseMetadata(
            event_metadata=json.loads(
                base64.b64decode(self.event_metadata).decode()
            )
            if self.event_metadata
            else {},
        )
    resources = None
    if include_resources:
        resources = TriggerExecutionResponseResources(
            trigger=self.trigger.to_model(),
        )

    return TriggerExecutionResponse(
        id=self.id, body=body, metadata=metadata, resources=resources
    )

TriggerSchema

Bases: NamedSchema

SQL Model for triggers.

Functions

from_request(request: TriggerRequest) -> TriggerSchema classmethod

Convert a TriggerRequest to a TriggerSchema.

Parameters:

Name	Type	Description	Default
request	TriggerRequest	The request model to convert.	required

Returns:

Type	Description
TriggerSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

@classmethod
def from_request(cls, request: "TriggerRequest") -> "TriggerSchema":
    """Convert a `TriggerRequest` to a `TriggerSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        action_id=request.action_id,
        event_source_id=request.event_source_id,
        event_filter=base64.b64encode(
            json.dumps(
                request.event_filter, default=pydantic_encoder
            ).encode("utf-8")
        ),
        schedule=base64.b64encode(request.schedule.json().encode("utf-8"))
        if request.schedule
        else None,
        description=request.description,
        is_active=True,  # Makes no sense for it to be created inactive
    )

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 = [
        joinedload(jl_arg(TriggerSchema.action), innerjoin=True),
        joinedload(jl_arg(TriggerSchema.event_source), innerjoin=True),
    ]

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(TriggerSchema.user)),
                # joinedload(jl_arg(TriggerSchema.executions)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TriggerResponse

Converts the schema to a model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
include_resources	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
TriggerResponse	The converted model.

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,
) -> "TriggerResponse":
    """Converts the schema to a model.

    Args:
        include_metadata: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        include_resources: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The converted model.
    """
    from zenml.models import TriggerExecutionResponse

    body = TriggerResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        action_flavor=self.action.flavor,
        action_subtype=self.action.plugin_subtype,
        event_source_flavor=self.event_source.flavor
        if self.event_source
        else None,
        event_source_subtype=self.event_source.plugin_subtype
        if self.event_source
        else None,
        is_active=self.is_active,
    )
    metadata = None
    if include_metadata:
        metadata = TriggerResponseMetadata(
            event_filter=json.loads(
                base64.b64decode(self.event_filter).decode()
            ),
            schedule=Schedule.parse_raw(
                base64.b64decode(self.schedule).decode()
            )
            if self.schedule
            else None,
            description=self.description,
        )
    resources = None
    if include_resources:
        executions = cast(
            Page[TriggerExecutionResponse],
            get_page_from_list(
                items_list=self.executions,
                response_model=TriggerExecutionResponse,
                include_resources=False,
                include_metadata=False,
            ),
        )
        resources = TriggerResponseResources(
            user=self.user.to_model() if self.user else None,
            action=self.action.to_model(),
            event_source=self.event_source.to_model()
            if self.event_source
            else None,
            executions=executions,
        )
    return TriggerResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(trigger_update: TriggerUpdate) -> TriggerSchema

Updates a trigger schema with a trigger update model.

Parameters:

Name	Type	Description	Default
trigger_update	TriggerUpdate	TriggerUpdate to update the trigger with.	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":
    """Updates a trigger schema with a trigger update model.

    Args:
        trigger_update: `TriggerUpdate` to update the trigger with.

    Returns:
        The updated TriggerSchema.
    """
    for field, value in trigger_update.model_dump(
        exclude_unset=True,
        exclude_none=True,
    ).items():
        if field == "event_filter":
            self.event_filter = base64.b64encode(
                json.dumps(
                    trigger_update.event_filter, default=pydantic_encoder
                ).encode("utf-8")
            )
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

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

action_schemas

SQL Model Implementations for Actions.

Classes

ActionSchema

Bases: NamedSchema

SQL Model for actions.

Functions

from_request(request: ActionRequest) -> ActionSchema classmethod

Convert a ActionRequest to a ActionSchema.

Parameters:

Name	Type	Description	Default
request	ActionRequest	The request model to convert.	required

Returns:

Type	Description
ActionSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/action_schemas.py

@classmethod
def from_request(cls, request: "ActionRequest") -> "ActionSchema":
    """Convert a `ActionRequest` to a `ActionSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        configuration=base64.b64encode(
            json.dumps(
                request.configuration, default=pydantic_encoder
            ).encode("utf-8"),
        ),
        flavor=request.flavor,
        plugin_subtype=request.plugin_subtype,
        description=request.description,
        service_account_id=request.service_account_id,
        auth_window=request.auth_window,
    )

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/action_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(ActionSchema.user)),
                joinedload(
                    jl_arg(ActionSchema.service_account), innerjoin=True
                ),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> ActionResponse

Converts the action schema to a model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
include_resources	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
ActionResponse	The converted model.

Source code in src/zenml/zen_stores/schemas/action_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> "ActionResponse":
    """Converts the action schema to a model.

    Args:
        include_metadata: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        include_resources: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The converted model.
    """
    body = ActionResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        flavor=self.flavor,
        plugin_subtype=self.plugin_subtype,
    )
    metadata = None
    if include_metadata:
        metadata = ActionResponseMetadata(
            configuration=json.loads(
                base64.b64decode(self.configuration).decode()
            ),
            description=self.description,
            auth_window=self.auth_window,
        )
    resources = None
    if include_resources:
        resources = ActionResponseResources(
            user=self.user.to_model() if self.user else None,
            service_account=self.service_account.to_model(),
        )
    return ActionResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(action_update: ActionUpdate) -> ActionSchema

Updates a action schema with a action update model.

Parameters:

Name	Type	Description	Default
action_update	ActionUpdate	ActionUpdate to update the action with.	required

Returns:

Type	Description
ActionSchema	The updated ActionSchema.

Source code in src/zenml/zen_stores/schemas/action_schemas.py

def update(self, action_update: "ActionUpdate") -> "ActionSchema":
    """Updates a action schema with a action update model.

    Args:
        action_update: `ActionUpdate` to update the action with.

    Returns:
        The updated ActionSchema.
    """
    for field, value in action_update.dict(
        exclude_unset=True,
        exclude_none=True,
    ).items():
        if field == "configuration":
            self.configuration = base64.b64encode(
                json.dumps(
                    action_update.configuration, default=pydantic_encoder
                ).encode("utf-8")
            )
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    return self

Functions

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	{}
**kwargs	Any	Keyword arguments to filter models.	{}

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

        **kwargs: Keyword arguments to filter models.

    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

ApiTransactionSchema

Bases: BaseSchema

SQL Model for API transactions.

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,
        ),
    )
    if self.result is not None:
        response.set_result(self.result)
    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.
    """
    if update.result is not None:
        self.result = update.get_result()
    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,
    )

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,
    )

    # 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	{}

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

    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	{}
kwargs	Any	Additional keyword arguments.	{}

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

        kwargs: Additional keyword arguments.

    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"},
    ).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.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,
            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

event_source_schemas

SQL Model Implementations for event sources.

Classes

EventSourceSchema

Bases: NamedSchema

SQL Model for tag.

Functions

from_request(request: EventSourceRequest) -> EventSourceSchema classmethod

Convert an EventSourceRequest to an EventSourceSchema.

Parameters:

Name	Type	Description	Default
request	EventSourceRequest	The request model to convert.	required

Returns:

Type	Description
EventSourceSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/event_source_schemas.py

@classmethod
def from_request(cls, request: EventSourceRequest) -> "EventSourceSchema":
    """Convert an `EventSourceRequest` to an `EventSourceSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        project_id=request.project,
        user_id=request.user,
        flavor=request.flavor,
        plugin_subtype=request.plugin_subtype,
        name=request.name,
        description=request.description,
        configuration=base64.b64encode(
            json.dumps(
                request.configuration,
                sort_keys=False,
                default=pydantic_encoder,
            ).encode("utf-8")
        ),
        is_active=True,  # Makes no sense to create an inactive event source
    )

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/event_source_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(EventSourceSchema.user)),
                # joinedload(jl_arg(EventSourceSchema.triggers)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> EventSourceResponse

Convert an EventSourceSchema to an EventSourceResponse.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
include_resources	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
EventSourceResponse	The created EventSourceResponse.

Source code in src/zenml/zen_stores/schemas/event_source_schemas.py

def to_model(
    self,
    include_metadata: bool = False,
    include_resources: bool = False,
    **kwargs: Any,
) -> EventSourceResponse:
    """Convert an `EventSourceSchema` to an `EventSourceResponse`.

    Args:
        include_metadata: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        include_resources: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The created `EventSourceResponse`.
    """
    from zenml.models import TriggerResponse

    body = EventSourceResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        flavor=self.flavor,
        plugin_subtype=self.plugin_subtype,
        is_active=self.is_active,
    )
    resources = None
    if include_resources:
        triggers = cast(
            Page[TriggerResponse],
            get_page_from_list(
                items_list=self.triggers,
                response_model=TriggerResponse,
                include_resources=include_resources,
                include_metadata=include_metadata,
            ),
        )
        resources = EventSourceResponseResources(
            user=self.user.to_model() if self.user else None,
            triggers=triggers,
        )
    metadata = None
    if include_metadata:
        metadata = EventSourceResponseMetadata(
            description=self.description,
            configuration=json.loads(
                base64.b64decode(self.configuration).decode()
            ),
        )
    return EventSourceResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(update: EventSourceUpdate) -> EventSourceSchema

Updates a EventSourceSchema from a EventSourceUpdate.

Parameters:

Name	Type	Description	Default
update	EventSourceUpdate	The EventSourceUpdate to update from.	required

Returns:

Type	Description
EventSourceSchema	The updated EventSourceSchema.

Source code in src/zenml/zen_stores/schemas/event_source_schemas.py

def update(self, update: EventSourceUpdate) -> "EventSourceSchema":
    """Updates a `EventSourceSchema` from a `EventSourceUpdate`.

    Args:
        update: The `EventSourceUpdate` to update from.

    Returns:
        The updated `EventSourceSchema`.
    """
    for field, value in update.model_dump(
        exclude_unset=True, exclude_none=True
    ).items():
        if field == "configuration":
            self.configuration = base64.b64encode(
                json.dumps(
                    update.configuration, default=pydantic_encoder
                ).encode("utf-8")
            )
        else:
            setattr(self, field, 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),
        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

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,
    )
    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,
        )
    return LogsResponse(
        id=self.id,
        body=body,
        metadata=metadata,
    )

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

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) -> PipelineRunSchema classmethod

Convert a PipelineRunRequest to a PipelineRunSchema.

Parameters:

Name	Type	Description	Default
request	PipelineRunRequest	The request to convert.	required

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"
) -> "PipelineRunSchema":
    """Convert a `PipelineRunRequest` to a `PipelineRunSchema`.

    Args:
        request: The request to convert.

    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,
        status=request.status.value,
        in_progress=not request.status.is_finished,
        status_reason=request.status_reason,
        pipeline_id=request.pipeline,
        snapshot_id=request.snapshot,
        trigger_execution_id=request.trigger_execution_id,
        triggered_by=triggered_by,
        triggered_by_type=triggered_by_type,
    )

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(
    #         [
    #             joinedload(jl_arg(PipelineRunSchema.run_metadata)),
    #         ]
    #     )

    if include_resources:
        options.extend(
            [
                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.user)),
                selectinload(jl_arg(PipelineRunSchema.tags)),
                selectinload(jl_arg(PipelineRunSchema.visualizations)),
            ]
        )

    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,
    )
    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)

        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,
        )

    resources = None
    if include_resources:
        # Add the client logs as "logs" if they exist, for backwards compatibility
        # TODO: This will be safe to remove in future releases (>0.84.0).
        client_logs = [
            log_entry
            for log_entry in self.logs
            if log_entry.source == "client"
        ]

        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,
            trigger_execution=(
                self.trigger_execution.to_model()
                if self.trigger_execution
                else None
            ),
            model_version=self.model_version.to_model()
            if self.model_version
            else None,
            tags=[tag.to_model() for tag in self.tags],
            logs=client_logs[0].to_model() if client_logs else None,
            log_collection=[log.to_model() for log in self.logs],
            visualizations=[
                visualization.to_model(
                    include_metadata=False,
                    include_resources=False,
                )
                for visualization in self.visualizations
            ],
        )

    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

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.

    Returns:
        The updated `PipelineRunSchema`.
    """
    if run_update.status:
        if (
            run_update.status == ExecutionStatus.PROVISIONING
            and self.status != ExecutionStatus.INITIALIZING.value
        ):
            # This run is already past the provisioning status, so we ignore
            # the update.
            pass
        else:
            self.status = run_update.status.value
            self.end_time = run_update.end_time

            if run_update.status_reason:
                self.status_reason = run_update.status_reason

        self.in_progress = self._check_if_run_in_progress()

    if run_update.orchestrator_run_id:
        self.orchestrator_run_id = run_update.orchestrator_run_id

    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.pipeline_id != request.pipeline
        or self.project_id != request.project
    ):
        raise ValueError(
            "Snapshot, project or pipeline 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

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,
    )

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

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,
        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.
    """
    runnable = False
    if self.build and not self.build.is_local and self.build.stack_id:
        runnable = True

    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=runnable,
        deployable=deployable,
    )
    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,
                step_configurations=all_step_configurations,
            )

        metadata = PipelineSnapshotResponseMetadata(
            description=self.description,
            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

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

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,
    )
    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

    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.
ValueError	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.
        ValueError: 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, components: Sequence[StackComponentSchema]) -> 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
components	Sequence[StackComponentSchema]	List of components to link to 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",
    components: Sequence["StackComponentSchema"],
) -> "StackSchema":
    """Create a stack schema from a request.

    Args:
        request: The request from which to create the stack.
        components: List of components to link to 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,
        components=components,
        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.components)).joinedload(
    #                 jl_arg(StackComponentSchema.flavor_schema)
    #             ),
    #         ]
    #     )

    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()
            )
        metadata = StackResponseMetadata(
            components={c.type: [c.to_model()] for c in self.components},
            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, components: List[StackComponentSchema]) -> 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
components	List[StackComponentSchema]	List of StackComponentSchema 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",
    components: List["StackComponentSchema"],
) -> "StackSchema":
    """Updates a stack schema with a stack update model.

    Args:
        stack_update: `StackUpdate` to update the stack with.
        components: List of `StackComponentSchema` to update the stack with.

    Returns:
        The updated StackSchema.
    """
    for field, value in stack_update.model_dump(
        exclude_unset=True,
        exclude={"user", "add_secrets", "remove_secrets"},
    ).items():
        if field == "components":
            self.components = components
        elif 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.step_configuration_schema)),
    ]

    if include_metadata:
        options.extend(
            [
                selectinload(jl_arg(StepRunSchema.logs)),
                # 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
                ),
            ]
        )

    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 self.step_configuration_schema:
            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(self.step_configuration_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,
        status=ExecutionStatus(self.status),
        version=self.version,
        is_retriable=self.is_retriable,
        start_time=self.start_time,
        end_time=self.end_time,
        created=self.created,
        updated=self.updated,
        model_version_id=self.model_version_id,
        substitutions=step.config.substitutions,
    )
    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,
            logs=self.logs.to_model() if self.logs 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),
                **input_artifact.artifact_version.to_model().model_dump(),
            )
            input_artifacts[input_artifact.name].append(step_run_input)

        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,
            inputs=input_artifacts,
            outputs=output_artifacts,
        )

    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_schemas

SQL Model Implementations for Triggers.

Classes

TriggerExecutionSchema

Bases: BaseSchema

SQL Model for trigger executions.

Functions

from_request(request: TriggerExecutionRequest) -> TriggerExecutionSchema classmethod

Convert a TriggerExecutionRequest to a TriggerExecutionSchema.

Parameters:

Name	Type	Description	Default
request	TriggerExecutionRequest	The request model to convert.	required

Returns:

Type	Description
TriggerExecutionSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

@classmethod
def from_request(
    cls, request: "TriggerExecutionRequest"
) -> "TriggerExecutionSchema":
    """Convert a `TriggerExecutionRequest` to a `TriggerExecutionSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        trigger_id=request.trigger,
        event_metadata=base64.b64encode(
            json.dumps(request.event_metadata).encode("utf-8")
        ),
    )

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TriggerExecutionResponse

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
TriggerExecutionResponse	The converted model.

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,
) -> "TriggerExecutionResponse":
    """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 = TriggerExecutionResponseBody(
        created=self.created,
        updated=self.updated,
    )
    metadata = None
    if include_metadata:
        metadata = TriggerExecutionResponseMetadata(
            event_metadata=json.loads(
                base64.b64decode(self.event_metadata).decode()
            )
            if self.event_metadata
            else {},
        )
    resources = None
    if include_resources:
        resources = TriggerExecutionResponseResources(
            trigger=self.trigger.to_model(),
        )

    return TriggerExecutionResponse(
        id=self.id, body=body, metadata=metadata, resources=resources
    )

TriggerSchema

Bases: NamedSchema

SQL Model for triggers.

Functions

from_request(request: TriggerRequest) -> TriggerSchema classmethod

Convert a TriggerRequest to a TriggerSchema.

Parameters:

Name	Type	Description	Default
request	TriggerRequest	The request model to convert.	required

Returns:

Type	Description
TriggerSchema	The converted schema.

Source code in src/zenml/zen_stores/schemas/trigger_schemas.py

@classmethod
def from_request(cls, request: "TriggerRequest") -> "TriggerSchema":
    """Convert a `TriggerRequest` to a `TriggerSchema`.

    Args:
        request: The request model to convert.

    Returns:
        The converted schema.
    """
    return cls(
        name=request.name,
        project_id=request.project,
        user_id=request.user,
        action_id=request.action_id,
        event_source_id=request.event_source_id,
        event_filter=base64.b64encode(
            json.dumps(
                request.event_filter, default=pydantic_encoder
            ).encode("utf-8")
        ),
        schedule=base64.b64encode(request.schedule.json().encode("utf-8"))
        if request.schedule
        else None,
        description=request.description,
        is_active=True,  # Makes no sense for it to be created inactive
    )

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 = [
        joinedload(jl_arg(TriggerSchema.action), innerjoin=True),
        joinedload(jl_arg(TriggerSchema.event_source), innerjoin=True),
    ]

    if include_resources:
        options.extend(
            [
                joinedload(jl_arg(TriggerSchema.user)),
                # joinedload(jl_arg(TriggerSchema.executions)),
            ]
        )

    return options

to_model(include_metadata: bool = False, include_resources: bool = False, **kwargs: Any) -> TriggerResponse

Converts the schema to a model.

Parameters:

Name	Type	Description	Default
include_metadata	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
include_resources	bool	Flag deciding whether to include the output model(s) metadata fields in the response.	False
**kwargs	Any	Keyword arguments to allow schema specific logic	{}

Returns:

Type	Description
TriggerResponse	The converted model.

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,
) -> "TriggerResponse":
    """Converts the schema to a model.

    Args:
        include_metadata: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        include_resources: Flag deciding whether to include the output model(s)
            metadata fields in the response.
        **kwargs: Keyword arguments to allow schema specific logic

    Returns:
        The converted model.
    """
    from zenml.models import TriggerExecutionResponse

    body = TriggerResponseBody(
        user_id=self.user_id,
        project_id=self.project_id,
        created=self.created,
        updated=self.updated,
        action_flavor=self.action.flavor,
        action_subtype=self.action.plugin_subtype,
        event_source_flavor=self.event_source.flavor
        if self.event_source
        else None,
        event_source_subtype=self.event_source.plugin_subtype
        if self.event_source
        else None,
        is_active=self.is_active,
    )
    metadata = None
    if include_metadata:
        metadata = TriggerResponseMetadata(
            event_filter=json.loads(
                base64.b64decode(self.event_filter).decode()
            ),
            schedule=Schedule.parse_raw(
                base64.b64decode(self.schedule).decode()
            )
            if self.schedule
            else None,
            description=self.description,
        )
    resources = None
    if include_resources:
        executions = cast(
            Page[TriggerExecutionResponse],
            get_page_from_list(
                items_list=self.executions,
                response_model=TriggerExecutionResponse,
                include_resources=False,
                include_metadata=False,
            ),
        )
        resources = TriggerResponseResources(
            user=self.user.to_model() if self.user else None,
            action=self.action.to_model(),
            event_source=self.event_source.to_model()
            if self.event_source
            else None,
            executions=executions,
        )
    return TriggerResponse(
        id=self.id,
        name=self.name,
        body=body,
        metadata=metadata,
        resources=resources,
    )

update(trigger_update: TriggerUpdate) -> TriggerSchema

Updates a trigger schema with a trigger update model.

Parameters:

Name	Type	Description	Default
trigger_update	TriggerUpdate	TriggerUpdate to update the trigger with.	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":
    """Updates a trigger schema with a trigger update model.

    Args:
        trigger_update: `TriggerUpdate` to update the trigger with.

    Returns:
        The updated TriggerSchema.
    """
    for field, value in trigger_update.model_dump(
        exclude_unset=True,
        exclude_none=True,
    ).items():
        if field == "event_filter":
            self.event_filter = base64.b64encode(
                json.dumps(
                    trigger_update.event_filter, default=pydantic_encoder
                ).encode("utf-8")
            )
        else:
            setattr(self, field, value)

    self.updated = utc_now()
    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)
    return {
        k: sorted(v, key=lambda x: x.created, reverse=True)[0].value
        for k, v in metadata_collection.items()
    }

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

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.

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_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.
vault_namespace	Optional[str]	The Vault Enterprise namespace.
mount_point	Optional[str]	The mount point to use for all secrets.
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_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.

engine: Engine property

The SQLAlchemy engine.

Returns:

Type	Description
Engine	The SQLAlchemy engine.

Raises:

Type	Description
ValueError	If the store is not initialized.

migration_utils: MigrationUtils property

The migration utils.

Returns:

Type	Description
MigrationUtils	The migration utils.

Raises:

Type	Description
ValueError	If the store is not initialized.

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

backup_database(strategy: Optional[DatabaseBackupStrategy] = None, location: Optional[str] = None, overwrite: bool = False) -> Tuple[str, Any]

Backup the database.

Parameters:

Name	Type	Description	Default
strategy	Optional[DatabaseBackupStrategy]	Custom backup strategy to use. If not set, the backup strategy from the store configuration will be used.	None
location	Optional[str]	Custom target location to backup the database to. If not set, the configured backup location will be used. Depending on the backup strategy, this can be a file path or a database name.	None
overwrite	bool	Whether to overwrite an existing backup if it exists. If set to False, the existing backup will be reused.	False

Returns:

Type	Description
str	The location where the database was backed up to and an accompanying
Any	user-friendly message that describes the backup location, or None
Tuple[str, Any]	if no backup was created (i.e. because the backup already exists).

Raises:

Type	Description
ValueError	If the backup database name is not set when the backup database is requested or if the backup strategy is invalid.

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

def backup_database(
    self,
    strategy: Optional[DatabaseBackupStrategy] = None,
    location: Optional[str] = None,
    overwrite: bool = False,
) -> Tuple[str, Any]:
    """Backup the database.

    Args:
        strategy: Custom backup strategy to use. If not set, the backup
            strategy from the store configuration will be used.
        location: Custom target location to backup the database to. If not
            set, the configured backup location will be used. Depending on
            the backup strategy, this can be a file path or a database name.
        overwrite: Whether to overwrite an existing backup if it exists.
            If set to False, the existing backup will be reused.

    Returns:
        The location where the database was backed up to and an accompanying
        user-friendly message that describes the backup location, or None
        if no backup was created (i.e. because the backup already exists).

    Raises:
        ValueError: If the backup database name is not set when the backup
            database is requested or if the backup strategy is invalid.
    """
    strategy = strategy or self.config.backup_strategy

    if (
        strategy == DatabaseBackupStrategy.DUMP_FILE
        or self.config.driver == SQLDatabaseDriver.SQLITE
    ):
        dump_file = location or self._get_db_backup_file_path()

        if not overwrite and os.path.isfile(dump_file):
            logger.warning(
                f"A previous backup file already exists at '{dump_file}'. "
                "Reusing the existing backup."
            )
        else:
            self.migration_utils.backup_database_to_file(
                dump_file=dump_file
            )
        return f"the '{dump_file}' backup file", dump_file
    elif strategy == DatabaseBackupStrategy.DATABASE:
        backup_db_name = location or self.config.backup_database
        if not backup_db_name:
            raise ValueError(
                "The backup database name must be set in the store "
                "configuration to use the backup database strategy."
            )

        if not overwrite and self.migration_utils.database_exists(
            backup_db_name
        ):
            logger.warning(
                "A previous backup database already exists at "
                f"'{backup_db_name}'. Reusing the existing backup."
            )
        else:
            self.migration_utils.backup_database_to_db(
                backup_db_name=backup_db_name
            )
        return f"the '{backup_db_name}' backup database", backup_db_name
    elif strategy == DatabaseBackupStrategy.IN_MEMORY:
        return (
            "memory",
            self.migration_utils.backup_database_to_memory(),
        )

    else:
        raise ValueError(f"Invalid backup strategy: {strategy}.")

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

noqa: DAR401

Raises: BackupSecretsStoreNotConfiguredError: if no backup secrets store is configured.

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.

    # noqa: DAR401
    Raises:
        BackupSecretsStoreNotConfiguredError: if no backup secrets store is
            configured.
    """
    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_database_backup(strategy: Optional[DatabaseBackupStrategy] = None, location: Optional[Any] = None) -> None

Delete the database backup.

Parameters:

Name	Type	Description	Default
strategy	Optional[DatabaseBackupStrategy]	Custom backup strategy to use. If not set, the backup strategy from the store configuration will be used.	None
location	Optional[Any]	Custom target location to delete the database backup from. If not set, the configured backup location will be used. Depending on the backup strategy, this can be a file path or a database name.	None

Raises:

Type	Description
ValueError	If the backup database name is not set when the backup database is requested.

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

def cleanup_database_backup(
    self,
    strategy: Optional[DatabaseBackupStrategy] = None,
    location: Optional[Any] = None,
) -> None:
    """Delete the database backup.

    Args:
        strategy: Custom backup strategy to use. If not set, the backup
            strategy from the store configuration will be used.
        location: Custom target location to delete the database backup
            from. If not set, the configured backup location will be used.
            Depending on the backup strategy, this can be a file path or a
            database name.

    Raises:
        ValueError: If the backup database name is not set when the backup
            database is requested.
    """
    strategy = strategy or self.config.backup_strategy

    if (
        strategy == DatabaseBackupStrategy.DUMP_FILE
        or self.config.driver == SQLDatabaseDriver.SQLITE
    ):
        dump_file = location or self._get_db_backup_file_path()
        if dump_file is not None and os.path.isfile(dump_file):
            try:
                os.remove(dump_file)
            except OSError:
                logger.warning(
                    f"Failed to cleanup database dump file {dump_file}."
                )
            else:
                logger.info(
                    f"Successfully cleaned up database dump file "
                    f"{dump_file}."
                )
    elif strategy == DatabaseBackupStrategy.DATABASE:
        backup_db_name = location or self.config.backup_database

        if not backup_db_name:
            raise ValueError(
                "The backup database name must be set in the store "
                "configuration to use the backup database strategy."
            )
        if self.migration_utils.database_exists(backup_db_name):
            # Drop the backup database
            self.migration_utils.drop_database(
                database=backup_db_name,
            )
            logger.info(
                f"Successfully cleaned up backup database "
                f"{backup_db_name}."
            )

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_action(action: ActionRequest) -> ActionResponse

Create an action.

Parameters:

Name	Type	Description	Default
action	ActionRequest	The action to create.	required

Returns:

Type	Description
ActionResponse	The created action.

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

def create_action(self, action: ActionRequest) -> ActionResponse:
    """Create an action.

    Args:
        action: The action to create.

    Returns:
        The created action.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=action, session=session)

        self._verify_name_uniqueness(
            resource=action,
            schema=ActionSchema,
            session=session,
        )

        # Verify that the given service account exists
        self._get_account_schema(
            account_name_or_id=action.service_account_id,
            session=session,
            service_account=True,
        )

        new_action = ActionSchema.from_request(action)
        session.add(new_action)
        session.commit()
        session.refresh(new_action)

        return new_action.to_model(
            include_metadata=True, include_resources=True
        )

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 = session.exec(
            select(schema_class).where(
                schema_class.id == visualization.resource_id
            )
        ).first()

        if not resource_schema:
            raise KeyError(
                f"Resource of type '{visualization.resource_type.value}' "
                f"with ID {visualization.resource_id} not found."
            )

        if hasattr(resource_schema, "project_id"):
            resource_project_id = resource_schema.project_id
            if resource_project_id and resource_project_id != project_id:
                raise IllegalOperationError(
                    f"Resource {visualization.resource_type.value} with ID "
                    f"{visualization.resource_id} belongs to a different project than "
                    f"the curated visualization (project ID: {project_id})."
                )

        self._assert_curated_visualization_duplicate(
            session=session,
            artifact_visualization_id=visualization.artifact_visualization_id,
            resource_id=visualization.resource_id,
            resource_type=visualization.resource_type,
        )
        if visualization.display_order is not None:
            self._assert_curated_visualization_display_order_unique(
                session=session,
                resource_id=visualization.resource_id,
                resource_type=visualization.resource_type,
                display_order=visualization.display_order,
            )

        schema = CuratedVisualizationSchema.from_request(visualization)

        session.add(schema)
        session.commit()
        session.refresh(schema)

        return schema.to_model(
            include_metadata=True,
            include_resources=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/sql_zen_store.py

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

    Args:
        deployment: The deployment to create.

    Returns:
        The newly created deployment.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=deployment, session=session
        )
        self._verify_name_uniqueness(
            resource=deployment,
            schema=DeploymentSchema,
            session=session,
        )
        self._get_reference_schema_by_id(
            resource=deployment,
            reference_schema=PipelineSnapshotSchema,
            reference_id=deployment.snapshot_id,
            session=session,
        )
        self._get_reference_schema_by_id(
            resource=deployment,
            reference_schema=StackComponentSchema,
            reference_id=deployment.deployer_id,
            session=session,
            reference_type="deployer",
        )
        deployment_schema = DeploymentSchema.from_request(deployment)
        session.add(deployment_schema)
        session.commit()

        self._attach_tags_to_resources(
            tags=deployment.tags,
            resources=deployment_schema,
            session=session,
        )

        session.refresh(deployment_schema)
        return deployment_schema.to_model(
            include_metadata=True, include_resources=True
        )

create_event_source(event_source: EventSourceRequest) -> EventSourceResponse

Create an event_source.

Parameters:

Name	Type	Description	Default
event_source	EventSourceRequest	The event_source to create.	required

Returns:

Type	Description
EventSourceResponse	The created event_source.

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

def create_event_source(
    self, event_source: EventSourceRequest
) -> EventSourceResponse:
    """Create an event_source.

    Args:
        event_source: The event_source to create.

    Returns:
        The created event_source.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=event_source, session=session
        )

        self._verify_name_uniqueness(
            resource=event_source,
            schema=EventSourceSchema,
            session=session,
        )

        new_event_source = EventSourceSchema.from_request(event_source)
        session.add(new_event_source)
        session.commit()
        session.refresh(new_event_source)

        return new_event_source.to_model(
            include_metadata=True, include_resources=True
        )

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.

Raises:

Type	Description
EntityExistsError	If a flavor with the same name and type is already owned by this user.
ValueError	In case the config_schema string exceeds the max length.

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

@track_decorator(AnalyticsEvent.CREATED_FLAVOR)
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.

    Raises:
        EntityExistsError: If a flavor with the same name and type
            is already owned by this user.
        ValueError: In case the config_schema string exceeds the max length.
    """
    with Session(self.engine) as session:
        if flavor.is_custom is False:
            # Set the user to None for built-in flavors
            flavor.user = None
        else:
            self._set_request_user_id(
                request_model=flavor, session=session
            )
        # Check if flavor with the same domain key (name, type) already
        # exists
        existing_flavor = session.exec(
            select(FlavorSchema)
            .where(FlavorSchema.name == flavor.name)
            .where(FlavorSchema.type == flavor.type)
        ).first()

        if existing_flavor is not None:
            raise EntityExistsError(
                f"Unable to register '{flavor.type.value}' flavor "
                f"with name '{flavor.name}' and type '{flavor.type}': "
                "Found an existing flavor with the same name and type."
            )

        config_schema = json.dumps(flavor.config_schema)

        if len(config_schema) > TEXT_FIELD_MAX_LENGTH:
            raise ValueError(
                "Json representation of configuration schema"
                "exceeds max length."
            )

        else:
            new_flavor = FlavorSchema(
                name=flavor.name,
                type=flavor.type,
                source=flavor.source,
                config_schema=config_schema,
                integration=flavor.integration,
                connector_type=flavor.connector_type,
                connector_resource_type=flavor.connector_resource_type,
                connector_resource_id_attr=flavor.connector_resource_id_attr,
                user_id=flavor.user,
                logo_url=flavor.logo_url,
                docs_url=flavor.docs_url,
                sdk_docs_url=flavor.sdk_docs_url,
                is_custom=flavor.is_custom,
            )
            session.add(new_flavor)
            session.commit()

            return new_flavor.to_model(
                include_metadata=True, include_resources=True
            )

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.

Raises:

Type	Description
EntityExistsError	If a model with the given name already exists.

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

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

    Args:
        model: the Model to be created.

    Returns:
        The newly created model.

    Raises:
        EntityExistsError: If a model with the given name already exists.
    """
    validate_name(model)
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=model, session=session)

        self._verify_name_uniqueness(
            resource=model,
            schema=ModelSchema,
            session=session,
        )

        model_schema = ModelSchema.from_request(model)
        session.add(model_schema)

        try:
            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 model {model.name}: "
                "A model with this name already exists."
            )

        self._attach_tags_to_resources(
            tags=model.tags,
            resources=model_schema,
            session=session,
        )

        session.refresh(model_schema)

        return model_schema.to_model(
            include_metadata=True, include_resources=True
        )

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/sql_zen_store.py

@track_decorator(AnalyticsEvent.CREATED_MODEL_VERSION)
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_model_version(model_version=model_version)

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/sql_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.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=model_version_artifact_link, session=session
        )

        # If the link already exists, return it
        existing_model_version_artifact_link = session.exec(
            select(ModelVersionArtifactSchema)
            .where(
                ModelVersionArtifactSchema.model_version_id
                == model_version_artifact_link.model_version
            )
            .where(
                ModelVersionArtifactSchema.artifact_version_id
                == model_version_artifact_link.artifact_version,
            )
        ).first()
        if existing_model_version_artifact_link is not None:
            return existing_model_version_artifact_link.to_model()

        model_version_artifact_link_schema = (
            ModelVersionArtifactSchema.from_request(
                model_version_artifact_request=model_version_artifact_link,
            )
        )
        session.add(model_version_artifact_link_schema)
        session.commit()
        return model_version_artifact_link_schema.to_model(
            include_metadata=True, include_resources=True
        )

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/sql_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.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=model_version_pipeline_run_link, session=session
        )

        # If the link already exists, return it
        existing_model_version_pipeline_run_link = session.exec(
            select(ModelVersionPipelineRunSchema)
            .where(
                ModelVersionPipelineRunSchema.model_version_id
                == model_version_pipeline_run_link.model_version
            )
            .where(
                ModelVersionPipelineRunSchema.pipeline_run_id
                == model_version_pipeline_run_link.pipeline_run,
            )
        ).first()
        if existing_model_version_pipeline_run_link is not None:
            return existing_model_version_pipeline_run_link.to_model()

        # Otherwise, create a new link
        model_version_pipeline_run_link_schema = (
            ModelVersionPipelineRunSchema.from_request(
                model_version_pipeline_run_link
            )
        )
        session.add(model_version_pipeline_run_link_schema)
        session.commit()
        return model_version_pipeline_run_link_schema.to_model(
            include_metadata=True, include_resources=True
        )

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.

Raises:

Type	Description
EntityExistsError	If an identical pipeline already exists.

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

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

    Args:
        pipeline: The pipeline to create.

    Returns:
        The newly created pipeline.

    Raises:
        EntityExistsError: If an identical pipeline already exists.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=pipeline, session=session)

        new_pipeline = PipelineSchema.from_request(pipeline)

        session.add(new_pipeline)
        try:
            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 pipeline in project "
                f"'{pipeline.project}': A pipeline with the name "
                f"{pipeline.name} already exists."
            )
        session.refresh(new_pipeline)

        self._attach_tags_to_resources(
            tags=pipeline.tags,
            resources=new_pipeline,
            session=session,
        )

        session.refresh(new_pipeline)

        return new_pipeline.to_model(
            include_metadata=True, include_resources=True
        )

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/sql_zen_store.py

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

    Args:
        project: The project to create.

    Returns:
        The newly created project.
    """
    with Session(self.engine) as session:
        # Check if project with the given name already exists
        self._verify_name_uniqueness(
            resource=project,
            schema=ProjectSchema,
            session=session,
        )

        # Create the project
        new_project = ProjectSchema.from_request(project)
        session.add(new_project)
        session.commit()

        # Explicitly refresh the new_project schema
        session.refresh(new_project)

        project_model = new_project.to_model(
            include_metadata=True, include_resources=True
        )

        self._update_onboarding_state(
            completed_steps={OnboardingStep.PROJECT_CREATED},
            session=session,
        )

    return project_model

create_run_metadata(run_metadata: RunMetadataRequest) -> None

Creates run metadata.

Parameters:

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

Returns:

Type	Description
None	The created run metadata.

Raises:

Type	Description
RuntimeError	If the resource type is not supported.

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

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

    Args:
        run_metadata: The run metadata to create.

    Returns:
        The created run metadata.

    Raises:
        RuntimeError: If the resource type is not supported.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=run_metadata, session=session
        )

        self._get_reference_schema_by_id(
            resource=run_metadata,
            reference_schema=StackComponentSchema,
            reference_id=run_metadata.stack_component_id,
            session=session,
        )

        for resource in run_metadata.resources:
            reference_schema: Type[BaseSchema]
            if resource.type == MetadataResourceTypes.PIPELINE_RUN:
                reference_schema = PipelineRunSchema
            elif resource.type == MetadataResourceTypes.STEP_RUN:
                reference_schema = StepRunSchema
            elif resource.type == MetadataResourceTypes.ARTIFACT_VERSION:
                reference_schema = ArtifactVersionSchema
            elif resource.type == MetadataResourceTypes.MODEL_VERSION:
                reference_schema = ModelVersionSchema
            elif resource.type == MetadataResourceTypes.SCHEDULE:
                reference_schema = ScheduleSchema
            else:
                raise RuntimeError(
                    f"Unknown resource type: {resource.type}"
                )

            self._get_reference_schema_by_id(
                resource=run_metadata,
                reference_schema=reference_schema,
                reference_id=resource.id,
                session=session,
            )

        if run_metadata.resources:
            for key, value in run_metadata.values.items():
                type_ = run_metadata.types[key]

                run_metadata_schema = RunMetadataSchema(
                    project_id=run_metadata.project,
                    user_id=run_metadata.user,
                    stack_component_id=run_metadata.stack_component_id,
                    key=key,
                    value=json.dumps(value),
                    type=type_,
                    publisher_step_id=run_metadata.publisher_step_id,
                )

                session.add(run_metadata_schema)
                session.commit()

                for resource in run_metadata.resources:
                    rm_resource_link = RunMetadataResourceSchema(
                        resource_id=resource.id,
                        resource_type=resource.type.value,
                        run_metadata_id=run_metadata_schema.id,
                    )
                    session.add(rm_resource_link)
                    session.commit()
    return None

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.

Raises:

Type	Description
ValueError	If trying to create a step run with a retried status.
EntityExistsError	if the step run already exists or a log entry with the same source already exists within the scope of the same step.
IllegalOperationError	if the pipeline run is stopped or stopping.

Source code in src/zenml/zen_stores/sql_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.

    Raises:
        ValueError: If trying to create a step run with a retried status.
        EntityExistsError: if the step run already exists or a log entry
            with the same source already exists within the scope of the
            same step.
        IllegalOperationError: if the pipeline run is stopped or stopping.
    """
    if step_run.status in {
        ExecutionStatus.RETRIED,
        ExecutionStatus.RETRYING,
    }:
        raise ValueError(
            "Retrying/retried status can not be set manually."
        )

    with Session(self.engine) as session:
        self._set_request_user_id(request_model=step_run, session=session)

        # Check if the pipeline run exists
        run = self._get_reference_schema_by_id(
            resource=step_run,
            reference_schema=PipelineRunSchema,
            reference_id=step_run.pipeline_run_id,
            session=session,
        )

        # Validate pipeline status before creating step
        if run.status in [
            ExecutionStatus.STOPPING,
            ExecutionStatus.STOPPED,
        ]:
            raise IllegalOperationError(
                f"Cannot create step '{step_run.name}' for pipeline in "
                f"{run.status} state. Pipeline run ID: {step_run.pipeline_run_id}"
            )

        if run.status == ExecutionStatus.FAILED:
            execution_mode = (
                run.get_pipeline_configuration().execution_mode
            )

            if execution_mode != ExecutionMode.CONTINUE_ON_FAILURE:
                raise IllegalOperationError(
                    f"Cannot creat step '{step_run.name}' for the run '{run.name}'"
                    "because the run is in a FAILED state and the execution mode is"
                    f"{execution_mode}."
                )

        self._get_reference_schema_by_id(
            resource=step_run,
            reference_schema=StepRunSchema,
            reference_id=step_run.original_step_run_id,
            session=session,
            reference_type="original step run",
        )
        step_config = run.get_step_configuration(step_name=step_run.name)

        # Release the read locks of the previous two queries before we
        # try to acquire more exclusive locks
        session.commit()

        # Acquire exclusive lock on the pipeline run to prevent deadlocks
        # during insertion
        session.exec(
            select(PipelineRunSchema.id)
            .with_for_update()
            .where(PipelineRunSchema.id == step_run.pipeline_run_id)
        )

        existing_step_runs = session.exec(
            select(StepRunSchema)
            .options(
                load_only(
                    jl_arg(StepRunSchema.status),
                    jl_arg(StepRunSchema.model_version_id),
                )
            )
            .where(
                col(StepRunSchema.pipeline_run_id)
                == step_run.pipeline_run_id
            )
            .where(col(StepRunSchema.name) == step_run.name)
            .order_by(desc(StepRunSchema.version))
        ).all()

        if any(
            ExecutionStatus(sr.status).is_successful
            for sr in existing_step_runs
        ):
            # This step already completed successfully
            raise EntityExistsError(
                f"Unable to create step `{step_run.name}`: A successful "
                f"step with this name already exists in the pipeline run "
                f"with ID '{step_run.pipeline_run_id}'."
            )

        retry_config = step_config.config.retry
        max_retries = retry_config.max_retries if retry_config else 0

        if len(existing_step_runs) > max_retries:
            raise EntityExistsError(
                f"Unable to create step `{step_run.name}`: The step has "
                f"exceeded the maximum number of retries."
            )

        if existing_step_runs:
            # Update all existing step runs to retried.
            # TODO: Once we have the health check, this should probably also
            # cancel the existing step runs in case they're still running?
            session.execute(
                update(StepRunSchema)
                .where(
                    col(StepRunSchema.pipeline_run_id)
                    == step_run.pipeline_run_id
                )
                .where(col(StepRunSchema.name) == step_run.name)
                .values(
                    status=ExecutionStatus.RETRIED.value,
                    end_time=func.coalesce(
                        StepRunSchema.end_time,
                        func.now(),
                    ),
                )
            )

        is_retriable = len(existing_step_runs) < max_retries
        if is_retriable and step_run.status == ExecutionStatus.FAILED:
            # This step will be retried by the orchestrator, so we
            # set its status accordingly.
            step_run.status = ExecutionStatus.RETRYING

        step_schema = StepRunSchema.from_request(
            step_run,
            snapshot_id=run.snapshot_id,
            version=len(existing_step_runs) + 1,
            # TODO: This isn't actually guaranteed to be correct, how
            # do we handle these cases? E.g. if the step on kubernetes
            # is retried during startup, it will not actually create X
            # step runs. Or if it doesn't reach the point in code where
            # the step run is created.
            is_retriable=is_retriable,
        )

        session.add(step_schema)
        try:
            session.commit()
        except IntegrityError:
            raise EntityExistsError(
                f"Unable to create step `{step_run.name}`: A step with "
                f"this name already exists in the pipeline run with ID "
                f"'{step_run.pipeline_run_id}'."
            )

        # Add logs entry for the step if exists
        if step_run.logs is not None:
            self._get_reference_schema_by_id(
                resource=step_run,
                reference_schema=StackComponentSchema,
                reference_id=step_run.logs.artifact_store_id,
                session=session,
                reference_type="logs artifact store",
            )

            log_entry = LogsSchema(
                uri=step_run.logs.uri,
                # TODO: Remove fallback when not supporting
                # clients <0.84.0 anymore
                source=step_run.logs.source or "execution",
                step_run_id=step_schema.id,
                artifact_store_id=step_run.logs.artifact_store_id,
            )
            try:
                session.add(log_entry)
                session.commit()
            except IntegrityError:
                session.rollback()
                raise EntityExistsError(
                    "Unable to create log entry: A log entry with this "
                    f"source '{step_run.logs.source}' already exists "
                    f"within the scope of the same step '{step_schema.id}'."
                )
        # If cached, attach metadata of the original step
        if (
            step_run.status == ExecutionStatus.CACHED
            and step_run.original_step_run_id is not None
        ):
            original_metadata_links = session.exec(
                select(RunMetadataResourceSchema)
                .where(
                    RunMetadataResourceSchema.run_metadata_id
                    == RunMetadataSchema.id
                )
                .where(
                    RunMetadataResourceSchema.resource_id
                    == step_run.original_step_run_id
                )
                .where(
                    RunMetadataResourceSchema.resource_type
                    == MetadataResourceTypes.STEP_RUN
                )
                .where(
                    RunMetadataSchema.publisher_step_id
                    == step_run.original_step_run_id
                )
            ).all()

            # Create new links in a batch
            new_links = [
                RunMetadataResourceSchema(
                    resource_id=step_schema.id,
                    resource_type=link.resource_type,
                    run_metadata_id=link.run_metadata_id,
                )
                for link in original_metadata_links
            ]

            if new_links:
                session.add_all(new_links)
                session.commit()
                session.refresh(step_schema, ["run_metadata"])

        if step_run.status == ExecutionStatus.CACHED:
            from zenml.utils.tag_utils import Tag

            cascading_tags = [
                tag
                for tag in run.get_pipeline_configuration().tags or []
                if isinstance(tag, Tag) and tag.cascade
            ]

            if cascading_tags:
                output_artifact_ids = [
                    id for ids in step_run.outputs.values() for id in ids
                ]
                output_artifacts = list(
                    session.exec(
                        select(ArtifactVersionSchema).where(
                            col(ArtifactVersionSchema.id).in_(
                                output_artifact_ids
                            )
                        )
                    ).all()
                )
                self._attach_tags_to_resources(
                    cascading_tags,
                    resources=output_artifacts,
                    session=session,
                )

        session.commit()

        for upstream_step in step_config.spec.upstream_steps:
            self._set_run_step_parent_step(
                child_step_run=step_schema,
                parent_step_name=upstream_step,
                session=session,
            )

        # Save input artifact IDs into the database.
        for input_name, artifact_version_ids in step_run.inputs.items():
            for artifact_version_id in artifact_version_ids:
                if step_run.original_step_run_id:
                    # This is a cached step run, for which the input
                    # artifacts might include manually loaded artifacts
                    # which can not be inferred from the step config. In
                    # this case, we check the input type of the artifact
                    # for the original step run.
                    input_type = self._get_step_run_input_type_from_cached_step_run(
                        input_name=input_name,
                        artifact_version_id=artifact_version_id,
                        cached_step_run_id=step_run.original_step_run_id,
                        session=session,
                    )
                else:
                    # This is a non-cached step run, which means all input
                    # artifacts we receive at creation time are inputs that
                    # are defined in the step config.
                    input_type = self._get_step_run_input_type_from_config(
                        input_name=input_name,
                        step_config=step_config.config,
                        step_spec=step_config.spec,
                    )
                self._set_run_step_input_artifact(
                    step_run=step_schema,
                    artifact_version_id=artifact_version_id,
                    name=input_name,
                    input_type=input_type,
                    session=session,
                )

        # Save output artifact IDs into the database.
        for name, artifact_version_ids in step_run.outputs.items():
            for artifact_version_id in artifact_version_ids:
                self._set_run_step_output_artifact(
                    step_run=step_schema,
                    artifact_version_id=artifact_version_id,
                    name=name,
                    session=session,
                )

        session.commit()

        if step_run.status != ExecutionStatus.RUNNING:
            self._update_pipeline_run_status(
                pipeline_run_id=step_run.pipeline_run_id, session=session
            )

        session.commit()
        session.refresh(
            step_schema, ["input_artifacts", "output_artifacts"]
        )

        if existing_step_runs:
            model_version_id = existing_step_runs[-1].model_version_id
        else:
            model_version_id = self._get_or_create_model_version_for_run(
                step_schema
            )

        if model_version_id:
            step_schema.model_version_id = model_version_id
            session.add(step_schema)
            session.commit()

            self.create_model_version_pipeline_run_link(
                ModelVersionPipelineRunRequest(
                    model_version=model_version_id,
                    pipeline_run=step_schema.pipeline_run_id,
                )
            )
            session.refresh(step_schema)

        return step_schema.to_model(
            include_metadata=True, include_resources=True
        )

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/sql_zen_store.py

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

    Args:
        template: The template to create.

    Returns:
        The newly created template.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=template, session=session)

        self._verify_name_uniqueness(
            resource=template,
            schema=RunTemplateSchema,
            session=session,
        )

        snapshot = self._get_reference_schema_by_id(
            resource=template,
            reference_schema=PipelineSnapshotSchema,
            reference_id=template.source_snapshot_id,
            session=session,
        )

        template_utils.validate_snapshot_is_templatable(snapshot)

        template_schema = RunTemplateSchema.from_request(request=template)

        if not template.hidden:
            # Also update the name and description of the underlying
            # snapshot
            snapshot.name = template.name
            snapshot.description = template.description
            session.add(snapshot)

        session.add(template_schema)
        session.commit()
        session.refresh(template_schema)

        self._attach_tags_to_resources(
            tags=template.tags,
            resources=template_schema,
            session=session,
        )

        session.refresh(template_schema)

        return template_schema.to_model(
            include_metadata=True, include_resources=True
        )

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/sql_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.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=schedule, session=session)

        self._verify_name_uniqueness(
            resource=schedule,
            schema=ScheduleSchema,
            session=session,
        )

        self._get_reference_schema_by_id(
            resource=schedule,
            reference_schema=StackComponentSchema,
            reference_id=schedule.orchestrator_id,
            session=session,
            reference_type="orchestrator",
        )

        self._get_reference_schema_by_id(
            resource=schedule,
            reference_schema=PipelineSchema,
            reference_id=schedule.pipeline_id,
            session=session,
        )

        new_schedule = ScheduleSchema.from_request(schedule)
        session.add(new_schedule)
        session.commit()
        return new_schedule.to_model(
            include_metadata=True, include_resources=True
        )

create_secret(secret: SecretRequest) -> SecretResponse

Creates a new secret.

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

• a user cannot own two private secrets with the same name
• two public secrets cannot have the same name

Parameters:

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

Returns:

Type	Description
SecretResponse	The newly created secret.

Raises:

Type	Description
EntityExistsError	If a secret with the same name already exists in the same scope.

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

@track_decorator(AnalyticsEvent.CREATED_SECRET)
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:

    - a user cannot own two private secrets with the same name
    - two public secrets cannot have the same name

    Args:
        secret: The secret to create.

    Returns:
        The newly created secret.

    Raises:
        EntityExistsError: If a secret with the same name already exists in
            the same scope.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=secret, session=session)
        assert secret.user is not None
        # Check if a secret with the same name already exists in the same
        # scope.
        secret_exists, msg = self._check_sql_secret_scope(
            session=session,
            secret_name=secret.name,
            private=secret.private,
            user=secret.user,
        )
        if secret_exists:
            raise EntityExistsError(msg)

        new_secret = self._create_secret_schema(
            secret=secret,
            session=session,
        )

        secret_model = new_secret.to_model(
            include_metadata=True, include_resources=True
        )

        secret_model.set_secrets(secret.secret_values)
        return secret_model

create_service(service: ServiceRequest) -> ServiceResponse

Create a new service.

Parameters:

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

Returns:

Type	Description
ServiceResponse	The newly created service.

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

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

    Args:
        service: The service to create.

    Returns:
        The newly created service.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=service, session=session)
        # Check if a service with the given name already exists
        self._fail_if_service_with_config_exists(
            service_request=service,
            session=session,
        )

        self._get_reference_schema_by_id(
            resource=service,
            reference_schema=PipelineRunSchema,
            reference_id=service.pipeline_run_id,
            session=session,
        )

        self._get_reference_schema_by_id(
            resource=service,
            reference_schema=ModelVersionSchema,
            reference_id=service.model_version_id,
            session=session,
        )

        service_schema = ServiceSchema.from_request(service)
        logger.debug("Creating service: %s", service_schema)
        session.add(service_schema)
        session.commit()

        return service_schema.to_model(
            include_metadata=True, include_resources=True
        )

create_service_account(service_account: Union[ServiceAccountRequest, ServiceAccountInternalRequest]) -> ServiceAccountResponse

Creates a new service account.

Parameters:

Name	Type	Description	Default
service_account	Union[ServiceAccountRequest, ServiceAccountInternalRequest]	Service account to be created.	required

Returns:

Type	Description
ServiceAccountResponse	The newly created service account.

Raises:

Type	Description
EntityExistsError	If a user or service account with the given name already exists.

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

@track_decorator(AnalyticsEvent.CREATED_SERVICE_ACCOUNT)
def create_service_account(
    self,
    service_account: Union[
        ServiceAccountRequest, ServiceAccountInternalRequest
    ],
) -> ServiceAccountResponse:
    """Creates a new service account.

    Args:
        service_account: Service account to be created.

    Returns:
        The newly created service account.

    Raises:
        EntityExistsError: If a user or service account with the given name
            already exists.
    """
    with Session(self.engine) as session:
        # Check if a service account with the given name already
        # exists
        err_msg = (
            f"Unable to create service account with name "
            f"'{service_account.name}': Found existing service "
            "account with this name."
        )
        try:
            self._get_account_schema(
                service_account.name, session=session, service_account=True
            )
            raise EntityExistsError(err_msg)
        except KeyError:
            pass

        # Create the service account
        new_account = UserSchema.from_service_account_request(
            service_account
        )
        session.add(new_account)
        # on commit an IntegrityError may arise we let it bubble up
        session.commit()

        return new_account.to_service_account_model(
            include_metadata=True, include_resources=True
        )

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.

Raises:

Type	Description
Exception	If anything goes wrong during the creation of the service connector.

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

@track_decorator(AnalyticsEvent.CREATED_SERVICE_CONNECTOR)
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.

    Raises:
        Exception: If anything goes wrong during the creation of the
            service connector.
    """
    # If the connector type is locally available, we validate the request
    # against the connector type schema before storing it in the database
    if service_connector_registry.is_registered(service_connector.type):
        connector_type = (
            service_connector_registry.get_service_connector_type(
                service_connector.type
            )
        )
        service_connector.validate_and_configure_resources(
            connector_type=connector_type,
            resource_types=service_connector.resource_types,
            resource_id=service_connector.resource_id,
            configuration=service_connector.configuration,
        )

    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=service_connector, session=session
        )
        assert service_connector.user is not None

        self._verify_name_uniqueness(
            resource=service_connector,
            schema=ServiceConnectorSchema,
            session=session,
        )

        # Create the secret
        secret_id = self._create_connector_secret(
            connector_name=service_connector.name,
            secrets=service_connector.configuration.secrets,
            session=session,
        )
        try:
            # Create the service connector
            new_service_connector = ServiceConnectorSchema.from_request(
                service_connector,
                secret_id=secret_id,
            )

            session.add(new_service_connector)
            session.commit()

            session.refresh(new_service_connector)
        except Exception:
            # Delete the secret if it was created
            if secret_id:
                try:
                    self.delete_secret(secret_id)
                except Exception:
                    # Ignore any errors that occur while deleting the
                    # secret
                    pass

            raise

        connector = new_service_connector.to_model(
            include_metadata=True, include_resources=True
        )
        if new_service_connector.secret_id:
            secrets = self._get_secret_values(
                secret_id=new_service_connector.secret_id
            )
            connector.add_secrets(secrets)
        self._populate_connector_type(connector)

        return connector

create_snapshot(snapshot: PipelineSnapshotRequest) -> PipelineSnapshotResponse

Creates a new snapshot.

Parameters:

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

Raises:

Type	Description
EntityExistsError	If a snapshot with the same name already exists for the same pipeline.
RuntimeError	If the snapshot creation fails.

Returns:

Type	Description
PipelineSnapshotResponse	The newly created snapshot.

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

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

    Args:
        snapshot: The snapshot to create.

    Raises:
        EntityExistsError: If a snapshot with the same name already
            exists for the same pipeline.
        RuntimeError: If the snapshot creation fails.

    Returns:
        The newly created snapshot.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=snapshot, session=session)
        self._get_reference_schema_by_id(
            resource=snapshot,
            reference_schema=StackSchema,
            reference_id=snapshot.stack,
            session=session,
        )

        self._get_reference_schema_by_id(
            resource=snapshot,
            reference_schema=PipelineSchema,
            reference_id=snapshot.pipeline,
            session=session,
        )

        self._get_reference_schema_by_id(
            resource=snapshot,
            reference_schema=PipelineBuildSchema,
            reference_id=snapshot.build,
            session=session,
        )

        self._get_reference_schema_by_id(
            resource=snapshot,
            reference_schema=ScheduleSchema,
            reference_id=snapshot.schedule,
            session=session,
        )

        if snapshot.code_reference:
            self._get_reference_schema_by_id(
                resource=snapshot,
                reference_schema=CodeRepositorySchema,
                reference_id=snapshot.code_reference.code_repository,
                session=session,
            )

        self._get_reference_schema_by_id(
            resource=snapshot,
            reference_schema=RunTemplateSchema,
            reference_id=snapshot.template,
            session=session,
        )

        self._get_reference_schema_by_id(
            resource=snapshot,
            reference_schema=PipelineSnapshotSchema,
            reference_id=snapshot.source_snapshot,
            session=session,
        )

        if isinstance(snapshot.name, str):
            validate_name(snapshot)

            if snapshot.replace:
                self._remove_name_from_snapshot(
                    session=session,
                    pipeline_id=snapshot.pipeline,
                    name=snapshot.name,
                )

        code_reference_id = self._create_or_reuse_code_reference(
            session=session,
            project_id=snapshot.project,
            code_reference=snapshot.code_reference,
        )

        new_snapshot = PipelineSnapshotSchema.from_request(
            snapshot, code_reference_id=code_reference_id
        )

        try:
            session.add(new_snapshot)
            session.commit()
        except IntegrityError as e:
            session.rollback()
            if new_snapshot.name and self._snapshot_exists(
                session=session,
                pipeline_id=snapshot.pipeline,
                name=new_snapshot.name,
            ):
                raise EntityExistsError(
                    f"Snapshot with name `{new_snapshot.name}` already "
                    f"exists for pipeline `{snapshot.pipeline}`. If you "
                    "want to replace the existing snapshot, set the "
                    "`replace` flag to `True`."
                )
            else:
                raise RuntimeError("Snapshot creation failed.") from e

        for index, (step_name, step_configuration) in enumerate(
            snapshot.step_configurations.items()
        ):
            step_configuration_schema = StepConfigurationSchema(
                index=index,
                name=step_name,
                # Don't include the merged config in the step
                # configurations, we reconstruct it in the `to_model` method
                # using the pipeline configuration.
                config=step_configuration.model_dump_json(
                    exclude={"config"}
                ),
                snapshot_id=new_snapshot.id,
            )
            session.add(step_configuration_schema)
        session.commit()

        self._attach_tags_to_resources(
            tags=snapshot.tags,
            resources=new_snapshot,
            session=session,
        )
        session.refresh(new_snapshot)

        return new_snapshot.to_model(
            include_metadata=True, include_resources=True
        )

create_stack(stack: StackRequest) -> StackResponse

Register a full stack.

Parameters:

Name	Type	Description	Default
stack	StackRequest	The full stack configuration.	required

Returns:

Type	Description
StackResponse	The registered stack.

Raises:

Type	Description
ValueError	If the full stack creation fails, due to the corrupted input.
Exception	If the full stack creation fails, due to unforeseen errors.

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

@track_decorator(AnalyticsEvent.REGISTERED_STACK)
def create_stack(self, stack: StackRequest) -> StackResponse:
    """Register a full stack.

    Args:
        stack: The full stack configuration.

    Returns:
        The registered stack.

    Raises:
        ValueError: If the full stack creation fails, due to the corrupted
            input.
        Exception: If the full stack creation fails, due to unforeseen
            errors.
    """
    with Session(self.engine) as session:
        if isinstance(stack, DefaultStackRequest):
            # Set the user to None for default stacks
            stack.user = None
        else:
            self._set_request_user_id(request_model=stack, session=session)

        # For clean-up purposes, each created entity is tracked here
        service_connectors_created_ids: List[UUID] = []
        components_created_ids: List[UUID] = []

        try:
            # Validate the name of the new stack
            validate_name(stack)

            if stack.labels is None:
                stack.labels = {}

            # Service Connectors
            service_connectors: List[ServiceConnectorResponse] = []

            orchestrator_components = stack.components[
                StackComponentType.ORCHESTRATOR
            ]
            for orchestrator_component in orchestrator_components:
                if isinstance(orchestrator_component, UUID):
                    orchestrator = self.get_stack_component(
                        orchestrator_component,
                        hydrate=False,
                    )
                    need_to_generate_permanent_tokens = (
                        orchestrator.flavor_name.startswith("vm_")
                    )
                else:
                    need_to_generate_permanent_tokens = (
                        orchestrator_component.flavor.startswith("vm_")
                    )

            for connector_id_or_info in stack.service_connectors:
                # Fetch an existing service connector
                if isinstance(connector_id_or_info, UUID):
                    existing_service_connector = (
                        self.get_service_connector(
                            connector_id_or_info, expand_secrets=True
                        )
                    )
                    if need_to_generate_permanent_tokens:
                        if (
                            existing_service_connector.configuration.get(
                                "generate_temporary_tokens", None
                            )
                            is not False
                        ):
                            connector_config = existing_service_connector.configuration.plain
                            connector_config[
                                "generate_temporary_tokens"
                            ] = False
                            self.update_service_connector(
                                existing_service_connector.id,
                                ServiceConnectorUpdate(
                                    configuration=ServiceConnectorConfiguration(
                                        **connector_config
                                    )
                                ),
                            )
                    service_connectors.append(
                        self.get_service_connector(
                            connector_id_or_info, expand_secrets=True
                        )
                    )
                # Create a new service connector
                else:
                    connector_name = stack.name
                    connector_config = connector_id_or_info.configuration
                    connector_config[
                        "generate_temporary_tokens"
                    ] = not need_to_generate_permanent_tokens

                    while True:
                        try:
                            service_connector_request = ServiceConnectorRequest(
                                name=connector_name,
                                connector_type=connector_id_or_info.type,
                                auth_method=connector_id_or_info.auth_method,
                                configuration=ServiceConnectorConfiguration(
                                    **connector_config
                                ),
                                labels={
                                    k: str(v)
                                    for k, v in stack.labels.items()
                                },
                            )
                            service_connector_response = self.create_service_connector(
                                service_connector=service_connector_request
                            )
                            service_connectors.append(
                                service_connector_response
                            )
                            service_connectors_created_ids.append(
                                service_connector_response.id
                            )
                            break
                        except EntityExistsError:
                            connector_name = (
                                f"{stack.name}-{random_str(4)}".lower()
                            )
                            continue

            # Stack Components
            components_mapping: Dict[StackComponentType, List[UUID]] = {}
            for (
                component_type,
                components,
            ) in stack.components.items():
                for component_info in components:
                    # Fetch an existing component
                    if isinstance(component_info, UUID):
                        component = self.get_stack_component(
                            component_id=component_info
                        )
                    # Create a new component
                    else:
                        flavor_list = self.list_flavors(
                            flavor_filter_model=FlavorFilter(
                                name=component_info.flavor,
                                type=component_type,
                            )
                        )
                        if not len(flavor_list):
                            raise ValueError(
                                f"Flavor '{component_info.flavor}' not found "
                                f"for component type '{component_type}'."
                            )

                        flavor_model = flavor_list[0]

                        component_name = stack.name
                        while True:
                            try:
                                component_request = ComponentRequest(
                                    name=component_name,
                                    type=component_type,
                                    flavor=component_info.flavor,
                                    configuration=component_info.configuration,
                                    labels=stack.labels,
                                )
                                component = self.create_stack_component(
                                    component=component_request
                                )
                                components_created_ids.append(component.id)
                                break
                            except EntityExistsError:
                                component_name = (
                                    f"{stack.name}-{random_str(4)}".lower()
                                )
                                continue

                        if (
                            component_info.service_connector_index
                            is not None
                        ):
                            service_connector = service_connectors[
                                component_info.service_connector_index
                            ]

                            requirements = (
                                flavor_model.connector_requirements
                            )

                            if not requirements:
                                raise ValueError(
                                    f"The '{flavor_model.name}' implementation "
                                    "does not support using a service "
                                    "connector to connect to resources."
                                )

                            if component_info.service_connector_resource_id:
                                resource_id = component_info.service_connector_resource_id
                            else:
                                resource_id = None
                                resource_type = requirements.resource_type
                                if (
                                    requirements.resource_id_attr
                                    is not None
                                ):
                                    resource_id = (
                                        component_info.configuration.get(
                                            requirements.resource_id_attr
                                        )
                                    )

                            satisfied, msg = requirements.is_satisfied_by(
                                connector=service_connector,
                                component=component,
                            )

                            if not satisfied:
                                raise ValueError(
                                    "Please pick a connector that is "
                                    "compatible with the component flavor and "
                                    "try again.."
                                )

                            if not resource_id:
                                if service_connector.resource_id:
                                    resource_id = (
                                        service_connector.resource_id
                                    )
                                elif service_connector.supports_instances:
                                    raise ValueError(
                                        f"Multiple {resource_type} resources "
                                        "are available for the selected "
                                        "connector. Please use a `resource_id` "
                                        "to configure a "
                                        f"{resource_type} resource."
                                    )

                            component_update = ComponentUpdate(
                                connector=service_connector.id,
                                connector_resource_id=resource_id,
                            )
                            self.update_stack_component(
                                component_id=component.id,
                                component_update=component_update,
                            )

                    components_mapping[component_type] = [
                        component.id,
                    ]

            # Stack
            self._verify_name_uniqueness(
                resource=stack,
                schema=StackSchema,
                session=session,
            )

            component_ids = (
                [
                    component_id
                    for list_of_component_ids in components_mapping.values()
                    for component_id in list_of_component_ids
                ]
                if stack.components is not None
                else []
            )
            filters = [
                (StackComponentSchema.id == component_id)
                for component_id in component_ids
            ]

            defined_components = session.exec(
                select(StackComponentSchema).where(or_(*filters))
            ).all()

            new_stack_schema = StackSchema.from_request(
                request=stack,
                components=defined_components,
            )

            self._link_secrets_to_resource(
                resource=new_stack_schema,
                secrets=stack.secrets,
                session=session,
            )

            session.add(new_stack_schema)
            session.commit()
            session.refresh(new_stack_schema)

            for defined_component in defined_components:
                if (
                    defined_component.type
                    == StackComponentType.ORCHESTRATOR
                ):
                    if defined_component.flavor not in {
                        "local",
                        "local_docker",
                    }:
                        self._update_onboarding_state(
                            completed_steps={
                                OnboardingStep.STACK_WITH_REMOTE_ORCHESTRATOR_CREATED
                            },
                            session=session,
                        )
                if (
                    defined_component.type
                    == StackComponentType.ARTIFACT_STORE
                ):
                    if defined_component.flavor != "local":
                        self._update_onboarding_state(
                            completed_steps={
                                OnboardingStep.STACK_WITH_REMOTE_ARTIFACT_STORE_CREATED
                            },
                            session=session,
                        )

            return new_stack_schema.to_model(
                include_metadata=True, include_resources=True
            )

        except Exception:
            for component_id in components_created_ids:
                self.delete_stack_component(component_id=component_id)
            for service_connector_id in service_connectors_created_ids:
                self.delete_service_connector(
                    service_connector_id=service_connector_id
                )
            logger.error(
                "Stack creation has failed. Cleaned up the entities "
                "that are created in the process."
            )
            raise

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/sql_zen_store.py

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

    Args:
        component: The stack component to create.

    Returns:
        The created stack component.
    """
    validate_name(component)
    with Session(self.engine) as session:
        if isinstance(component, DefaultComponentRequest):
            # Set the user to None for default components
            component.user = None
        else:
            self._set_request_user_id(
                request_model=component, session=session
            )

        self._fail_if_component_with_name_type_exists(
            name=component.name,
            component_type=component.type,
            session=session,
        )

        is_default_stack_component = (
            component.name == DEFAULT_STACK_AND_COMPONENT_NAME
            and component.type
            in {
                StackComponentType.ORCHESTRATOR,
                StackComponentType.DEPLOYER,
                StackComponentType.ARTIFACT_STORE,
            }
        )
        # We have to skip the validation of the default components
        # as it creates a loop of initialization.
        if not is_default_stack_component:
            from zenml.stack.utils import validate_stack_component_config

            validate_stack_component_config(
                configuration_dict=component.configuration,
                flavor=component.flavor,
                component_type=component.type,
                zen_store=self,
                validate_custom_flavors=False,
            )

        service_connector = self._get_reference_schema_by_id(
            resource=component,
            reference_schema=ServiceConnectorSchema,
            reference_id=component.connector,
            session=session,
        )

        # warn about skypilot regions, if needed
        # TODO: this sooo does not belong here!
        if component.flavor in {"vm_gcp", "vm_azure"}:
            stack_deployment_class = get_stack_deployment_class(
                StackDeploymentProvider.GCP
                if component.flavor == "vm_gcp"
                else StackDeploymentProvider.AZURE
            )
            skypilot_regions = (
                stack_deployment_class.skypilot_default_regions().values()
            )
            if (
                component.configuration.get("region", None)
                and component.configuration["region"]
                not in skypilot_regions
            ):
                logger.warning(
                    f"Region `{component.configuration['region']}` is "
                    "not enabled in Skypilot by default. Supported regions "
                    f"by default are: {skypilot_regions}. Check the "
                    "Skypilot documentation to learn how to enable "
                    "regions rather than default ones. (If you have "
                    "already extended your configuration - "
                    "simply ignore this warning)"
                )

        # Create the component
        new_component = StackComponentSchema.from_request(
            request=component, service_connector=service_connector
        )

        self._link_secrets_to_resource(
            resource=new_component,
            secrets=component.secrets,
            session=session,
        )

        session.add(new_component)
        session.commit()

        session.refresh(new_component)

        return new_component.to_model(
            include_metadata=True, include_resources=True
        )

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/sql_zen_store.py

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

    Args:
        tag: the tag to be created.

    Returns:
        The newly created tag.
    """
    with Session(self.engine) as session:
        tag_schema = self._create_tag_schema(tag=tag, session=session)
        return tag_schema.to_model(
            include_metadata=True, include_resources=True
        )

create_tag_resource(tag_resource: TagResourceRequest) -> TagResourceResponse

Creates a new tag resource relationship.

Parameters:

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

Returns:

Type	Description
TagResourceResponse	The newly created tag resource relationship.

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

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

    Args:
        tag_resource: the tag resource relationship to be created.

    Returns:
        The newly created tag resource relationship.
    """
    return self.batch_create_tag_resource(tag_resources=[tag_resource])[0]

create_trigger(trigger: TriggerRequest) -> TriggerResponse

Creates a new trigger.

Parameters:

Name	Type	Description	Default
trigger	TriggerRequest	Trigger to be created.	required

Returns:

Type	Description
TriggerResponse	The newly created trigger.

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

@track_decorator(AnalyticsEvent.CREATED_TRIGGER)
def create_trigger(self, trigger: TriggerRequest) -> TriggerResponse:
    """Creates a new trigger.

    Args:
        trigger: Trigger to be created.

    Returns:
        The newly created trigger.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(request_model=trigger, session=session)

        # Verify that the trigger name is unique
        self._verify_name_uniqueness(
            resource=trigger,
            schema=TriggerSchema,
            session=session,
        )

        # Verify that the given action exists
        self._get_reference_schema_by_id(
            resource=trigger,
            reference_schema=ActionSchema,
            reference_id=trigger.action_id,
            session=session,
        )

        self._get_reference_schema_by_id(
            resource=trigger,
            reference_schema=EventSourceSchema,
            reference_id=trigger.event_source_id,
            session=session,
        )

        new_trigger = TriggerSchema.from_request(trigger)
        session.add(new_trigger)
        session.commit()
        session.refresh(new_trigger)

        return new_trigger.to_model(
            include_metadata=True, include_resources=True
        )

create_trigger_execution(trigger_execution: TriggerExecutionRequest) -> TriggerExecutionResponse

Create a trigger execution.

Parameters:

Name	Type	Description	Default
trigger_execution	TriggerExecutionRequest	The trigger execution to create.	required

Returns:

Type	Description
TriggerExecutionResponse	The created trigger execution.

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

def create_trigger_execution(
    self, trigger_execution: TriggerExecutionRequest
) -> TriggerExecutionResponse:
    """Create a trigger execution.

    Args:
        trigger_execution: The trigger execution to create.

    Returns:
        The created trigger execution.
    """
    with Session(self.engine) as session:
        self._set_request_user_id(
            request_model=trigger_execution, session=session
        )
        self._get_reference_schema_by_id(
            resource=trigger_execution,
            reference_schema=TriggerSchema,
            reference_id=trigger_execution.trigger,
            session=session,
        )
        new_execution = TriggerExecutionSchema.from_request(
            trigger_execution
        )
        session.add(new_execution)
        session.commit()
        session.refresh(new_execution)

        return new_execution.to_model(
            include_metadata=True, include_resources=True
        )

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.

Raises:

Type	Description
EntityExistsError	If a user or service account with the given name already exists.

Source code in src/zenml/zen_stores/sql_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.

    Raises:
        EntityExistsError: If a user or service account with the given name
            already exists.
    """
    with Session(self.engine) as session:
        # Check if a user account with the given name already exists
        err_msg = (
            f"Unable to create user with name '{user.name}': "
            f"Found an existing user account with this name."
        )
        try:
            self._get_account_schema(
                user.name,
                session=session,
                # Filter out service accounts
                service_account=False,
            )
            raise EntityExistsError(err_msg)
        except KeyError:
            pass

        # Create the user
        new_user = UserSchema.from_user_request(user)
        session.add(new_user)
        # on commit an IntegrityError may arise we let it bubble up
        session.commit()

        server_info = self.get_store_info()
        with AnalyticsContext() as context:
            context.user_id = new_user.id

            context.group(
                group_id=server_info.id,
                traits={
                    "server_id": server_info.id,
                    "version": server_info.version,
                    "deployment_type": str(server_info.deployment_type),
                    "database_type": str(server_info.database_type),
                },
            )

        return new_user.to_model(
            include_metadata=True, include_resources=True
        )

delete_action(action_id: UUID) -> None

Delete an action.

Parameters:

Name	Type	Description	Default
action_id	UUID	The ID of the action to delete.	required

Raises:

Type	Description
IllegalOperationError	If the action can't be deleted because it's used by triggers.

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

def delete_action(self, action_id: UUID) -> None:
    """Delete an action.

    Args:
        action_id: The ID of the action to delete.

    Raises:
        IllegalOperationError: If the action can't be deleted
            because it's used by triggers.
    """
    with Session(self.engine) as session:
        action = self._get_schema_by_id(
            resource_id=action_id,
            schema_class=ActionSchema,
            session=session,
        )

        # Prevent deletion of action if it is used by a trigger
        if action.triggers:
            raise IllegalOperationError(
                f"Unable to delete action with ID `{action_id}` "
                f"as it is used by {len(action.triggers)} triggers."
            )

        session.delete(action)
        session.commit()

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

Deletes all model version to artifact links.

Parameters:

Name	Type	Description	Default
model_version_id	UUID	ID of the model version containing the link.	required
only_links	bool	Whether to only delete the link to the artifact.	True

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

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

    Args:
        model_version_id: ID of the model version containing the link.
        only_links: Whether to only delete the link to the artifact.
    """
    with Session(self.engine) as session:
        if not only_links:
            artifact_version_ids = session.execute(
                select(
                    ModelVersionArtifactSchema.artifact_version_id
                ).where(
                    ModelVersionArtifactSchema.model_version_id
                    == model_version_id
                )
            ).fetchall()
            session.execute(
                delete(ArtifactVersionSchema).where(
                    col(ArtifactVersionSchema.id).in_(
                        [a[0] for a in artifact_version_ids]
                    )
                ),
            )
        session.execute(
            delete(ModelVersionArtifactSchema).where(
                ModelVersionArtifactSchema.model_version_id  # type: ignore[arg-type]
                == model_version_id
            )
        )

        session.commit()

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/sql_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.
    """
    with Session(self.engine) as session:
        api_key = self._get_api_key(
            service_account_id=service_account_id,
            api_key_name_or_id=api_key_name_or_id,
            session=session,
        )

        session.delete(api_key)
        session.commit()

delete_api_transaction(api_transaction_id: UUID) -> None

Delete an API transaction.

Parameters:

Name	Type	Description	Default
api_transaction_id	UUID	The ID of the API transaction to delete.	required

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

def delete_api_transaction(self, api_transaction_id: UUID) -> None:
    """Delete an API transaction.

    Args:
        api_transaction_id: The ID of the API transaction to delete.
    """
    with Session(self.engine) as session:
        session.execute(
            delete(ApiTransactionSchema).where(
                col(ApiTransactionSchema.id) == api_transaction_id
            )
        )
        session.commit()

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/sql_zen_store.py

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

    Args:
        artifact_id: The ID of the artifact to delete.
    """
    with Session(self.engine) as session:
        existing_artifact = self._get_schema_by_id(
            resource_id=artifact_id,
            schema_class=ArtifactSchema,
            session=session,
        )
        session.delete(existing_artifact)
        session.commit()

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/sql_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.
    """
    with Session(self.engine) as session:
        artifact_version = self._get_schema_by_id(
            resource_id=artifact_version_id,
            schema_class=ArtifactVersionSchema,
            session=session,
        )
        session.delete(artifact_version)
        session.commit()

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/sql_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.
    """
    with Session(self.engine) as session:
        existing_device = self._get_schema_by_id(
            resource_id=device_id,
            schema_class=OAuthDeviceSchema,
            session=session,
            resource_type="authorized device",
        )

        session.delete(existing_device)
        session.commit()

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/sql_zen_store.py

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

    Args:
        build_id: The ID of the build to delete.
    """
    with Session(self.engine) as session:
        # Check if build with the given ID exists
        build = self._get_schema_by_id(
            resource_id=build_id,
            schema_class=PipelineBuildSchema,
            session=session,
        )

        session.delete(build)
        session.commit()

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/sql_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.
    """
    with Session(self.engine) as session:
        existing_repo = self._get_schema_by_id(
            resource_id=code_repository_id,
            schema_class=CodeRepositorySchema,
            session=session,
        )

        session.delete(existing_repo)
        session.commit()

delete_curated_visualization(visualization_id: UUID) -> None

Delete a curated visualization.

Parameters:

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

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

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

    Args:
        visualization_id: The ID of the curated visualization to delete.
    """
    with Session(self.engine) as session:
        schema = self._get_schema_by_id(
            resource_id=visualization_id,
            schema_class=CuratedVisualizationSchema,
            session=session,
        )
        session.delete(schema)
        session.commit()

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/sql_zen_store.py

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

    Args:
        deployment_id: The ID of the deployment to delete.
    """
    with Session(self.engine) as session:
        deployment = self._get_schema_by_id(
            resource_id=deployment_id,
            schema_class=DeploymentSchema,
            session=session,
        )

        session.delete(deployment)
        session.commit()

delete_event_source(event_source_id: UUID) -> None

Delete an event_source.

Parameters:

Name	Type	Description	Default
event_source_id	UUID	The ID of the event_source to delete.	required

Raises:

Type	Description
IllegalOperationError	If the event source can't be deleted because it's used by triggers.

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

def delete_event_source(self, event_source_id: UUID) -> None:
    """Delete an event_source.

    Args:
        event_source_id: The ID of the event_source to delete.

    Raises:
        IllegalOperationError: If the event source can't be deleted
            because it's used by triggers.
    """
    with Session(self.engine) as session:
        event_source = self._get_schema_by_id(
            resource_id=event_source_id,
            schema_class=EventSourceSchema,
            session=session,
        )

        # Prevent deletion of event source if it is used by a trigger
        if event_source.triggers:
            raise IllegalOperationError(
                f"Unable to delete event_source with ID `{event_source_id}`"
                f" as it is used by {len(event_source.triggers)} triggers."
            )

        session.delete(event_source)
        session.commit()

delete_expired_authorized_devices() -> None

Deletes all expired OAuth 2.0 authorized devices.

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

def delete_expired_authorized_devices(self) -> None:
    """Deletes all expired OAuth 2.0 authorized devices."""
    with Session(self.engine) as session:
        expired_devices = session.exec(
            select(OAuthDeviceSchema).where(OAuthDeviceSchema.user is None)
        ).all()
        for device in expired_devices:
            # Delete devices that have expired
            if (
                device.expires is not None
                and device.expires < utc_now()
                and device.user_id is None
            ):
                session.delete(device)
        session.commit()

delete_flavor(flavor_id: UUID) -> None

Delete a flavor.

Parameters:

Name	Type	Description	Default
flavor_id	UUID	The id of the flavor to delete.	required

Raises:

Type	Description
IllegalOperationError	if the flavor is used by a stack component.

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

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

    Args:
        flavor_id: The id of the flavor to delete.

    Raises:
        IllegalOperationError: if the flavor is used by a stack component.
    """
    with Session(self.engine) as session:
        flavor_in_db = self._get_schema_by_id(
            resource_id=flavor_id,
            schema_class=FlavorSchema,
            session=session,
        )
        components_of_flavor = session.exec(
            select(StackComponentSchema).where(
                StackComponentSchema.flavor == flavor_in_db.name
            )
        ).all()
        if len(components_of_flavor) > 0:
            raise IllegalOperationError(
                f"Stack Component `{flavor_in_db.name}` of type "
                f"`{flavor_in_db.type} cannot be "
                f"deleted as it is used by "
                f"{len(components_of_flavor)} "
                f"components. Before deleting this "
                f"flavor, make sure to delete all "
                f"associated components."
            )
        else:
            session.delete(flavor_in_db)
            session.commit()

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/sql_zen_store.py

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

    Args:
        model_id: id of the model to be deleted.
    """
    with Session(self.engine) as session:
        model = self._get_schema_by_id(
            resource_id=model_id,
            schema_class=ModelSchema,
            session=session,
        )

        session.delete(model)
        session.commit()

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

Raises:

Type	Description
KeyError	specified ID or name not found.

Source code in src/zenml/zen_stores/sql_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.

    Raises:
        KeyError: specified ID or name not found.
    """
    with Session(self.engine) as session:
        query = select(ModelVersionSchema).where(
            ModelVersionSchema.id == model_version_id
        )
        model_version = session.exec(query).first()
        if model_version is None:
            raise KeyError(
                "Unable to delete model version with id "
                f"`{model_version_id}`: "
                "No model version with this id found."
            )
        session.delete(model_version)
        session.commit()

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

Raises:

Type	Description
KeyError	specified ID or name not found.

Source code in src/zenml/zen_stores/sql_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.

    Raises:
        KeyError: specified ID or name not found.
    """
    with Session(self.engine) as session:
        model_version = self.get_model_version(model_version_id)
        query = select(ModelVersionArtifactSchema).where(
            ModelVersionArtifactSchema.model_version_id == model_version.id
        )
        try:
            UUID(str(model_version_artifact_link_name_or_id))
            query = query.where(
                ModelVersionArtifactSchema.id
                == model_version_artifact_link_name_or_id
            )
        except ValueError:
            query = (
                query.where(
                    ModelVersionArtifactSchema.artifact_version_id
                    == ArtifactVersionSchema.id
                )
                .where(
                    ArtifactVersionSchema.artifact_id == ArtifactSchema.id
                )
                .where(
                    ArtifactSchema.name
                    == model_version_artifact_link_name_or_id
                )
            )

        model_version_artifact_link = session.exec(query).first()
        if model_version_artifact_link is None:
            raise KeyError(
                f"Unable to delete model version link with name or ID "
                f"`{model_version_artifact_link_name_or_id}`: "
                f"No model version link with this name found."
            )

        session.delete(model_version_artifact_link)
        session.commit()

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	name or 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

Raises:

Type	Description
KeyError	specified ID not found.

Source code in src/zenml/zen_stores/sql_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: name or 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.

    Raises:
        KeyError: specified ID not found.
    """
    with Session(self.engine) as session:
        model_version = self.get_model_version(
            model_version_id=model_version_id
        )
        query = select(ModelVersionPipelineRunSchema).where(
            ModelVersionPipelineRunSchema.model_version_id
            == model_version.id
        )
        try:
            UUID(str(model_version_pipeline_run_link_name_or_id))
            query = query.where(
                ModelVersionPipelineRunSchema.id
                == model_version_pipeline_run_link_name_or_id
            )
        except ValueError:
            query = query.where(
                ModelVersionPipelineRunSchema.pipeline_run_id
                == PipelineRunSchema.id
            ).where(
                PipelineRunSchema.name
                == model_version_pipeline_run_link_name_or_id
            )

        model_version_pipeline_run_link = session.exec(query).first()
        if model_version_pipeline_run_link is None:
            raise KeyError(
                f"Unable to delete model version link with name "
                f"`{model_version_pipeline_run_link_name_or_id}`: "
                f"No model version link with this name found."
            )

        session.delete(model_version_pipeline_run_link)
        session.commit()

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/sql_zen_store.py

def delete_pipeline(self, pipeline_id: UUID) -> None:
    """Deletes a pipeline.

    Args:
        pipeline_id: The ID of the pipeline to delete.
    """
    with Session(self.engine) as session:
        # Check if pipeline with the given ID exists
        pipeline = self._get_schema_by_id(
            resource_id=pipeline_id,
            schema_class=PipelineSchema,
            session=session,
        )

        session.delete(pipeline)
        session.commit()

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

Raises:

Type	Description
IllegalOperationError	If the project is the default project.

Source code in src/zenml/zen_stores/sql_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.

    Raises:
        IllegalOperationError: If the project is the default project.
    """
    with Session(self.engine) as session:
        # Check if project with the given name exists
        project = self._get_schema_by_name_or_id(
            object_name_or_id=project_name_or_id,
            schema_class=ProjectSchema,
            session=session,
        )
        if (
            self._default_project_enabled
            and project.name == self._default_project_name
        ):
            raise IllegalOperationError(
                "The default project cannot be deleted."
            )

        session.delete(project)
        session.commit()

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/sql_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.
    """
    with Session(self.engine) as session:
        # Check if pipeline run with the given ID exists
        existing_run = self._get_schema_by_id(
            resource_id=run_id,
            schema_class=PipelineRunSchema,
            session=session,
        )

        # Delete the pipeline run
        session.delete(existing_run)
        session.commit()

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/sql_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.
    """
    with Session(self.engine) as session:
        template = self._get_schema_by_id(
            resource_id=template_id,
            schema_class=RunTemplateSchema,
            session=session,
        )

        session.delete(template)
        # We set the reference of all snapshots to this template to null
        # manually as we can't have a foreign key there to avoid a cycle
        snapshots = session.exec(
            select(PipelineSnapshotSchema).where(
                PipelineSnapshotSchema.template_id == template_id
            )
        ).all()
        for snapshot in snapshots:
            snapshot.template_id