zenml-nightly 0.75.0.dev20250312__py3-none-any.whl → 0.75.0.dev20250313__py3-none-any.whl

zenml/zen_stores/schemas/stack_schemas.py

@@ -18,7 +18,8 @@ import json
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
-from sqlmodel import Relationship, SQLModel
+from sqlalchemy import UniqueConstraint
+from sqlmodel import Field, Relationship, SQLModel
 
 from zenml.models import (
     StackResponse,
@@ -30,7 +31,6 @@ from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
-from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas.component_schemas import (
@@ -76,19 +76,17 @@ class StackSchema(NamedSchema, table=True):
     """SQL Model for stacks."""
 
     __tablename__ = "stack"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            name="unique_stack_name",
+        ),
+    )
+
+    description: Optional[str] = Field(default=None)
     stack_spec_path: Optional[str]
     labels: Optional[bytes]
 
-    workspace_id: UUID = build_foreign_key_field(
-        source=__tablename__,
-        target=WorkspaceSchema.__tablename__,
-        source_column="workspace_id",
-        target_column="id",
-        ondelete="CASCADE",
-        nullable=False,
-    )
-    workspace: "WorkspaceSchema" = Relationship(back_populates="stacks")
-
     user_id: Optional[UUID] = build_foreign_key_field(
         source=__tablename__,
         target=UserSchema.__tablename__,
@@ -123,7 +121,7 @@ class StackSchema(NamedSchema, table=True):
             The updated StackSchema.
         """
         for field, value in stack_update.model_dump(
-            exclude_unset=True, exclude={"workspace", "user"}
+            exclude_unset=True, exclude={"user"}
         ).items():
             if field == "components":
                 self.components = components
@@ -150,7 +148,6 @@ class StackSchema(NamedSchema, table=True):
             include_resources: Whether the resources will be filled.
             **kwargs: Keyword arguments to allow schema specific logic
 
-
         Returns:
             The converted model.
         """
@@ -162,12 +159,12 @@ class StackSchema(NamedSchema, table=True):
         metadata = None
         if include_metadata:
             metadata = StackResponseMetadata(
-                workspace=self.workspace.to_model(),
                 components={c.type: [c.to_model()] for c in self.components},
                 stack_spec_path=self.stack_spec_path,
                 labels=json.loads(base64.b64decode(self.labels).decode())
                 if self.labels
                 else None,
+                description=self.description,
             )
 
         return StackResponse(

zenml/zen_stores/schemas/step_run_schemas.py

@@ -102,7 +102,7 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
         ondelete="SET NULL",
         nullable=True,
     )
-    deployment_id: UUID = build_foreign_key_field(
+    deployment_id: Optional[UUID] = build_foreign_key_field(
         source=__tablename__,
         target=PipelineDeploymentSchema.__tablename__,
         source_column="deployment_id",
@@ -134,7 +134,7 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
         ondelete="CASCADE",
         nullable=False,
     )
-    model_version_id: UUID = build_foreign_key_field(
+    model_version_id: Optional[UUID] = build_foreign_key_field(
         source=__tablename__,
         target=MODEL_VERSION_TABLENAME,
         source_column="model_version_id",
@@ -186,11 +186,14 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
     model_config = ConfigDict(protected_namespaces=())  # type: ignore[assignment]
 
     @classmethod
-    def from_request(cls, request: StepRunRequest) -> "StepRunSchema":
+    def from_request(
+        cls, request: StepRunRequest, deployment_id: Optional[UUID]
+    ) -> "StepRunSchema":
         """Create a step run schema from a step run request model.
 
         Args:
             request: The step run request model.
+            deployment_id: The deployment ID.
 
         Returns:
             The step run schema.
@@ -202,14 +205,13 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
             start_time=request.start_time,
             end_time=request.end_time,
             status=request.status.value,
+            deployment_id=deployment_id,
             original_step_run_id=request.original_step_run_id,
             pipeline_run_id=request.pipeline_run_id,
-            deployment_id=request.deployment,
             docstring=request.docstring,
             cache_key=request.cache_key,
             code_hash=request.code_hash,
             source_code=request.source_code,
-            model_version_id=request.model_version_id,
         )
 
     def to_model(
@@ -355,9 +357,6 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
                 self.status = value.value
             if key == "end_time":
                 self.end_time = value
-            if key == "model_version_id":
-                if value and self.model_version_id is None:
-                    self.model_version_id = value
 
         self.updated = utc_now()
 

zenml/zen_stores/schemas/tag_schemas.py

@@ -13,10 +13,10 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of tag tables."""
 
-from typing import Any, List
+from typing import Any, List, Optional
 from uuid import UUID
 
-from sqlalchemy import VARCHAR, Column
+from sqlalchemy import VARCHAR, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.enums import ColorVariants, TaggableResourceTypes
@@ -27,6 +27,7 @@ from zenml.models import (
     TagResourceResponseBody,
     TagResponse,
     TagResponseBody,
+    TagResponseMetadata,
     TagUpdate,
 )
 from zenml.utils.time_utils import utc_now
@@ -35,14 +36,33 @@ from zenml.zen_stores.schemas.schema_utils import (
     build_foreign_key_field,
     build_index,
 )
+from zenml.zen_stores.schemas.user_schemas import UserSchema
 
 
 class TagSchema(NamedSchema, table=True):
     """SQL Model for tag."""
 
     __tablename__ = "tag"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            name="unique_tag_name",
+        ),
+    )
+
+    user_id: Optional[UUID] = build_foreign_key_field(
+        source=__tablename__,
+        target=UserSchema.__tablename__,
+        source_column="user_id",
+        target_column="id",
+        ondelete="SET NULL",
+        nullable=True,
+    )
+    user: Optional["UserSchema"] = Relationship(back_populates="tags")
 
     color: str = Field(sa_column=Column(VARCHAR(255), nullable=False))
+    exclusive: bool = Field(default=False)
+
     links: List["TagResourceSchema"] = Relationship(
         back_populates="tag",
         sa_relationship_kwargs={"overlaps": "tags", "cascade": "delete"},
@@ -60,7 +80,9 @@ class TagSchema(NamedSchema, table=True):
         """
         return cls(
             name=request.name,
+            exclusive=request.exclusive,
             color=request.color.value,
+            user_id=request.user,
         )
 
     def to_model(
@@ -80,15 +102,21 @@ class TagSchema(NamedSchema, table=True):
         Returns:
             The created `TagResponse`.
         """
+        metadata = None
+        if include_metadata:
+            metadata = TagResponseMetadata()
         return TagResponse(
             id=self.id,
             name=self.name,
             body=TagResponseBody(
+                user=self.user.to_model() if self.user else None,
                 created=self.created,
                 updated=self.updated,
                 color=ColorVariants(self.color),
+                exclusive=self.exclusive,
                 tagged_count=len(self.links),
             ),
+            metadata=metadata,
         )
 
     def update(self, update: TagUpdate) -> "TagSchema":

zenml/zen_stores/schemas/trigger_schemas.py

@@ -18,7 +18,7 @@ import json
 from typing import Any, List, Optional, cast
 from uuid import UUID
 
-from sqlalchemy import TEXT, Column
+from sqlalchemy import TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.config.schedule import Schedule
@@ -51,6 +51,13 @@ class TriggerSchema(NamedSchema, table=True):
     """SQL Model for triggers."""
 
     __tablename__ = "trigger"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_trigger_name_in_workspace",
+        ),
+    )
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,

zenml/zen_stores/schemas/user_schemas.py

@@ -34,11 +34,14 @@ from zenml.models import (
 )
 from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
+from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
+from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas import (
         ActionSchema,
         APIKeySchema,
+        ArtifactSchema,
         ArtifactVersionSchema,
         CodeRepositorySchema,
         EventSourceSchema,
@@ -51,6 +54,7 @@ if TYPE_CHECKING:
         PipelineRunSchema,
         PipelineSchema,
         RunMetadataSchema,
+        RunTemplateSchema,
         ScheduleSchema,
         SecretSchema,
         ServiceConnectorSchema,
@@ -58,6 +62,7 @@ if TYPE_CHECKING:
         StackComponentSchema,
         StackSchema,
         StepRunSchema,
+        TagSchema,
         TriggerSchema,
     )
 
@@ -79,6 +84,14 @@ class UserSchema(NamedSchema, table=True):
     external_user_id: Optional[UUID] = Field(nullable=True)
     is_admin: bool = Field(default=False)
     user_metadata: Optional[str] = Field(nullable=True)
+    default_workspace_id: Optional[UUID] = build_foreign_key_field(
+        source=__tablename__,
+        target=WorkspaceSchema.__tablename__,
+        source_column="default_workspace_id",
+        target_column="id",
+        ondelete="SET NULL",
+        nullable=True,
+    )
 
     stacks: List["StackSchema"] = Relationship(back_populates="user")
     components: List["StackComponentSchema"] = Relationship(
@@ -100,8 +113,12 @@ class UserSchema(NamedSchema, table=True):
         back_populates="user",
     )
     runs: List["PipelineRunSchema"] = Relationship(back_populates="user")
+    run_templates: List["RunTemplateSchema"] = Relationship(
+        back_populates="user",
+    )
     step_runs: List["StepRunSchema"] = Relationship(back_populates="user")
     builds: List["PipelineBuildSchema"] = Relationship(back_populates="user")
+    artifacts: List["ArtifactSchema"] = Relationship(back_populates="user")
     artifact_versions: List["ArtifactVersionSchema"] = Relationship(
         back_populates="user"
     )
@@ -150,6 +167,9 @@ class UserSchema(NamedSchema, table=True):
         back_populates="service_account",
         sa_relationship_kwargs={"cascade": "delete"},
     )
