Skip to content

Zen Stores

zenml.zen_stores special

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

base_zen_store

BaseZenStore (ABC)

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

Source code in zenml/zen_stores/base_zen_store.py
class BaseZenStore(ABC):
    """Base class for accessing data in ZenML Repository and new Service."""

    def initialize(
        self,
        url: str,
        skip_default_registrations: bool = False,
        *args: Any,
        **kwargs: Any,
    ) -> "BaseZenStore":
        """Initialize the store.

        Args:
            url: The URL of the store.
            skip_default_registrations: If `True`, the creation of the default
                stack and user 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_registrations:
            logger.info("Registering default stack and user...")
            if self.is_empty:
                self.register_default_stack()

            self.create_default_user()

        return self

    # Static methods:

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

        Args:
            path: the path string to build a URL out of.

        Returns:
            Url pointing to the path for the store type.
        """

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

    # Public Interface:

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

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

    @property
    @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 zen 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 update_stack_component(
        self,
        name: str,
        component_type: StackComponentType,
        component: StackComponentWrapper,
    ) -> Dict[str, str]:
        """Update a stack component.

        Args:
            name: The original name of the stack component.
            component_type: The type of the stack component to update.
            component: The new component to update with.

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

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

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

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

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

    @abstractmethod
    def _save_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, configuration) 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.

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

    # User, project and role management

    @property
    @abstractmethod
    def users(self) -> List[User]:
        """All registered users.

        Returns:
            A list of all registered users.
        """

    @abstractmethod
    def get_user(self, user_name: str) -> User:
        """Gets a specific user.

        Args:
            user_name: Name of the user to get.

        Returns:
            The requested user.

        Raises:
            KeyError: If no user with the given name exists.
        """

    @abstractmethod
    def create_user(self, user_name: str) -> User:
        """Creates a new user.

        Args:
            user_name: Unique username.

        Returns:
             The newly created user.

        Raises:
            EntityExistsError: If a user with the given name already exists.
        """

    @abstractmethod
    def delete_user(self, user_name: str) -> None:
        """Deletes a user.

        Args:
            user_name: Name of the user to delete.

        Raises:
            KeyError: If no user with the given name exists.
        """

    @property
    @abstractmethod
    def teams(self) -> List[Team]:
        """All registered teams.

        Returns:
            A list of all registered teams.
        """

    @abstractmethod
    def create_team(self, team_name: str) -> Team:
        """Creates a new team.

        Args:
            team_name: Unique team name.

        Returns:
             The newly created team.

        Raises:
            EntityExistsError: If a team with the given name already exists.
        """

    @abstractmethod
    def get_team(self, team_name: str) -> Team:
        """Gets a specific team.

        Args:
            team_name: Name of the team to get.

        Returns:
            The requested team.

        Raises:
            KeyError: If no team with the given name exists.
        """

    @abstractmethod
    def delete_team(self, team_name: str) -> None:
        """Deletes a team.

        Args:
            team_name: Name of the team to delete.

        Raises:
            KeyError: If no team with the given name exists.
        """

    @abstractmethod
    def add_user_to_team(self, team_name: str, user_name: str) -> None:
        """Adds a user to a team.

        Args:
            team_name: Name of the team.
            user_name: Name of the user.

        Raises:
            KeyError: If no user and team with the given names exists.
        """

    @abstractmethod
    def remove_user_from_team(self, team_name: str, user_name: str) -> None:
        """Removes a user from a team.

        Args:
            team_name: Name of the team.
            user_name: Name of the user.

        Raises:
            KeyError: If no user and team with the given names exists.
        """

    @property
    @abstractmethod
    def projects(self) -> List[Project]:
        """All registered projects.

        Returns:
            A list of all registered projects.
        """

    @abstractmethod
    def get_project(self, project_name: str) -> Project:
        """Gets a specific project.

        Args:
            project_name: Name of the project to get.

        Returns:
            The requested project.

        Raises:
            KeyError: If no project with the given name exists.
        """

    @abstractmethod
    def create_project(
        self, project_name: str, description: Optional[str] = None
    ) -> Project:
        """Creates a new project.

        Args:
            project_name: Unique project name.
            description: Optional project description.

        Returns:
             The newly created project.

        Raises:
            EntityExistsError: If a project with the given name already exists.
        """

    @abstractmethod
    def delete_project(self, project_name: str) -> None:
        """Deletes a project.

        Args:
            project_name: Name of the project to delete.

        Raises:
            KeyError: If no project with the given name exists.
        """

    @property
    @abstractmethod
    def roles(self) -> List[Role]:
        """All registered roles.

        Returns:
            A list of all registered roles.
        """

    @property
    @abstractmethod
    def role_assignments(self) -> List[RoleAssignment]:
        """All registered role assignments.

        Returns:
            A list of all registered role assignments.
        """

    @abstractmethod
    def get_role(self, role_name: str) -> Role:
        """Gets a specific role.

        Args:
            role_name: Name of the role to get.

        Returns:
            The requested role.

        Raises:
            KeyError: If no role with the given name exists.
        """

    @abstractmethod
    def create_role(self, role_name: str) -> Role:
        """Creates a new role.

        Args:
            role_name: Unique role name.

        Returns:
             The newly created role.

        Raises:
            EntityExistsError: If a role with the given name already exists.
        """

    @abstractmethod
    def delete_role(self, role_name: str) -> None:
        """Deletes a role.

        Args:
            role_name: Name of the role to delete.

        Raises:
            KeyError: If no role with the given name exists.
        """

    @abstractmethod
    def assign_role(
        self,
        role_name: str,
        entity_name: str,
        project_name: Optional[str] = None,
        is_user: bool = True,
    ) -> None:
        """Assigns a role to a user or team.

        Args:
            role_name: Name of the role to assign.
            entity_name: User or team name.
            project_name: Optional project name.
            is_user: Boolean indicating whether the given `entity_name` refers
                to a user.

        Raises:
            KeyError: If no role, entity or project with the given names exists.
        """

    @abstractmethod
    def revoke_role(
        self,
        role_name: str,
        entity_name: str,
        project_name: Optional[str] = None,
        is_user: bool = True,
    ) -> None:
        """Revokes a role from a user or team.

        Args:
            role_name: Name of the role to revoke.
            entity_name: User or team name.
            project_name: Optional project name.
            is_user: Boolean indicating whether the given `entity_name` refers
                to a user.

        Raises:
            KeyError: If no role, entity or project with the given names exists.
        """

    @abstractmethod
    def get_users_for_team(self, team_name: str) -> List[User]:
        """Fetches all users of a team.

        Args:
            team_name: Name of the team.

        Returns:
            List of users that are part of the team.

        Raises:
            KeyError: If no team with the given name exists.
        """

    @abstractmethod
    def get_teams_for_user(self, user_name: str) -> List[Team]:
        """Fetches all teams for a user.

        Args:
            user_name: Name of the user.

        Returns:
            List of teams that the user is part of.

        Raises:
            KeyError: If no user with the given name exists.
        """

    @abstractmethod
    def get_role_assignments_for_user(
        self,
        user_name: str,
        project_name: Optional[str] = None,
        include_team_roles: bool = True,
    ) -> List[RoleAssignment]:
        """Fetches all role assignments for a user.

        Args:
            user_name: Name of the user.
            project_name: Optional filter to only return roles assigned for
                this project.
            include_team_roles: If `True`, includes roles for all teams that
                the user is part of.

        Returns:
            List of role assignments for this user.

        Raises:
            KeyError: If no user or project with the given names exists.
        """

    @abstractmethod
    def get_role_assignments_for_team(
        self,
        team_name: str,
        project_name: Optional[str] = None,
    ) -> List[RoleAssignment]:
        """Fetches all role assignments for a team.

        Args:
            team_name: Name of the user.
            project_name: Optional filter to only return roles assigned for
                this project.

        Returns:
            List of role assignments for this team.

        Raises:
            KeyError: If no team or project with the given names exists.
        """

    # Common code (user facing):

    @property
    def stacks(self) -> List[StackWrapper]:
        """All stacks registered in this zen 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 stack's components aren't registered in the zen 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 to register a stack component, if it doesn't exist.

            Args:
                component: StackComponentWrapper to register.

            Returns:
                metadata key value pair for telemetry.

            Raises:
                StackComponentExistsError: If a component with same name exists.
            """
            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._save_stack(stack.name, stack_configuration)
        logger.info("Registered stack with name '%s'.", stack.name)
        return metadata

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

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

        Args:
            name: The original name of the stack.
            stack: The new stack to use in the update.

        Returns:
            metadata dict for telemetry or logging.

        Raises:
            DoesNotExistException: If no stack exists with the given name.
        """
        try:
            self.get_stack(name)
        except KeyError:
            raise KeyError(
                f"Unable to update stack with name '{stack.name}': No existing "
                f"stack found with this name."
            )

        try:
            renamed_stack = self.get_stack(stack.name)
            if (name != stack.name) and renamed_stack:
                raise StackExistsError(
                    f"Unable to update stack with name '{stack.name}': Found "
                    f"existing stack with this name."
                )
        except KeyError:
            pass

        def __check_component(
            component: StackComponentWrapper,
        ) -> Tuple[StackComponentType, str]:
            try:
                _ = self.get_stack_component(
                    component_type=component.type, name=component.name
                )
            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._save_stack(stack.name, stack_configuration)

        logger.info("Updated stack with name '%s'.", name)
        if name != stack.name:
            self.deregister_stack(name)

        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)

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

    def create_default_user(self) -> None:
        """Creates a default user."""
        try:
            self.get_user(user_name=DEFAULT_USERNAME)
        except KeyError:
            self.create_user(user_name=DEFAULT_USERNAME)

    # 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)
is_empty: bool property readonly

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.

projects: List[zenml.zen_stores.models.user_management_models.Project] property readonly

All registered projects.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Project]

A list of all registered projects.

role_assignments: List[zenml.zen_stores.models.user_management_models.RoleAssignment] property readonly

All registered role assignments.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

A list of all registered role assignments.

roles: List[zenml.zen_stores.models.user_management_models.Role] property readonly

All registered roles.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Role]

A list of all registered roles.

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

Configurations for all stacks registered in this zen store.

Returns:

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

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

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

All stacks registered in this zen store.

teams: List[zenml.zen_stores.models.user_management_models.Team] property readonly

All registered teams.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Team]

A list of all registered teams.

type: StoreType property readonly

The type of zen store.

url: str property readonly

Get the repository URL.

users: List[zenml.zen_stores.models.user_management_models.User] property readonly

All registered users.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.User]

A list of all registered users.

add_user_to_team(self, team_name, user_name)

Adds a user to a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required
user_name str

Name of the user.

required

Exceptions:

Type Description
KeyError

If no user and team with the given names exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def add_user_to_team(self, team_name: str, user_name: str) -> None:
    """Adds a user to a team.

    Args:
        team_name: Name of the team.
        user_name: Name of the user.

    Raises:
        KeyError: If no user and team with the given names exists.
    """
assign_role(self, role_name, entity_name, project_name=None, is_user=True)

Assigns a role to a user or team.

Parameters:

Name Type Description Default
role_name str

Name of the role to assign.

required
entity_name str

User or team name.

required
project_name Optional[str]

Optional project name.

None
is_user bool

Boolean indicating whether the given entity_name refers to a user.

True

Exceptions:

Type Description
KeyError

If no role, entity or project with the given names exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def assign_role(
    self,
    role_name: str,
    entity_name: str,
    project_name: Optional[str] = None,
    is_user: bool = True,
) -> None:
    """Assigns a role to a user or team.

    Args:
        role_name: Name of the role to assign.
        entity_name: User or team name.
        project_name: Optional project name.
        is_user: Boolean indicating whether the given `entity_name` refers
            to a user.

    Raises:
        KeyError: If no role, entity or project with the given names exists.
    """
create_default_user(self)

Creates a default user.

Source code in zenml/zen_stores/base_zen_store.py
def create_default_user(self) -> None:
    """Creates a default user."""
    try:
        self.get_user(user_name=DEFAULT_USERNAME)
    except KeyError:
        self.create_user(user_name=DEFAULT_USERNAME)
create_project(self, project_name, description=None)

Creates a new project.

Parameters:

Name Type Description Default
project_name str

Unique project name.

required
description Optional[str]

Optional project description.

None

Returns:

Type Description
Project

The newly created project.

Exceptions:

Type Description
EntityExistsError

If a project with the given name already exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def create_project(
    self, project_name: str, description: Optional[str] = None
) -> Project:
    """Creates a new project.

    Args:
        project_name: Unique project name.
        description: Optional project description.

    Returns:
         The newly created project.

    Raises:
        EntityExistsError: If a project with the given name already exists.
    """
create_role(self, role_name)

Creates a new role.

Parameters:

Name Type Description Default
role_name str

Unique role name.

required

Returns:

Type Description
Role

The newly created role.

Exceptions:

Type Description
EntityExistsError

If a role with the given name already exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def create_role(self, role_name: str) -> Role:
    """Creates a new role.

    Args:
        role_name: Unique role name.

    Returns:
         The newly created role.

    Raises:
        EntityExistsError: If a role with the given name already exists.
    """
create_team(self, team_name)

Creates a new team.

Parameters:

Name Type Description Default
team_name str

Unique team name.

required

Returns:

Type Description
Team

The newly created team.

Exceptions:

Type Description
EntityExistsError

If a team with the given name already exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def create_team(self, team_name: str) -> Team:
    """Creates a new team.

    Args:
        team_name: Unique team name.

    Returns:
         The newly created team.

    Raises:
        EntityExistsError: If a team with the given name already exists.
    """
create_user(self, user_name)

Creates a new user.

Parameters:

Name Type Description Default
user_name str

Unique username.

required

Returns:

Type Description
User

The newly created user.

Exceptions:

Type Description
EntityExistsError

If a user with the given name already exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def create_user(self, user_name: str) -> User:
    """Creates a new user.

    Args:
        user_name: Unique username.

    Returns:
         The newly created user.

    Raises:
        EntityExistsError: If a user with the given name already exists.
    """
delete_project(self, project_name)

Deletes a project.

Parameters:

Name Type Description Default
project_name str

Name of the project to delete.

required

Exceptions:

Type Description
KeyError

If no project with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def delete_project(self, project_name: str) -> None:
    """Deletes a project.

    Args:
        project_name: Name of the project to delete.

    Raises:
        KeyError: If no project with the given name exists.
    """
delete_role(self, role_name)

Deletes a role.

Parameters:

Name Type Description Default
role_name str

Name of the role to delete.

required

Exceptions:

Type Description
KeyError

If no role with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def delete_role(self, role_name: str) -> None:
    """Deletes a role.

    Args:
        role_name: Name of the role to delete.

    Raises:
        KeyError: If no role with the given name exists.
    """
delete_team(self, team_name)

Deletes a team.

Parameters:

Name Type Description Default
team_name str

Name of the team to delete.

required

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def delete_team(self, team_name: str) -> None:
    """Deletes a team.

    Args:
        team_name: Name of the team to delete.

    Raises:
        KeyError: If no team with the given name exists.
    """
delete_user(self, user_name)

Deletes a user.

Parameters:

Name Type Description Default
user_name str

Name of the user to delete.

required

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def delete_user(self, user_name: str) -> None:
    """Deletes a user.

    Args:
        user_name: Name of the user to delete.

    Raises:
        KeyError: If no user with the given name exists.
    """
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

Exceptions:

Type Description
KeyError

If no stack exists for the given name.

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

    Raises:
        KeyError: If no stack exists for the given name.
    """
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/zen_stores/base_zen_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.

Parameters:

Name Type Description Default
path str

the path string to build a URL out of.

required

Returns:

Type Description
str

Url pointing to the path for the store type.

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

    Args:
        path: the path string to build a URL out of.

    Returns:
        Url pointing to the path for the store type.
    """
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/zen_stores/base_zen_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_project(self, project_name)

Gets a specific project.

Parameters:

Name Type Description Default
project_name str

Name of the project to get.

required

Returns:

Type Description
Project

The requested project.

Exceptions:

Type Description
KeyError

If no project with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def get_project(self, project_name: str) -> Project:
    """Gets a specific project.

    Args:
        project_name: Name of the project to get.

    Returns:
        The requested project.

    Raises:
        KeyError: If no project with the given name exists.
    """
get_role(self, role_name)

Gets a specific role.

Parameters:

Name Type Description Default
role_name str

Name of the role to get.

required

Returns:

Type Description
Role

The requested role.

Exceptions:

Type Description
KeyError

If no role with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def get_role(self, role_name: str) -> Role:
    """Gets a specific role.

    Args:
        role_name: Name of the role to get.

    Returns:
        The requested role.

    Raises:
        KeyError: If no role with the given name exists.
    """
get_role_assignments_for_team(self, team_name, project_name=None)

Fetches all role assignments for a team.

Parameters:

Name Type Description Default
team_name str

Name of the user.

required
project_name Optional[str]

Optional filter to only return roles assigned for this project.

None

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

List of role assignments for this team.

Exceptions:

Type Description
KeyError

If no team or project with the given names exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def get_role_assignments_for_team(
    self,
    team_name: str,
    project_name: Optional[str] = None,
) -> List[RoleAssignment]:
    """Fetches all role assignments for a team.

    Args:
        team_name: Name of the user.
        project_name: Optional filter to only return roles assigned for
            this project.

    Returns:
        List of role assignments for this team.

    Raises:
        KeyError: If no team or project with the given names exists.
    """
get_role_assignments_for_user(self, user_name, project_name=None, include_team_roles=True)

Fetches all role assignments for a user.

Parameters:

Name Type Description Default
user_name str

Name of the user.

required
project_name Optional[str]

Optional filter to only return roles assigned for this project.

None
include_team_roles bool

If True, includes roles for all teams that the user is part of.

True

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

List of role assignments for this user.

Exceptions:

Type Description
KeyError

If no user or project with the given names exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def get_role_assignments_for_user(
    self,
    user_name: str,
    project_name: Optional[str] = None,
    include_team_roles: bool = True,
) -> List[RoleAssignment]:
    """Fetches all role assignments for a user.

    Args:
        user_name: Name of the user.
        project_name: Optional filter to only return roles assigned for
            this project.
        include_team_roles: If `True`, includes roles for all teams that
            the user is part of.

    Returns:
        List of role assignments for this user.

    Raises:
        KeyError: If no user or project with the given names exists.
    """
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/zen_stores/base_zen_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/zen_stores/base_zen_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.zen_stores.models.stack_component_wrapper.StackComponentWrapper]

A list of StackComponentConfiguration instances.

