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

zenml/zen_stores/schemas/action_schemas.py

@@ -19,7 +19,7 @@ from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
 from pydantic.json import pydantic_encoder
-from sqlalchemy import TEXT, Column
+from sqlalchemy import TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.models import (
@@ -44,6 +44,13 @@ class ActionSchema(NamedSchema, table=True):
     """SQL Model for actions."""
 
     __tablename__ = "action"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_action_name_in_workspace",
+        ),
+    )
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,

zenml/zen_stores/schemas/api_key_schemas.py

@@ -19,7 +19,7 @@ from typing import Any, Optional, Tuple
 from uuid import UUID
 
 from passlib.context import CryptContext
-from sqlalchemy import TEXT, Column
+from sqlalchemy import TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.models import (
@@ -42,6 +42,13 @@ class APIKeySchema(NamedSchema, table=True):
     """SQL Model for API keys."""
 
     __tablename__ = "api_key"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "service_account_id",
+            name="unique_api_key_name_in_service_account",
+        ),
+    )
 
     description: str = Field(sa_column=Column(TEXT))
     key: str

zenml/zen_stores/schemas/artifact_schemas.py

@@ -71,7 +71,8 @@ class ArtifactSchema(NamedSchema, table=True):
     __table_args__ = (
         UniqueConstraint(
             "name",
-            name="unique_artifact_name",
+            "workspace_id",
+            name="unique_artifact_name_in_workspace",
         ),
     )
 
@@ -91,6 +92,28 @@ class ArtifactSchema(NamedSchema, table=True):
         ),
     )
 
+    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()
+
+    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="artifacts",
+    )
+
     @property
     def latest_version(self) -> Optional["ArtifactVersionSchema"]:
         """Fetch the latest version for this artifact.
@@ -133,6 +156,8 @@ class ArtifactSchema(NamedSchema, table=True):
         return cls(
             name=artifact_request.name,
             has_custom_name=artifact_request.has_custom_name,
+            workspace_id=artifact_request.workspace,
+            user_id=artifact_request.user,
         )
 
     def to_model(
@@ -165,6 +190,7 @@ class ArtifactSchema(NamedSchema, table=True):
             tags=[tag.to_model() for tag in self.tags],
             latest_version_name=latest_name,
             latest_version_id=latest_id,
+            user=self.user.to_model() if self.user else None,
         )
 
         # Create the metadata of the model
@@ -172,6 +198,7 @@ class ArtifactSchema(NamedSchema, table=True):
         if include_metadata:
             metadata = ArtifactResponseMetadata(
                 has_custom_name=self.has_custom_name,
+                workspace=self.workspace.to_model(),
             )
 
         return ArtifactResponse(

zenml/zen_stores/schemas/code_repository_schemas.py

@@ -17,7 +17,7 @@ import json
 from typing import Any, Optional
 from uuid import UUID
 
-from sqlalchemy import TEXT, Column
+from sqlalchemy import TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.models import (
@@ -42,6 +42,13 @@ class CodeRepositorySchema(NamedSchema, table=True):
     """SQL Model for code repositories."""
 
     __tablename__ = "code_repository"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_code_repository_name_in_workspace",
+        ),
+    )
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,

zenml/zen_stores/schemas/component_schemas.py

@@ -18,6 +18,7 @@ import json
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
+from sqlalchemy import UniqueConstraint
 from sqlmodel import Relationship
 
 from zenml.enums import StackComponentType
@@ -37,7 +38,6 @@ from zenml.zen_stores.schemas.service_connector_schemas import (
 )
 from zenml.zen_stores.schemas.stack_schemas import StackCompositionSchema
 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.flavor_schemas import FlavorSchema
@@ -51,22 +51,19 @@ class StackComponentSchema(NamedSchema, table=True):
     """SQL Model for stack components."""
 
     __tablename__ = "stack_component"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "type",
+            name="unique_component_name_and_type",
+        ),
+    )
 
     type: str
     flavor: str
     configuration: bytes
     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="components")
-
     user_id: Optional[UUID] = build_foreign_key_field(
         source=__tablename__,
         target=UserSchema.__tablename__,
@@ -132,7 +129,6 @@ class StackComponentSchema(NamedSchema, table=True):
         """
         return cls(
             name=request.name,
-            workspace_id=request.workspace,
             user_id=request.user,
             type=request.type,
             flavor=request.flavor,
@@ -158,7 +154,7 @@ class StackComponentSchema(NamedSchema, table=True):
             The updated `StackComponentSchema`.
         """
         for field, value in component_update.model_dump(
-            exclude_unset=True, exclude={"workspace", "user", "connector"}
+            exclude_unset=True, exclude={"user", "connector"}
         ).items():
             if field == "configuration":
                 self.configuration = base64.b64encode(
@@ -209,7 +205,6 @@ class StackComponentSchema(NamedSchema, table=True):
         metadata = None
         if include_metadata:
             metadata = ComponentResponseMetadata(
-                workspace=self.workspace.to_model(),
                 configuration=json.loads(
                     base64.b64decode(self.configuration).decode()
                 ),

zenml/zen_stores/schemas/event_source_schemas.py

@@ -18,7 +18,7 @@ import json
 from typing import TYPE_CHECKING, 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 import EventSourceResponseMetadata
@@ -46,6 +46,13 @@ class EventSourceSchema(NamedSchema, table=True):
     """SQL Model for tag."""
 
     __tablename__ = "event_source"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_event_source_name_in_workspace",
+        ),
+    )
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,

zenml/zen_stores/schemas/flavor_schemas.py

@@ -17,7 +17,7 @@ import json
 from typing import Any, Optional
 from uuid import UUID
 
-from sqlalchemy import TEXT, Column
+from sqlalchemy import TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.enums import StackComponentType
@@ -31,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
 
 
 class FlavorSchema(NamedSchema, table=True):
@@ -45,6 +44,13 @@ class FlavorSchema(NamedSchema, table=True):
     """
 
     __tablename__ = "flavor"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "type",
+            name="unique_flavor_name_and_type",
+        ),
+    )
 
     type: str
     source: str