+    tags: List["TagSchema"] = Relationship(
+        back_populates="user",
+    )
 
     @classmethod
     def from_user_request(cls, model: UserRequest) -> "UserSchema":
@@ -175,6 +195,7 @@ class UserSchema(NamedSchema, table=True):
             user_metadata=json.dumps(model.user_metadata)
             if model.user_metadata
             else None,
+            default_workspace_id=model.default_workspace_id,
         )
 
     @classmethod
@@ -262,8 +283,8 @@ class UserSchema(NamedSchema, table=True):
             include_resources: Whether the resources will be filled.
             **kwargs: Keyword arguments to allow schema specific logic
             include_private: Whether to include the user private information
-                             this is to limit the amount of data one can get
-                             about other users
+                this is to limit the amount of data one can get about other
+                users.
 
         Returns:
             The converted `UserResponse`.
@@ -289,6 +310,7 @@ class UserSchema(NamedSchema, table=True):
                 created=self.created,
                 updated=self.updated,
                 is_admin=self.is_admin,
+                default_workspace_id=self.default_workspace_id,
             ),
             metadata=metadata,
         )

zenml/zen_stores/schemas/utils.py

@@ -110,3 +110,19 @@ class RunMetadataInterface:
             k: sorted(v, key=lambda x: x.created, reverse=True)[0].value
             for k, v in metadata_collection.items()
         }