Source code in zenml/zen_stores/base_zen_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/zen_stores/base_zen_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.
    """
get_team(self, team_name)

Gets a specific team.

Parameters:

Name Type Description Default
team_name str

Name of the team to get.

required

Returns:

Type Description
Team

The requested team.

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def get_team(self, team_name: str) -> Team:
    """Gets a specific team.

    Args:
        team_name: Name of the team to get.

    Returns:
        The requested team.

    Raises:
        KeyError: If no team with the given name exists.
    """
get_teams_for_user(self, user_name)

Fetches all teams for a user.

Parameters:

Name Type Description Default
user_name str

Name of the user.

required

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Team]

List of teams that the user is part of.

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def get_teams_for_user(self, user_name: str) -> List[Team]:
    """Fetches all teams for a user.

    Args:
        user_name: Name of the user.

    Returns:
        List of teams that the user is part of.

    Raises:
        KeyError: If no user with the given name exists.
    """
get_user(self, user_name)

Gets a specific user.

Parameters:

Name Type Description Default
user_name str

Name of the user to get.

required

Returns:

Type Description
User

The requested user.

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def get_user(self, user_name: str) -> User:
    """Gets a specific user.

    Args:
        user_name: Name of the user to get.

    Returns:
        The requested user.

    Raises:
        KeyError: If no user with the given name exists.
    """
get_users_for_team(self, team_name)

Fetches all users of a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.User]

List of users that are part of the team.

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def get_users_for_team(self, team_name: str) -> List[User]:
    """Fetches all users of a team.

    Args:
        team_name: Name of the team.

    Returns:
        List of users that are part of the team.

    Raises:
        KeyError: If no team with the given name exists.
    """
initialize(self, url, skip_default_registrations=False, *args, **kwargs)

Initialize the store.

Parameters:

Name Type Description Default
url str

The URL of the store.

required
skip_default_registrations bool

If True, the creation of the default stack and user 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
BaseZenStore

The initialized concrete store instance.

Source code in zenml/zen_stores/base_zen_store.py
def initialize(
    self,
    url: str,
    skip_default_registrations: bool = False,
    *args: Any,
    **kwargs: Any,
) -> "BaseZenStore":
    """Initialize the store.

    Args:
        url: The URL of the store.
        skip_default_registrations: If `True`, the creation of the default
            stack and user 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_registrations:
        logger.info("Registering default stack and user...")
        if self.is_empty:
            self.register_default_stack()

        self.create_default_user()

    return self
is_valid_url(url) staticmethod

Check if the given url is valid.

Source code in zenml/zen_stores/base_zen_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/zen_stores/base_zen_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.
    """
    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 stack's components aren't registered in the zen 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/zen_stores/base_zen_store.py
def register_stack(self, stack: StackWrapper) -> Dict[str, str]:
    """Register a stack and its components.

    If any of the stack's components aren't registered in the zen 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 to register a stack component, if it doesn't exist.

        Args:
            component: StackComponentWrapper to register.

        Returns:
            metadata key value pair for telemetry.

        Raises:
            StackComponentExistsError: If a component with same name exists.
        """
        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._save_stack(stack.name, stack_configuration)
    logger.info("Registered stack with name '%s'.", stack.name)
    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/zen_stores/base_zen_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.
    """
remove_user_from_team(self, team_name, user_name)

Removes a user from a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required
user_name str

Name of the user.

required

Exceptions:

Type Description
KeyError

If no user and team with the given names exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def remove_user_from_team(self, team_name: str, user_name: str) -> None:
    """Removes a user from a team.

    Args:
        team_name: Name of the team.
        user_name: Name of the user.

    Raises:
        KeyError: If no user and team with the given names exists.
    """
revoke_role(self, role_name, entity_name, project_name=None, is_user=True)

Revokes a role from a user or team.

Parameters:

Name Type Description Default
role_name str

Name of the role to revoke.

required
entity_name str

User or team name.

required
project_name Optional[str]

Optional project name.

None
is_user bool

Boolean indicating whether the given entity_name refers to a user.

True

Exceptions:

Type Description
KeyError

If no role, entity or project with the given names exists.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def revoke_role(
    self,
    role_name: str,
    entity_name: str,
    project_name: Optional[str] = None,
    is_user: bool = True,
) -> None:
    """Revokes a role from a user or team.

    Args:
        role_name: Name of the role to revoke.
        entity_name: User or team name.
        project_name: Optional project name.
        is_user: Boolean indicating whether the given `entity_name` refers
            to a user.

    Raises:
        KeyError: If no role, entity or project with the given names exists.
    """
update_stack(self, name, stack)

Update a stack and its components.

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

Parameters:

Name Type Description Default
name str

The original name of the stack.

required
stack StackWrapper

The new stack to use in the update.

required

Returns:

Type Description
Dict[str, str]

metadata dict for telemetry or logging.

Exceptions:

Type Description
DoesNotExistException

If no stack exists with the given name.

Source code in zenml/zen_stores/base_zen_store.py
def update_stack(self, name: str, stack: StackWrapper) -> Dict[str, str]:
    """Update a stack and its components.

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

    Args:
        name: The original name of the stack.
        stack: The new stack to use in the update.

    Returns:
        metadata dict for telemetry or logging.

    Raises:
        DoesNotExistException: If no stack exists with the given name.
    """
    try:
        self.get_stack(name)
    except KeyError:
        raise KeyError(
            f"Unable to update stack with name '{stack.name}': No existing "
            f"stack found with this name."
        )

    try:
        renamed_stack = self.get_stack(stack.name)
        if (name != stack.name) and renamed_stack:
            raise StackExistsError(
                f"Unable to update stack with name '{stack.name}': Found "
                f"existing stack with this name."
            )
    except KeyError:
        pass

    def __check_component(
        component: StackComponentWrapper,
    ) -> Tuple[StackComponentType, str]:
        try:
            _ = self.get_stack_component(
                component_type=component.type, name=component.name
            )
        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._save_stack(stack.name, stack_configuration)

    logger.info("Updated stack with name '%s'.", name)
    if name != stack.name:
        self.deregister_stack(name)

    return metadata
update_stack_component(self, name, component_type, component)

Update a stack component.

Parameters:

Name Type Description Default
name str

The original name of the stack component.

required
component_type StackComponentType

The type of the stack component to update.

required
component StackComponentWrapper

The new component to update with.

required

Exceptions:

Type Description
KeyError

If no stack component exists with the given name.

Source code in zenml/zen_stores/base_zen_store.py
@abstractmethod
def update_stack_component(
    self,
    name: str,
    component_type: StackComponentType,
    component: StackComponentWrapper,
) -> Dict[str, str]:
    """Update a stack component.

    Args:
        name: The original name of the stack component.
        component_type: The type of the stack component to update.
        component: The new component to update with.

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

local_zen_store

LocalZenStore (BaseZenStore)

Source code in zenml/zen_stores/local_zen_store.py
class LocalZenStore(BaseZenStore):
    def initialize(
        self,
        url: str,
        *args: Any,
        store_data: Optional[ZenStoreModel] = None,
        **kwargs: Any,
    ) -> "LocalZenStore":
        """Initializes a local ZenStore instance.

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

        Returns:
            The initialized ZenStore 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}"
        utils.create_dir_recursive_if_not_exists(str(self._root))

        if store_data is not None:
            self.__store = store_data
            self._write_store()
        elif fileio.exists(self._store_path()):
            config_dict = yaml_utils.read_yaml(self._store_path())
            self.__store = ZenStoreModel.parse_obj(config_dict)
        else:
            self.__store = ZenStoreModel.empty_store()
            self._write_store()

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

    # Public interface implementations:

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

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

    # Static methods:

    @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 LocalZenStore.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://"

    @property
    def is_empty(self) -> bool:
        """Check if the zen 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 zen 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
        )
        utils.create_dir_recursive_if_not_exists(
            os.path.dirname(component_config_path)
        )
        utils.write_file_contents_as_string(
            component_config_path,
            base64.b64decode(component.config).decode(),
        )

        # add the component to the zen 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 update_stack_component(
        self,
        name: str,
        component_type: StackComponentType,
        component: StackComponentWrapper,
    ) -> Dict[str, str]:
        """Update a stack component.

        Args:
            name: The original name of the stack component.
            component_type: The type of the stack component to update.
            component: The new component to update with.

        Raises:
            KeyError: If no stack component exists with the given name.
        """
        components = self.__store.stack_components[component_type]
        if name not in components:
            raise KeyError(
                f"Unable to update stack component (type: {component_type}) "
                f"with name '{name}': No existing stack component "
                f"found with this name."
            )
        elif name != component.name and component.name in components:
            raise StackComponentExistsError(
                f"Unable to update stack component (type: {component_type}) "
                f"with name '{component.name}': a stack component already "
                f"is registered with this name."
            )
        component_config_path = self._get_stack_component_config_path(
            component_type=component.type, name=component.name
        )
        utils.create_dir_recursive_if_not_exists(
            os.path.dirname(component_config_path)
        )
        utils.write_file_contents_as_string(
            component_config_path,
            base64.b64decode(component.config).decode(),
        )
        if name != component.name:
            self._delete_stack_component(component_type, name)

        # add the component to the stack store dict and write it to disk
        components[component.name] = component.flavor

        for _, conf in self.stack_configurations.items():
            for component_type, component_name in conf.items():
                if component_name == name and component_type == component.type:
                    conf[component_type] = component.name
        self._write_store()

        logger.info(
            "Updated stack component with type '%s' and name '%s'.",
            component_type,
            component.name,
        )
        return {component.type.value: component.flavor}

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

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

        Raises:
            KeyError: If no stack exists for the given name.
        """
        del self.__store.stacks[name]
        self._write_store()

    # Private interface implementations:

    def _save_stack(
        self,
        name: str,
        stack_configuration: Dict[StackComponentType, str],
    ) -> None:
        """Save a stack.

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

    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, configuration) 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(
            utils.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.

        Raises:
            KeyError: If no component exists for given type and name.
        """
        component_config_path = self._get_stack_component_config_path(
            component_type=component_type, name=name
        )

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

        components = self.__store.stack_components[component_type]
        del components[name]
        self._write_store()

    # User, project and role management

    @property
    def users(self) -> List[User]:
        """All registered users.

        Returns:
            A list of all registered users.
        """
        return self.__store.users

    def get_user(self, user_name: str) -> User:
        """Gets a specific user.

        Args:
            user_name: Name of the user to get.

        Returns:
            The requested user.

        Raises:
            KeyError: If no user with the given name exists.
        """
        return _get_unique_entity(user_name, collection=self.__store.users)

    def create_user(self, user_name: str) -> User:
        """Creates a new user.

        Args:
            user_name: Unique username.

        Returns:
             The newly created user.

        Raises:
            EntityExistsError: If a user with the given name already exists.
        """
        if _get_unique_entity(
            user_name, collection=self.__store.users, ensure_exists=False
        ):
            raise EntityExistsError(
                f"User with name '{user_name}' already exists."
            )

        user = User(name=user_name)
        self.__store.users.append(user)
        self._write_store()
        return user

    def delete_user(self, user_name: str) -> None:
        """Deletes a user.

        Args:
            user_name: Name of the user to delete.

        Raises:
            KeyError: If no user with the given name exists.
        """
        user = _get_unique_entity(user_name, collection=self.__store.users)
        self.__store.users.remove(user)
        for user_names in self.__store.team_assignments.values():
            user_names.discard(user.name)

        self.__store.role_assignments = [
            assignment
            for assignment in self.__store.role_assignments
            if assignment.user_id != user.id
        ]
        self._write_store()
        logger.info("Deleted user %s.", user)

    @property
    def teams(self) -> List[Team]:
        """All registered teams.

        Returns:
            A list of all registered teams.
        """
        return self.__store.teams

    def get_team(self, team_name: str) -> Team:
        """Gets a specific team.

        Args:
            team_name: Name of the team to get.

        Returns:
            The requested team.

        Raises:
            KeyError: If no team with the given name exists.
        """
        return _get_unique_entity(team_name, collection=self.__store.teams)

    def create_team(self, team_name: str) -> Team:
        """Creates a new team.

        Args:
            team_name: Unique team name.

        Returns:
             The newly created team.

        Raises:
            EntityExistsError: If a team with the given name already exists.
        """
        if _get_unique_entity(
            team_name, collection=self.__store.teams, ensure_exists=False
        ):
            raise EntityExistsError(
                f"Team with name '{team_name}' already exists."
            )

        team = Team(name=team_name)
        self.__store.teams.append(team)
        self._write_store()
        return team

    def delete_team(self, team_name: str) -> None:
        """Deletes a team.

        Args:
            team_name: Name of the team to delete.

        Raises:
            KeyError: If no team with the given name exists.
        """
        team = _get_unique_entity(team_name, collection=self.__store.teams)
        self.__store.teams.remove(team)
        self.__store.team_assignments.pop(team.name, None)
        self.__store.role_assignments = [
            assignment
            for assignment in self.__store.role_assignments
            if assignment.team_id != team.id
        ]
        self._write_store()
        logger.info("Deleted team %s.", team)

    def add_user_to_team(self, team_name: str, user_name: str) -> None:
        """Adds a user to a team.

        Args:
            team_name: Name of the team.
            user_name: Name of the user.

        Raises:
            KeyError: If no user and team with the given names exists.
        """
        team = _get_unique_entity(team_name, self.__store.teams)
        user = _get_unique_entity(user_name, self.__store.users)
        self.__store.team_assignments[team.name].add(user.name)
        self._write_store()

    def remove_user_from_team(self, team_name: str, user_name: str) -> None:
        """Removes a user from a team.

        Args:
            team_name: Name of the team.
            user_name: Name of the user.

        Raises:
            KeyError: If no user and team with the given names exists.
        """
        team = _get_unique_entity(team_name, self.__store.teams)
        user = _get_unique_entity(user_name, self.__store.users)
        self.__store.team_assignments[team.name].remove(user.name)
        self._write_store()

    @property
    def projects(self) -> List[Project]:
        """All registered projects.

        Returns:
            A list of all registered projects.
        """
        return self.__store.projects

    def get_project(self, project_name: str) -> Project:
        """Gets a specific project.

        Args:
            project_name: Name of the project to get.

        Returns:
            The requested project.

        Raises:
            KeyError: If no project with the given name exists.
        """
        return _get_unique_entity(
            project_name, collection=self.__store.projects
        )

    def create_project(
        self, project_name: str, description: Optional[str] = None
    ) -> Project:
        """Creates a new project.

        Args:
            project_name: Unique project name.
            description: Optional project description.

        Returns:
             The newly created project.

        Raises:
            EntityExistsError: If a project with the given name already exists.
        """
        if _get_unique_entity(
            project_name, collection=self.__store.projects, ensure_exists=False
        ):
            raise EntityExistsError(
                f"Project with name '{project_name}' already exists."
            )

        project = Project(name=project_name, description=description)
        self.__store.projects.append(project)
        self._write_store()
        return project

    def delete_project(self, project_name: str) -> None:
        """Deletes a project.

        Args:
            project_name: Name of the project to delete.

        Raises:
            KeyError: If no project with the given name exists.
        """
        project = _get_unique_entity(
            project_name, collection=self.__store.projects
        )
        self.__store.projects.remove(project)
        self.__store.role_assignments = [
            assignment
            for assignment in self.__store.role_assignments
            if assignment.project_id != project.id
        ]

        self._write_store()
        logger.info("Deleted project %s.", project)

    @property
    def roles(self) -> List[Role]:
        """All registered roles.

        Returns:
            A list of all registered roles.
        """
        return self.__store.roles

    @property
    def role_assignments(self) -> List[RoleAssignment]:
        """All registered role assignments.

        Returns:
            A list of all registered role assignments.
        """
        return self.__store.role_assignments

    def get_role(self, role_name: str) -> Role:
        """Gets a specific role.

        Args:
            role_name: Name of the role to get.

        Returns:
            The requested role.

        Raises:
            KeyError: If no role with the given name exists.
        """
        return _get_unique_entity(role_name, collection=self.__store.roles)

    def create_role(self, role_name: str) -> Role:
        """Creates a new role.

        Args:
            role_name: Unique role name.

        Returns:
             The newly created role.

        Raises:
            EntityExistsError: If a role with the given name already exists.
        """
        if _get_unique_entity(
            role_name, collection=self.__store.roles, ensure_exists=False
        ):
            raise EntityExistsError(
                f"Role with name '{role_name}' already exists."
            )

        role = Role(name=role_name)
        self.__store.roles.append(role)
        self._write_store()
        return role

    def delete_role(self, role_name: str) -> None:
        """Deletes a role.

        Args:
            role_name: Name of the role to delete.

        Raises:
            KeyError: If no role with the given name exists.
        """
        role = _get_unique_entity(role_name, collection=self.__store.roles)
        self.__store.roles.remove(role)
        self.__store.role_assignments = [
            assignment
            for assignment in self.__store.role_assignments
            if assignment.role_id != role.id
        ]

        self._write_store()
        logger.info("Deleted role %s.", role)

    def assign_role(
        self,
        role_name: str,
        entity_name: str,
        project_name: Optional[str] = None,
        is_user: bool = True,
    ) -> None:
        """Assigns a role to a user or team.

        Args:
            role_name: Name of the role to assign.
            entity_name: User or team name.
            project_name: Optional project name.
            is_user: Boolean indicating whether the given `entity_name` refers
                to a user.

        Raises:
            KeyError: If no role, entity or project with the given names exists.
        """
        role = _get_unique_entity(role_name, collection=self.__store.roles)
        project_id: Optional[UUID] = None
        if project_name:
            project_id = _get_unique_entity(
                project_name, collection=self.__store.projects
            ).id

        if is_user:
            user = _get_unique_entity(entity_name, self.__store.users)
            assignment = RoleAssignment(
                role_id=role.id, project_id=project_id, user_id=user.id
            )
        else:
            team = _get_unique_entity(entity_name, self.__store.teams)
            assignment = RoleAssignment(
                role_id=role.id, project_id=project_id, team_id=team.id
            )

        self.__store.role_assignments.append(assignment)
        self._write_store()

    def revoke_role(
        self,
        role_name: str,
        entity_name: str,
        project_name: Optional[str] = None,
        is_user: bool = True,
    ) -> None:
        """Revokes a role from a user or team.

        Args:
            role_name: Name of the role to revoke.
            entity_name: User or team name.
            project_name: Optional project name.
            is_user: Boolean indicating whether the given `entity_name` refers
                to a user.

        Raises:
            KeyError: If no role, entity or project with the given names exists.
        """
        role = _get_unique_entity(role_name, collection=self.__store.roles)

        user_id: Optional[UUID] = None
        team_id: Optional[UUID] = None
        project_id: Optional[UUID] = None

        if is_user:
            user_id = _get_unique_entity(entity_name, self.__store.users).id
        else:
            team_id = _get_unique_entity(entity_name, self.__store.teams).id

        if project_name:
            project_id = _get_unique_entity(
                project_name, collection=self.__store.projects
            ).id

        assignments = self._get_role_assignments(
            role_id=role.id,
            user_id=user_id,
            team_id=team_id,
            project_id=project_id,
        )
        if assignments:
            self.__store.role_assignments.remove(
                assignments[0]
            )  # there should only be one
            self._write_store()

    def get_users_for_team(self, team_name: str) -> List[User]:
        """Fetches all users of a team.

        Args:
            team_name: Name of the team.

        Returns:
            List of users that are part of the team.

        Raises:
            KeyError: If no team with the given name exists.
        """
        team = _get_unique_entity(team_name, collection=self.__store.teams)
        user_names = self.__store.team_assignments[team.name]
        return [user for user in self.users if user.name in user_names]

    def get_teams_for_user(self, user_name: str) -> List[Team]:
        """Fetches all teams for a user.

        Args:
            user_name: Name of the user.

        Returns:
            List of teams that the user is part of.

        Raises:
            KeyError: If no user with the given name exists.
        """
        user = _get_unique_entity(user_name, collection=self.__store.users)
        team_names = [
            team_name
            for team_name, user_names in self.__store.team_assignments.items()
            if user.name in user_names
        ]
        return [team for team in self.teams if team.name in team_names]

    def get_role_assignments_for_user(
        self,
        user_name: str,
        project_name: Optional[str] = None,
        include_team_roles: bool = True,
    ) -> List[RoleAssignment]:
        """Fetches all role assignments for a user.

        Args:
            user_name: Name of the user.
            project_name: Optional filter to only return roles assigned for
                this project.
            include_team_roles: If `True`, includes roles for all teams that
                the user is part of.

        Returns:
            List of role assignments for this user.

        Raises:
            KeyError: If no user or project with the given names exists.
        """
        user = _get_unique_entity(user_name, collection=self.__store.users)
        project_id = (
            _get_unique_entity(
                project_name, collection=self.__store.projects
            ).id
            if project_name
            else None
        )
        assignments = self._get_role_assignments(
            user_id=user.id, project_id=project_id
        )

        if include_team_roles:
            for team in self.get_teams_for_user(user_name):
                assignments += self.get_role_assignments_for_team(
                    team.name, project_name=project_name
                )
        return assignments

    def get_role_assignments_for_team(
        self,
        team_name: str,
        project_name: Optional[str] = None,
    ) -> List[RoleAssignment]:
        """Fetches all role assignments for a team.

        Args:
            team_name: Name of the user.
            project_name: Optional filter to only return roles assigned for
                this project.

        Returns:
            List of role assignments for this team.

        Raises:
            KeyError: If no team or project with the given names exists.
        """
        team = _get_unique_entity(team_name, collection=self.__store.teams)
        project_id = (
            _get_unique_entity(
                project_name, collection=self.__store.projects
            ).id
            if project_name
            else None
        )
        return self._get_role_assignments(
            team_id=team.id, project_id=project_id
        )

    # Implementation-specific internal methods:

    @property
    def root(self) -> Path:
        """The root directory of the zen store."""
        if not self._root:
            raise RuntimeError(
                "Local zen 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 zen store yaml file."""
        return str(self.root / "stacks.yaml")

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

    def _get_role_assignments(
        self,
        role_id: Optional[UUID] = None,
        project_id: Optional[UUID] = None,
        user_id: Optional[UUID] = None,
        team_id: Optional[UUID] = None,
    ) -> List[RoleAssignment]:
        """Gets all role assignments that match the criteria.

        Args:
            role_id: Only include role assignments associated with this role id.
            project_id: Only include role assignments associated with this
                project id.
            user_id: Only include role assignments associated with this user id.
            team_id: Only include role assignments associated with this team id.

        Returns:
            List of role assignments.
        """
        return [
            assignment
            for assignment in self.__store.role_assignments
            if not (
                (role_id and assignment.role_id != role_id)
                or (project_id and project_id != assignment.project_id)
                or (user_id and user_id != assignment.user_id)
                or (team_id and team_id != assignment.team_id)
            )
        ]
is_empty: bool property readonly

Check if the zen store is empty.

projects: List[zenml.zen_stores.models.user_management_models.Project] property readonly

All registered projects.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Project]