@@ -54,18 +60,6 @@ class FlavorSchema(NamedSchema, table=True):
     connector_resource_type: Optional[str]
     connector_resource_id_attr: Optional[str]
 
-    workspace_id: Optional[UUID] = build_foreign_key_field(
-        source=__tablename__,
-        target=WorkspaceSchema.__tablename__,
-        source_column="workspace_id",
-        target_column="id",
-        ondelete="CASCADE",
-        nullable=True,
-    )
-    workspace: Optional["WorkspaceSchema"] = Relationship(
-        back_populates="flavors"
-    )
-
     user_id: Optional[UUID] = build_foreign_key_field(
         source=__tablename__,
         target=UserSchema.__tablename__,
@@ -84,7 +78,10 @@ class FlavorSchema(NamedSchema, table=True):
 
     is_custom: bool = Field(default=True)
 
-    def update(self, flavor_update: "FlavorUpdate") -> "FlavorSchema":
+    def update(
+        self,
+        flavor_update: "FlavorUpdate",
+    ) -> "FlavorSchema":
         """Update a `FlavorSchema` from a `FlavorUpdate`.
 
         Args:
@@ -94,7 +91,7 @@ class FlavorSchema(NamedSchema, table=True):
             The updated `FlavorSchema`.
         """
         for field, value in flavor_update.model_dump(
-            exclude_unset=True, exclude={"workspace", "user"}
+            exclude_unset=True, exclude={"user"}
         ).items():
             if field == "config_schema":
                 setattr(self, field, json.dumps(value))
@@ -129,22 +126,19 @@ class FlavorSchema(NamedSchema, table=True):
             integration=self.integration,
             source=self.source,
             logo_url=self.logo_url,
+            is_custom=self.is_custom,
             created=self.created,
             updated=self.updated,
         )
         metadata = None
         if include_metadata:
             metadata = FlavorResponseMetadata(
-                workspace=self.workspace.to_model()
-                if self.workspace
-                else None,
                 config_schema=json.loads(self.config_schema),
                 connector_type=self.connector_type,
                 connector_resource_type=self.connector_resource_type,
                 connector_resource_id_attr=self.connector_resource_id_attr,
                 docs_url=self.docs_url,
                 sdk_docs_url=self.sdk_docs_url,
-                is_custom=self.is_custom,
             )
         return FlavorResponse(
             id=self.id,

zenml/zen_stores/schemas/model_schemas.py

@@ -248,6 +248,9 @@ class ModelSchema(NamedSchema, table=True):
         for field, value in model_update.model_dump(
             exclude_unset=True, exclude_none=True
         ).items():
+            if field in ["add_tags", "remove_tags"]:
+                # Tags are handled separately
+                continue
             setattr(self, field, value)
         self.updated = utc_now()
         return self

zenml/zen_stores/schemas/pipeline_build_schemas.py

@@ -100,6 +100,8 @@ class PipelineBuildSchema(BaseSchema, table=True):
     python_version: Optional[str]
     checksum: Optional[str]
     stack_checksum: Optional[str]
+    # Build duration in seconds
+    duration: Optional[int] = None
 
     @classmethod
     def from_request(
@@ -125,6 +127,7 @@ class PipelineBuildSchema(BaseSchema, table=True):
             python_version=request.python_version,
             checksum=request.checksum,
             stack_checksum=request.stack_checksum,
+            duration=request.duration,
         )
 
     def to_model(
@@ -162,6 +165,7 @@ class PipelineBuildSchema(BaseSchema, table=True):
                 stack_checksum=self.stack_checksum,
                 is_local=self.is_local,
                 contains_code=self.contains_code,
+                duration=self.duration,
             )
         return PipelineBuildResponse(
             id=self.id,

zenml/zen_stores/schemas/pipeline_deployment_schemas.py

@@ -148,7 +148,9 @@ class PipelineDeploymentSchema(BaseSchema, table=True):
     template_id: Optional[UUID] = None
 
     # SQLModel Relationships
-    user: Optional["UserSchema"] = Relationship()
+    user: Optional["UserSchema"] = Relationship(
+        back_populates="deployments",
+    )
     workspace: "WorkspaceSchema" = Relationship()
     stack: Optional["StackSchema"] = Relationship()
     pipeline: Optional["PipelineSchema"] = Relationship()

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -251,7 +251,6 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
             pipeline_id=request.pipeline,
             deployment_id=request.deployment,
             trigger_execution_id=request.trigger_execution_id,
-            model_version_id=request.model_version_id,
         )
 
     def fetch_metadata_collection(self) -> Dict[str, List[RunMetadataEntry]]:
@@ -435,8 +434,6 @@ class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
         if run_update.status:
             self.status = run_update.status.value
             self.end_time = run_update.end_time
-        if run_update.model_version_id and self.model_version_id is None:
-            self.model_version_id = run_update.model_version_id
 
         self.updated = utc_now()
         return self

zenml/zen_stores/schemas/run_template_schemas.py

@@ -32,7 +32,7 @@ from zenml.models import (
     RunTemplateUpdate,
 )
 from zenml.utils.time_utils import utc_now
-from zenml.zen_stores.schemas.base_schemas import BaseSchema
+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
@@ -45,7 +45,7 @@ if TYPE_CHECKING:
     from zenml.zen_stores.schemas.tag_schemas import TagSchema
 
 
-class RunTemplateSchema(BaseSchema, table=True):
+class RunTemplateSchema(NamedSchema, table=True):
     """SQL Model for run templates."""
 
     __tablename__ = "run_template"
@@ -57,7 +57,6 @@ class RunTemplateSchema(BaseSchema, table=True):
         ),
     )
 
