Skip to content

Stack Stores

zenml.stack_stores special

The stack store defines exactly where and how stacks are persisted across their life.

base_stack_store

BaseStackStore (ABC)

Base class for accessing data in ZenML Repository and new Service.

Source code in zenml/stack_stores/base_stack_store.py
class BaseStackStore(ABC):
    """Base class for accessing data in ZenML Repository and new Service."""

    def register_default_stack(self) -> None:
        """Populates the store with the default Stack.

        The default stack contains a local orchestrator,
        a local artifact store and a local SQLite metadata store.
        """

        # register the default stack
        stack = Stack.default_local_stack()
        metadata = self.register_stack(StackWrapper.from_stack(stack))
        metadata["store_type"] = self.type.value
        track_event(AnalyticsEvent.REGISTERED_STACK, metadata=metadata)

    # Public Interface:

    @property
    @abstractmethod
    def type(self) -> StoreType:
        """The type of stack store."""

    @abstractmethod
    def initialize(
        self,
        url: str,
        skip_default_stack: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> "BaseStackStore":
        """Initialize the store.

        Args:
            url: The URL of the store.
            skip_default_stack: If True, the creation of the default stack will
                be skipped.
            *args: Additional arguments to pass to the concrete store
                implementation.
            **kwargs: Additional keyword arguments to pass to the concrete
                store implementation.

        Returns:
            The initialized concrete store instance.
        """

        if not skip_default_stack and self.is_empty():

            logger.info("Initializing store...")
            self.register_default_stack()

        return self

    @property
    @abstractmethod
    def url(self) -> str:
        """Get the repository URL."""

    @staticmethod
    @abstractmethod
    def get_path_from_url(url: str) -> Optional[Path]:
        """Get the path from a URL, if it points or is backed by a local file.

        Args:
            url: The URL to get the path from.

        Returns:
            The local path backed by the URL, or None if the URL is not backed
            by a local file or directory
        """

    @staticmethod
    @abstractmethod
    def get_local_url(path: str) -> str:
        """Get a local URL for a given local path."""

    @staticmethod
    @abstractmethod
    def is_valid_url(url: str) -> bool:
        """Check if the given url is valid."""

    @abstractmethod
    def is_empty(self) -> bool:
        """Check if the store is empty (no stacks are configured).

        The implementation of this method should check if the store is empty
        without having to load all the stacks from the persistent storage.
        """

    @abstractmethod
    def get_stack_configuration(
        self, name: str
    ) -> Dict[StackComponentType, str]:
        """Fetches a stack configuration by name.

        Args:
            name: The name of the stack to fetch.

        Returns:
            Dict[StackComponentType, str] for the requested stack name.

        Raises:
            KeyError: If no stack exists for the given name.
        """

    @property
    @abstractmethod
    def stack_configurations(self) -> Dict[str, Dict[StackComponentType, str]]:
        """Configurations for all stacks registered in this stack store.

        Returns:
            Dictionary mapping stack names to Dict[StackComponentType, str]'s
        """

    @abstractmethod
    def register_stack_component(
        self,
        component: StackComponentWrapper,
    ) -> None:
        """Register a stack component.

        Args:
            component: The component to register.

        Raises:
            StackComponentExistsError: If a stack component with the same type
                and name already exists.
        """

    @abstractmethod
    def deregister_stack(self, name: str) -> None:
        """Delete a stack from storage.

        Args:
            name: The name of the stack to be deleted.
        """

    # Private interface (must be implemented, not to be called by user):

    @abstractmethod
    def _create_stack(
        self, name: str, stack_configuration: Dict[StackComponentType, str]
    ) -> None:
        """Add a stack to storage.

        Args:
            name: The name to save the stack as.
            stack_configuration: Dict[StackComponentType, str] to persist.
        """

    @abstractmethod
    def _get_component_flavor_and_config(
        self, component_type: StackComponentType, name: str
    ) -> Tuple[str, bytes]:
        """Fetch the flavor and configuration for a stack component.

        Args:
            component_type: The type of the component to fetch.
            name: The name of the component to fetch.

        Returns:
            Pair of (flavor, congfiguration) for stack component, as string and
            base64-encoded yaml document, respectively

        Raises:
            KeyError: If no stack component exists for the given type and name.
        """

    @abstractmethod
    def _get_stack_component_names(
        self, component_type: StackComponentType
    ) -> List[str]:
        """Get names of all registered stack components of a given type.

        Args:
            component_type: The type of the component to list names for.

        Returns:
            A list of names as strings.
        """

    @abstractmethod
    def _delete_stack_component(
        self, component_type: StackComponentType, name: str
    ) -> None:
        """Remove a StackComponent from storage.

        Args:
            component_type: The type of component to delete.
            name: Then name of the component to delete.
        """

    # Common code (user facing):

    @property
    def stacks(self) -> List[StackWrapper]:
        """All stacks registered in this stack store."""
        return [
            self._stack_from_dict(name, conf)
            for name, conf in self.stack_configurations.items()
        ]

    def get_stack(self, name: str) -> StackWrapper:
        """Fetch a stack by name.

        Args:
            name: The name of the stack to retrieve.

        Returns:
            StackWrapper instance if the stack exists.

        Raises:
            KeyError: If no stack exists for the given name.
        """
        return self._stack_from_dict(name, self.get_stack_configuration(name))

    def register_stack(self, stack: StackWrapper) -> Dict[str, str]:
        """Register a stack and its components.

        If any of the stacks' components aren't registered in the stack store
        yet, this method will try to register them as well.

        Args:
            stack: The stack to register.

        Returns:
            metadata dict for telemetry or logging.

        Raises:
            StackExistsError: If a stack with the same name already exists.
            StackComponentExistsError: If a component of the stack wasn't
                registered and a different component with the same name
                already exists.
        """
        try:
            self.get_stack(stack.name)
        except KeyError:
            pass
        else:
            raise StackExistsError(
                f"Unable to register stack with name '{stack.name}': Found "
                f"existing stack with this name."
            )

        def __check_component(
            component: StackComponentWrapper,
        ) -> Tuple[StackComponentType, str]:
            try:
                existing_component = self.get_stack_component(
                    component_type=component.type, name=component.name
                )
                if existing_component.uuid != component.uuid:
                    raise StackComponentExistsError(
                        f"Unable to register one of the stacks components: "
                        f"A component of type '{component.type}' and name "
                        f"'{component.name}' already exists."
                    )
            except KeyError:
                self.register_stack_component(component)
            return component.type, component.name

        stack_configuration = {
            typ: name for typ, name in map(__check_component, stack.components)
        }
        metadata = {c.type.value: c.flavor for c in stack.components}
        self._create_stack(stack.name, stack_configuration)
        return metadata

    def get_stack_component(
        self, component_type: StackComponentType, name: str
    ) -> StackComponentWrapper:
        """Get a registered stack component.

        Raises:
            KeyError: If no component with the requested type and name exists.
        """
        flavor, config = self._get_component_flavor_and_config(
            component_type, name=name
        )
        uuid = yaml.safe_load(base64.b64decode(config).decode())["uuid"]
        return StackComponentWrapper(
            type=component_type,
            flavor=flavor,
            name=name,
            uuid=uuid,
            config=config,
        )

    def get_stack_components(
        self, component_type: StackComponentType
    ) -> List[StackComponentWrapper]:
        """Fetches all registered stack components of the given type.

        Args:
            component_type: StackComponentType to list members of

        Returns:
            A list of StackComponentConfiguration instances.
        """
        return [
            self.get_stack_component(component_type=component_type, name=name)
            for name in self._get_stack_component_names(component_type)
        ]

    def deregister_stack_component(
        self, component_type: StackComponentType, name: str
    ) -> None:
        """Deregisters a stack component.

        Args:
            component_type: The type of the component to deregister.
            name: The name of the component to deregister.

        Raises:
            ValueError: if trying to deregister a component that's part
                of a stack.
        """
        for stack_name, stack_config in self.stack_configurations.items():
            if stack_config.get(component_type) == name:
                raise ValueError(
                    f"Unable to deregister stack component (type: "
                    f"{component_type}, name: {name}) that is part of a "
                    f"registered stack (stack name: '{stack_name}')."
                )
        self._delete_stack_component(component_type, name=name)

    # Common code (internal implementations, private):

    def _stack_from_dict(
        self, name: str, stack_configuration: Dict[StackComponentType, str]
    ) -> StackWrapper:
        """Build a StackWrapper from stored configurations"""
        stack_components = [
            self.get_stack_component(
                component_type=component_type, name=component_name
            )
            for component_type, component_name in stack_configuration.items()
        ]
        return StackWrapper(name=name, components=stack_components)
stack_configurations: Dict[str, Dict[zenml.enums.StackComponentType, str]] property readonly

Configurations for all stacks registered in this stack store.

Returns:

Type Description
Dict[str, Dict[zenml.enums.StackComponentType, str]]

Dictionary mapping stack names to Dict[StackComponentType, str]'s

stacks: List[zenml.stack_stores.models.stack_wrapper.StackWrapper] property readonly

All stacks registered in this stack store.

type: StoreType property readonly

The type of stack store.

url: str property readonly

Get the repository URL.

deregister_stack(self, name)

Delete a stack from storage.

Parameters:

Name Type Description Default
name str

The name of the stack to be deleted.

required
Source code in zenml/stack_stores/base_stack_store.py
@abstractmethod
def deregister_stack(self, name: str) -> None:
    """Delete a stack from storage.

    Args:
        name: The name of the stack to be deleted.
    """
deregister_stack_component(self, component_type, name)

Deregisters a stack component.

Parameters:

Name Type Description Default
component_type StackComponentType

The type of the component to deregister.

required
name str

The name of the component to deregister.

required

Exceptions:

Type Description
ValueError

if trying to deregister a component that's part of a stack.

Source code in zenml/stack_stores/base_stack_store.py
def deregister_stack_component(
    self, component_type: StackComponentType, name: str
) -> None:
    """Deregisters a stack component.

    Args:
        component_type: The type of the component to deregister.
        name: The name of the component to deregister.

    Raises:
        ValueError: if trying to deregister a component that's part
            of a stack.
    """
    for stack_name, stack_config in self.stack_configurations.items():
        if stack_config.get(component_type) == name:
            raise ValueError(
                f"Unable to deregister stack component (type: "
                f"{component_type}, name: {name}) that is part of a "
                f"registered stack (stack name: '{stack_name}')."
            )
    self._delete_stack_component(component_type, name=name)
get_local_url(path) staticmethod

Get a local URL for a given local path.

Source code in zenml/stack_stores/base_stack_store.py
@staticmethod
@abstractmethod
def get_local_url(path: str) -> str:
    """Get a local URL for a given local path."""
get_path_from_url(url) staticmethod

Get the path from a URL, if it points or is backed by a local file.

Parameters:

Name Type Description Default
url str

The URL to get the path from.

required

Returns:

Type Description
Optional[pathlib.Path]

The local path backed by the URL, or None if the URL is not backed by a local file or directory

Source code in zenml/stack_stores/base_stack_store.py
@staticmethod
@abstractmethod
def get_path_from_url(url: str) -> Optional[Path]:
    """Get the path from a URL, if it points or is backed by a local file.

    Args:
        url: The URL to get the path from.

    Returns:
        The local path backed by the URL, or None if the URL is not backed
        by a local file or directory
    """
get_stack(self, name)

Fetch a stack by name.

Parameters:

Name Type Description Default
name str

The name of the stack to retrieve.

required

Returns:

Type Description
StackWrapper

StackWrapper instance if the stack exists.

Exceptions:

Type Description
KeyError

If no stack exists for the given name.

Source code in zenml/stack_stores/base_stack_store.py
def get_stack(self, name: str) -> StackWrapper:
    """Fetch a stack by name.

    Args:
        name: The name of the stack to retrieve.

    Returns:
        StackWrapper instance if the stack exists.

    Raises:
        KeyError: If no stack exists for the given name.
    """
    return self._stack_from_dict(name, self.get_stack_configuration(name))
get_stack_component(self, component_type, name)

Get a registered stack component.

Exceptions:

Type Description
KeyError

If no component with the requested type and name exists.

Source code in zenml/stack_stores/base_stack_store.py
def get_stack_component(
    self, component_type: StackComponentType, name: str
) -> StackComponentWrapper:
    """Get a registered stack component.

    Raises:
        KeyError: If no component with the requested type and name exists.
    """
    flavor, config = self._get_component_flavor_and_config(
        component_type, name=name
    )
    uuid = yaml.safe_load(base64.b64decode(config).decode())["uuid"]
    return StackComponentWrapper(
        type=component_type,
        flavor=flavor,
        name=name,
        uuid=uuid,
        config=config,
    )
get_stack_components(self, component_type)

Fetches all registered stack components of the given type.

Parameters:

Name Type Description Default
component_type StackComponentType

StackComponentType to list members of

required

Returns:

Type Description
List[zenml.stack_stores.models.stack_component_wrapper.StackComponentWrapper]

A list of StackComponentConfiguration instances.

Source code in zenml/stack_stores/base_stack_store.py
def get_stack_components(
    self, component_type: StackComponentType
) -> List[StackComponentWrapper]:
    """Fetches all registered stack components of the given type.

    Args:
        component_type: StackComponentType to list members of

    Returns:
        A list of StackComponentConfiguration instances.
    """
    return [
        self.get_stack_component(component_type=component_type, name=name)
        for name in self._get_stack_component_names(component_type)
    ]
get_stack_configuration(self, name)

Fetches a stack configuration by name.

Parameters:

Name Type Description Default
name str

The name of the stack to fetch.

required

Returns:

Type Description
Dict[zenml.enums.StackComponentType, str]

Dict[StackComponentType, str] for the requested stack name.

Exceptions:

Type Description
KeyError

If no stack exists for the given name.

Source code in zenml/stack_stores/base_stack_store.py
@abstractmethod
def get_stack_configuration(
    self, name: str
) -> Dict[StackComponentType, str]:
    """Fetches a stack configuration by name.

    Args:
        name: The name of the stack to fetch.

    Returns:
        Dict[StackComponentType, str] for the requested stack name.

    Raises:
        KeyError: If no stack exists for the given name.
    """
initialize(self, url, skip_default_stack=False, *args, **kwargs)

Initialize the store.

Parameters:

Name Type Description Default
url str

The URL of the store.

required
skip_default_stack bool

If True, the creation of the default stack will be skipped.

False
*args Any

Additional arguments to pass to the concrete store implementation.

()
**kwargs Any

Additional keyword arguments to pass to the concrete store implementation.

{}

Returns:

Type Description
BaseStackStore

The initialized concrete store instance.

Source code in zenml/stack_stores/base_stack_store.py
@abstractmethod
def initialize(
    self,
    url: str,
    skip_default_stack: bool = False,
    *args: Any,
    **kwargs: Any,
) -> "BaseStackStore":
    """Initialize the store.

    Args:
        url: The URL of the store.
        skip_default_stack: If True, the creation of the default stack will
            be skipped.
        *args: Additional arguments to pass to the concrete store
            implementation.
        **kwargs: Additional keyword arguments to pass to the concrete
            store implementation.

    Returns:
        The initialized concrete store instance.
    """

    if not skip_default_stack and self.is_empty():

        logger.info("Initializing store...")
        self.register_default_stack()

    return self
is_empty(self)

Check if the store is empty (no stacks are configured).

The implementation of this method should check if the store is empty without having to load all the stacks from the persistent storage.

Source code in zenml/stack_stores/base_stack_store.py
@abstractmethod
def is_empty(self) -> bool:
    """Check if the store is empty (no stacks are configured).

    The implementation of this method should check if the store is empty
    without having to load all the stacks from the persistent storage.
    """
is_valid_url(url) staticmethod

Check if the given url is valid.

Source code in zenml/stack_stores/base_stack_store.py
@staticmethod
@abstractmethod
def is_valid_url(url: str) -> bool:
    """Check if the given url is valid."""
register_default_stack(self)

Populates the store with the default Stack.

The default stack contains a local orchestrator, a local artifact store and a local SQLite metadata store.

Source code in zenml/stack_stores/base_stack_store.py
def register_default_stack(self) -> None:
    """Populates the store with the default Stack.

    The default stack contains a local orchestrator,
    a local artifact store and a local SQLite metadata store.
    """

    # register the default stack
    stack = Stack.default_local_stack()
    metadata = self.register_stack(StackWrapper.from_stack(stack))
    metadata["store_type"] = self.type.value
    track_event(AnalyticsEvent.REGISTERED_STACK, metadata=metadata)
register_stack(self, stack)

Register a stack and its components.

If any of the stacks' components aren't registered in the stack store yet, this method will try to register them as well.

Parameters:

Name Type Description Default
stack StackWrapper

The stack to register.

required

Returns:

Type Description
Dict[str, str]

metadata dict for telemetry or logging.

Exceptions:

Type Description
StackExistsError

If a stack with the same name already exists.

StackComponentExistsError

If a component of the stack wasn't registered and a different component with the same name already exists.

Source code in zenml/stack_stores/base_stack_store.py
def register_stack(self, stack: StackWrapper) -> Dict[str, str]:
    """Register a stack and its components.

    If any of the stacks' components aren't registered in the stack store
    yet, this method will try to register them as well.

    Args:
        stack: The stack to register.

    Returns:
        metadata dict for telemetry or logging.

    Raises:
        StackExistsError: If a stack with the same name already exists.
        StackComponentExistsError: If a component of the stack wasn't
            registered and a different component with the same name
            already exists.
    """
    try:
        self.get_stack(stack.name)
    except KeyError:
        pass
    else:
        raise StackExistsError(
            f"Unable to register stack with name '{stack.name}': Found "
            f"existing stack with this name."
        )

    def __check_component(
        component: StackComponentWrapper,
    ) -> Tuple[StackComponentType, str]:
        try:
            existing_component = self.get_stack_component(
                component_type=component.type, name=component.name
            )
            if existing_component.uuid != component.uuid:
                raise StackComponentExistsError(
                    f"Unable to register one of the stacks components: "
                    f"A component of type '{component.type}' and name "
                    f"'{component.name}' already exists."
                )
        except KeyError:
            self.register_stack_component(component)
        return component.type, component.name

    stack_configuration = {
        typ: name for typ, name in map(__check_component, stack.components)
    }
    metadata = {c.type.value: c.flavor for c in stack.components}
    self._create_stack(stack.name, stack_configuration)
    return metadata
register_stack_component(self, component)

Register a stack component.

Parameters:

Name Type Description Default
component StackComponentWrapper

The component to register.

required

Exceptions:

Type Description
StackComponentExistsError

If a stack component with the same type and name already exists.

Source code in zenml/stack_stores/base_stack_store.py
@abstractmethod
def register_stack_component(
    self,
    component: StackComponentWrapper,
) -> None:
    """Register a stack component.

    Args:
        component: The component to register.

    Raises:
        StackComponentExistsError: If a stack component with the same type
            and name already exists.
    """

local_stack_store

LocalStackStore (BaseStackStore)

Source code in zenml/stack_stores/local_stack_store.py
class LocalStackStore(BaseStackStore):
    def initialize(
        self,
        url: str,
        *args: Any,
        stack_data: Optional[StackStoreModel] = None,
        **kwargs: Any,
    ) -> "LocalStackStore":
        """Initializes a local stack store instance.

        Args:
            url: URL of local directory of the repository to use for
                stack storage.
            stack_data: optional stack data store object to pre-populate the
                stack store with.
            args: additional positional arguments (ignored).
            kwargs: additional keyword arguments (ignored).

        Returns:
            The initialized stack store instance.
        """
        if not self.is_valid_url(url):
            raise ValueError(f"Invalid URL for local store: {url}")

        self._root = self.get_path_from_url(url)
        self._url = f"file://{self._root}"
        fileio.create_dir_recursive_if_not_exists(str(self._root))

        if stack_data is not None:
            self.__store = stack_data
            self._write_store()
        elif fileio.file_exists(self._store_path()):
            config_dict = yaml_utils.read_yaml(self._store_path())
            self.__store = StackStoreModel.parse_obj(config_dict)
        else:
            self.__store = StackStoreModel.empty_store()
            self._write_store()

        super().initialize(url, *args, **kwargs)
        return self

    # Public interface implementations:

    @property
    def type(self) -> StoreType:
        """The type of stack store."""
        return StoreType.LOCAL

    @property
    def url(self) -> str:
        """URL of the repository."""
        return self._url

    @staticmethod
    def get_path_from_url(url: str) -> Optional[Path]:
        """Get the path from a URL.

        Args:
            url: The URL to get the path from.

        Returns:
            The path from the URL.
        """
        if not LocalStackStore.is_valid_url(url):
            raise ValueError(f"Invalid URL for local store: {url}")
        url = url.replace("file://", "")
        return Path(url)

    @staticmethod
    def get_local_url(path: str) -> str:
        """Get a local URL for a given local path."""
        return f"file://{path}"

    @staticmethod
    def is_valid_url(url: str) -> bool:
        """Check if the given url is a valid local path."""
        scheme = re.search("^([a-z0-9]+://)", url)
        return not scheme or scheme.group() == "file://"

    def is_empty(self) -> bool:
        """Check if the stack store is empty."""
        return len(self.__store.stacks) == 0

    def get_stack_configuration(
        self, name: str
    ) -> Dict[StackComponentType, str]:
        """Fetches a stack configuration by name.

        Args:
            name: The name of the stack to fetch.

        Returns:
            Dict[StackComponentType, str] for the requested stack name.

        Raises:
            KeyError: If no stack exists for the given name.
        """
        logger.debug("Fetching stack with name '%s'.", name)
        if name not in self.__store.stacks:
            raise KeyError(
                f"Unable to find stack with name '{name}'. Available names: "
                f"{set(self.__store.stacks)}."
            )

        return self.__store.stacks[name]

    @property
    def stack_configurations(self) -> Dict[str, Dict[StackComponentType, str]]:
        """Configuration for all stacks registered in this stack store.

        Returns:
            Dictionary mapping stack names to Dict[StackComponentType, str]
        """
        return self.__store.stacks.copy()

    def register_stack_component(
        self,
        component: StackComponentWrapper,
    ) -> None:
        """Register a stack component.

        Args:
            component: The component to register.

        Raises:
            StackComponentExistsError: If a stack component with the same type
                and name already exists.
        """
        components = self.__store.stack_components[component.type]
        if component.name in components:
            raise StackComponentExistsError(
                f"Unable to register stack component (type: {component.type}) "
                f"with name '{component.name}': Found existing stack component "
                f"with this name."
            )

        # write the component configuration file
        component_config_path = self._get_stack_component_config_path(
            component_type=component.type, name=component.name
        )
        fileio.create_dir_recursive_if_not_exists(
            os.path.dirname(component_config_path)
        )
        write_file_contents_as_string(
            component_config_path,
            base64.b64decode(component.config).decode(),
        )

        # add the component to the stack store dict and write it to disk
        components[component.name] = component.flavor
        self._write_store()
        logger.info(
            "Registered stack component with type '%s' and name '%s'.",
            component.type,
            component.name,
        )

    def deregister_stack(self, name: str) -> None:
        """Remove a stack from storage.

        Args:
            name: The name of the stack to be deleted.
        """
        try:
            del self.__store.stacks[name]
            self._write_store()
            logger.info("Deregistered stack with name '%s'.", name)
        except KeyError:
            logger.warning(
                "Unable to deregister stack with name '%s': No stack exists "
                "with this name.",
                name,
            )

    # Private interface implementations:

    def _create_stack(
        self, name: str, stack_configuration: Dict[StackComponentType, str]
    ) -> None:
        """Add a stack to storage.

        Args:
            name: The name to save the stack as.
            stack_configuration: Dict[StackComponentType, str] to persist.
        """
        self.__store.stacks[name] = stack_configuration
        self._write_store()
        logger.info("Registered stack with name '%s'.", name)

    def _get_component_flavor_and_config(
        self, component_type: StackComponentType, name: str
    ) -> Tuple[str, bytes]:
        """Fetch the flavor and configuration for a stack component.

        Args:
            component_type: The type of the component to fetch.
            name: The name of the component to fetch.

        Returns:
            Pair of (flavor, congfiguration) for stack component, as string and
            base64-encoded yaml document, respectively

        Raises:
            KeyError: If no stack component exists for the given type and name.
        """
        components: Dict[str, str] = self.__store.stack_components[
            component_type
        ]
        if name not in components:
            raise KeyError(
                f"Unable to find stack component (type: {component_type}) "
                f"with name '{name}'. Available names: {set(components)}."
            )

        component_config_path = self._get_stack_component_config_path(
            component_type=component_type, name=name
        )
        flavor = components[name]
        config = base64.b64encode(
            read_file_contents_as_string(component_config_path).encode()
        )
        return flavor, config

    def _get_stack_component_names(
        self, component_type: StackComponentType
    ) -> List[str]:
        """Get names of all registered stack components of a given type."""
        return list(self.__store.stack_components[component_type])

    def _delete_stack_component(
        self, component_type: StackComponentType, name: str
    ) -> None:
        """Remove a StackComponent from storage.

        Args:
            component_type: The type of component to delete.
            name: Then name of the component to delete.
        """
        components = self.__store.stack_components[component_type]
        try:
            del components[name]
            self._write_store()
            logger.info(
                "Deregistered stack component (type: %s) with name '%s'.",
                component_type.value,
                name,
            )
        except KeyError:
            logger.warning(
                "Unable to deregister stack component (type: %s) with name "
                "'%s': No stack component exists with this name.",
                component_type.value,
                name,
            )
        component_config_path = self._get_stack_component_config_path(
            component_type=component_type, name=name
        )

        if fileio.file_exists(component_config_path):
            fileio.remove(component_config_path)

    # Implementation-specific internal methods:

    @property
    def root(self) -> Path:
        """The root directory of the stack store."""
        if not self._root:
            raise RuntimeError(
                "Local stack store has not been initialized. Call `initialize` "
                "before using the store."
            )
        return self._root

    def _get_stack_component_config_path(
        self, component_type: StackComponentType, name: str
    ) -> str:
        """Path to the configuration file of a stack component."""
        path = self.root / component_type.plural / f"{name}.yaml"
        return str(path)

    def _store_path(self) -> str:
        """Path to the repository configuration file."""
        return str(self.root / "stacks.yaml")

    def _write_store(self) -> None:
        """Writes the stack store yaml file."""
        config_dict = json.loads(self.__store.json())
        yaml_utils.write_yaml(self._store_path(), config_dict)
root: Path property readonly

The root directory of the stack store.

stack_configurations: Dict[str, Dict[zenml.enums.StackComponentType, str]] property readonly

Configuration for all stacks registered in this stack store.

Returns:

Type Description
Dict[str, Dict[zenml.enums.StackComponentType, str]]

Dictionary mapping stack names to Dict[StackComponentType, str]

type: StoreType property readonly

The type of stack store.

url: str property readonly

URL of the repository.

deregister_stack(self, name)

Remove a stack from storage.

Parameters:

Name Type Description Default
name str

The name of the stack to be deleted.

required
Source code in zenml/stack_stores/local_stack_store.py
def deregister_stack(self, name: str) -> None:
    """Remove a stack from storage.

    Args:
        name: The name of the stack to be deleted.
    """
    try:
        del self.__store.stacks[name]
        self._write_store()
        logger.info("Deregistered stack with name '%s'.", name)
    except KeyError:
        logger.warning(
            "Unable to deregister stack with name '%s': No stack exists "
            "with this name.",
            name,
        )
get_local_url(path) staticmethod

Get a local URL for a given local path.

Source code in zenml/stack_stores/local_stack_store.py
@staticmethod
def get_local_url(path: str) -> str:
    """Get a local URL for a given local path."""
    return f"file://{path}"
get_path_from_url(url) staticmethod

Get the path from a URL.

Parameters:

Name Type Description Default
url str

The URL to get the path from.

required

Returns:

Type Description
Optional[pathlib.Path]

The path from the URL.

Source code in zenml/stack_stores/local_stack_store.py
@staticmethod
def get_path_from_url(url: str) -> Optional[Path]:
    """Get the path from a URL.

    Args:
        url: The URL to get the path from.

    Returns:
        The path from the URL.
    """
    if not LocalStackStore.is_valid_url(url):
        raise ValueError(f"Invalid URL for local store: {url}")
    url = url.replace("file://", "")
    return Path(url)
get_stack_configuration(self, name)

Fetches a stack configuration by name.

Parameters:

Name Type Description Default
name str

The name of the stack to fetch.

required

Returns:

Type Description
Dict[zenml.enums.StackComponentType, str]

Dict[StackComponentType, str] for the requested stack name.

Exceptions:

Type Description
KeyError

If no stack exists for the given name.

Source code in zenml/stack_stores/local_stack_store.py
def get_stack_configuration(
    self, name: str
) -> Dict[StackComponentType, str]:
    """Fetches a stack configuration by name.

    Args:
        name: The name of the stack to fetch.

    Returns:
        Dict[StackComponentType, str] for the requested stack name.

    Raises:
        KeyError: If no stack exists for the given name.
    """
    logger.debug("Fetching stack with name '%s'.", name)
    if name not in self.__store.stacks:
        raise KeyError(
            f"Unable to find stack with name '{name}'. Available names: "
            f"{set(self.__store.stacks)}."
        )

    return self.__store.stacks[name]
initialize(self, url, *args, *, stack_data=None, **kwargs)

Initializes a local stack store instance.

Parameters:

Name Type Description Default
url str

URL of local directory of the repository to use for stack storage.

required
stack_data Optional[zenml.stack_stores.models.stack_store_model.StackStoreModel]

optional stack data store object to pre-populate the stack store with.

None
args Any

additional positional arguments (ignored).

()
kwargs Any

additional keyword arguments (ignored).

{}

Returns:

Type Description
LocalStackStore

The initialized stack store instance.

Source code in zenml/stack_stores/local_stack_store.py
def initialize(
    self,
    url: str,
    *args: Any,
    stack_data: Optional[StackStoreModel] = None,
    **kwargs: Any,
) -> "LocalStackStore":
    """Initializes a local stack store instance.

    Args:
        url: URL of local directory of the repository to use for
            stack storage.
        stack_data: optional stack data store object to pre-populate the
            stack store with.
        args: additional positional arguments (ignored).
        kwargs: additional keyword arguments (ignored).

    Returns:
        The initialized stack store instance.
    """
    if not self.is_valid_url(url):
        raise ValueError(f"Invalid URL for local store: {url}")

    self._root = self.get_path_from_url(url)
    self._url = f"file://{self._root}"
    fileio.create_dir_recursive_if_not_exists(str(self._root))

    if stack_data is not None:
        self.__store = stack_data
        self._write_store()
    elif fileio.file_exists(self._store_path()):
        config_dict = yaml_utils.read_yaml(self._store_path())
        self.__store = StackStoreModel.parse_obj(config_dict)
    else:
        self.__store = StackStoreModel.empty_store()
        self._write_store()

    super().initialize(url, *args, **kwargs)
    return self
is_empty(self)

Check if the stack store is empty.

Source code in zenml/stack_stores/local_stack_store.py
def is_empty(self) -> bool:
    """Check if the stack store is empty."""
    return len(self.__store.stacks) == 0
is_valid_url(url) staticmethod

Check if the given url is a valid local path.

Source code in zenml/stack_stores/local_stack_store.py
@staticmethod
def is_valid_url(url: str) -> bool:
    """Check if the given url is a valid local path."""
    scheme = re.search("^([a-z0-9]+://)", url)
    return not scheme or scheme.group() == "file://"
register_stack_component(self, component)

Register a stack component.

Parameters:

Name Type Description Default
component StackComponentWrapper

The component to register.

required

Exceptions:

Type Description
StackComponentExistsError

If a stack component with the same type and name already exists.

Source code in zenml/stack_stores/local_stack_store.py
def register_stack_component(
    self,
    component: StackComponentWrapper,
) -> None:
    """Register a stack component.

    Args:
        component: The component to register.

    Raises:
        StackComponentExistsError: If a stack component with the same type
            and name already exists.
    """
    components = self.__store.stack_components[component.type]
    if component.name in components:
        raise StackComponentExistsError(
            f"Unable to register stack component (type: {component.type}) "
            f"with name '{component.name}': Found existing stack component "
            f"with this name."
        )

    # write the component configuration file
    component_config_path = self._get_stack_component_config_path(
        component_type=component.type, name=component.name
    )
    fileio.create_dir_recursive_if_not_exists(
        os.path.dirname(component_config_path)
    )
    write_file_contents_as_string(
        component_config_path,
        base64.b64decode(component.config).decode(),
    )

    # add the component to the stack store dict and write it to disk
    components[component.name] = component.flavor
    self._write_store()
    logger.info(
        "Registered stack component with type '%s' and name '%s'.",
        component.type,
        component.name,
    )

models special

stack_component_wrapper

StackComponentWrapper (BaseModel) pydantic-model

Serializable Configuration of a StackComponent

Source code in zenml/stack_stores/models/stack_component_wrapper.py
class StackComponentWrapper(BaseModel):
    """Serializable Configuration of a StackComponent"""

    type: StackComponentType
    flavor: str  # due to subclassing, can't properly use enum type here
    name: str
    uuid: UUID
    config: bytes  # b64 encoded yaml config

    @classmethod
    def from_component(
        cls, component: StackComponent
    ) -> "StackComponentWrapper":
        return cls(
            type=component.type,
            flavor=component.flavor.value,
            name=component.name,
            uuid=component.uuid,
            config=base64.b64encode(
                yaml.dump(json.loads(component.json())).encode()
            ),
        )

stack_store_model

StackStoreModel (BaseModel) pydantic-model

Pydantic object used for serializing a ZenML Stack Store.

Attributes:

Name Type Description
stacks Dict[str, Dict[zenml.enums.StackComponentType, str]]

Maps stack names to a configuration object containing the names and flavors of all stack components.

stack_components DefaultDict[zenml.enums.StackComponentType, Dict[str, str]]

Contains names and flavors of all registered stack components.

Source code in zenml/stack_stores/models/stack_store_model.py
class StackStoreModel(BaseModel):
    """Pydantic object used for serializing a ZenML Stack Store.

    Attributes:
        stacks: Maps stack names to a configuration object containing the
            names and flavors of all stack components.
        stack_components: Contains names and flavors of all registered stack
            components.
    """

    stacks: Dict[str, Dict[StackComponentType, str]]
    stack_components: DefaultDict[StackComponentType, Dict[str, str]]

    @validator("stack_components")
    def _construct_defaultdict(
        cls, stack_components: Dict[StackComponentType, Dict[str, str]]
    ) -> DefaultDict[StackComponentType, Dict[str, str]]:
        """Ensures that `stack_components` is a defaultdict so stack
        components of a new component type can be added without issues."""
        return defaultdict(dict, stack_components)

    @classmethod
    def empty_store(cls) -> "StackStoreModel":
        """Initialize a new empty stack store with current zen version."""
        return cls(stacks={}, stack_components={})

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

        # Validate attributes when assigning them. We need to set this in order
        # to have a mix of mutable and immutable attributes
        validate_assignment = True
        # Ignore extra attributes from configs of previous ZenML versions
        extra = "ignore"
Config

Pydantic configuration class.

Source code in zenml/stack_stores/models/stack_store_model.py
class Config:
    """Pydantic configuration class."""

    # Validate attributes when assigning them. We need to set this in order
    # to have a mix of mutable and immutable attributes
    validate_assignment = True
    # Ignore extra attributes from configs of previous ZenML versions
    extra = "ignore"
empty_store() classmethod

Initialize a new empty stack store with current zen version.

Source code in zenml/stack_stores/models/stack_store_model.py
@classmethod
def empty_store(cls) -> "StackStoreModel":
    """Initialize a new empty stack store with current zen version."""
    return cls(stacks={}, stack_components={})

stack_wrapper

StackWrapper (BaseModel) pydantic-model

Network Serializable Wrapper describing a Stack.

Source code in zenml/stack_stores/models/stack_wrapper.py
class StackWrapper(BaseModel):
    """Network Serializable Wrapper describing a Stack."""

    name: str
    components: List[StackComponentWrapper]

    @classmethod
    def from_stack(cls, stack: Stack) -> "StackWrapper":
        return cls(
            name=stack.name,
            components=[
                StackComponentWrapper.from_component(component)
                for t, component in stack.components.items()
            ],
        )

sql_stack_store

SqlStackStore (BaseStackStore)

Repository Implementation that uses SQL database backend

Source code in zenml/stack_stores/sql_stack_store.py
class SqlStackStore(BaseStackStore):
    """Repository Implementation that uses SQL database backend"""

    def initialize(
        self,
        url: str,
        *args: Any,
        **kwargs: Any,
    ) -> "SqlStackStore":
        """Initialize a new SqlStackStore.

        Args:
            url: odbc path to a database.
            args, kwargs: additional parameters for SQLModel.

        Returns:
            The initialized stack store instance.
        """
        if not self.is_valid_url(url):
            raise ValueError(f"Invalid URL for SQL store: {url}")

        logger.debug("Initializing SqlStackStore at %s", url)
        self._url = url

        local_path = self.get_path_from_url(url)
        if local_path:
            fileio.create_dir_recursive_if_not_exists(str(local_path.parent))

        self.engine = create_engine(url, *args, **kwargs)
        SQLModel.metadata.create_all(self.engine)
        with Session(self.engine) as session:
            if not session.exec(select(ZenUser)).first():
                session.add(ZenUser(id=1, name="LocalZenUser"))
            session.commit()

        super().initialize(url, *args, **kwargs)

        return self

    # Public interface implementations:

    @property
    def type(self) -> StoreType:
        """The type of stack store."""
        return StoreType.SQL

    @property
    def url(self) -> str:
        """URL of the repository."""
        if not self._url:
            raise RuntimeError(
                "SQL stack store has not been initialized. Call `initialize` "
                "before using the store."
            )
        return self._url

    @staticmethod
    def get_path_from_url(url: str) -> Optional[Path]:
        """Get the local path from a URL, if it points to a local sqlite file.

        This method first checks that the URL is a valid SQLite URL, which is
        backed by a file in the local filesystem. All other types of supported
        SQLAlchemy connection URLs are considered non-local and won't return
        a valid local path.

        Args:
            url: The URL to get the path from.

        Returns:
            The path extracted from the URL, or None, if the URL does not
            point to a local sqlite file.
        """
        if not SqlStackStore.is_valid_url(url):
            raise ValueError(f"Invalid URL for SQL store: {url}")
        if not url.startswith("sqlite:///"):
            return None
        url = url.replace("sqlite:///", "")
        return Path(url)

    @staticmethod
    def get_local_url(path: str) -> str:
        """Get a local SQL url for a given local path."""
        return f"sqlite:///{path}/zenml.db"

    @staticmethod
    def is_valid_url(url: str) -> bool:
        """Check if the given url is a valid SQL url."""
        try:
            make_url(url)
        except ArgumentError:
            logger.debug("Invalid SQL URL: %s", url)
            return False

        return True

    def is_empty(self) -> bool:
        """Check if the stack store is empty."""
        with Session(self.engine) as session:
            return session.exec(select(ZenStack)).first() is None

    def get_stack_configuration(
        self, name: str
    ) -> Dict[StackComponentType, str]:
        """Fetches a stack configuration by name.

        Args:
            name: The name of the stack to fetch.

        Returns:
            Dict[StackComponentType, str] for the requested stack name.

        Raises:
            KeyError: If no stack exists for the given name.
        """
        logger.debug("Fetching stack with name '%s'.", name)
        # first check that the stack exists
        with Session(self.engine) as session:
            maybe_stack = session.exec(
                select(ZenStack).where(ZenStack.name == name)
            ).first()
        if maybe_stack is None:
            raise KeyError(
                f"Unable to find stack with name '{name}'. Available names: "
                f"{set(self.stack_names)}."
            )
        # then get all components assigned to that stack
        with Session(self.engine) as session:
            definitions_and_components = session.exec(
                select(ZenStackDefinition, ZenStackComponent)
                .where(
                    ZenStackDefinition.component_type
                    == ZenStackComponent.component_type
                )
                .where(
                    ZenStackDefinition.component_name == ZenStackComponent.name
                )
                .where(ZenStackDefinition.stack_name == name)
            )
            params = {
                component.component_type: component.name
                for _, component in definitions_and_components
            }
        return {StackComponentType(typ): name for typ, name in params.items()}

    @property
    def stack_configurations(self) -> Dict[str, Dict[StackComponentType, str]]:
        """Configuration for all stacks registered in this stack store.

        Returns:
            Dictionary mapping stack names to Dict[StackComponentType, str]
        """
        return {n: self.get_stack_configuration(n) for n in self.stack_names}

    def register_stack_component(
        self,
        component: StackComponentWrapper,
    ) -> None:
        """Register a stack component.

        Args:
            component: The component to register.

        Raises:
            StackComponentExistsError: If a stack component with the same type
                and name already exists.
        """
        with Session(self.engine) as session:
            existing_component = session.exec(
                select(ZenStackComponent)
                .where(ZenStackComponent.name == component.name)
                .where(ZenStackComponent.component_type == component.type)
            ).first()
            if existing_component is not None:
                raise StackComponentExistsError(
                    f"Unable to register stack component (type: "
                    f"{component.type}) with name '{component.name}': Found "
                    f"existing stack component with this name."
                )
            new_component = ZenStackComponent(
                component_type=component.type,
                name=component.name,
                component_flavor=component.flavor,
                configuration=component.config,
            )
            session.add(new_component)
            session.commit()

    def deregister_stack(self, name: str) -> None:
        """Delete a stack from storage.

        Args:
            name: The name of the stack to be deleted.
        """
        with Session(self.engine) as session:
            stack = session.exec(
                select(ZenStack).where(ZenStack.name == name)
            ).one()
            session.delete(stack)
            definitions = session.exec(
                select(ZenStackDefinition).where(
                    ZenStackDefinition.stack_name == name
                )
            ).all()
            for definition in definitions:
                session.delete(definition)
            session.commit()

    # Private interface implementations:

    def _create_stack(
        self, name: str, stack_configuration: Dict[StackComponentType, str]
    ) -> None:
        """Add a stack to storage.

        Args:
            name: The name to save the stack as.
            stack_configuration: Dict[StackComponentType, str] to persist.
        """
        with Session(self.engine) as session:
            stack = ZenStack(name=name, created_by=1)
            session.add(stack)
            for ctype, cname in stack_configuration.items():
                if cname is not None:
                    session.add(
                        ZenStackDefinition(
                            stack_name=name,
                            component_type=ctype,
                            component_name=cname,
                        )
                    )
            session.commit()

    def _get_component_flavor_and_config(
        self, component_type: StackComponentType, name: str
    ) -> Tuple[str, bytes]:
        """Fetch the flavor and configuration for a stack component.

        Args:
            component_type: The type of the component to fetch.
            name: The name of the component to fetch.

        Returns:
            Pair of (flavor, congfiguration) for stack component, as string and
            base64-encoded yaml document, respectively

        Raises:
            KeyError: If no stack component exists for the given type and name.
        """
        with Session(self.engine) as session:
            component = session.exec(
                select(ZenStackComponent)
                .where(ZenStackComponent.component_type == component_type)
                .where(ZenStackComponent.name == name)
            ).one_or_none()
            if component is None:
                raise KeyError(
                    f"Unable to find stack component (type: {component_type}) "
                    f"with name '{name}'."
                )
        return component.component_flavor, component.configuration

    def _get_stack_component_names(
        self, component_type: StackComponentType
    ) -> List[str]:
        """Get names of all registered stack components of a given type.

        Args:
            component_type: The type of the component to list names for.

        Returns:
            A list of names as strings.
        """
        with Session(self.engine) as session:
            statement = select(ZenStackComponent).where(
                ZenStackComponent.component_type == component_type
            )
            return [component.name for component in session.exec(statement)]

    def _delete_stack_component(
        self, component_type: StackComponentType, name: str
    ) -> None:
        """Remove a StackComponent from storage.

        Args:
            component_type: The type of component to delete.
            name: Then name of the component to delete.
        """
        with Session(self.engine) as session:
            component = session.exec(
                select(ZenStackComponent)
                .where(ZenStackComponent.component_type == component_type)
                .where(ZenStackComponent.name == name)
            ).first()
            if component is not None:
                session.delete(component)
                session.commit()
                logger.info(
                    "Deregistered stack component (type: %s) with name '%s'.",
                    component_type.value,
                    name,
                )
            else:
                logger.warning(
                    "Unable to deregister stack component (type: %s) with name "
                    "'%s': No stack component exists with this name.",
                    component_type.value,
                    name,
                )

    # Implementation-specific internal methods:

    @property
    def stack_names(self) -> List[str]:
        """Names of all stacks registered in this StackStore."""
        with Session(self.engine) as session:
            return [s.name for s in session.exec(select(ZenStack))]
stack_configurations: Dict[str, Dict[zenml.enums.StackComponentType, str]] property readonly

Configuration for all stacks registered in this stack store.

Returns:

Type Description
Dict[str, Dict[zenml.enums.StackComponentType, str]]

Dictionary mapping stack names to Dict[StackComponentType, str]

stack_names: List[str] property readonly

Names of all stacks registered in this StackStore.

type: StoreType property readonly

The type of stack store.

url: str property readonly

URL of the repository.

deregister_stack(self, name)

Delete a stack from storage.

Parameters:

Name Type Description Default
name str

The name of the stack to be deleted.

required
Source code in zenml/stack_stores/sql_stack_store.py
def deregister_stack(self, name: str) -> None:
    """Delete a stack from storage.

    Args:
        name: The name of the stack to be deleted.
    """
    with Session(self.engine) as session:
        stack = session.exec(
            select(ZenStack).where(ZenStack.name == name)
        ).one()
        session.delete(stack)
        definitions = session.exec(
            select(ZenStackDefinition).where(
                ZenStackDefinition.stack_name == name
            )
        ).all()
        for definition in definitions:
            session.delete(definition)
        session.commit()
get_local_url(path) staticmethod

Get a local SQL url for a given local path.

Source code in zenml/stack_stores/sql_stack_store.py
@staticmethod
def get_local_url(path: str) -> str:
    """Get a local SQL url for a given local path."""
    return f"sqlite:///{path}/zenml.db"
get_path_from_url(url) staticmethod

Get the local path from a URL, if it points to a local sqlite file.

This method first checks that the URL is a valid SQLite URL, which is backed by a file in the local filesystem. All other types of supported SQLAlchemy connection URLs are considered non-local and won't return a valid local path.

Parameters:

Name Type Description Default
url str

The URL to get the path from.

required

Returns:

Type Description
Optional[pathlib.Path]

The path extracted from the URL, or None, if the URL does not point to a local sqlite file.

Source code in zenml/stack_stores/sql_stack_store.py
@staticmethod
def get_path_from_url(url: str) -> Optional[Path]:
    """Get the local path from a URL, if it points to a local sqlite file.

    This method first checks that the URL is a valid SQLite URL, which is
    backed by a file in the local filesystem. All other types of supported
    SQLAlchemy connection URLs are considered non-local and won't return
    a valid local path.

    Args:
        url: The URL to get the path from.

    Returns:
        The path extracted from the URL, or None, if the URL does not
        point to a local sqlite file.
    """
    if not SqlStackStore.is_valid_url(url):
        raise ValueError(f"Invalid URL for SQL store: {url}")
    if not url.startswith("sqlite:///"):
        return None
    url = url.replace("sqlite:///", "")
    return Path(url)
get_stack_configuration(self, name)

Fetches a stack configuration by name.

Parameters:

Name Type Description Default
name str

The name of the stack to fetch.

required

Returns:

Type Description
Dict[zenml.enums.StackComponentType, str]

Dict[StackComponentType, str] for the requested stack name.

Exceptions:

Type Description
KeyError

If no stack exists for the given name.

Source code in zenml/stack_stores/sql_stack_store.py
def get_stack_configuration(
    self, name: str
) -> Dict[StackComponentType, str]:
    """Fetches a stack configuration by name.

    Args:
        name: The name of the stack to fetch.

    Returns:
        Dict[StackComponentType, str] for the requested stack name.

    Raises:
        KeyError: If no stack exists for the given name.
    """
    logger.debug("Fetching stack with name '%s'.", name)
    # first check that the stack exists
    with Session(self.engine) as session:
        maybe_stack = session.exec(
            select(ZenStack).where(ZenStack.name == name)
        ).first()
    if maybe_stack is None:
        raise KeyError(
            f"Unable to find stack with name '{name}'. Available names: "
            f"{set(self.stack_names)}."
        )
    # then get all components assigned to that stack
    with Session(self.engine) as session:
        definitions_and_components = session.exec(
            select(ZenStackDefinition, ZenStackComponent)
            .where(
                ZenStackDefinition.component_type
                == ZenStackComponent.component_type
            )
            .where(
                ZenStackDefinition.component_name == ZenStackComponent.name
            )
            .where(ZenStackDefinition.stack_name == name)
        )
        params = {
            component.component_type: component.name
            for _, component in definitions_and_components
        }
    return {StackComponentType(typ): name for typ, name in params.items()}
initialize(self, url, *args, **kwargs)

Initialize a new SqlStackStore.

Parameters:

Name Type Description Default
url str

odbc path to a database.

required
args, kwargs

additional parameters for SQLModel.

required

Returns:

Type Description
SqlStackStore

The initialized stack store instance.

Source code in zenml/stack_stores/sql_stack_store.py
def initialize(
    self,
    url: str,
    *args: Any,
    **kwargs: Any,
) -> "SqlStackStore":
    """Initialize a new SqlStackStore.

    Args:
        url: odbc path to a database.
        args, kwargs: additional parameters for SQLModel.

    Returns:
        The initialized stack store instance.
    """
    if not self.is_valid_url(url):
        raise ValueError(f"Invalid URL for SQL store: {url}")

    logger.debug("Initializing SqlStackStore at %s", url)
    self._url = url

    local_path = self.get_path_from_url(url)
    if local_path:
        fileio.create_dir_recursive_if_not_exists(str(local_path.parent))

    self.engine = create_engine(url, *args, **kwargs)
    SQLModel.metadata.create_all(self.engine)
    with Session(self.engine) as session:
        if not session.exec(select(ZenUser)).first():
            session.add(ZenUser(id=1, name="LocalZenUser"))
        session.commit()

    super().initialize(url, *args, **kwargs)

    return self
is_empty(self)

Check if the stack store is empty.

Source code in zenml/stack_stores/sql_stack_store.py
def is_empty(self) -> bool:
    """Check if the stack store is empty."""
    with Session(self.engine) as session:
        return session.exec(select(ZenStack)).first() is None
is_valid_url(url) staticmethod

Check if the given url is a valid SQL url.

Source code in zenml/stack_stores/sql_stack_store.py
@staticmethod
def is_valid_url(url: str) -> bool:
    """Check if the given url is a valid SQL url."""
    try:
        make_url(url)
    except ArgumentError:
        logger.debug("Invalid SQL URL: %s", url)
        return False

    return True
register_stack_component(self, component)

Register a stack component.

Parameters:

Name Type Description Default
component StackComponentWrapper

The component to register.

required

Exceptions:

Type Description
StackComponentExistsError

If a stack component with the same type and name already exists.

Source code in zenml/stack_stores/sql_stack_store.py
def register_stack_component(
    self,
    component: StackComponentWrapper,
) -> None:
    """Register a stack component.

    Args:
        component: The component to register.

    Raises:
        StackComponentExistsError: If a stack component with the same type
            and name already exists.
    """
    with Session(self.engine) as session:
        existing_component = session.exec(
            select(ZenStackComponent)
            .where(ZenStackComponent.name == component.name)
            .where(ZenStackComponent.component_type == component.type)
        ).first()
        if existing_component is not None:
            raise StackComponentExistsError(
                f"Unable to register stack component (type: "
                f"{component.type}) with name '{component.name}': Found "
                f"existing stack component with this name."
            )
        new_component = ZenStackComponent(
            component_type=component.type,
            name=component.name,
            component_flavor=component.flavor,
            configuration=component.config,
        )
        session.add(new_component)
        session.commit()

ZenStackDefinition (SQLModel) pydantic-model

Join table between Stacks and StackComponents

Source code in zenml/stack_stores/sql_stack_store.py
class ZenStackDefinition(SQLModel, table=True):
    """Join table between Stacks and StackComponents"""

    stack_name: str = Field(primary_key=True, foreign_key="zenstack.name")
    component_type: StackComponentType = Field(
        primary_key=True, foreign_key="zenstackcomponent.component_type"
    )
    component_name: str = Field(
        primary_key=True, foreign_key="zenstackcomponent.name"
    )