A list of all registered projects.

role_assignments: List[zenml.zen_stores.models.user_management_models.RoleAssignment] property readonly

All registered role assignments.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

A list of all registered role assignments.

roles: List[zenml.zen_stores.models.user_management_models.Role] property readonly

All registered roles.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Role]

A list of all registered roles.

root: Path property readonly

The root directory of the zen store.

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

Configuration for all stacks registered in this zen store.

Returns:

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

Dictionary mapping stack names to Dict[StackComponentType, str]

teams: List[zenml.zen_stores.models.user_management_models.Team] property readonly

All registered teams.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Team]

A list of all registered teams.

type: StoreType property readonly

The type of zen store.

url: str property readonly

URL of the repository.

users: List[zenml.zen_stores.models.user_management_models.User] property readonly

All registered users.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.User]

A list of all registered users.

add_user_to_team(self, team_name, user_name)

Adds a user to a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required
user_name str

Name of the user.

required

Exceptions:

Type Description
KeyError

If no user and team with the given names exists.

Source code in zenml/zen_stores/local_zen_store.py
def add_user_to_team(self, team_name: str, user_name: str) -> None:
    """Adds a user to a team.

    Args:
        team_name: Name of the team.
        user_name: Name of the user.

    Raises:
        KeyError: If no user and team with the given names exists.
    """
    team = _get_unique_entity(team_name, self.__store.teams)
    user = _get_unique_entity(user_name, self.__store.users)
    self.__store.team_assignments[team.name].add(user.name)
    self._write_store()
assign_role(self, role_name, entity_name, project_name=None, is_user=True)

Assigns a role to a user or team.

Parameters:

Name Type Description Default
role_name str

Name of the role to assign.

required
entity_name str

User or team name.

required
project_name Optional[str]

Optional project name.

None
is_user bool

Boolean indicating whether the given entity_name refers to a user.

True

Exceptions:

Type Description
KeyError

If no role, entity or project with the given names exists.

Source code in zenml/zen_stores/local_zen_store.py
def assign_role(
    self,
    role_name: str,
    entity_name: str,
    project_name: Optional[str] = None,
    is_user: bool = True,
) -> None:
    """Assigns a role to a user or team.

    Args:
        role_name: Name of the role to assign.
        entity_name: User or team name.
        project_name: Optional project name.
        is_user: Boolean indicating whether the given `entity_name` refers
            to a user.

    Raises:
        KeyError: If no role, entity or project with the given names exists.
    """
    role = _get_unique_entity(role_name, collection=self.__store.roles)
    project_id: Optional[UUID] = None
    if project_name:
        project_id = _get_unique_entity(
            project_name, collection=self.__store.projects
        ).id

    if is_user:
        user = _get_unique_entity(entity_name, self.__store.users)
        assignment = RoleAssignment(
            role_id=role.id, project_id=project_id, user_id=user.id
        )
    else:
        team = _get_unique_entity(entity_name, self.__store.teams)
        assignment = RoleAssignment(
            role_id=role.id, project_id=project_id, team_id=team.id
        )

    self.__store.role_assignments.append(assignment)
    self._write_store()
create_project(self, project_name, description=None)

Creates a new project.

Parameters:

Name Type Description Default
project_name str

Unique project name.

required
description Optional[str]

Optional project description.

None

Returns:

Type Description
Project

The newly created project.

Exceptions:

Type Description
EntityExistsError

If a project with the given name already exists.

Source code in zenml/zen_stores/local_zen_store.py
def create_project(
    self, project_name: str, description: Optional[str] = None
) -> Project:
    """Creates a new project.

    Args:
        project_name: Unique project name.
        description: Optional project description.

    Returns:
         The newly created project.

    Raises:
        EntityExistsError: If a project with the given name already exists.
    """
    if _get_unique_entity(
        project_name, collection=self.__store.projects, ensure_exists=False
    ):
        raise EntityExistsError(
            f"Project with name '{project_name}' already exists."
        )

    project = Project(name=project_name, description=description)
    self.__store.projects.append(project)
    self._write_store()
    return project
create_role(self, role_name)

Creates a new role.

Parameters:

Name Type Description Default
role_name str

Unique role name.

required

Returns:

Type Description
Role

The newly created role.

Exceptions:

Type Description
EntityExistsError

If a role with the given name already exists.

Source code in zenml/zen_stores/local_zen_store.py
def create_role(self, role_name: str) -> Role:
    """Creates a new role.

    Args:
        role_name: Unique role name.

    Returns:
         The newly created role.

    Raises:
        EntityExistsError: If a role with the given name already exists.
    """
    if _get_unique_entity(
        role_name, collection=self.__store.roles, ensure_exists=False
    ):
        raise EntityExistsError(
            f"Role with name '{role_name}' already exists."
        )

    role = Role(name=role_name)
    self.__store.roles.append(role)
    self._write_store()
    return role
create_team(self, team_name)

Creates a new team.

Parameters:

Name Type Description Default
team_name str

Unique team name.

required

Returns:

Type Description
Team

The newly created team.

Exceptions:

Type Description
EntityExistsError

If a team with the given name already exists.

Source code in zenml/zen_stores/local_zen_store.py
def create_team(self, team_name: str) -> Team:
    """Creates a new team.

    Args:
        team_name: Unique team name.

    Returns:
         The newly created team.

    Raises:
        EntityExistsError: If a team with the given name already exists.
    """
    if _get_unique_entity(
        team_name, collection=self.__store.teams, ensure_exists=False
    ):
        raise EntityExistsError(
            f"Team with name '{team_name}' already exists."
        )

    team = Team(name=team_name)
    self.__store.teams.append(team)
    self._write_store()
    return team
create_user(self, user_name)

Creates a new user.

Parameters:

Name Type Description Default
user_name str

Unique username.

required

Returns:

Type Description
User

The newly created user.

Exceptions:

Type Description
EntityExistsError

If a user with the given name already exists.

Source code in zenml/zen_stores/local_zen_store.py
def create_user(self, user_name: str) -> User:
    """Creates a new user.

    Args:
        user_name: Unique username.

    Returns:
         The newly created user.

    Raises:
        EntityExistsError: If a user with the given name already exists.
    """
    if _get_unique_entity(
        user_name, collection=self.__store.users, ensure_exists=False
    ):
        raise EntityExistsError(
            f"User with name '{user_name}' already exists."
        )

    user = User(name=user_name)
    self.__store.users.append(user)
    self._write_store()
    return user
delete_project(self, project_name)

Deletes a project.

Parameters:

Name Type Description Default
project_name str

Name of the project to delete.

required

Exceptions:

Type Description
KeyError

If no project with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def delete_project(self, project_name: str) -> None:
    """Deletes a project.

    Args:
        project_name: Name of the project to delete.

    Raises:
        KeyError: If no project with the given name exists.
    """
    project = _get_unique_entity(
        project_name, collection=self.__store.projects
    )
    self.__store.projects.remove(project)
    self.__store.role_assignments = [
        assignment
        for assignment in self.__store.role_assignments
        if assignment.project_id != project.id
    ]

    self._write_store()
    logger.info("Deleted project %s.", project)
delete_role(self, role_name)

Deletes a role.

Parameters:

Name Type Description Default
role_name str

Name of the role to delete.

required

Exceptions:

Type Description
KeyError

If no role with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def delete_role(self, role_name: str) -> None:
    """Deletes a role.

    Args:
        role_name: Name of the role to delete.

    Raises:
        KeyError: If no role with the given name exists.
    """
    role = _get_unique_entity(role_name, collection=self.__store.roles)
    self.__store.roles.remove(role)
    self.__store.role_assignments = [
        assignment
        for assignment in self.__store.role_assignments
        if assignment.role_id != role.id
    ]

    self._write_store()
    logger.info("Deleted role %s.", role)
delete_team(self, team_name)

Deletes a team.

Parameters:

Name Type Description Default
team_name str

Name of the team to delete.

required

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def delete_team(self, team_name: str) -> None:
    """Deletes a team.

    Args:
        team_name: Name of the team to delete.

    Raises:
        KeyError: If no team with the given name exists.
    """
    team = _get_unique_entity(team_name, collection=self.__store.teams)
    self.__store.teams.remove(team)
    self.__store.team_assignments.pop(team.name, None)
    self.__store.role_assignments = [
        assignment
        for assignment in self.__store.role_assignments
        if assignment.team_id != team.id
    ]
    self._write_store()
    logger.info("Deleted team %s.", team)
delete_user(self, user_name)

Deletes a user.

Parameters:

Name Type Description Default
user_name str

Name of the user to delete.

required

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def delete_user(self, user_name: str) -> None:
    """Deletes a user.

    Args:
        user_name: Name of the user to delete.

    Raises:
        KeyError: If no user with the given name exists.
    """
    user = _get_unique_entity(user_name, collection=self.__store.users)
    self.__store.users.remove(user)
    for user_names in self.__store.team_assignments.values():
        user_names.discard(user.name)

    self.__store.role_assignments = [
        assignment
        for assignment in self.__store.role_assignments
        if assignment.user_id != user.id
    ]
    self._write_store()
    logger.info("Deleted user %s.", user)
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

Exceptions:

Type Description
KeyError

If no stack exists for the given name.

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

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

    Raises:
        KeyError: If no stack exists for the given name.
    """
    del self.__store.stacks[name]
    self._write_store()
get_local_url(path) staticmethod

Get a local URL for a given local path.

Source code in zenml/zen_stores/local_zen_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/zen_stores/local_zen_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 LocalZenStore.is_valid_url(url):
        raise ValueError(f"Invalid URL for local store: {url}")
    url = url.replace("file://", "")
    return Path(url)
get_project(self, project_name)

Gets a specific project.

Parameters:

Name Type Description Default
project_name str

Name of the project to get.

required

Returns:

Type Description
Project

The requested project.

Exceptions:

Type Description
KeyError

If no project with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def get_project(self, project_name: str) -> Project:
    """Gets a specific project.

    Args:
        project_name: Name of the project to get.

    Returns:
        The requested project.

    Raises:
        KeyError: If no project with the given name exists.
    """
    return _get_unique_entity(
        project_name, collection=self.__store.projects
    )
get_role(self, role_name)

Gets a specific role.

Parameters:

Name Type Description Default
role_name str

Name of the role to get.

required

Returns:

Type Description
Role

The requested role.

Exceptions:

Type Description
KeyError

If no role with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def get_role(self, role_name: str) -> Role:
    """Gets a specific role.

    Args:
        role_name: Name of the role to get.

    Returns:
        The requested role.

    Raises:
        KeyError: If no role with the given name exists.
    """
    return _get_unique_entity(role_name, collection=self.__store.roles)
get_role_assignments_for_team(self, team_name, project_name=None)

Fetches all role assignments for a team.

Parameters:

Name Type Description Default
team_name str

Name of the user.

required
project_name Optional[str]

Optional filter to only return roles assigned for this project.

None

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

List of role assignments for this team.

Exceptions:

Type Description
KeyError

If no team or project with the given names exists.

Source code in zenml/zen_stores/local_zen_store.py
def get_role_assignments_for_team(
    self,
    team_name: str,
    project_name: Optional[str] = None,
) -> List[RoleAssignment]:
    """Fetches all role assignments for a team.

    Args:
        team_name: Name of the user.
        project_name: Optional filter to only return roles assigned for
            this project.

    Returns:
        List of role assignments for this team.

    Raises:
        KeyError: If no team or project with the given names exists.
    """
    team = _get_unique_entity(team_name, collection=self.__store.teams)
    project_id = (
        _get_unique_entity(
            project_name, collection=self.__store.projects
        ).id
        if project_name
        else None
    )
    return self._get_role_assignments(
        team_id=team.id, project_id=project_id
    )
get_role_assignments_for_user(self, user_name, project_name=None, include_team_roles=True)

Fetches all role assignments for a user.

Parameters:

Name Type Description Default
user_name str

Name of the user.

required
project_name Optional[str]

Optional filter to only return roles assigned for this project.

None
include_team_roles bool

If True, includes roles for all teams that the user is part of.

True

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

List of role assignments for this user.

Exceptions:

Type Description
KeyError

If no user or project with the given names exists.

Source code in zenml/zen_stores/local_zen_store.py
def get_role_assignments_for_user(
    self,
    user_name: str,
    project_name: Optional[str] = None,
    include_team_roles: bool = True,
) -> List[RoleAssignment]:
    """Fetches all role assignments for a user.

    Args:
        user_name: Name of the user.
        project_name: Optional filter to only return roles assigned for
            this project.
        include_team_roles: If `True`, includes roles for all teams that
            the user is part of.

    Returns:
        List of role assignments for this user.

    Raises:
        KeyError: If no user or project with the given names exists.
    """
    user = _get_unique_entity(user_name, collection=self.__store.users)
    project_id = (
        _get_unique_entity(
            project_name, collection=self.__store.projects
        ).id
        if project_name
        else None
    )
    assignments = self._get_role_assignments(
        user_id=user.id, project_id=project_id
    )

    if include_team_roles:
        for team in self.get_teams_for_user(user_name):
            assignments += self.get_role_assignments_for_team(
                team.name, project_name=project_name
            )
    return assignments
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/zen_stores/local_zen_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]
get_team(self, team_name)

Gets a specific team.

Parameters:

Name Type Description Default
team_name str

Name of the team to get.

required

Returns:

Type Description
Team

The requested team.

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def get_team(self, team_name: str) -> Team:
    """Gets a specific team.

    Args:
        team_name: Name of the team to get.

    Returns:
        The requested team.

    Raises:
        KeyError: If no team with the given name exists.
    """
    return _get_unique_entity(team_name, collection=self.__store.teams)
get_teams_for_user(self, user_name)

Fetches all teams for a user.

Parameters:

Name Type Description Default
user_name str

Name of the user.

required

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Team]

List of teams that the user is part of.

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def get_teams_for_user(self, user_name: str) -> List[Team]:
    """Fetches all teams for a user.

    Args:
        user_name: Name of the user.

    Returns:
        List of teams that the user is part of.

    Raises:
        KeyError: If no user with the given name exists.
    """
    user = _get_unique_entity(user_name, collection=self.__store.users)
    team_names = [
        team_name
        for team_name, user_names in self.__store.team_assignments.items()
        if user.name in user_names
    ]
    return [team for team in self.teams if team.name in team_names]
get_user(self, user_name)

Gets a specific user.

Parameters:

Name Type Description Default
user_name str

Name of the user to get.

required

Returns:

Type Description
User

The requested user.

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def get_user(self, user_name: str) -> User:
    """Gets a specific user.

    Args:
        user_name: Name of the user to get.

    Returns:
        The requested user.

    Raises:
        KeyError: If no user with the given name exists.
    """
    return _get_unique_entity(user_name, collection=self.__store.users)
get_users_for_team(self, team_name)

Fetches all users of a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.User]

List of users that are part of the team.

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/local_zen_store.py
def get_users_for_team(self, team_name: str) -> List[User]:
    """Fetches all users of a team.

    Args:
        team_name: Name of the team.

    Returns:
        List of users that are part of the team.

    Raises:
        KeyError: If no team with the given name exists.
    """
    team = _get_unique_entity(team_name, collection=self.__store.teams)
    user_names = self.__store.team_assignments[team.name]
    return [user for user in self.users if user.name in user_names]
initialize(self, url, *args, *, store_data=None, **kwargs)

Initializes a local ZenStore instance.

Parameters:

Name Type Description Default
url str

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

required
store_data Optional[zenml.zen_stores.models.zen_store_model.ZenStoreModel]

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

None
args Any

additional positional arguments (ignored).

()
kwargs Any

additional keyword arguments (ignored).

{}

Returns:

Type Description
LocalZenStore

The initialized ZenStore instance.

Source code in zenml/zen_stores/local_zen_store.py
def initialize(
    self,
    url: str,
    *args: Any,
    store_data: Optional[ZenStoreModel] = None,
    **kwargs: Any,
) -> "LocalZenStore":
    """Initializes a local ZenStore instance.

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

    Returns:
        The initialized ZenStore 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}"
    utils.create_dir_recursive_if_not_exists(str(self._root))

    if store_data is not None:
        self.__store = store_data
        self._write_store()
    elif fileio.exists(self._store_path()):
        config_dict = yaml_utils.read_yaml(self._store_path())
        self.__store = ZenStoreModel.parse_obj(config_dict)
    else:
        self.__store = ZenStoreModel.empty_store()
        self._write_store()

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

Check if the given url is a valid local path.