-    name: str = Field(nullable=False)
     description: Optional[str] = Field(
         sa_column=Column(
             String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
@@ -92,7 +91,9 @@ class RunTemplateSchema(BaseSchema, table=True):
         nullable=True,
     )
 
-    user: Optional["UserSchema"] = Relationship()
+    user: Optional["UserSchema"] = Relationship(
+        back_populates="run_templates",
+    )
     workspace: "WorkspaceSchema" = Relationship()
     source_deployment: Optional["PipelineDeploymentSchema"] = Relationship(
         sa_relationship_kwargs={
@@ -180,6 +181,9 @@ class RunTemplateSchema(BaseSchema, table=True):
         for field, value in update.model_dump(
             exclude_unset=True, exclude_none=True
         ).items():
+            if field in ["add_tags", "remove_tags"]:
+                # Tags are handled separately
+                continue
             setattr(self, field, value)
 
         self.updated = utc_now()

zenml/zen_stores/schemas/schedule_schema.py

@@ -17,6 +17,7 @@ from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
 
+from sqlalchemy import UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.enums import MetadataResourceTypes
@@ -49,6 +50,13 @@ class ScheduleSchema(NamedSchema, RunMetadataInterface, table=True):
     """SQL Model for schedules."""
 
     __tablename__ = "schedule"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_schedule_name_in_workspace",
+        ),
+    )
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,
@@ -154,20 +162,7 @@ class ScheduleSchema(NamedSchema, RunMetadataInterface, table=True):
         """
         if schedule_update.name is not None:
             self.name = schedule_update.name
-        if schedule_update.active is not None:
-            self.active = schedule_update.active
-        if schedule_update.cron_expression is not None:
-            self.cron_expression = schedule_update.cron_expression
-        if schedule_update.start_time is not None:
-            self.start_time = schedule_update.start_time
-        if schedule_update.end_time is not None:
-            self.end_time = schedule_update.end_time
-        if schedule_update.interval_second is not None:
-            self.interval_second = (
-                schedule_update.interval_second.total_seconds()
-            )
-        if schedule_update.catchup is not None:
-            self.catchup = schedule_update.catchup
+
         self.updated = utc_now()
         return self
 

zenml/zen_stores/schemas/secret_schemas.py

@@ -18,7 +18,7 @@ import json
 from typing import Any, Dict, Optional, cast
 from uuid import UUID
 
-from sqlalchemy import TEXT, Column
+from sqlalchemy import TEXT, Column, UniqueConstraint
 from sqlalchemy_utils.types.encrypted.encrypted_type import (
     AesGcmEngine,
     InvalidCiphertextError,
@@ -26,7 +26,6 @@ from sqlalchemy_utils.types.encrypted.encrypted_type import (
 from sqlmodel import Field, Relationship
 
 from zenml.constants import TEXT_FIELD_MAX_LENGTH
-from zenml.enums import SecretScope
 from zenml.models import (
     SecretRequest,
     SecretResponse,
@@ -38,7 +37,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
 
 
 class SecretDecodeError(Exception):
@@ -54,21 +52,19 @@ class SecretSchema(NamedSchema, table=True):
     """
 
     __tablename__ = "secret"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "private",
+            "user_id",
+            name="unique_secret_name_private_scope_user",
+        ),
+    )
 
