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