Source code in zenml/zen_stores/local_zen_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/zen_stores/local_zen_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
    )
    utils.create_dir_recursive_if_not_exists(
        os.path.dirname(component_config_path)
    )
    utils.write_file_contents_as_string(
        component_config_path,
        base64.b64decode(component.config).decode(),
    )

    # add the component to the zen 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,
    )
remove_user_from_team(self, team_name, user_name)

Removes a user from a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required
user_name str

Name of the user.

required

Exceptions:

Type Description
KeyError

If no user and team with the given names exists.

Source code in zenml/zen_stores/local_zen_store.py
def remove_user_from_team(self, team_name: str, user_name: str) -> None:
    """Removes a user from a team.

    Args:
        team_name: Name of the team.
        user_name: Name of the user.

    Raises:
        KeyError: If no user and team with the given names exists.
    """
    team = _get_unique_entity(team_name, self.__store.teams)
    user = _get_unique_entity(user_name, self.__store.users)
    self.__store.team_assignments[team.name].remove(user.name)
    self._write_store()
revoke_role(self, role_name, entity_name, project_name=None, is_user=True)

Revokes a role from a user or team.

Parameters:

Name Type Description Default
role_name str

Name of the role to revoke.

required
entity_name str

User or team name.

required
project_name Optional[str]

Optional project name.

None
is_user bool

Boolean indicating whether the given entity_name refers to a user.

True

Exceptions:

Type Description
KeyError

If no role, entity or project with the given names exists.

Source code in zenml/zen_stores/local_zen_store.py
def revoke_role(
    self,
    role_name: str,
    entity_name: str,
    project_name: Optional[str] = None,
    is_user: bool = True,
) -> None:
    """Revokes a role from a user or team.

    Args:
        role_name: Name of the role to revoke.
        entity_name: User or team name.
        project_name: Optional project name.
        is_user: Boolean indicating whether the given `entity_name` refers
            to a user.

    Raises:
        KeyError: If no role, entity or project with the given names exists.
    """
    role = _get_unique_entity(role_name, collection=self.__store.roles)

    user_id: Optional[UUID] = None
    team_id: Optional[UUID] = None
    project_id: Optional[UUID] = None

    if is_user:
        user_id = _get_unique_entity(entity_name, self.__store.users).id
    else:
        team_id = _get_unique_entity(entity_name, self.__store.teams).id

    if project_name:
        project_id = _get_unique_entity(
            project_name, collection=self.__store.projects
        ).id

    assignments = self._get_role_assignments(
        role_id=role.id,
        user_id=user_id,
        team_id=team_id,
        project_id=project_id,
    )
    if assignments:
        self.__store.role_assignments.remove(
            assignments[0]
        )  # there should only be one
        self._write_store()
update_stack_component(self, name, component_type, component)

Update a stack component.

Parameters:

Name Type Description Default
name str

The original name of the stack component.

required
component_type StackComponentType

The type of the stack component to update.

required
component StackComponentWrapper

The new component to update with.

required

Exceptions:

Type Description
KeyError

If no stack component exists with the given name.

Source code in zenml/zen_stores/local_zen_store.py
def update_stack_component(
    self,
    name: str,
    component_type: StackComponentType,
    component: StackComponentWrapper,
) -> Dict[str, str]:
    """Update a stack component.

    Args:
        name: The original name of the stack component.
        component_type: The type of the stack component to update.
        component: The new component to update with.

    Raises:
        KeyError: If no stack component exists with the given name.
    """
    components = self.__store.stack_components[component_type]
    if name not in components:
        raise KeyError(
            f"Unable to update stack component (type: {component_type}) "
            f"with name '{name}': No existing stack component "
            f"found with this name."
        )
    elif name != component.name and component.name in components:
        raise StackComponentExistsError(
            f"Unable to update stack component (type: {component_type}) "
            f"with name '{component.name}': a stack component already "
            f"is registered with this name."
        )
    component_config_path = self._get_stack_component_config_path(
        component_type=component.type, name=component.name
    )
    utils.create_dir_recursive_if_not_exists(
        os.path.dirname(component_config_path)
    )
    utils.write_file_contents_as_string(
        component_config_path,
        base64.b64decode(component.config).decode(),
    )
    if name != component.name:
        self._delete_stack_component(component_type, name)

    # add the component to the stack store dict and write it to disk
    components[component.name] = component.flavor

    for _, conf in self.stack_configurations.items():
        for component_type, component_name in conf.items():
            if component_name == name and component_type == component.type:
                conf[component_type] = component.name
    self._write_store()

    logger.info(
        "Updated stack component with type '%s' and name '%s'.",
        component_type,
        component.name,
    )
    return {component.type.value: component.flavor}

models special

stack_component_wrapper

StackComponentWrapper (BaseModel) pydantic-model

Serializable Configuration of a StackComponent

Source code in zenml/zen_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,
            name=component.name,
            uuid=component.uuid,
            config=base64.b64encode(
                yaml.dump(json.loads(component.json())).encode()
            ),
        )

stack_wrapper

StackWrapper (BaseModel) pydantic-model

Network Serializable Wrapper describing a Stack.

Source code in zenml/zen_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()
            ],
        )

user_management_models

Operation (BaseModel) pydantic-model

Pydantic object representing an operation that requires permission.

Attributes:

Name Type Description
id int

Operation id.

name str

Operation name.

Source code in zenml/zen_stores/models/user_management_models.py
class Operation(BaseModel):
    """Pydantic object representing an operation that requires permission.

    Attributes:
        id: Operation id.
        name: Operation name.
    """

    id: int
    name: str
Permission (BaseModel) pydantic-model

Pydantic object representing permissions on a specific resource.

Attributes:

Name Type Description
operation Operation

The operation for which the permissions are.

types Set[zenml.zen_stores.models.user_management_models.PermissionType]

Types of permissions.

Source code in zenml/zen_stores/models/user_management_models.py
class Permission(BaseModel):
    """Pydantic object representing permissions on a specific resource.

    Attributes:
        operation: The operation for which the permissions are.
        types: Types of permissions.
    """

    operation: Operation
    types: Set[PermissionType]

    class Config:
        # similar to non-mutable but also makes the object hashable
        frozen = True
PermissionType (Enum)

All permission types.

Source code in zenml/zen_stores/models/user_management_models.py
class PermissionType(Enum):
    """All permission types."""

    CREATE = "create"
    READ = "read"
    UPDATE = "update"
    DELETE = "delete"
Project (BaseModel) pydantic-model

Pydantic object representing a project.

Attributes:

Name Type Description
id UUID

Id of the project.

creation_date datetime

Date when the project was created.

name str

Name of the project.

description Optional[str]

Optional project description.

Source code in zenml/zen_stores/models/user_management_models.py
class Project(BaseModel):
    """Pydantic object representing a project.

    Attributes:
        id: Id of the project.
        creation_date: Date when the project was created.
        name: Name of the project.
        description: Optional project description.
    """

    id: UUID = Field(default_factory=uuid4)
    creation_date: datetime = Field(default_factory=datetime.now)
    name: str
    description: Optional[str] = None
Role (BaseModel) pydantic-model

Pydantic object representing a role.

Attributes:

Name Type Description
id UUID

Id of the role.

creation_date datetime

Date when the role was created.

name str

Name of the role.

permissions Set[zenml.zen_stores.models.user_management_models.Permission]

Set of permissions allowed by this role.

Source code in zenml/zen_stores/models/user_management_models.py
class Role(BaseModel):
    """Pydantic object representing a role.

    Attributes:
        id: Id of the role.
        creation_date: Date when the role was created.
        name: Name of the role.
        permissions: Set of permissions allowed by this role.
    """

    id: UUID = Field(default_factory=uuid4)
    creation_date: datetime = Field(default_factory=datetime.now)
    name: str
    permissions: Set[Permission] = set()
RoleAssignment (BaseModel) pydantic-model

Pydantic object representing a role assignment.

Attributes:

Name Type Description
id UUID

Id of the role assignment.

creation_date datetime

Date when the role was assigned.

role_id UUID

Id of the role.

project_id Optional[uuid.UUID]

Optional ID of a project that the role is limited to.

team_id Optional[uuid.UUID]

Id of a team to which the role is assigned.

user_id Optional[uuid.UUID]

Id of a user to which the role is assigned.