-    scope: str
+    private: bool
 
     values: Optional[bytes] = Field(sa_column=Column(TEXT, nullable=True))
 
-    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="secrets")
-
     user_id: UUID = build_foreign_key_field(
         source=__tablename__,
         target=UserSchema.__tablename__,
@@ -177,8 +173,7 @@ class SecretSchema(NamedSchema, table=True):
         assert secret.user is not None, "User must be set for secret creation."
         return cls(
             name=secret.name,
-            scope=secret.scope.value,
-            workspace_id=secret.workspace,
+            private=secret.private,
             user_id=secret.user,
             # Don't store secret values implicitly in the secret. The
             # SQL secret store will call `store_secret_values` to store the
@@ -202,12 +197,9 @@ class SecretSchema(NamedSchema, table=True):
         # SQL secret store will call `set_secret_values` to update the
         # values separately if SQL is used as the secrets store.
         for field, value in secret_update.model_dump(
-            exclude_unset=True, exclude={"workspace", "user", "values"}
+            exclude_unset=True, exclude={"user", "values"}
         ).items():
-            if field == "scope":
-                setattr(self, field, value.value)
-            else:
-                setattr(self, field, value)
+            setattr(self, field, value)
 
         self.updated = utc_now()
         return self
@@ -231,9 +223,7 @@ class SecretSchema(NamedSchema, table=True):
         """
         metadata = None
         if include_metadata:
-            metadata = SecretResponseMetadata(
-                workspace=self.workspace.to_model(),
-            )
+            metadata = SecretResponseMetadata()
 
         # Don't load the secret values implicitly in the secret. The
         # SQL secret store will call `get_secret_values` to load the
@@ -242,7 +232,7 @@ class SecretSchema(NamedSchema, table=True):
             user=self.user.to_model() if self.user else None,
             created=self.created,
             updated=self.updated,
-            scope=SecretScope(self.scope),
+            private=self.private,
         )
         return SecretResponse(
             id=self.id,

zenml/zen_stores/schemas/service_connector_schemas.py

@@ -19,7 +19,7 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, 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.models import (
@@ -33,7 +33,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 StackComponentSchema
@@ -43,6 +42,12 @@ class ServiceConnectorSchema(NamedSchema, table=True):
     """SQL Model for service connectors."""
 
     __tablename__ = "service_connector"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            name="unique_service_connector_name",
+        ),
+    )
 
     connector_type: str = Field(sa_column=Column(TEXT))
     description: str
@@ -57,18 +62,6 @@ class ServiceConnectorSchema(NamedSchema, table=True):
     expiration_seconds: Optional[int]
     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="service_connectors"
-    )
-
     user_id: Optional[UUID] = build_foreign_key_field(
         source=__tablename__,
         target=UserSchema.__tablename__,
@@ -146,7 +139,6 @@ class ServiceConnectorSchema(NamedSchema, table=True):
         """
         assert connector_request.user is not None, "User must be set."
         return cls(
-            workspace_id=connector_request.workspace,
             user_id=connector_request.user,
             name=connector_request.name,
             description=connector_request.description,
@@ -189,7 +181,7 @@ class ServiceConnectorSchema(NamedSchema, table=True):
         """
         for field, value in connector_update.model_dump(
             exclude_unset=False,
-            exclude={"workspace", "user", "secrets"},
+            exclude={"user", "secrets"},
         ).items():
             if value is None:
                 if field == "resource_id":
@@ -264,7 +256,6 @@ class ServiceConnectorSchema(NamedSchema, table=True):
         metadata = None
         if include_metadata:
             metadata = ServiceConnectorResponseMetadata(
-                workspace=self.workspace.to_model(),
                 configuration=json.loads(
                     base64.b64decode(self.configuration).decode()
                 )

zenml/zen_stores/schemas/service_schemas.py

@@ -135,7 +135,6 @@ class ServiceSchema(NamedSchema, table=True):
         """
         body = ServiceResponseBody(
             user=self.user.to_model() if self.user else None,
-            workspace=self.workspace.to_model(),
             created=self.created,
             updated=self.updated,
             service_type=json.loads(self.service_type),