+
+
+def get_resource_type_name(schema_class: Type[BaseSchema]) -> str:
+    """Get the name of a resource from a schema class.
+
+    Args:
+        schema_class: The schema class to get the name of.
+
+    Returns:
+        The name of the resource.
+    """
+    entity_name = schema_class.__tablename__
+    assert isinstance(entity_name, str)
+    # Some entities are plural, some are singular, some have multiple words
+    # in their table name connected by underscores (e.g. pipeline_run)
+    return entity_name.replace("_", " ").rstrip("s")

zenml/zen_stores/schemas/workspace_schemas.py

@@ -15,6 +15,7 @@
 
 from typing import TYPE_CHECKING, Any, List
 
+from sqlalchemy import UniqueConstraint
 from sqlmodel import Relationship
 
 from zenml.models import (
@@ -33,7 +34,6 @@ if TYPE_CHECKING:
         ArtifactVersionSchema,
         CodeRepositorySchema,
         EventSourceSchema,
-        FlavorSchema,
         ModelSchema,
         ModelVersionSchema,
         PipelineBuildSchema,
@@ -42,11 +42,7 @@ if TYPE_CHECKING:
         PipelineSchema,
         RunMetadataSchema,
         ScheduleSchema,
-        SecretSchema,
-        ServiceConnectorSchema,
         ServiceSchema,
-        StackComponentSchema,
-        StackSchema,
         StepRunSchema,
         TriggerSchema,
     )
@@ -56,21 +52,15 @@ class WorkspaceSchema(NamedSchema, table=True):
     """SQL Model for workspaces."""
 
     __tablename__ = "workspace"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            name="unique_workspace_name",
+        ),
+    )
 
     description: str
 
-    stacks: List["StackSchema"] = Relationship(
-        back_populates="workspace",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
-    components: List["StackComponentSchema"] = Relationship(
-        back_populates="workspace",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
-    flavors: List["FlavorSchema"] = Relationship(
-        back_populates="workspace",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
     pipelines: List["PipelineSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
@@ -99,10 +89,6 @@ class WorkspaceSchema(NamedSchema, table=True):
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    secrets: List["SecretSchema"] = Relationship(
-        back_populates="workspace",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
     actions: List["ActionSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
@@ -128,10 +114,6 @@ class WorkspaceSchema(NamedSchema, table=True):
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},
     )
-    service_connectors: List["ServiceConnectorSchema"] = Relationship(
-        back_populates="workspace",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
     models: List["ModelSchema"] = Relationship(
         back_populates="workspace",
         sa_relationship_kwargs={"cascade": "delete"},

zenml/zen_stores/secrets_stores/service_connector_secrets_store.py

@@ -23,7 +23,6 @@ from typing import (
     Optional,
     Type,
 )
-from uuid import uuid4
 
 from pydantic import Field, model_validator
 
@@ -133,8 +132,6 @@ class ServiceConnectorSecretsStore(BaseSecretsStore):
                 name="secrets-store",
                 connector_type=self.SERVICE_CONNECTOR_TYPE,
                 resource_types=[self.SERVICE_CONNECTOR_RESOURCE_TYPE],
-                user=uuid4(),  # Use a fake user ID
-                workspace=uuid4(),  # Use a fake workspace ID
                 auth_method=self.config.auth_method,
                 configuration=self.config.auth_config,
             )