Source code in zenml/zen_stores/models/user_management_models.py
class RoleAssignment(BaseModel):
    """Pydantic object representing a role assignment.

    Attributes:
        id: Id of the role assignment.
        creation_date: Date when the role was assigned.
        role_id: Id of the role.
        project_id: Optional ID of a project that the role is limited to.
        team_id: Id of a team to which the role is assigned.
        user_id: Id of a user to which the role is assigned.
    """

    id: UUID = Field(default_factory=uuid4)
    creation_date: datetime = Field(default_factory=datetime.now)
    role_id: UUID
    project_id: Optional[UUID] = None
    team_id: Optional[UUID] = None
    user_id: Optional[UUID] = None

    @root_validator
    def ensure_single_entity(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        """Validates that either `user_id` or `team_id` is set."""

        user_id = values.get("user_id", None)
        team_id = values.get("team_id", None)
        if user_id and team_id:
            raise ValueError("Only `user_id` or `team_id` is allowed.")

        if not (user_id or team_id):
            raise ValueError(
                "Missing `user_id` or `team_id` for role assignment."
            )

        return values
ensure_single_entity(values) classmethod

Validates that either user_id or team_id is set.

Source code in zenml/zen_stores/models/user_management_models.py
@root_validator
def ensure_single_entity(cls, values: Dict[str, Any]) -> Dict[str, Any]:
    """Validates that either `user_id` or `team_id` is set."""

    user_id = values.get("user_id", None)
    team_id = values.get("team_id", None)
    if user_id and team_id:
        raise ValueError("Only `user_id` or `team_id` is allowed.")

    if not (user_id or team_id):
        raise ValueError(
            "Missing `user_id` or `team_id` for role assignment."
        )

    return values
Team (BaseModel) pydantic-model

Pydantic object representing a team.

Attributes:

Name Type Description
id UUID

Id of the team.

creation_date datetime

Date when the team was created.

name str

Name of the team.

Source code in zenml/zen_stores/models/user_management_models.py
class Team(BaseModel):
    """Pydantic object representing a team.

    Attributes:
        id: Id of the team.
        creation_date: Date when the team was created.
        name: Name of the team.
    """

    id: UUID = Field(default_factory=uuid4)
    creation_date: datetime = Field(default_factory=datetime.now)
    name: str
User (BaseModel) pydantic-model

Pydantic object representing a user.

Attributes:

Name Type Description
id UUID

Id of the user.

creation_date datetime

Date when the user was created.

name str

Name of the user.

Source code in zenml/zen_stores/models/user_management_models.py
class User(BaseModel):
    """Pydantic object representing a user.

    Attributes:
        id: Id of the user.
        creation_date: Date when the user was created.
        name: Name of the user.
    """

    id: UUID = Field(default_factory=uuid4)
    creation_date: datetime = Field(default_factory=datetime.now)
    name: str
    # email: str
    # password: str

zen_store_model

ZenStoreModel (BaseModel) pydantic-model

Pydantic object used for serializing a ZenStore.

Attributes:

Name Type Description
version

zenml version number

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.

users List[zenml.zen_stores.models.user_management_models.User]

All registered users.

teams List[zenml.zen_stores.models.user_management_models.Team]

All registered teams.

projects List[zenml.zen_stores.models.user_management_models.Project]

All registered projects.

roles List[zenml.zen_stores.models.user_management_models.Role]

All registered roles.

role_assignments List[zenml.zen_stores.models.user_management_models.RoleAssignment]

All role assignments.

team_assignments DefaultDict[str, Set[str]]

Maps team names to names of users that are part of the team.

Source code in zenml/zen_stores/models/zen_store_model.py
class ZenStoreModel(BaseModel):
    """Pydantic object used for serializing a ZenStore.

    Attributes:
        version: zenml version number
        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.
        users: All registered users.
        teams: All registered teams.
        projects: All registered projects.
        roles: All registered roles.
        role_assignments: All role assignments.
        team_assignments: Maps team names to names of users that are part of
            the team.
    """

    stacks: Dict[str, Dict[StackComponentType, str]]
    stack_components: DefaultDict[StackComponentType, Dict[str, str]]
    users: List[User] = []
    teams: List[Team] = []
    projects: List[Project] = []
    roles: List[Role] = []
    role_assignments: List[RoleAssignment] = []
    team_assignments: DefaultDict[str, Set[str]] = defaultdict(set)

    @validator("stack_components")
    def _construct_stack_components_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)

    @validator("team_assignments")
    def _construct_team_assignments_defaultdict(
        cls, team_assignments: Dict[str, Set[str]]
    ) -> DefaultDict[str, Set[str]]:
        """Ensures that `team_assignments` is a defaultdict so users
        of a new teams can be added without issues."""
        return defaultdict(set, team_assignments)

    @classmethod
    def empty_store(cls) -> "ZenStoreModel":
        """Initialize a new empty zen 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/zen_stores/models/zen_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 zen store with current zen version.

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

rest_zen_store

RestZenStore (BaseZenStore)

ZenStore implementation for accessing data from a REST api.

Source code in zenml/zen_stores/rest_zen_store.py
class RestZenStore(BaseZenStore):
    """ZenStore implementation for accessing data from a REST api."""

    def initialize(
        self,
        url: str,
        *args: Any,
        **kwargs: Any,
    ) -> "RestZenStore":
        """Initializes a rest zen store instance.

        Args:
            url: Endpoint URL of the service for zen storage.
            args: additional positional arguments (ignored).
            kwargs: additional keyword arguments (ignored).

        Returns:
            The initialized zen store instance.
        """
        if not self.is_valid_url(url.strip("/")):
            raise ValueError("Invalid URL for REST store: {url}")
        self._url = url.strip("/")
        if "skip_default_stack" not in kwargs:
            kwargs["skip_default_stack"] = True
        super().initialize(url, *args, **kwargs)
        return self

    # Static methods:

    @staticmethod
    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:
            None, because there are no local paths from REST urls.
        """
        return None

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

        Args:
             path: the path string to build a URL out of.

        Returns:
            Url pointing to the path for the store type.

        Raises:
            NotImplementedError: always
        """
        raise NotImplementedError("Cannot build a REST url from a 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 (
            scheme is not None
            and scheme.group() in ("https://", "http://")
            and url[-1] != "/"
        )

    # Public Interface:

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

    @property
    def url(self) -> str:
        """Get the stack store URL."""
        return self._url

    @property
    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.
        """
        empty = self.get(IS_EMPTY)
        if not isinstance(empty, bool):
            raise ValueError(
                f"Bad API Response. Expected boolean, got:\n{empty}"
            )
        return empty

    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.
        """
        return self._parse_stack_configuration(
            self.get(f"{STACK_CONFIGURATIONS}/{name}")
        )

    @property
    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
        """
        body = self.get(STACK_CONFIGURATIONS)
        if not isinstance(body, dict):
            raise ValueError(
                f"Bad API Response. Expected dict, got {type(body)}"
            )
        return {
            key: self._parse_stack_configuration(value)
            for key, value in body.items()
        }

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

        Args:
            component: The component to register.

        Raises:
            KeyError: If a stack component with the same type
                and name already exists.
        """
        self.post(STACK_COMPONENTS, body=component)

    def update_stack_component(
        self,
        name: str,
        component_type: StackComponentType,
        component: StackComponentWrapper,
    ) -> Dict[str, str]:
        """Update a stack component.

        Args:
            name: The original name of the stack component.
            component_type: The type of the stack component to update.
            component: The new component to update with.

        Raises:
            KeyError: If no stack component exists with the given name.
        """
        body = self.put(
            f"{STACK_COMPONENTS}/{component_type}/{name}", body=component
        )
        if isinstance(body, dict):
            return cast(Dict[str, str], body)
        else:
            raise ValueError(
                f"Bad API Response. Expected dict, got {type(body)}"
            )

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

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

        Raises:
            KeyError: If no stack exists for the given name.
        """
        self.delete(f"{STACKS}/{name}")

    def _save_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.
        """
        raise NotImplementedError

    # Custom implementations:

    @property
    def stacks(self) -> List[StackWrapper]:
        """All stacks registered in this repository."""
        body = self.get(STACKS)
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [StackWrapper.parse_obj(s) for s in body]

    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 StackWrapper.parse_obj(self.get(f"{STACKS}/{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.
        """
        body = self.post(STACKS, stack)
        if isinstance(body, dict):
            return cast(Dict[str, str], body)
        else:
            raise ValueError(
                f"Bad API Response. Expected dict, got {type(body)}"
            )

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

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

        Args:
            name: The original name of the stack.
            stack: The new stack to use in the update.

        Returns:
            metadata dict for telemetry or logging.

        Raises:
            ValueError: If a dict is not returned from the API.
        """
        body = self.put(f"{STACKS}/{name}", body=stack)
        if name != stack.name:
            self.deregister_stack(name)
        if isinstance(body, dict):
            return cast(Dict[str, str], body)
        else:
            raise ValueError(
                f"Bad API Response. Expected dict, got {type(body)}"
            )

    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.
        """
        return StackComponentWrapper.parse_obj(
            self.get(f"{STACK_COMPONENTS}/{component_type}/{name}")
        )

    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.
        """
        body = self.get(f"{STACK_COMPONENTS}/{component_type}")
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [StackComponentWrapper.parse_obj(c) for c in body]

    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.
        """
        self.delete(f"{STACK_COMPONENTS}/{component_type}/{name}")

    # User, project and role management

    @property
    def users(self) -> List[User]:
        """All registered users.

        Returns:
            A list of all registered users.

        Raises:
            ValueError: In case of a bad API response.
        """
        body = self.get(USERS)
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [User.parse_obj(user_dict) for user_dict in body]

    def get_user(self, user_name: str) -> User:
        """Gets a specific user.

        Args:
            user_name: Name of the user to get.

        Returns:
            The requested user.

        Raises:
            KeyError: If no user with the given name exists.
        """
        return User.parse_obj(self.get(f"{USERS}/{user_name}"))

    def create_user(self, user_name: str) -> User:
        """Creates a new user.

        Args:
            user_name: Unique username.

        Returns:
             The newly created user.

        Raises:
            EntityExistsError: If a user with the given name already exists.
        """
        user = User(name=user_name)
        return User.parse_obj(self.post(USERS, body=user))

    def delete_user(self, user_name: str) -> None:
        """Deletes a user.

        Args:
            user_name: Name of the user to delete.

        Raises:
            KeyError: If no user with the given name exists.
        """
        self.delete(f"{USERS}/{user_name}")

    @property
    def teams(self) -> List[Team]:
        """All registered teams.

        Returns:
            A list of all registered teams.

        Raises:
            ValueError: In case of a bad API response.
        """
        body = self.get(TEAMS)
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [Team.parse_obj(team_dict) for team_dict in body]

    def get_team(self, team_name: str) -> Team:
        """Gets a specific team.

        Args:
            team_name: Name of the team to get.

        Returns:
            The requested team.

        Raises:
            KeyError: If no team with the given name exists.
        """
        return Team.parse_obj(self.get(f"{TEAMS}/{team_name}"))

    def create_team(self, team_name: str) -> Team:
        """Creates a new team.

        Args:
            team_name: Unique team name.

        Returns:
             The newly created team.

        Raises:
            EntityExistsError: If a team with the given name already exists.
        """
        team = Team(name=team_name)
        return Team.parse_obj(self.post(TEAMS, body=team))

    def delete_team(self, team_name: str) -> None:
        """Deletes a team.

        Args:
            team_name: Name of the team to delete.

        Raises:
            KeyError: If no team with the given name exists.
        """
        self.delete(f"{TEAMS}/{team_name}")

    def add_user_to_team(self, team_name: str, user_name: str) -> None:
        """Adds a user to a team.

        Args:
            team_name: Name of the team.
            user_name: Name of the user.

        Raises:
            KeyError: If no user and team with the given names exists.
        """
        user = User(name=user_name)
        self.post(f"{TEAMS}/{team_name}/users", user)

    def remove_user_from_team(self, team_name: str, user_name: str) -> None:
        """Removes a user from a team.

        Args:
            team_name: Name of the team.
            user_name: Name of the user.

        Raises:
            KeyError: If no user and team with the given names exists.
        """
        self.delete(f"{TEAMS}/{team_name}/users/{user_name}")

    @property
    def projects(self) -> List[Project]:
        """All registered projects.

        Returns:
            A list of all registered projects.

        Raises:
            ValueError: In case of a bad API response.
        """
        body = self.get(PROJECTS)
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [Project.parse_obj(project_dict) for project_dict in body]

    def get_project(self, project_name: str) -> Project:
        """Gets a specific project.

        Args:
            project_name: Name of the project to get.

        Returns:
            The requested project.

        Raises:
            KeyError: If no project with the given name exists.
        """
        return Project.parse_obj(self.get(f"{PROJECTS}/{project_name}"))

    def create_project(
        self, project_name: str, description: Optional[str] = None
    ) -> Project:
        """Creates a new project.

        Args:
            project_name: Unique project name.
            description: Optional project description.

        Returns:
             The newly created project.

        Raises:
            EntityExistsError: If a project with the given name already exists.
        """
        project = Project(name=project_name, description=description)
        return Project.parse_obj(self.post(PROJECTS, body=project))

    def delete_project(self, project_name: str) -> None:
        """Deletes a project.

        Args:
            project_name: Name of the project to delete.

        Raises:
            KeyError: If no project with the given name exists.
        """
        self.delete(f"{PROJECTS}/{project_name}")

    @property
    def roles(self) -> List[Role]:
        """All registered roles.

        Returns:
            A list of all registered roles.

        Raises:
            ValueError: In case of a bad API response.
        """
        body = self.get(ROLES)
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [Role.parse_obj(role_dict) for role_dict in body]

    @property
    def role_assignments(self) -> List[RoleAssignment]:
        """All registered role assignments.

        Returns:
            A list of all registered role assignments.

        Raises:
            ValueError: In case of a bad API response.
        """
        body = self.get(ROLE_ASSIGNMENTS)
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [
            RoleAssignment.parse_obj(assignment_dict)
            for assignment_dict in body
        ]

    def get_role(self, role_name: str) -> Role:
        """Gets a specific role.

        Args:
            role_name: Name of the role to get.

        Returns:
            The requested role.

        Raises:
            KeyError: If no role with the given name exists.
        """
        return Role.parse_obj(self.get(f"{ROLES}/{role_name}"))

    def create_role(self, role_name: str) -> Role:
        """Creates a new role.

        Args:
            role_name: Unique role name.

        Returns:
             The newly created role.

        Raises:
            EntityExistsError: If a role with the given name already exists.
        """
        role = Role(name=role_name)
        return Role.parse_obj(self.post(ROLES, body=role))

    def delete_role(self, role_name: str) -> None:
        """Deletes a role.

        Args:
            role_name: Name of the role to delete.

        Raises:
            KeyError: If no role with the given name exists.
        """
        self.delete(f"{ROLES}/{role_name}")

    def assign_role(
        self,
        role_name: str,
        entity_name: str,
        project_name: Optional[str] = None,
        is_user: bool = True,
    ) -> None:
        """Assigns a role to a user or team.

        Args:
            role_name: Name of the role to assign.
            entity_name: User or team name.
            project_name: Optional project name.
            is_user: Boolean indicating whether the given `entity_name` refers
                to a user.

        Raises:
            KeyError: If no role, entity or project with the given names exists.
        """
        data = {
            "role_name": role_name,
            "entity_name": entity_name,
            "project_name": project_name,
            "is_user": is_user,
        }
        self._handle_response(
            requests.post(
                self.url + ROLE_ASSIGNMENTS,
                json=data,
                auth=self._get_authentication(),
            )
        )

    def revoke_role(
        self,
        role_name: str,
        entity_name: str,
        project_name: Optional[str] = None,
        is_user: bool = True,
    ) -> None:
        """Revokes a role from a user or team.

        Args:
            role_name: Name of the role to revoke.
            entity_name: User or team name.
            project_name: Optional project name.
            is_user: Boolean indicating whether the given `entity_name` refers
                to a user.

        Raises:
            KeyError: If no role, entity or project with the given names exists.
        """
        data = {
            "role_name": role_name,
            "entity_name": entity_name,
            "project_name": project_name,
            "is_user": is_user,
        }
        self._handle_response(
            requests.delete(
                self.url + ROLE_ASSIGNMENTS,
                json=data,
                auth=self._get_authentication(),
            )
        )

    def get_users_for_team(self, team_name: str) -> List[User]:
        """Fetches all users of a team.

        Args:
            team_name: Name of the team.

        Returns:
            List of users that are part of the team.

        Raises:
            KeyError: If no team with the given name exists.
            ValueError: In case of a bad API response.
        """
        body = self.get(f"{TEAMS}/{team_name}/users")
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [User.parse_obj(user_dict) for user_dict in body]

    def get_teams_for_user(self, user_name: str) -> List[Team]:
        """Fetches all teams for a user.

        Args:
            user_name: Name of the user.

        Returns:
            List of teams that the user is part of.

        Raises:
            KeyError: If no user with the given name exists.
            ValueError: In case of a bad API response.
        """
        body = self.get(f"{USERS}/{user_name}/teams")
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [Team.parse_obj(team_dict) for team_dict in body]

    def get_role_assignments_for_user(
        self,
        user_name: str,
        project_name: Optional[str] = None,
        include_team_roles: bool = True,
    ) -> List[RoleAssignment]:
        """Fetches all role assignments for a user.

        Args:
            user_name: Name of the user.
            project_name: Optional filter to only return roles assigned for
                this project.
            include_team_roles: If `True`, includes roles for all teams that
                the user is part of.

        Returns:
            List of role assignments for this user.

        Raises:
            KeyError: If no user or project with the given names exists.
            ValueError: In case of a bad API response.
        """
        path = f"{USERS}/{user_name}/role_assignments"
        if project_name:
            path += f"?project_name={project_name}"

        body = self.get(path)
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        assignments = [
            RoleAssignment.parse_obj(assignment_dict)
            for assignment_dict in body
        ]
        if include_team_roles:
            for team in self.get_teams_for_user(user_name):
                assignments += self.get_role_assignments_for_team(
                    team.name, project_name=project_name
                )
        return assignments

    def get_role_assignments_for_team(
        self,
        team_name: str,
        project_name: Optional[str] = None,
    ) -> List[RoleAssignment]:
        """Fetches all role assignments for a team.

        Args:
            team_name: Name of the user.
            project_name: Optional filter to only return roles assigned for
                this project.

        Returns:
            List of role assignments for this team.

        Raises:
            KeyError: If no user or project with the given names exists.
            ValueError: In case of a bad API response.
        """
        path = f"{TEAMS}/{team_name}/role_assignments"
        if project_name:
            path += f"?project_name={project_name}"

        body = self.get(path)
        if not isinstance(body, list):
            raise ValueError(
                f"Bad API Response. Expected list, got {type(body)}"
            )
        return [
            RoleAssignment.parse_obj(assignment_dict)
            for assignment_dict in body
        ]

    # Private interface shall not be implemented for REST store, instead the
    # API only provides all public methods, including the ones that would
    # otherwise be inherited from the BaseZenStore in other implementations.
    # Don't call these! ABC complains that they aren't implemented, but they
    # aren't needed with the custom implementations of base methods.

    def _create_stack(
        self, name: str, stack_configuration: Dict[StackComponentType, str]
    ) -> None:
        """Add a stack to storage"""
        raise NotImplementedError("Not to be accessed directly in client!")

    def _get_component_flavor_and_config(
        self, component_type: StackComponentType, name: str
    ) -> Tuple[str, bytes]:
        """Fetch the flavor and configuration for a stack component."""
        raise NotImplementedError("Not to be accessed directly in client!")

    def _get_stack_component_names(
        self, component_type: StackComponentType
    ) -> List[str]:
        """Get names of all registered stack components of a given type."""
        raise NotImplementedError("Not to be accessed directly in client!")

    def _delete_stack_component(
        self, component_type: StackComponentType, name: str
    ) -> None:
        """Remove a StackComponent from storage."""
        raise NotImplementedError("Not to be accessed directly in client!")

    # Implementation specific methods:

    def _parse_stack_configuration(
        self, to_parse: Json
    ) -> Dict[StackComponentType, str]:
        """Parse an API response into `Dict[StackComponentType, str]`."""
        if not isinstance(to_parse, dict):
            raise ValueError(
                f"Bad API Response. Expected dict, got {type(to_parse)}."
            )
        return {
            StackComponentType(typ): component_name
            for typ, component_name in to_parse.items()
        }

    def _handle_response(self, response: requests.Response) -> Json:
        """Handle API response, translating http status codes to Exception."""
        if response.status_code >= 200 and response.status_code < 300:
            try:
                payload: Json = response.json()
                return payload
            except requests.exceptions.JSONDecodeError:
                raise ValueError(
                    "Bad response from API. Expected json, got\n"
                    f"{response.text}"
                )
        elif response.status_code == 401:
            raise requests.HTTPError(
                f"{response.status_code} Client Error: Unauthorized request to URL {response.url}: {response.json().get('detail')}"
            )
        elif response.status_code == 404:
            if "DoesNotExistException" not in response.text:
                raise KeyError(*response.json().get("detail", (response.text,)))
            message = ": ".join(response.json().get("detail", (response.text,)))
            raise DoesNotExistException(message)
        elif response.status_code == 409:
            if "StackComponentExistsError" in response.text:
                raise StackComponentExistsError(
                    *response.json().get("detail", (response.text,))
                )
            elif "StackExistsError" in response.text:
                raise StackExistsError(
                    *response.json().get("detail", (response.text,))
                )
            elif "EntityExistsError" in response.text:
                raise EntityExistsError(
                    *response.json().get("detail", (response.text,))
                )
            else:
                raise ValueError(
                    *response.json().get("detail", (response.text,))
                )
        elif response.status_code == 422:
            raise RuntimeError(*response.json().get("detail", (response.text,)))
        elif response.status_code == 500:
            raise KeyError(response.text)
        else:
            raise RuntimeError(
                "Error retrieving from API. Got response "
                f"{response.status_code} with body:\n{response.text}"
            )

    @staticmethod
    def _get_authentication() -> Tuple[str, str]:
        """Gets HTTP basic auth credentials."""
        from zenml.repository import Repository

        return Repository().active_user_name, ""

    def get(self, path: str) -> Json:
        """Make a GET request to the given endpoint path."""
        return self._handle_response(
            requests.get(self.url + path, auth=self._get_authentication())
        )

    def delete(self, path: str) -> Json:
        """Make a DELETE request to the given endpoint path."""
        return self._handle_response(
            requests.delete(self.url + path, auth=self._get_authentication())
        )

    def post(self, path: str, body: BaseModel) -> Json:
        """Make a POST request to the given endpoint path."""
        endpoint = self.url + path
        return self._handle_response(
            requests.post(
                endpoint, data=body.json(), auth=self._get_authentication()
            )
        )

    def put(self, path: str, body: BaseModel) -> Json:
        """Make a PUT request to the given endpoint path."""
        endpoint = self.url + path
        return self._handle_response(
            requests.put(
                endpoint, data=body.json(), auth=self._get_authentication()
            )
        )
is_empty: bool property readonly

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.

projects: List[zenml.zen_stores.models.user_management_models.Project] property readonly

All registered projects.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Project]

A list of all registered projects.

Exceptions:

Type Description
ValueError

In case of a bad API response.

role_assignments: List[zenml.zen_stores.models.user_management_models.RoleAssignment] property readonly

All registered role assignments.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

A list of all registered role assignments.

Exceptions:

Type Description
ValueError

In case of a bad API response.

roles: List[zenml.zen_stores.models.user_management_models.Role] property readonly

All registered roles.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Role]

A list of all registered roles.

Exceptions:

Type Description
ValueError

In case of a bad API response.

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.zen_stores.models.stack_wrapper.StackWrapper] property readonly

All stacks registered in this repository.

teams: List[zenml.zen_stores.models.user_management_models.Team] property readonly

All registered teams.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Team]

A list of all registered teams.

Exceptions:

Type Description
ValueError

In case of a bad API response.

type: StoreType property readonly

The type of stack store.

url: str property readonly

Get the stack store URL.

users: List[zenml.zen_stores.models.user_management_models.User] property readonly

All registered users.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.User]

A list of all registered users.

Exceptions:

Type Description
ValueError

In case of a bad API response.

add_user_to_team(self, team_name, user_name)

Adds a user to a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required
user_name str

Name of the user.

required

Exceptions:

Type Description
KeyError

If no user and team with the given names exists.

Source code in zenml/zen_stores/rest_zen_store.py
def add_user_to_team(self, team_name: str, user_name: str) -> None:
    """Adds a user to a team.

    Args:
        team_name: Name of the team.
        user_name: Name of the user.

    Raises:
        KeyError: If no user and team with the given names exists.
    """
    user = User(name=user_name)
    self.post(f"{TEAMS}/{team_name}/users", user)
assign_role(self, role_name, entity_name, project_name=None, is_user=True)

Assigns a role to a user or team.

Parameters:

Name Type Description Default
role_name str

Name of the role to assign.

required
entity_name str

User or team name.

required
project_name Optional[str]

Optional project name.

None
is_user bool

Boolean indicating whether the given entity_name refers to a user.

True

Exceptions:

Type Description
KeyError

If no role, entity or project with the given names exists.

Source code in zenml/zen_stores/rest_zen_store.py
def assign_role(
    self,
    role_name: str,
    entity_name: str,
    project_name: Optional[str] = None,
    is_user: bool = True,
) -> None:
    """Assigns a role to a user or team.

    Args:
        role_name: Name of the role to assign.
        entity_name: User or team name.
        project_name: Optional project name.
        is_user: Boolean indicating whether the given `entity_name` refers
            to a user.

    Raises:
        KeyError: If no role, entity or project with the given names exists.
    """
    data = {
        "role_name": role_name,
        "entity_name": entity_name,
        "project_name": project_name,
        "is_user": is_user,
    }
    self._handle_response(
        requests.post(
            self.url + ROLE_ASSIGNMENTS,
            json=data,
            auth=self._get_authentication(),
        )
    )
create_project(self, project_name, description=None)

Creates a new project.

Parameters:

Name Type Description Default
project_name str

Unique project name.

required
description Optional[str]

Optional project description.

None

Returns:

Type Description
Project

The newly created project.

Exceptions:

Type Description
EntityExistsError

If a project with the given name already exists.

Source code in zenml/zen_stores/rest_zen_store.py
def create_project(
    self, project_name: str, description: Optional[str] = None
) -> Project:
    """Creates a new project.

    Args:
        project_name: Unique project name.
        description: Optional project description.

    Returns:
         The newly created project.

    Raises:
        EntityExistsError: If a project with the given name already exists.
    """
    project = Project(name=project_name, description=description)
    return Project.parse_obj(self.post(PROJECTS, body=project))
create_role(self, role_name)

Creates a new role.

Parameters:

Name Type Description Default
role_name str

Unique role name.

required

Returns:

Type Description
Role

The newly created role.

Exceptions:

Type Description
EntityExistsError

If a role with the given name already exists.

Source code in zenml/zen_stores/rest_zen_store.py
def create_role(self, role_name: str) -> Role:
    """Creates a new role.

    Args:
        role_name: Unique role name.

    Returns:
         The newly created role.

    Raises:
        EntityExistsError: If a role with the given name already exists.
    """
    role = Role(name=role_name)
    return Role.parse_obj(self.post(ROLES, body=role))
create_team(self, team_name)

Creates a new team.

Parameters:

Name Type Description Default
team_name str

Unique team name.

required

Returns:

Type Description
Team

The newly created team.

Exceptions:

Type Description
EntityExistsError

If a team with the given name already exists.

Source code in zenml/zen_stores/rest_zen_store.py
def create_team(self, team_name: str) -> Team:
    """Creates a new team.

    Args:
        team_name: Unique team name.

    Returns:
         The newly created team.

    Raises:
        EntityExistsError: If a team with the given name already exists.
    """
    team = Team(name=team_name)
    return Team.parse_obj(self.post(TEAMS, body=team))
create_user(self, user_name)

Creates a new user.

Parameters:

Name Type Description Default
user_name str

Unique username.

required

Returns:

Type Description
User

The newly created user.

Exceptions:

Type Description
EntityExistsError

If a user with the given name already exists.

Source code in zenml/zen_stores/rest_zen_store.py
def create_user(self, user_name: str) -> User:
    """Creates a new user.

    Args:
        user_name: Unique username.

    Returns:
         The newly created user.

    Raises:
        EntityExistsError: If a user with the given name already exists.
    """
    user = User(name=user_name)
    return User.parse_obj(self.post(USERS, body=user))
delete(self, path)

Make a DELETE request to the given endpoint path.

Source code in zenml/zen_stores/rest_zen_store.py
def delete(self, path: str) -> Json:
    """Make a DELETE request to the given endpoint path."""
    return self._handle_response(
        requests.delete(self.url + path, auth=self._get_authentication())
    )
delete_project(self, project_name)

Deletes a project.

Parameters:

Name Type Description Default
project_name str

Name of the project to delete.

required

Exceptions:

Type Description
KeyError

If no project with the given name exists.

Source code in zenml/zen_stores/rest_zen_store.py
def delete_project(self, project_name: str) -> None:
    """Deletes a project.

    Args:
        project_name: Name of the project to delete.

    Raises:
        KeyError: If no project with the given name exists.
    """
    self.delete(f"{PROJECTS}/{project_name}")
delete_role(self, role_name)

Deletes a role.

Parameters:

Name Type Description Default
role_name str

Name of the role to delete.

required

Exceptions:

Type Description
KeyError

If no role with the given name exists.

Source code in zenml/zen_stores/rest_zen_store.py
def delete_role(self, role_name: str) -> None:
    """Deletes a role.

    Args:
        role_name: Name of the role to delete.

    Raises:
        KeyError: If no role with the given name exists.
    """
    self.delete(f"{ROLES}/{role_name}")
delete_team(self, team_name)

Deletes a team.

Parameters:

Name Type Description Default
team_name str

Name of the team to delete.

required

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/rest_zen_store.py
def delete_team(self, team_name: str) -> None:
    """Deletes a team.

    Args:
        team_name: Name of the team to delete.

    Raises:
        KeyError: If no team with the given name exists.
    """
    self.delete(f"{TEAMS}/{team_name}")
delete_user(self, user_name)

Deletes a user.

Parameters:

Name Type Description Default
user_name str

Name of the user to delete.

