zenml-nightly 0.68.1.dev20241111__py3-none-any.whl → 0.70.0.dev20241114__py3-none-any.whl

zenml/zen_stores/migrations/versions/904464ea4041_add_pipeline_model_run_unique_constraints.py

@@ -0,0 +1,192 @@
+"""Add pipeline, model and run unique constraints [904464ea4041].
+
+Revision ID: 904464ea4041
+Revises: b557b2871693
+Create Date: 2024-11-04 10:27:05.450092
+
+"""
+
+from collections import defaultdict
+from typing import Any, Dict, Set
+
+import sqlalchemy as sa
+from alembic import op
+
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
+
+# revision identifiers, used by Alembic.
+revision = "904464ea4041"
+down_revision = "b557b2871693"
+branch_labels = None
+depends_on = None
+
+
+def resolve_duplicate_entities() -> None:
+    """Resolve duplicate entities."""
+    connection = op.get_bind()
+    meta = sa.MetaData()
+    meta.reflect(
+        bind=connection,
+        only=("pipeline_run", "pipeline", "model", "model_version"),
+    )
+
+    # Remove duplicate names for runs, pipelines and models
+    for table_name in ["pipeline_run", "pipeline", "model"]:
+        table = sa.Table(table_name, meta)
+        result = connection.execute(
+            sa.select(table.c.id, table.c.name, table.c.workspace_id)
+        ).all()
+        existing: Dict[str, Set[str]] = defaultdict(set)
+
+        for id_, name, workspace_id in result:
+            names_in_workspace = existing[workspace_id]
+
+            if name in names_in_workspace:
+                new_name = f"{name}_{id_[:6]}"
+                logger.warning(
+                    "Migrating %s name from %s to %s to resolve duplicate name.",
+                    table_name,
+                    name,
+                    new_name,
+                )
+                connection.execute(
+                    sa.update(table)
+                    .where(table.c.id == id_)
+                    .values(name=new_name)
+                )
+                names_in_workspace.add(new_name)
+            else:
+                names_in_workspace.add(name)
+
+    # Remove duplicate names and version numbers for model versions
+    model_version_table = sa.Table("model_version", meta)
+    result = connection.execute(
+        sa.select(
+            model_version_table.c.id,
+            model_version_table.c.name,
+            model_version_table.c.number,
+            model_version_table.c.model_id,
+        )
+    ).all()
+
+    existing_names: Dict[str, Set[str]] = defaultdict(set)
+    existing_numbers: Dict[str, Set[int]] = defaultdict(set)
+
+    needs_update = []
+
+    for id_, name, number, model_id in result:
+        names_for_model = existing_names[model_id]
+        numbers_for_model = existing_numbers[model_id]
+
+        needs_new_name = name in names_for_model
+        needs_new_number = number in numbers_for_model
+
+        if needs_new_name or needs_new_number:
+            needs_update.append(
+                (id_, name, number, model_id, needs_new_name, needs_new_number)
+            )
+
+        names_for_model.add(name)
+        numbers_for_model.add(number)
+
+    for (
+        id_,
+        name,
+        number,
+        model_id,
+        needs_new_name,
+        needs_new_number,
+    ) in needs_update:
+        values: Dict[str, Any] = {}
+
+        is_numeric_version = str(number) == name
+        next_numeric_version = max(existing_numbers[model_id]) + 1
+
+        if is_numeric_version:
+            # No matter if the name or number clashes, we need to update both
+            values["number"] = next_numeric_version
+            values["name"] = str(next_numeric_version)
+            existing_numbers[model_id].add(next_numeric_version)
+            logger.warning(
+                "Migrating model version %s to %s to resolve duplicate name.",
+                name,
+                values["name"],
+            )
+        else:
+            if needs_new_name:
+                values["name"] = f"{name}_{id_[:6]}"
+                logger.warning(
+                    "Migrating model version %s to %s to resolve duplicate name.",
+                    name,
+                    values["name"],
+                )
+
+            if needs_new_number:
+                values["number"] = next_numeric_version
+                existing_numbers[model_id].add(next_numeric_version)
+
+        connection.execute(
+            sa.update(model_version_table)
+            .where(model_version_table.c.id == id_)
+            .values(**values)
+        )
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    resolve_duplicate_entities()
+
+    with op.batch_alter_table("pipeline", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_pipeline_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    with op.batch_alter_table("pipeline_run", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_run_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    with op.batch_alter_table("model", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_model_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    with op.batch_alter_table("model_version", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_version_for_model_id", ["name", "model_id"]
+        )
+        batch_op.create_unique_constraint(
+            "unique_version_number_for_model_id", ["number", "model_id"]
+        )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("model_version", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_version_number_for_model_id", type_="unique"
+        )
+        batch_op.drop_constraint("unique_version_for_model_id", type_="unique")
+
+    with op.batch_alter_table("model", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_model_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("pipeline_run", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_run_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("pipeline", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_pipeline_name_in_workspace", type_="unique"
+        )
+
+    # ### end Alembic commands ###

zenml/zen_stores/rest_zen_store.py

@@ -110,6 +110,7 @@ from zenml.constants import (
     WORKSPACES,
 )
 from zenml.enums import (
+    APITokenType,
     OAuthGrantTypes,
     StackDeploymentProvider,
     StoreType,
@@ -3872,17 +3873,16 @@ class RestZenStore(BaseZenStore):
 
     def get_api_token(
         self,
-        pipeline_id: Optional[UUID] = None,
         schedule_id: Optional[UUID] = None,
-        expires_minutes: Optional[int] = None,
+        pipeline_run_id: Optional[UUID] = None,
+        step_run_id: Optional[UUID] = None,
     ) -> str:
         """Get an API token for a workload.
 
         Args:
-            pipeline_id: The ID of the pipeline to get a token for.
             schedule_id: The ID of the schedule to get a token for.
-            expires_minutes: The number of minutes for which the token should
-                be valid. If not provided, the token will be valid indefinitely.
+            pipeline_run_id: The ID of the pipeline run to get a token for.
+            step_run_id: The ID of the step run to get a token for.
 
         Returns:
             The API token.
@@ -3890,13 +3890,16 @@ class RestZenStore(BaseZenStore):
         Raises:
             ValueError: if the server response is not valid.
         """
-        params: Dict[str, Any] = {}
-        if pipeline_id:
-            params["pipeline_id"] = pipeline_id
+        params: Dict[str, Any] = {
+            # Python clients may only request workload tokens.
+            "token_type": APITokenType.WORKLOAD.value,
+        }
         if schedule_id:
             params["schedule_id"] = schedule_id
-        if expires_minutes:
-            params["expires_minutes"] = expires_minutes
+        if pipeline_run_id:
+            params["pipeline_run_id"] = pipeline_run_id
+        if step_run_id:
+            params["step_run_id"] = step_run_id
         response_body = self.get(API_TOKEN, params=params)
         if not isinstance(response_body, str):
             raise ValueError(

zenml/zen_stores/schemas/model_schemas.py

@@ -19,7 +19,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 from uuid import UUID
 
 from pydantic import ConfigDict
-from sqlalchemy import BOOLEAN, INTEGER, TEXT, Column
+from sqlalchemy import BOOLEAN, INTEGER, TEXT, Column, UniqueConstraint
 from sqlmodel import Field, Relationship
 
 from zenml.enums import MetadataResourceTypes, TaggableResourceTypes
@@ -62,6 +62,13 @@ class ModelSchema(NamedSchema, table=True):
     """SQL Model for model."""
 
     __tablename__ = "model"
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_model_name_in_workspace",
+        ),
+    )
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,
@@ -220,6 +227,23 @@ class ModelVersionSchema(NamedSchema, table=True):
     """SQL Model for model version."""
 
     __tablename__ = MODEL_VERSION_TABLENAME
+    __table_args__ = (
+        # We need two unique constraints here:
+        # - The first to ensure that each model version for a
+        #   model has a unique version number
+        # - The second one to ensure that explicit names given by
+        #   users are unique
+        UniqueConstraint(
+            "number",
+            "model_id",
+            name="unique_version_number_for_model_id",
+        ),
+        UniqueConstraint(
+            "name",
+            "model_id",
+            name="unique_version_for_model_id",
+        ),
+    )
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -72,6 +72,11 @@ class PipelineRunSchema(NamedSchema, table=True):
             "orchestrator_run_id",
             name="unique_orchestrator_run_id_for_deployment_id",
         ),
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_run_name_in_workspace",
+        ),
     )
 
     # Fields

zenml/zen_stores/schemas/pipeline_schemas.py

@@ -17,7 +17,7 @@ from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, 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 TaggableResourceTypes
@@ -50,7 +50,13 @@ class PipelineSchema(NamedSchema, table=True):
     """SQL Model for pipelines."""
 
     __tablename__ = "pipeline"
-
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "workspace_id",
+            name="unique_pipeline_name_in_workspace",
+        ),
+    )
     # Fields
     description: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))