required

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/rest_zen_store.py
def delete_user(self, user_name: str) -> None:
    """Deletes a user.

    Args:
        user_name: Name of the user to delete.

    Raises:
        KeyError: If no user with the given name exists.
    """
    self.delete(f"{USERS}/{user_name}")
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

Exceptions:

Type Description
KeyError

If no stack exists for the given name.

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

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

    Raises:
        KeyError: If no stack exists for the given name.
    """
    self.delete(f"{STACKS}/{name}")
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/zen_stores/rest_zen_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.
    """
    self.delete(f"{STACK_COMPONENTS}/{component_type}/{name}")
get(self, path)

Make a GET request to the given endpoint path.

Source code in zenml/zen_stores/rest_zen_store.py
def get(self, path: str) -> Json:
    """Make a GET request to the given endpoint path."""
    return self._handle_response(
        requests.get(self.url + path, auth=self._get_authentication())
    )
get_local_url(path) staticmethod

Get a local URL for a given local path.

Parameters:

Name Type Description Default
path str

the path string to build a URL out of.

required

Returns:

Type Description
str

Url pointing to the path for the store type.

Exceptions:

Type Description
NotImplementedError

always

Source code in zenml/zen_stores/rest_zen_store.py
@staticmethod
def get_local_url(path: str) -> str:
    """Get a local URL for a given local path.

    Args:
         path: the path string to build a URL out of.

    Returns:
        Url pointing to the path for the store type.

    Raises:
        NotImplementedError: always
    """
    raise NotImplementedError("Cannot build a REST url from a 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]

None, because there are no local paths from REST urls.

Source code in zenml/zen_stores/rest_zen_store.py
@staticmethod
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:
        None, because there are no local paths from REST urls.
    """
    return None
get_project(self, project_name)

Gets a specific project.

Parameters:

Name Type Description Default
project_name str

Name of the project to get.

required

Returns:

Type Description
Project

The requested project.

Exceptions:

Type Description
KeyError

If no project with the given name exists.

Source code in zenml/zen_stores/rest_zen_store.py
def get_project(self, project_name: str) -> Project:
    """Gets a specific project.

    Args:
        project_name: Name of the project to get.

    Returns:
        The requested project.

    Raises:
        KeyError: If no project with the given name exists.
    """
    return Project.parse_obj(self.get(f"{PROJECTS}/{project_name}"))
get_role(self, role_name)

Gets a specific role.

Parameters:

Name Type Description Default
role_name str

Name of the role to get.

required

Returns:

Type Description
Role

The requested role.

Exceptions:

Type Description
KeyError

If no role with the given name exists.

Source code in zenml/zen_stores/rest_zen_store.py
def get_role(self, role_name: str) -> Role:
    """Gets a specific role.

    Args:
        role_name: Name of the role to get.

    Returns:
        The requested role.

    Raises:
        KeyError: If no role with the given name exists.
    """
    return Role.parse_obj(self.get(f"{ROLES}/{role_name}"))
get_role_assignments_for_team(self, team_name, project_name=None)

Fetches all role assignments for a team.

Parameters:

Name Type Description Default
team_name str

Name of the user.

required
project_name Optional[str]

Optional filter to only return roles assigned for this project.

None

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

List of role assignments for this team.

Exceptions:

Type Description
KeyError

If no user or project with the given names exists.

ValueError

In case of a bad API response.

Source code in zenml/zen_stores/rest_zen_store.py
def get_role_assignments_for_team(
    self,
    team_name: str,
    project_name: Optional[str] = None,
) -> List[RoleAssignment]:
    """Fetches all role assignments for a team.

    Args:
        team_name: Name of the user.
        project_name: Optional filter to only return roles assigned for
            this project.

    Returns:
        List of role assignments for this team.

    Raises:
        KeyError: If no user or project with the given names exists.
        ValueError: In case of a bad API response.
    """
    path = f"{TEAMS}/{team_name}/role_assignments"
    if project_name:
        path += f"?project_name={project_name}"

    body = self.get(path)
    if not isinstance(body, list):
        raise ValueError(
            f"Bad API Response. Expected list, got {type(body)}"
        )
    return [
        RoleAssignment.parse_obj(assignment_dict)
        for assignment_dict in body
    ]
get_role_assignments_for_user(self, user_name, project_name=None, include_team_roles=True)

Fetches all role assignments for a user.

Parameters:

Name Type Description Default
user_name str

Name of the user.

required
project_name Optional[str]

Optional filter to only return roles assigned for this project.

None
include_team_roles bool

If True, includes roles for all teams that the user is part of.

True

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

List of role assignments for this user.

Exceptions:

Type Description
KeyError

If no user or project with the given names exists.

ValueError

In case of a bad API response.

Source code in zenml/zen_stores/rest_zen_store.py
def get_role_assignments_for_user(
    self,
    user_name: str,
    project_name: Optional[str] = None,
    include_team_roles: bool = True,
) -> List[RoleAssignment]:
    """Fetches all role assignments for a user.

    Args:
        user_name: Name of the user.
        project_name: Optional filter to only return roles assigned for
            this project.
        include_team_roles: If `True`, includes roles for all teams that
            the user is part of.

    Returns:
        List of role assignments for this user.

    Raises:
        KeyError: If no user or project with the given names exists.
        ValueError: In case of a bad API response.
    """
    path = f"{USERS}/{user_name}/role_assignments"
    if project_name:
        path += f"?project_name={project_name}"

    body = self.get(path)
    if not isinstance(body, list):
        raise ValueError(
            f"Bad API Response. Expected list, got {type(body)}"
        )
    assignments = [
        RoleAssignment.parse_obj(assignment_dict)
        for assignment_dict in body
    ]
    if include_team_roles:
        for team in self.get_teams_for_user(user_name):
            assignments += self.get_role_assignments_for_team(
                team.name, project_name=project_name
            )
    return assignments
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/zen_stores/rest_zen_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 StackWrapper.parse_obj(self.get(f"{STACKS}/{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/zen_stores/rest_zen_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.
    """
    return StackComponentWrapper.parse_obj(
        self.get(f"{STACK_COMPONENTS}/{component_type}/{name}")
    )
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.zen_stores.models.stack_component_wrapper.StackComponentWrapper]

A list of StackComponentConfiguration instances.

Source code in zenml/zen_stores/rest_zen_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.
    """
    body = self.get(f"{STACK_COMPONENTS}/{component_type}")
    if not isinstance(body, list):
        raise ValueError(
            f"Bad API Response. Expected list, got {type(body)}"
        )
    return [StackComponentWrapper.parse_obj(c) for c in body]
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/zen_stores/rest_zen_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.
    """
    return self._parse_stack_configuration(
        self.get(f"{STACK_CONFIGURATIONS}/{name}")
    )
get_team(self, team_name)

Gets a specific team.

Parameters:

Name Type Description Default
team_name str

Name of the team to get.

required

Returns:

Type Description
Team

The requested team.

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/rest_zen_store.py
def get_team(self, team_name: str) -> Team:
    """Gets a specific team.

    Args:
        team_name: Name of the team to get.

    Returns:
        The requested team.

    Raises:
        KeyError: If no team with the given name exists.
    """
    return Team.parse_obj(self.get(f"{TEAMS}/{team_name}"))
get_teams_for_user(self, user_name)

Fetches all teams for a user.

Parameters:

Name Type Description Default
user_name str

Name of the user.

required

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Team]

List of teams that the user is part of.

Exceptions:

Type Description
KeyError

If no user with the given name exists.

ValueError

In case of a bad API response.

Source code in zenml/zen_stores/rest_zen_store.py
def get_teams_for_user(self, user_name: str) -> List[Team]:
    """Fetches all teams for a user.

    Args:
        user_name: Name of the user.

    Returns:
        List of teams that the user is part of.

    Raises:
        KeyError: If no user with the given name exists.
        ValueError: In case of a bad API response.
    """
    body = self.get(f"{USERS}/{user_name}/teams")
    if not isinstance(body, list):
        raise ValueError(
            f"Bad API Response. Expected list, got {type(body)}"
        )
    return [Team.parse_obj(team_dict) for team_dict in body]
get_user(self, user_name)

Gets a specific user.

Parameters:

Name Type Description Default
user_name str

Name of the user to get.

required

Returns:

Type Description
User

The requested user.

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/rest_zen_store.py
def get_user(self, user_name: str) -> User:
    """Gets a specific user.

    Args:
        user_name: Name of the user to get.

    Returns:
        The requested user.

    Raises:
        KeyError: If no user with the given name exists.
    """
    return User.parse_obj(self.get(f"{USERS}/{user_name}"))
get_users_for_team(self, team_name)

Fetches all users of a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.User]

List of users that are part of the team.

Exceptions:

Type Description
KeyError

If no team with the given name exists.

ValueError

In case of a bad API response.

Source code in zenml/zen_stores/rest_zen_store.py
def get_users_for_team(self, team_name: str) -> List[User]:
    """Fetches all users of a team.

    Args:
        team_name: Name of the team.

    Returns:
        List of users that are part of the team.

    Raises:
        KeyError: If no team with the given name exists.
        ValueError: In case of a bad API response.
    """
    body = self.get(f"{TEAMS}/{team_name}/users")
    if not isinstance(body, list):
        raise ValueError(
            f"Bad API Response. Expected list, got {type(body)}"
        )
    return [User.parse_obj(user_dict) for user_dict in body]
initialize(self, url, *args, **kwargs)

Initializes a rest zen store instance.

Parameters:

Name Type Description Default
url str

Endpoint URL of the service for zen storage.

required
args Any

additional positional arguments (ignored).

()
kwargs Any

additional keyword arguments (ignored).

{}

Returns:

Type Description
RestZenStore

The initialized zen store instance.

Source code in zenml/zen_stores/rest_zen_store.py
def initialize(
    self,
    url: str,
    *args: Any,
    **kwargs: Any,
) -> "RestZenStore":
    """Initializes a rest zen store instance.

    Args:
        url: Endpoint URL of the service for zen storage.
        args: additional positional arguments (ignored).
        kwargs: additional keyword arguments (ignored).

    Returns:
        The initialized zen store instance.
    """
    if not self.is_valid_url(url.strip("/")):
        raise ValueError("Invalid URL for REST store: {url}")
    self._url = url.strip("/")
    if "skip_default_stack" not in kwargs:
        kwargs["skip_default_stack"] = True
    super().initialize(url, *args, **kwargs)
    return self
is_valid_url(url) staticmethod

Check if the given url is a valid local path.

Source code in zenml/zen_stores/rest_zen_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 (
        scheme is not None
        and scheme.group() in ("https://", "http://")
        and url[-1] != "/"
    )
post(self, path, body)

Make a POST request to the given endpoint path.

Source code in zenml/zen_stores/rest_zen_store.py
def post(self, path: str, body: BaseModel) -> Json:
    """Make a POST request to the given endpoint path."""
    endpoint = self.url + path
    return self._handle_response(
        requests.post(
            endpoint, data=body.json(), auth=self._get_authentication()
        )
    )
put(self, path, body)

Make a PUT request to the given endpoint path.

Source code in zenml/zen_stores/rest_zen_store.py
def put(self, path: str, body: BaseModel) -> Json:
    """Make a PUT request to the given endpoint path."""
    endpoint = self.url + path
    return self._handle_response(
        requests.put(
            endpoint, data=body.json(), auth=self._get_authentication()
        )
    )
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/zen_stores/rest_zen_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.
    """
    body = self.post(STACKS, stack)
    if isinstance(body, dict):
        return cast(Dict[str, str], body)
    else:
        raise ValueError(
            f"Bad API Response. Expected dict, got {type(body)}"
        )
register_stack_component(self, component)

Register a stack component.

Parameters:

Name Type Description Default
component StackComponentWrapper

The component to register.

required

Exceptions:

Type Description
KeyError

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

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

    Args:
        component: The component to register.

    Raises:
        KeyError: If a stack component with the same type
            and name already exists.
    """
    self.post(STACK_COMPONENTS, body=component)
remove_user_from_team(self, team_name, user_name)

Removes a user from a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required
user_name str

Name of the user.

required

Exceptions:

Type Description
KeyError

If no user and team with the given names exists.

Source code in zenml/zen_stores/rest_zen_store.py
def remove_user_from_team(self, team_name: str, user_name: str) -> None:
    """Removes a user from a team.

    Args:
        team_name: Name of the team.
        user_name: Name of the user.

    Raises:
        KeyError: If no user and team with the given names exists.
    """
    self.delete(f"{TEAMS}/{team_name}/users/{user_name}")
revoke_role(self, role_name, entity_name, project_name=None, is_user=True)

Revokes a role from a user or team.

Parameters:

Name Type Description Default
role_name str

Name of the role to revoke.

required
entity_name str

User or team name.

required
project_name Optional[str]

Optional project name.

None
is_user bool

Boolean indicating whether the given entity_name refers to a user.

True

Exceptions:

Type Description
KeyError

If no role, entity or project with the given names exists.

Source code in zenml/zen_stores/rest_zen_store.py
def revoke_role(
    self,
    role_name: str,
    entity_name: str,
    project_name: Optional[str] = None,
    is_user: bool = True,
) -> None:
    """Revokes a role from a user or team.

    Args:
        role_name: Name of the role to revoke.
        entity_name: User or team name.
        project_name: Optional project name.
        is_user: Boolean indicating whether the given `entity_name` refers
            to a user.

    Raises:
        KeyError: If no role, entity or project with the given names exists.
    """
    data = {
        "role_name": role_name,
        "entity_name": entity_name,
        "project_name": project_name,
        "is_user": is_user,
    }
    self._handle_response(
        requests.delete(
            self.url + ROLE_ASSIGNMENTS,
            json=data,
            auth=self._get_authentication(),
        )
    )
update_stack(self, name, stack)

Update a stack and its components.

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

Parameters:

Name Type Description Default
name str

The original name of the stack.

required
stack StackWrapper

The new stack to use in the update.

required

Returns:

Type Description
Dict[str, str]

metadata dict for telemetry or logging.

Exceptions:

Type Description
ValueError

If a dict is not returned from the API.

Source code in zenml/zen_stores/rest_zen_store.py
def update_stack(self, name: str, stack: StackWrapper) -> Dict[str, str]:
    """Update a stack and its components.

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

    Args:
        name: The original name of the stack.
        stack: The new stack to use in the update.

    Returns:
        metadata dict for telemetry or logging.

    Raises:
        ValueError: If a dict is not returned from the API.
    """
    body = self.put(f"{STACKS}/{name}", body=stack)
    if name != stack.name:
        self.deregister_stack(name)
    if isinstance(body, dict):
        return cast(Dict[str, str], body)
    else:
        raise ValueError(
            f"Bad API Response. Expected dict, got {type(body)}"
        )
update_stack_component(self, name, component_type, component)

Update a stack component.

Parameters:

Name Type Description Default
name str

The original name of the stack component.

required
component_type StackComponentType

The type of the stack component to update.

required
component StackComponentWrapper

The new component to update with.

required

Exceptions:

Type Description
KeyError

If no stack component exists with the given name.

Source code in zenml/zen_stores/rest_zen_store.py
def update_stack_component(
    self,
    name: str,
    component_type: StackComponentType,
    component: StackComponentWrapper,
) -> Dict[str, str]:
    """Update a stack component.

    Args:
        name: The original name of the stack component.
        component_type: The type of the stack component to update.
        component: The new component to update with.

    Raises:
        KeyError: If no stack component exists with the given name.
    """
    body = self.put(
        f"{STACK_COMPONENTS}/{component_type}/{name}", body=component
    )
    if isinstance(body, dict):
        return cast(Dict[str, str], body)
    else:
        raise ValueError(
            f"Bad API Response. Expected dict, got {type(body)}"
        )

sql_zen_store

SqlZenStore (BaseZenStore)

Repository Implementation that uses SQL database backend

Source code in zenml/zen_stores/sql_zen_store.py
class SqlZenStore(BaseZenStore):
    """Repository Implementation that uses SQL database backend"""

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

        Args:
            url: odbc path to a database.
            args, kwargs: additional parameters for SQLModel.
        Returns:
            The initialized zen store instance.
        """
        if not self.is_valid_url(url):
            raise ValueError(f"Invalid URL for SQL store: {url}")

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

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

        # we need to remove `skip_default_registrations` from the kwargs,
        # because SQLModel will raise an error if it is present
        sql_kwargs = kwargs.copy()
        sql_kwargs.pop("skip_default_registrations", False)
        self.engine = create_engine(url, *args, **sql_kwargs)
        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 zen store."""
        return StoreType.SQL

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

    # Static methods:

    @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 SqlZenStore.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

    @property
    def is_empty(self) -> bool:
        """Check if the zen 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 zen 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 update_stack_component(
        self,
        name: str,
        component_type: StackComponentType,
        component: StackComponentWrapper,
    ) -> Dict[str, str]:
        """Update a stack component.

        Args:
            name: The original name of the stack component.
            component_type: The type of the stack component to update.
            component: The new component to update with.

        Raises:
            KeyError: If no stack component exists with the given name.
        """
        with Session(self.engine) as session:
            updated_component = session.exec(
                select(ZenStackComponent)
                .where(ZenStackComponent.component_type == component_type)
                .where(ZenStackComponent.name == name)
            ).first()

            if not updated_component:
                raise KeyError(
                    f"Unable to update stack component (type: {component.type}) "
                    f"with name '{component.name}': No existing stack component "
                    f"found with this name."
                )

            new_name_component = session.exec(
                select(ZenStackComponent)
                .where(ZenStackComponent.component_type == component_type)
                .where(ZenStackComponent.name == component.name)
            ).first()
            if (name != component.name) and new_name_component is not None:
                raise StackComponentExistsError(
                    f"Unable to update stack component (type: "
                    f"{component.type}) with name '{component.name}': Found "
                    f"existing stack component with this name."
                )

            updated_component.configuration = component.config

            # handle any potential renamed component
            updated_component.name = component.name

            # rename components inside stacks
            updated_stack_definitions = session.exec(
                select(ZenStackDefinition)
                .where(ZenStackDefinition.component_type == component_type)
                .where(ZenStackDefinition.component_name == name)
            ).all()
            for stack_definition in updated_stack_definitions:
                stack_definition.component_name = component.name
                session.add(stack_definition)

            session.add(updated_component)
            session.commit()
        logger.info(
            "Updated stack component with type '%s' and name '%s'.",
            component_type,
            component.name,
        )
        return {component.type.value: component.flavor}

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

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

        Raises:
            KeyError: If no stack exists for the given name.
        """
        with Session(self.engine) as session:
            try:
                stack = session.exec(
                    select(ZenStack).where(ZenStack.name == name)
                ).one()
                session.delete(stack)
            except NoResultFound as error:
                raise KeyError from error
            definitions = session.exec(
                select(ZenStackDefinition).where(
                    ZenStackDefinition.stack_name == name
                )
            ).all()
            for definition in definitions:
                session.delete(definition)
            session.commit()

    # Private interface implementations:

    def _save_stack(
        self,
        name: str,
        stack_configuration: Dict[StackComponentType, str],
    ) -> None:
        """Save a stack.

        Args:
            name: The name to save the stack as.
            stack_configuration: Dict[StackComponentType, str] to persist.
        """
        with Session(self.engine) as session:
            stack = session.exec(
                select(ZenStack).where(ZenStack.name == name)
            ).first()
            if stack is None:
                stack = ZenStack(name=name, created_by=1)
                session.add(stack)
            for ctype, cname in stack_configuration.items():
                statement = (
                    select(ZenStackDefinition)
                    .where(ZenStackDefinition.stack_name == name)
                    .where(ZenStackDefinition.component_type == ctype)
                )
                results = session.exec(statement)
                component = results.one_or_none()
                if component is None:
                    session.add(
                        ZenStackDefinition(
                            stack_name=name,
                            component_type=ctype,
                            component_name=cname,
                        )
                    )
                else:
                    component.component_name = cname
                    component.component_type = ctype
                    session.add(component)
            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, configuration) 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.

        Raises:
            KeyError: If no component exists for 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)
            ).first()
            if component is not None:
                session.delete(component)
                session.commit()
            else:
                raise KeyError(
                    "Unable to deregister stack component (type: "
                    f"{component_type.value}) with name '{name}': No stack "
                    "component exists with this name."
                )

    # User, project and role management

    @property
    def users(self) -> List[User]:
        """All registered users.

        Returns:
            A list of all registered users.
        """
        with Session(self.engine) as session:
            return [
                User(**user.dict())
                for user in session.exec(select(UserTable)).all()
            ]

    def get_user(self, user_name: str) -> User:
        """Gets a specific user.

        Args:
            user_name: Name of the user to get.

        Returns:
            The requested user.

        Raises:
            KeyError: If no user with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                user = session.exec(
                    select(UserTable).where(UserTable.name == user_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            return User(**user.dict())

    def create_user(self, user_name: str) -> User:
        """Creates a new user.

        Args:
            user_name: Unique username.

        Returns:
             The newly created user.

        Raises:
            EntityExistsError: If a user with the given name already exists.
        """
        with Session(self.engine) as session:
            existing_user = session.exec(
                select(UserTable).where(UserTable.name == user_name)
            ).first()
            if existing_user:
                raise EntityExistsError(
                    f"User with name '{user_name}' already exists."
                )
            user = UserTable(name=user_name)
            session.add(user)
            session.commit()
        return user

    def delete_user(self, user_name: str) -> None:
        """Deletes a user.

        Args:
            user_name: Name of the user to delete.

        Raises:
            KeyError: If no user with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                user = session.exec(
                    select(UserTable).where(UserTable.name == user_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            session.delete(user)
            session.commit()
            self._delete_query_results(
                select(RoleAssignmentTable).where(
                    RoleAssignmentTable.user_id == user.id
                )
            )
            self._delete_query_results(
                select(TeamAssignmentTable).where(
                    TeamAssignmentTable.user_id == user.id
                )
            )

    @property
    def teams(self) -> List[Team]:
        """All registered teams.

        Returns:
            A list of all registered teams.
        """
        with Session(self.engine) as session:
            return [
                Team(**team.dict())
                for team in session.exec(select(TeamTable)).all()
            ]

    def get_team(self, team_name: str) -> Team:
        """Gets a specific team.

        Args:
            team_name: Name of the team to get.

        Returns:
            The requested team.

        Raises:
            KeyError: If no team with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                team = session.exec(
                    select(TeamTable).where(TeamTable.name == team_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            return Team(**team.dict())

    def create_team(self, team_name: str) -> Team:
        """Creates a new team.

        Args:
            team_name: Unique team name.

        Returns:
             The newly created team.

        Raises:
            EntityExistsError: If a team with the given name already exists.
        """
        with Session(self.engine) as session:
            existing_team = session.exec(
                select(TeamTable).where(TeamTable.name == team_name)
            ).first()
            if existing_team:
                raise EntityExistsError(
                    f"Team with name '{team_name}' already exists."
                )
            sql_team = TeamTable(name=team_name)
            team = Team(**sql_team.dict())
            session.add(sql_team)
            session.commit()
        return team

    def delete_team(self, team_name: str) -> None:
        """Deletes a team.

        Args:
            team_name: Name of the team to delete.

        Raises:
            KeyError: If no team with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                team = session.exec(
                    select(TeamTable).where(TeamTable.name == team_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            session.delete(team)
            session.commit()
            self._delete_query_results(
                select(RoleAssignmentTable).where(
                    RoleAssignmentTable.team_id == team.id
                )
            )
            self._delete_query_results(
                select(TeamAssignmentTable).where(
                    TeamAssignmentTable.team_id == team.id
                )
            )

    def add_user_to_team(self, team_name: str, user_name: str) -> None:
        """Adds a user to a team.

        Args:
            team_name: Name of the team.
            user_name: Name of the user.

        Raises:
            KeyError: If no user and team with the given names exists.
        """
        with Session(self.engine) as session:
            try:
                team = session.exec(
                    select(TeamTable).where(TeamTable.name == team_name)
                ).one()
                user = session.exec(
                    select(UserTable).where(UserTable.name == user_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            assignment = TeamAssignmentTable(user_id=user.id, team_id=team.id)
            session.add(assignment)
            session.commit()

    def remove_user_from_team(self, team_name: str, user_name: str) -> None:
        """Removes a user from a team.

        Args:
            team_name: Name of the team.
            user_name: Name of the user.

        Raises:
            KeyError: If no user and team with the given names exists.
        """
        with Session(self.engine) as session:
            try:
                assignment = session.exec(
                    select(TeamAssignmentTable)
                    .where(TeamAssignmentTable.team_id == TeamTable.id)
                    .where(TeamAssignmentTable.user_id == UserTable.id)
                    .where(UserTable.name == user_name)
                    .where(TeamTable.name == team_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            session.delete(assignment)
            session.commit()

    @property
    def projects(self) -> List[Project]:
        """All registered projects.

        Returns:
            A list of all registered projects.
        """
        with Session(self.engine) as session:
            return [
                Project(**project.dict())
                for project in session.exec(select(ProjectTable)).all()
            ]

    def get_project(self, project_name: str) -> Project:
        """Gets a specific project.

        Args:
            project_name: Name of the project to get.

        Returns:
            The requested project.

        Raises:
            KeyError: If no project with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                project = session.exec(
                    select(ProjectTable).where(
                        ProjectTable.name == project_name
                    )
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            return Project(**project.dict())

    def create_project(
        self, project_name: str, description: Optional[str] = None
    ) -> Project:
        """Creates a new project.

        Args:
            project_name: Unique project name.
            description: Optional project description.

        Returns:
             The newly created project.

        Raises:
            EntityExistsError: If a project with the given name already exists.
        """
        with Session(self.engine) as session:
            existing_project = session.exec(
                select(ProjectTable).where(ProjectTable.name == project_name)
            ).first()
            if existing_project:
                raise EntityExistsError(
                    f"Project with name '{project_name}' already exists."
                )
            sql_project = ProjectTable(name=project_name)
            project = Project(**sql_project.dict())
            session.add(sql_project)
            session.commit()
        return project

    def delete_project(self, project_name: str) -> None:
        """Deletes a project.

        Args:
            project_name: Name of the project to delete.

        Raises:
            KeyError: If no project with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                project = session.exec(
                    select(ProjectTable).where(
                        ProjectTable.name == project_name
                    )
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            session.delete(project)
            session.commit()
            self._delete_query_results(
                select(RoleAssignmentTable).where(
                    RoleAssignmentTable.project_id == project.id
                )
            )

    @property
    def roles(self) -> List[Role]:
        """All registered roles.

        Returns:
            A list of all registered roles.
        """
        with Session(self.engine) as session:
            return [
                Role(**role.dict())
                for role in session.exec(select(RoleTable)).all()
            ]

    @property
    def role_assignments(self) -> List[RoleAssignment]:
        """All registered role assignments.

        Returns:
            A list of all registered role assignments.
        """
        with Session(self.engine) as session:
            return [
                RoleAssignment(**assignment.dict())
                for assignment in session.exec(
                    select(RoleAssignmentTable)
                ).all()
            ]

    def get_role(self, role_name: str) -> Role:
        """Gets a specific role.

        Args:
            role_name: Name of the role to get.

        Returns:
            The requested role.

        Raises:
            KeyError: If no role with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                role = session.exec(
                    select(RoleTable).where(RoleTable.name == role_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            return Role(**role.dict())

    def create_role(self, role_name: str) -> Role:
        """Creates a new role.

        Args:
            role_name: Unique role name.

        Returns:
             The newly created role.

        Raises:
            EntityExistsError: If a role with the given name already exists.
        """
        with Session(self.engine) as session:
            existing_role = session.exec(
                select(RoleTable).where(RoleTable.name == role_name)
            ).first()
            if existing_role:
                raise EntityExistsError(
                    f"Role with name '{role_name}' already exists."
                )
            sql_role = RoleTable(name=role_name)
            role = Role(**sql_role.dict())
            session.add(sql_role)
            session.commit()
        return role

    def delete_role(self, role_name: str) -> None:
        """Deletes a role.

        Args:
            role_name: Name of the role to delete.

        Raises:
            KeyError: If no role with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                role = session.exec(
                    select(RoleTable).where(RoleTable.name == role_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            session.delete(role)
            session.commit()
            self._delete_query_results(
                select(RoleAssignmentTable).where(
                    RoleAssignmentTable.role_id == role.id
                )
            )

    def assign_role(
        self,
        role_name: str,
        entity_name: str,
        project_name: Optional[str] = None,
        is_user: bool = True,
    ) -> None:
        """Assigns a role to a user or team.

        Args:
            role_name: Name of the role to assign.
            entity_name: User or team name.
            project_name: Optional project name.
            is_user: Boolean indicating whether the given `entity_name` refers
                to a user.

        Raises:
            KeyError: If no role, entity or project with the given names exists.
        """
        with Session(self.engine) as session:
            user_id: Optional[UUID] = None
            team_id: Optional[UUID] = None
            project_id: Optional[UUID] = None

            try:
                role_id = session.exec(
                    select(RoleTable.id).where(RoleTable.name == role_name)
                ).one()

                if project_name:
                    project_id = session.exec(
                        select(ProjectTable.id).where(
                            ProjectTable.name == project_name
                        )
                    ).one()

                if is_user:
                    user_id = session.exec(
                        select(UserTable.id).where(
                            UserTable.name == entity_name
                        )
                    ).one()
                else:
                    team_id = session.exec(
                        select(TeamTable.id).where(
                            TeamTable.name == entity_name
                        )
                    ).one()
            except NoResultFound as error:
                raise KeyError from error

            assignment = RoleAssignmentTable(
                role_id=role_id,
                project_id=project_id,
                user_id=user_id,
                team_id=team_id,
            )
            session.add(assignment)
            session.commit()

    def revoke_role(
        self,
        role_name: str,
        entity_name: str,
        project_name: Optional[str] = None,
        is_user: bool = True,
    ) -> None:
        """Revokes a role from a user or team.

        Args:
            role_name: Name of the role to revoke.
            entity_name: User or team name.
            project_name: Optional project name.
            is_user: Boolean indicating whether the given `entity_name` refers
                to a user.

        Raises:
            KeyError: If no role, entity or project with the given names exists.
        """
        with Session(self.engine) as session:
            statement = (
                select(RoleAssignmentTable)
                .where(RoleAssignmentTable.role_id == RoleTable.id)
                .where(RoleTable.name == role_name)
            )

            if project_name:
                statement = statement.where(
                    RoleAssignmentTable.project_id == ProjectTable.id
                ).where(ProjectTable.name == project_name)

            if is_user:
                statement = statement.where(
                    RoleAssignmentTable.user_id == UserTable.id
                ).where(UserTable.name == entity_name)
            else:
                statement = statement.where(
                    RoleAssignmentTable.team_id == TeamTable.id
                ).where(TeamTable.name == entity_name)

            try:
                assignment = session.exec(statement).one()
            except NoResultFound as error:
                raise KeyError from error

            session.delete(assignment)
            session.commit()

    def get_users_for_team(self, team_name: str) -> List[User]:
        """Fetches all users of a team.

        Args:
            team_name: Name of the team.

        Returns:
            List of users that are part of the team.

        Raises:
            KeyError: If no team with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                team_id = session.exec(
                    select(TeamTable.id).where(TeamTable.name == team_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            users = session.exec(
                select(UserTable)
                .where(UserTable.id == TeamAssignmentTable.user_id)
                .where(TeamAssignmentTable.team_id == team_id)
            ).all()
            return [User(**user.dict()) for user in users]

    def get_teams_for_user(self, user_name: str) -> List[Team]:
        """Fetches all teams for a user.

        Args:
            user_name: Name of the user.

        Returns:
            List of teams that the user is part of.

        Raises:
            KeyError: If no user with the given name exists.
        """
        with Session(self.engine) as session:
            try:
                user_id = session.exec(
                    select(UserTable.id).where(UserTable.name == user_name)
                ).one()
            except NoResultFound as error:
                raise KeyError from error

            teams = session.exec(
                select(TeamTable)
                .where(TeamTable.id == TeamAssignmentTable.team_id)
                .where(TeamAssignmentTable.user_id == user_id)
            ).all()
            return [Team(**team.dict()) for team in teams]

    def get_role_assignments_for_user(
        self,
        user_name: str,
        project_name: Optional[str] = None,
        include_team_roles: bool = True,
    ) -> List[RoleAssignment]:
        """Fetches all role assignments for a user.

        Args:
            user_name: Name of the user.
            project_name: Optional filter to only return roles assigned for
                this project.
            include_team_roles: If `True`, includes roles for all teams that
                the user is part of.

        Returns:
            List of role assignments for this user.

        Raises:
            KeyError: If no user or project with the given names exists.
        """
        with Session(self.engine) as session:
            try:
                user_id = session.exec(
                    select(UserTable.id).where(UserTable.name == user_name)
                ).one()
                statement = select(RoleAssignmentTable).where(
                    RoleAssignmentTable.user_id == user_id
                )
                if project_name:
                    project_id = session.exec(
                        select(ProjectTable.id).where(
                            ProjectTable.name == project_name
                        )
                    ).one()
                    statement = statement.where(
                        RoleAssignmentTable.project_id == project_id
                    )
            except NoResultFound as error:
                raise KeyError from error

            assignments = [
                RoleAssignment(**assignment.dict())
                for assignment in session.exec(statement).all()
            ]
            if include_team_roles:
                for team in self.get_teams_for_user(user_name):
                    assignments += self.get_role_assignments_for_team(
                        team.name, project_name=project_name
                    )

            return assignments

    def get_role_assignments_for_team(
        self,
        team_name: str,
        project_name: Optional[str] = None,
    ) -> List[RoleAssignment]:
        """Fetches all role assignments for a team.

        Args:
            team_name: Name of the user.
            project_name: Optional filter to only return roles assigned for
                this project.

        Returns:
            List of role assignments for this team.

        Raises:
            KeyError: If no team or project with the given names exists.
        """
        with Session(self.engine) as session:
            try:
                team_id = session.exec(
                    select(TeamTable.id).where(TeamTable.name == team_name)
                ).one()

                statement = select(RoleAssignmentTable).where(
                    RoleAssignmentTable.team_id == team_id
                )
                if project_name:
                    project_id = session.exec(
                        select(ProjectTable.id).where(
                            ProjectTable.name == project_name
                        )
                    ).one()
                    statement = statement.where(
                        RoleAssignmentTable.project_id == project_id
                    )
            except NoResultFound as error:
                raise KeyError from error

            return [
                RoleAssignment(**assignment.dict())
                for assignment in session.exec(statement).all()
            ]

    # Implementation-specific internal methods:

    @property
    def stack_names(self) -> List[str]:
        """Names of all stacks registered in this ZenStore."""
        with Session(self.engine) as session:
            return [s.name for s in session.exec(select(ZenStack))]

    def _delete_query_results(self, query: Any) -> None:
        """Deletes all rows returned by the input query."""
        with Session(self.engine) as session:
            for result in session.exec(query).all():
                session.delete(result)
            session.commit()
is_empty: bool property readonly

Check if the zen store is empty.

projects: List[zenml.zen_stores.models.user_management_models.Project] property readonly

All registered projects.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Project]

A list of all registered projects.

role_assignments: List[zenml.zen_stores.models.user_management_models.RoleAssignment] property readonly

All registered role assignments.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

A list of all registered role assignments.

roles: List[zenml.zen_stores.models.user_management_models.Role] property readonly

All registered roles.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Role]

A list of all registered roles.

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

Configuration for all stacks registered in this zen 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 ZenStore.

teams: List[zenml.zen_stores.models.user_management_models.Team] property readonly

All registered teams.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Team]

A list of all registered teams.

type: StoreType property readonly

The type of zen store.

url: str property readonly

URL of the repository.

users: List[zenml.zen_stores.models.user_management_models.User] property readonly

All registered users.

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.User]

A list of all registered users.

add_user_to_team(self, team_name, user_name)

Adds a user to a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required
user_name str

Name of the user.

required

Exceptions:

Type Description
KeyError

If no user and team with the given names exists.

Source code in zenml/zen_stores/sql_zen_store.py
def add_user_to_team(self, team_name: str, user_name: str) -> None:
    """Adds a user to a team.

    Args:
        team_name: Name of the team.
        user_name: Name of the user.

    Raises:
        KeyError: If no user and team with the given names exists.
    """
    with Session(self.engine) as session:
        try:
            team = session.exec(
                select(TeamTable).where(TeamTable.name == team_name)
            ).one()
            user = session.exec(
                select(UserTable).where(UserTable.name == user_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        assignment = TeamAssignmentTable(user_id=user.id, team_id=team.id)
        session.add(assignment)
        session.commit()
assign_role(self, role_name, entity_name, project_name=None, is_user=True)

Assigns a role to a user or team.

Parameters:

Name Type Description Default
role_name str

Name of the role to assign.

required
entity_name str

User or team name.

required
project_name Optional[str]

Optional project name.

None
is_user bool

Boolean indicating whether the given entity_name refers to a user.

True

Exceptions:

Type Description
KeyError

If no role, entity or project with the given names exists.

Source code in zenml/zen_stores/sql_zen_store.py
def assign_role(
    self,
    role_name: str,
    entity_name: str,
    project_name: Optional[str] = None,
    is_user: bool = True,
) -> None:
    """Assigns a role to a user or team.

    Args:
        role_name: Name of the role to assign.
        entity_name: User or team name.
        project_name: Optional project name.
        is_user: Boolean indicating whether the given `entity_name` refers
            to a user.

    Raises:
        KeyError: If no role, entity or project with the given names exists.
    """
    with Session(self.engine) as session:
        user_id: Optional[UUID] = None
        team_id: Optional[UUID] = None
        project_id: Optional[UUID] = None

        try:
            role_id = session.exec(
                select(RoleTable.id).where(RoleTable.name == role_name)
            ).one()

            if project_name:
                project_id = session.exec(
                    select(ProjectTable.id).where(
                        ProjectTable.name == project_name
                    )
                ).one()

            if is_user:
                user_id = session.exec(
                    select(UserTable.id).where(
                        UserTable.name == entity_name
                    )
                ).one()
            else:
                team_id = session.exec(
                    select(TeamTable.id).where(
                        TeamTable.name == entity_name
                    )
                ).one()
        except NoResultFound as error:
            raise KeyError from error

        assignment = RoleAssignmentTable(
            role_id=role_id,
            project_id=project_id,
            user_id=user_id,
            team_id=team_id,
        )
        session.add(assignment)
        session.commit()
create_project(self, project_name, description=None)

Creates a new project.

Parameters:

Name Type Description Default
project_name str

Unique project name.

required
description Optional[str]

Optional project description.

None

Returns:

Type Description
Project

The newly created project.

Exceptions:

Type Description
EntityExistsError

If a project with the given name already exists.

Source code in zenml/zen_stores/sql_zen_store.py
def create_project(
    self, project_name: str, description: Optional[str] = None
) -> Project:
    """Creates a new project.

    Args:
        project_name: Unique project name.
        description: Optional project description.

    Returns:
         The newly created project.

    Raises:
        EntityExistsError: If a project with the given name already exists.
    """
    with Session(self.engine) as session:
        existing_project = session.exec(
            select(ProjectTable).where(ProjectTable.name == project_name)
        ).first()
        if existing_project:
            raise EntityExistsError(
                f"Project with name '{project_name}' already exists."
            )
        sql_project = ProjectTable(name=project_name)
        project = Project(**sql_project.dict())
        session.add(sql_project)
        session.commit()
    return project
create_role(self, role_name)

Creates a new role.

Parameters:

Name Type Description Default
role_name str

Unique role name.

required

Returns:

Type Description
Role

The newly created role.

Exceptions:

Type Description
EntityExistsError

If a role with the given name already exists.

Source code in zenml/zen_stores/sql_zen_store.py
def create_role(self, role_name: str) -> Role:
    """Creates a new role.

    Args:
        role_name: Unique role name.

    Returns:
         The newly created role.

    Raises:
        EntityExistsError: If a role with the given name already exists.
    """
    with Session(self.engine) as session:
        existing_role = session.exec(
            select(RoleTable).where(RoleTable.name == role_name)
        ).first()
        if existing_role:
            raise EntityExistsError(
                f"Role with name '{role_name}' already exists."
            )
        sql_role = RoleTable(name=role_name)
        role = Role(**sql_role.dict())
        session.add(sql_role)
        session.commit()
    return role
create_team(self, team_name)

Creates a new team.

Parameters:

Name Type Description Default
team_name str

Unique team name.

required

Returns:

Type Description
Team

The newly created team.

Exceptions:

Type Description
EntityExistsError

If a team with the given name already exists.

Source code in zenml/zen_stores/sql_zen_store.py
def create_team(self, team_name: str) -> Team:
    """Creates a new team.

    Args:
        team_name: Unique team name.

    Returns:
         The newly created team.

    Raises:
        EntityExistsError: If a team with the given name already exists.
    """
    with Session(self.engine) as session:
        existing_team = session.exec(
            select(TeamTable).where(TeamTable.name == team_name)
        ).first()
        if existing_team:
            raise EntityExistsError(
                f"Team with name '{team_name}' already exists."
            )
        sql_team = TeamTable(name=team_name)
        team = Team(**sql_team.dict())
        session.add(sql_team)
        session.commit()
    return team
create_user(self, user_name)

Creates a new user.

Parameters:

Name Type Description Default
user_name str

Unique username.

required

Returns:

Type Description
User

The newly created user.

Exceptions:

Type Description
EntityExistsError

If a user with the given name already exists.

Source code in zenml/zen_stores/sql_zen_store.py
def create_user(self, user_name: str) -> User:
    """Creates a new user.

    Args:
        user_name: Unique username.

    Returns:
         The newly created user.

    Raises:
        EntityExistsError: If a user with the given name already exists.
    """
    with Session(self.engine) as session:
        existing_user = session.exec(
            select(UserTable).where(UserTable.name == user_name)
        ).first()
        if existing_user:
            raise EntityExistsError(
                f"User with name '{user_name}' already exists."
            )
        user = UserTable(name=user_name)
        session.add(user)
        session.commit()
    return user
delete_project(self, project_name)

Deletes a project.

Parameters:

Name Type Description Default
project_name str

Name of the project to delete.

required

Exceptions:

Type Description
KeyError

If no project with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def delete_project(self, project_name: str) -> None:
    """Deletes a project.

    Args:
        project_name: Name of the project to delete.

    Raises:
        KeyError: If no project with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            project = session.exec(
                select(ProjectTable).where(
                    ProjectTable.name == project_name
                )
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        session.delete(project)
        session.commit()
        self._delete_query_results(
            select(RoleAssignmentTable).where(
                RoleAssignmentTable.project_id == project.id
            )
        )
delete_role(self, role_name)

Deletes a role.

Parameters:

Name Type Description Default
role_name str

Name of the role to delete.

required

Exceptions:

Type Description
KeyError

If no role with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def delete_role(self, role_name: str) -> None:
    """Deletes a role.

    Args:
        role_name: Name of the role to delete.

    Raises:
        KeyError: If no role with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            role = session.exec(
                select(RoleTable).where(RoleTable.name == role_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        session.delete(role)
        session.commit()
        self._delete_query_results(
            select(RoleAssignmentTable).where(
                RoleAssignmentTable.role_id == role.id
            )
        )
delete_team(self, team_name)

Deletes a team.

Parameters:

Name Type Description Default
team_name str

Name of the team to delete.

required

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def delete_team(self, team_name: str) -> None:
    """Deletes a team.

    Args:
        team_name: Name of the team to delete.

    Raises:
        KeyError: If no team with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            team = session.exec(
                select(TeamTable).where(TeamTable.name == team_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        session.delete(team)
        session.commit()
        self._delete_query_results(
            select(RoleAssignmentTable).where(
                RoleAssignmentTable.team_id == team.id
            )
        )
        self._delete_query_results(
            select(TeamAssignmentTable).where(
                TeamAssignmentTable.team_id == team.id
            )
        )
delete_user(self, user_name)

Deletes a user.

Parameters:

Name Type Description Default
user_name str

Name of the user to delete.

required

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def delete_user(self, user_name: str) -> None:
    """Deletes a user.

    Args:
        user_name: Name of the user to delete.

    Raises:
        KeyError: If no user with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            user = session.exec(
                select(UserTable).where(UserTable.name == user_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        session.delete(user)
        session.commit()
        self._delete_query_results(
            select(RoleAssignmentTable).where(
                RoleAssignmentTable.user_id == user.id
            )
        )
        self._delete_query_results(
            select(TeamAssignmentTable).where(
                TeamAssignmentTable.user_id == user.id
            )
        )
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

Exceptions:

Type Description
KeyError

If no stack exists for the given name.

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

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

    Raises:
        KeyError: If no stack exists for the given name.
    """
    with Session(self.engine) as session:
        try:
            stack = session.exec(
                select(ZenStack).where(ZenStack.name == name)
            ).one()
            session.delete(stack)
        except NoResultFound as error:
            raise KeyError from error
        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/zen_stores/sql_zen_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/zen_stores/sql_zen_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 SqlZenStore.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_project(self, project_name)

Gets a specific project.

Parameters:

Name Type Description Default
project_name str

Name of the project to get.

required

Returns:

Type Description
Project

The requested project.

Exceptions:

Type Description
KeyError

If no project with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def get_project(self, project_name: str) -> Project:
    """Gets a specific project.

    Args:
        project_name: Name of the project to get.

    Returns:
        The requested project.

    Raises:
        KeyError: If no project with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            project = session.exec(
                select(ProjectTable).where(
                    ProjectTable.name == project_name
                )
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        return Project(**project.dict())
get_role(self, role_name)

Gets a specific role.

Parameters:

Name Type Description Default
role_name str

Name of the role to get.

required

Returns:

Type Description
Role

The requested role.

Exceptions:

Type Description
KeyError

If no role with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def get_role(self, role_name: str) -> Role:
    """Gets a specific role.

    Args:
        role_name: Name of the role to get.

    Returns:
        The requested role.

    Raises:
        KeyError: If no role with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            role = session.exec(
                select(RoleTable).where(RoleTable.name == role_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        return Role(**role.dict())
get_role_assignments_for_team(self, team_name, project_name=None)

Fetches all role assignments for a team.

Parameters:

Name Type Description Default
team_name str

Name of the user.

required
project_name Optional[str]

Optional filter to only return roles assigned for this project.

None

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

List of role assignments for this team.

Exceptions:

Type Description
KeyError

If no team or project with the given names exists.

Source code in zenml/zen_stores/sql_zen_store.py
def get_role_assignments_for_team(
    self,
    team_name: str,
    project_name: Optional[str] = None,
) -> List[RoleAssignment]:
    """Fetches all role assignments for a team.

    Args:
        team_name: Name of the user.
        project_name: Optional filter to only return roles assigned for
            this project.

    Returns:
        List of role assignments for this team.

    Raises:
        KeyError: If no team or project with the given names exists.
    """
    with Session(self.engine) as session:
        try:
            team_id = session.exec(
                select(TeamTable.id).where(TeamTable.name == team_name)
            ).one()

            statement = select(RoleAssignmentTable).where(
                RoleAssignmentTable.team_id == team_id
            )
            if project_name:
                project_id = session.exec(
                    select(ProjectTable.id).where(
                        ProjectTable.name == project_name
                    )
                ).one()
                statement = statement.where(
                    RoleAssignmentTable.project_id == project_id
                )
        except NoResultFound as error:
            raise KeyError from error

        return [
            RoleAssignment(**assignment.dict())
            for assignment in session.exec(statement).all()
        ]
get_role_assignments_for_user(self, user_name, project_name=None, include_team_roles=True)

Fetches all role assignments for a user.

Parameters:

Name Type Description Default
user_name str

Name of the user.

required
project_name Optional[str]

Optional filter to only return roles assigned for this project.

None
include_team_roles bool

If True, includes roles for all teams that the user is part of.

True

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.RoleAssignment]

List of role assignments for this user.

Exceptions:

Type Description
KeyError

If no user or project with the given names exists.

Source code in zenml/zen_stores/sql_zen_store.py
def get_role_assignments_for_user(
    self,
    user_name: str,
    project_name: Optional[str] = None,
    include_team_roles: bool = True,
) -> List[RoleAssignment]:
    """Fetches all role assignments for a user.

    Args:
        user_name: Name of the user.
        project_name: Optional filter to only return roles assigned for
            this project.
        include_team_roles: If `True`, includes roles for all teams that
            the user is part of.

    Returns:
        List of role assignments for this user.

    Raises:
        KeyError: If no user or project with the given names exists.
    """
    with Session(self.engine) as session:
        try:
            user_id = session.exec(
                select(UserTable.id).where(UserTable.name == user_name)
            ).one()
            statement = select(RoleAssignmentTable).where(
                RoleAssignmentTable.user_id == user_id
            )
            if project_name:
                project_id = session.exec(
                    select(ProjectTable.id).where(
                        ProjectTable.name == project_name
                    )
                ).one()
                statement = statement.where(
                    RoleAssignmentTable.project_id == project_id
                )
        except NoResultFound as error:
            raise KeyError from error

        assignments = [
            RoleAssignment(**assignment.dict())
            for assignment in session.exec(statement).all()
        ]
        if include_team_roles:
            for team in self.get_teams_for_user(user_name):
                assignments += self.get_role_assignments_for_team(
                    team.name, project_name=project_name
                )

        return assignments
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/zen_stores/sql_zen_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()}
get_team(self, team_name)

Gets a specific team.

Parameters:

Name Type Description Default
team_name str

Name of the team to get.

required

Returns:

Type Description
Team

The requested team.

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def get_team(self, team_name: str) -> Team:
    """Gets a specific team.

    Args:
        team_name: Name of the team to get.

    Returns:
        The requested team.

    Raises:
        KeyError: If no team with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            team = session.exec(
                select(TeamTable).where(TeamTable.name == team_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        return Team(**team.dict())
get_teams_for_user(self, user_name)

Fetches all teams for a user.

Parameters:

Name Type Description Default
user_name str

Name of the user.

required

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.Team]

List of teams that the user is part of.

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def get_teams_for_user(self, user_name: str) -> List[Team]:
    """Fetches all teams for a user.

    Args:
        user_name: Name of the user.

    Returns:
        List of teams that the user is part of.

    Raises:
        KeyError: If no user with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            user_id = session.exec(
                select(UserTable.id).where(UserTable.name == user_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        teams = session.exec(
            select(TeamTable)
            .where(TeamTable.id == TeamAssignmentTable.team_id)
            .where(TeamAssignmentTable.user_id == user_id)
        ).all()
        return [Team(**team.dict()) for team in teams]
get_user(self, user_name)

Gets a specific user.

Parameters:

Name Type Description Default
user_name str

Name of the user to get.

required

Returns:

Type Description
User

The requested user.

Exceptions:

Type Description
KeyError

If no user with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def get_user(self, user_name: str) -> User:
    """Gets a specific user.

    Args:
        user_name: Name of the user to get.

    Returns:
        The requested user.

    Raises:
        KeyError: If no user with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            user = session.exec(
                select(UserTable).where(UserTable.name == user_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        return User(**user.dict())
get_users_for_team(self, team_name)

Fetches all users of a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required

Returns:

Type Description
List[zenml.zen_stores.models.user_management_models.User]

List of users that are part of the team.

Exceptions:

Type Description
KeyError

If no team with the given name exists.

Source code in zenml/zen_stores/sql_zen_store.py
def get_users_for_team(self, team_name: str) -> List[User]:
    """Fetches all users of a team.

    Args:
        team_name: Name of the team.

    Returns:
        List of users that are part of the team.

    Raises:
        KeyError: If no team with the given name exists.
    """
    with Session(self.engine) as session:
        try:
            team_id = session.exec(
                select(TeamTable.id).where(TeamTable.name == team_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        users = session.exec(
            select(UserTable)
            .where(UserTable.id == TeamAssignmentTable.user_id)
            .where(TeamAssignmentTable.team_id == team_id)
        ).all()
        return [User(**user.dict()) for user in users]
initialize(self, url, *args, **kwargs)

Initialize a new SqlZenStore.

Parameters:

Name Type Description Default
url str

odbc path to a database.

required
args, kwargs

additional parameters for SQLModel.

required

Returns:

Type Description
SqlZenStore

The initialized zen store instance.

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

    Args:
        url: odbc path to a database.
        args, kwargs: additional parameters for SQLModel.
    Returns:
        The initialized zen store instance.
    """
    if not self.is_valid_url(url):
        raise ValueError(f"Invalid URL for SQL store: {url}")

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

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

    # we need to remove `skip_default_registrations` from the kwargs,
    # because SQLModel will raise an error if it is present
    sql_kwargs = kwargs.copy()
    sql_kwargs.pop("skip_default_registrations", False)
    self.engine = create_engine(url, *args, **sql_kwargs)
    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_valid_url(url) staticmethod

Check if the given url is a valid SQL url.

Source code in zenml/zen_stores/sql_zen_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/zen_stores/sql_zen_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()
remove_user_from_team(self, team_name, user_name)

Removes a user from a team.

Parameters:

Name Type Description Default
team_name str

Name of the team.

required
user_name str

Name of the user.

required

Exceptions:

Type Description
KeyError

If no user and team with the given names exists.

Source code in zenml/zen_stores/sql_zen_store.py
def remove_user_from_team(self, team_name: str, user_name: str) -> None:
    """Removes a user from a team.

    Args:
        team_name: Name of the team.
        user_name: Name of the user.

    Raises:
        KeyError: If no user and team with the given names exists.
    """
    with Session(self.engine) as session:
        try:
            assignment = session.exec(
                select(TeamAssignmentTable)
                .where(TeamAssignmentTable.team_id == TeamTable.id)
                .where(TeamAssignmentTable.user_id == UserTable.id)
                .where(UserTable.name == user_name)
                .where(TeamTable.name == team_name)
            ).one()
        except NoResultFound as error:
            raise KeyError from error

        session.delete(assignment)
        session.commit()
revoke_role(self, role_name, entity_name, project_name=None, is_user=True)

Revokes a role from a user or team.

Parameters:

Name Type Description Default
role_name str

Name of the role to revoke.

required
entity_name str

User or team name.

required
project_name Optional[str]

Optional project name.

None
is_user bool

Boolean indicating whether the given entity_name refers to a user.

True

Exceptions:

Type Description
KeyError

If no role, entity or project with the given names exists.

Source code in zenml/zen_stores/sql_zen_store.py
def revoke_role(
    self,
    role_name: str,
    entity_name: str,
    project_name: Optional[str] = None,
    is_user: bool = True,
) -> None:
    """Revokes a role from a user or team.

    Args:
        role_name: Name of the role to revoke.
        entity_name: User or team name.
        project_name: Optional project name.
        is_user: Boolean indicating whether the given `entity_name` refers
            to a user.

    Raises:
        KeyError: If no role, entity or project with the given names exists.
    """
    with Session(self.engine) as session:
        statement = (
            select(RoleAssignmentTable)
            .where(RoleAssignmentTable.role_id == RoleTable.id)
            .where(RoleTable.name == role_name)
        )

        if project_name:
            statement = statement.where(
                RoleAssignmentTable.project_id == ProjectTable.id
            ).where(ProjectTable.name == project_name)

        if is_user:
            statement = statement.where(
                RoleAssignmentTable.user_id == UserTable.id
            ).where(UserTable.name == entity_name)
        else:
            statement = statement.where(
                RoleAssignmentTable.team_id == TeamTable.id
            ).where(TeamTable.name == entity_name)

        try:
            assignment = session.exec(statement).one()
        except NoResultFound as error:
            raise KeyError from error

        session.delete(assignment)
        session.commit()
update_stack_component(self, name, component_type, component)

Update a stack component.

Parameters:

Name Type Description Default
name str

The original name of the stack component.

required
component_type StackComponentType

The type of the stack component to update.

required
component StackComponentWrapper

The new component to update with.

required

Exceptions:

Type Description
KeyError

If no stack component exists with the given name.

Source code in zenml/zen_stores/sql_zen_store.py
def update_stack_component(
    self,
    name: str,
    component_type: StackComponentType,
    component: StackComponentWrapper,
) -> Dict[str, str]:
    """Update a stack component.

    Args:
        name: The original name of the stack component.
        component_type: The type of the stack component to update.
        component: The new component to update with.

    Raises:
        KeyError: If no stack component exists with the given name.
    """
    with Session(self.engine) as session:
        updated_component = session.exec(
            select(ZenStackComponent)
            .where(ZenStackComponent.component_type == component_type)
            .where(ZenStackComponent.name == name)
        ).first()

        if not updated_component:
            raise KeyError(
                f"Unable to update stack component (type: {component.type}) "
                f"with name '{component.name}': No existing stack component "
                f"found with this name."
            )

        new_name_component = session.exec(
            select(ZenStackComponent)
            .where(ZenStackComponent.component_type == component_type)
            .where(ZenStackComponent.name == component.name)
        ).first()
        if (name != component.name) and new_name_component is not None:
            raise StackComponentExistsError(
                f"Unable to update stack component (type: "
                f"{component.type}) with name '{component.name}': Found "
                f"existing stack component with this name."
            )

        updated_component.configuration = component.config

        # handle any potential renamed component
        updated_component.name = component.name

        # rename components inside stacks
        updated_stack_definitions = session.exec(
            select(ZenStackDefinition)
            .where(ZenStackDefinition.component_type == component_type)
            .where(ZenStackDefinition.component_name == name)
        ).all()
        for stack_definition in updated_stack_definitions:
            stack_definition.component_name = component.name
            session.add(stack_definition)

        session.add(updated_component)
        session.commit()
    logger.info(
        "Updated stack component with type '%s' and name '%s'.",
        component_type,
        component.name,
    )
    return {component.type.value: component.flavor}

ZenStackDefinition (SQLModel) pydantic-model

Join table between Stacks and StackComponents

Source code in zenml/zen_stores/sql_zen_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"
    )