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

zenml/zen_stores/migrations/versions/1cb6477f72d6_move_artifact_save_type.py

@@ -23,7 +23,8 @@ def upgrade() -> None:
         batch_op.add_column(sa.Column("save_type", sa.TEXT(), nullable=True))
 
     # Step 2: Move data from step_run_output_artifact.type to artifact_version.save_type
-    op.execute("""
+    op.execute(
+        """
         UPDATE artifact_version
         SET save_type = (
             SELECT max(step_run_output_artifact.type)
@@ -31,17 +32,22 @@ def upgrade() -> None:
             WHERE step_run_output_artifact.artifact_id = artifact_version.id
             GROUP BY artifact_id
         )
-    """)
-    op.execute("""
+    """
+    )
+    op.execute(
+        """
         UPDATE artifact_version
         SET save_type = 'step_output'
         WHERE artifact_version.save_type = 'default'
-    """)
-    op.execute("""
+    """
+    )
+    op.execute(
+        """
         UPDATE artifact_version
         SET save_type = 'external'
         WHERE save_type is NULL
-    """)
+    """
+    )
 
     # # Step 3: Set save_type to non-nullable
     with op.batch_alter_table("artifact_version", schema=None) as batch_op:
@@ -69,7 +75,8 @@ def downgrade() -> None:
         )
 
     # Move data back from artifact_version.save_type to step_run_output_artifact.type
-    op.execute("""
+    op.execute(
+        """
         UPDATE step_run_output_artifact
         SET type = (
             SELECT max(artifact_version.save_type)
@@ -77,12 +84,15 @@ def downgrade() -> None:
             WHERE step_run_output_artifact.artifact_id = artifact_version.id
             GROUP BY artifact_id
         )
-    """)
-    op.execute("""
+    """
+    )
+    op.execute(
+        """
         UPDATE step_run_output_artifact
         SET type = 'default'
         WHERE step_run_output_artifact.type = 'step_output'
-    """)
+    """
+    )
 
     # Set type to non-nullable
     with op.batch_alter_table(

zenml/zen_stores/migrations/versions/1f9d1cd00b90_add_unique_name_constraints.py

@@ -0,0 +1,231 @@
+"""add unique name constraints [1f9d1cd00b90].
+
+Revision ID: 1f9d1cd00b90
+Revises: f76a368a25a5
+Create Date: 2025-02-22 20:18:34.258987
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.orm import Session
+
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
+
+# revision identifiers, used by Alembic.
+revision = "1f9d1cd00b90"
+down_revision = "f76a368a25a5"
+branch_labels = None
+depends_on = None
+
+
+def resolve_duplicate_names(
+    table_name: str, other_columns: list[str], session: Session
+) -> None:
+    """Resolve duplicate entities.
+
+    Args:
+        table_name: The name of the table to resolve duplicate entities for.
+        other_columns: The columns that are part of the unique constraint,
+            excluding the name column.
+        session: The SQLAlchemy session to use.
+    """
+    columns = ["name"] + other_columns
+    duplicates = session.execute(
+        sa.text(
+            f"""
+                SELECT id, name
+                FROM `{table_name}`
+                WHERE ({", ".join(columns)}) IN (
+                    SELECT {", ".join(columns)}
+                    FROM `{table_name}`
+                    GROUP BY {", ".join(columns)}
+                    HAVING COUNT(*) > 1
+                )
+            """  # nosec B608
+        )
+    )
+    for id_, name in list(duplicates)[1:]:
+        logger.warning(f"Duplicate {table_name}: {name} (id: {id_})")
+        session.execute(
+            sa.text(
+                f"""
+                    UPDATE {table_name}
+                    SET name = :new_name
+                    WHERE id = :id_
+                """  # nosec B608
+            ),
+            params={"id_": id_, "new_name": f"{name}_{id_[:6]}"},
+        )
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    bind = op.get_bind()
+    session = Session(bind=bind)
+
+    resolve_duplicate_names("action", ["workspace_id"], session)
+
+    with op.batch_alter_table("action", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_action_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("api_key", ["service_account_id"], session)
+
+    with op.batch_alter_table("api_key", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_api_key_name_in_service_account",
+            ["name", "service_account_id"],
+        )
+
+    resolve_duplicate_names("artifact", ["workspace_id"], session)
+
+    with op.batch_alter_table("artifact", schema=None) as batch_op:
+        batch_op.drop_constraint("unique_artifact_name", type_="unique")
+        batch_op.create_unique_constraint(
+            "unique_artifact_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("code_repository", ["workspace_id"], session)
+
+    with op.batch_alter_table("code_repository", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_code_repository_name_in_workspace",
+            ["name", "workspace_id"],
+        )
+
+    resolve_duplicate_names("event_source", ["workspace_id"], session)
+
+    with op.batch_alter_table("event_source", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_event_source_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("flavor", ["type"], session)
+
+    with op.batch_alter_table("flavor", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_flavor_name_and_type", ["name", "type"]
+        )
+
+    resolve_duplicate_names("schedule", ["workspace_id"], session)
+
+    with op.batch_alter_table("schedule", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_schedule_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("secret", ["private", "user_id"], session)
+
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_secret_name_private_scope_user",
+            ["name", "private", "user_id"],
+        )
+
+    resolve_duplicate_names("service_connector", [], session)
+
+    with op.batch_alter_table("service_connector", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_service_connector_name", ["name"]
+        )
+
+    resolve_duplicate_names("stack", [], session)
+
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.create_unique_constraint("unique_stack_name", ["name"])
+
+    resolve_duplicate_names("stack_component", ["type"], session)
+
+    with op.batch_alter_table("stack_component", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_component_name_and_type", ["name", "type"]
+        )
+
+    resolve_duplicate_names("tag", [], session)
+
+    with op.batch_alter_table("tag", schema=None) as batch_op:
+        batch_op.create_unique_constraint("unique_tag_name", ["name"])
+
+    resolve_duplicate_names("trigger", ["workspace_id"], session)
+
+    with op.batch_alter_table("trigger", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_trigger_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("workspace", [], session)
+
+    with op.batch_alter_table("workspace", schema=None) as batch_op:
+        batch_op.create_unique_constraint("unique_workspace_name", ["name"])
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("workspace", schema=None) as batch_op:
+        batch_op.drop_constraint("unique_workspace_name", type_="unique")
+
+    with op.batch_alter_table("trigger", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_trigger_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("tag", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_tag_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("stack_component", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_component_name_and_type", type_="unique"
+        )
+
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.drop_constraint("unique_stack_name", type_="unique")
+
+    with op.batch_alter_table("service_connector", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_service_connector_name", type_="unique"
+        )
+
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_secret_name_and_private_scope", type_="unique"
+        )
+
+    with op.batch_alter_table("schedule", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_schedule_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("flavor", schema=None) as batch_op:
+        batch_op.drop_constraint("unique_flavor_name_and_type", type_="unique")
+
+    with op.batch_alter_table("event_source", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_event_source_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("code_repository", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_code_repository_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("artifact", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_artifact_name_in_workspace", type_="unique"
+        )
+        batch_op.create_unique_constraint("unique_artifact_name", ["name"])
+
+    with op.batch_alter_table("api_key", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_api_key_name_in_service_account", type_="unique"
+        )
+
+    with op.batch_alter_table("action", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_action_name_in_workspace", type_="unique"
+        )

zenml/zen_stores/migrations/versions/288f4fb6e112_make_tags_user_scoped.py

@@ -0,0 +1,74 @@
+"""make tags user scoped [288f4fb6e112].
+
+Revision ID: 288f4fb6e112
+Revises: 3b1776345020
+Create Date: 2025-02-19 15:16:42.954792
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.orm import Session
+
+# revision identifiers, used by Alembic.
+revision = "288f4fb6e112"
+down_revision = "3b1776345020"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("tag", schema=None) as batch_op:
+        # First add columns as nullable
+        batch_op.add_column(
+            sa.Column("user_id", sqlmodel.sql.sqltypes.GUID(), nullable=True)
+        )
+
+        # Add foreign key constraints
+        batch_op.create_foreign_key(
+            "fk_tag_user_id_user",
+            "user",
+            ["user_id"],
+            ["id"],
+            ondelete="SET NULL",
+        )
+
+    bind = op.get_bind()
+    session = Session(bind=bind)
+
+    tags = session.execute(
+        sa.text("""
+            SELECT t.id, tr.resource_id, tr.resource_type
+            FROM tag t
+            JOIN tag_resource tr ON t.id = tr.tag_id
+        """)
+    )
+
+    tag_ids = []
+    for tag_id, resource_id, resource_type in tags:
+        if tag_id in tag_ids:
+            continue
+        tag_ids.append(tag_id)
+        session.execute(
+            sa.text(
+                f"""
+                    UPDATE tag
+                    SET user_id = (
+                        SELECT r.user_id
+                        FROM {resource_type} r
+                        WHERE r.id = :resource_id
+                    )
+                    WHERE id = :tag_id
+                """  # nosec B608
+            ),
+            params={"resource_id": resource_id, "tag_id": tag_id},
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("tag", schema=None) as batch_op:
+        batch_op.drop_constraint("fk_tag_user_id_user", type_="foreignkey")
+        batch_op.drop_column("user_id")

zenml/zen_stores/migrations/versions/2e695a26fe7a_add_user_default_workspace.py

@@ -0,0 +1,45 @@
+"""add user default workspace [2e695a26fe7a].
+
+Revision ID: 2e695a26fe7a
+Revises: 1f9d1cd00b90
+Create Date: 2025-02-24 18:19:43.121393
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2e695a26fe7a"
+down_revision = "1f9d1cd00b90"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "default_workspace_id",
+                sqlmodel.sql.sqltypes.GUID(),
+                nullable=True,
+            )
+        )
+        batch_op.create_foreign_key(
+            "fk_user_default_workspace_id_workspace",
+            "workspace",
+            ["default_workspace_id"],
+            ["id"],
+            ondelete="SET NULL",
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "fk_user_default_workspace_id_workspace", type_="foreignkey"
+        )
+        batch_op.drop_column("default_workspace_id")

zenml/zen_stores/migrations/versions/3b1776345020_remove_workspace_from_globals.py

@@ -0,0 +1,81 @@
+"""remove workspace from globals [3b1776345020].
+
+Revision ID: 3b1776345020
+Revises: 0392807467dc
+Create Date: 2025-02-13 15:57:38.255825
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "3b1776345020"
+down_revision = "0392807467dc"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision.
+
+    Raises:
+        RuntimeError: If more than one workspace exists.
+    """
+    # If more than one workspace exists, we fail the migration because it
+    # would mean merging together resources from different workspaces, which
+    # can lead to naming conflicts.
+    workspace_count_query = sa.text("SELECT COUNT(*) FROM workspace")
+    connection = op.get_bind()
+    workspace_count = connection.execute(
+        workspace_count_query,
+    ).scalar()
+    assert isinstance(workspace_count, int)
+    if workspace_count > 1:
+        raise RuntimeError(
+            "Your ZenML installation has more than just the default workspace "
+            "configured. This migration removes the workspace scopes of all "
+            "stacks, components, flavors, service connectors and secrets, "
+            "which may lead to naming conflicts if multiple workspaces are "
+            "present. Please delete all but the default workspace before "
+            "running this migration."
+        )
+
+    with op.batch_alter_table("flavor", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "fk_flavor_workspace_id_workspace", type_="foreignkey"
+        )
+        batch_op.drop_column("workspace_id")
+
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "fk_secret_workspace_id_workspace", type_="foreignkey"
+        )
+        batch_op.drop_column("workspace_id")
+
+    with op.batch_alter_table("service_connector", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "fk_service_connector_workspace_id_workspace", type_="foreignkey"
+        )
+        batch_op.drop_column("workspace_id")
+
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "fk_stack_workspace_id_workspace", type_="foreignkey"
+        )
+        batch_op.drop_column("workspace_id")
+
+    with op.batch_alter_table("stack_component", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "fk_stack_component_workspace_id_workspace", type_="foreignkey"
+        )
+        batch_op.drop_column("workspace_id")
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision.
+
+    Raises:
+        NotImplementedError: This migration is not reversible.
+    """
+    raise NotImplementedError("This migration is not reversible.")

zenml/zen_stores/migrations/versions/41b28cae31ce_make_artifacts_workspace_scoped.py

@@ -0,0 +1,136 @@
+"""make artifacts workspace scoped [41b28cae31ce].
+
+Revision ID: 41b28cae31ce
+Revises: 288f4fb6e112
+Create Date: 2025-02-19 23:23:08.133826
+
+"""
+
+import os
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.orm import Session
+from sqlalchemy.sql import column, table
+
+from zenml.constants import (
+    DEFAULT_PROJECT_NAME,
+    ENV_ZENML_DEFAULT_PROJECT_NAME,
+)
+
+# revision identifiers, used by Alembic.
+revision = "41b28cae31ce"
+down_revision = "288f4fb6e112"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision.
+
+    Raises:
+        Exception: If the default workspace is not found.
+    """
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("artifact", schema=None) as batch_op:
+        # First add columns as nullable
+        batch_op.add_column(
+            sa.Column(
+                "workspace_id", sqlmodel.sql.sqltypes.GUID(), nullable=True
+            )
+        )
+        batch_op.add_column(
+            sa.Column("user_id", sqlmodel.sql.sqltypes.GUID(), nullable=True)
+        )
+
+    # Create a temp table object for the update
+    artifact_table = table(
+        "artifact",
+        column("id", sqlmodel.sql.sqltypes.GUID()),
+        column("workspace_id", sqlmodel.sql.sqltypes.GUID()),
+    )
+
+    default_workspace_name = os.getenv(
+        ENV_ZENML_DEFAULT_PROJECT_NAME, DEFAULT_PROJECT_NAME
+    )
+
+    default_workspace_query = sa.text(
+        "SELECT id FROM workspace WHERE name = :default_workspace_name LIMIT 1"
+    )
+    connection = op.get_bind()
+    default_workspace_id = connection.execute(
+        default_workspace_query,
+        {"default_workspace_name": default_workspace_name},
+    ).scalar()
+
+    if default_workspace_id is None:
+        raise Exception(
+            "Default workspace not found. Cannot proceed with migration."
+        )
+
+    # Update existing records with the default workspace
+    op.execute(
+        artifact_table.update().values(workspace_id=default_workspace_id)
+    )
+
+    bind = op.get_bind()
+    session = Session(bind=bind)
+
+    # Set the artifact owner to the owner of the latest artifact version.
+    # NOTE: we skip this for SQLite because the subquery will fail.
+    if bind.dialect.name != "sqlite":
+        session.execute(
+            sa.text(
+                """
+                UPDATE artifact a
+                SET user_id = (
+                    SELECT v.user_id
+                    FROM `artifact_version` v
+                    WHERE v.artifact_id = a.id
+                    ORDER BY v.created DESC
+                    LIMIT 1
+                )
+            """
+            ),
+        )
+
+    # Now make workspace_id non-nullable
+    with op.batch_alter_table("artifact", schema=None) as batch_op:
+        batch_op.alter_column(
+            "workspace_id",
+            existing_type=sqlmodel.sql.sqltypes.GUID(),
+            nullable=False,
+        )
+
+        # Add foreign key constraints
+        batch_op.create_foreign_key(
+            "fk_artifact_workspace_id_workspace",
+            "workspace",
+            ["workspace_id"],
+            ["id"],
+            ondelete="CASCADE",
+        )
+        batch_op.create_foreign_key(
+            "fk_artifact_user_id_user",
+            "user",
+            ["user_id"],
+            ["id"],
+            ondelete="SET NULL",
+        )
+
+
+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("artifact", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "fk_artifact_user_id_user", type_="foreignkey"
+        )
+        batch_op.drop_constraint(
+            "fk_artifact_workspace_id_workspace", type_="foreignkey"
+        )
+        batch_op.drop_column("user_id")
+        batch_op.drop_column("workspace_id")
+
+    # ### end Alembic commands ###

zenml/zen_stores/migrations/versions/9e7bf0970266_adding_exclusive_attribute_to_tags.py

@@ -0,0 +1,47 @@
+"""adding exclusive attribute to tags [9e7bf0970266].
+
+Revision ID: 9e7bf0970266
+Revises: 2e695a26fe7a
+Create Date: 2025-03-03 15:17:49.341208
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "9e7bf0970266"
+down_revision = "2e695a26fe7a"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # Use batch_alter_table for safer schema modifications
+    with op.batch_alter_table("tag") as batch_op:
+        # First add the column as nullable
+        batch_op.add_column(
+            sa.Column(
+                "exclusive",
+                sa.Boolean(),
+                nullable=True,
+            ),
+        )
+
+    # Update existing rows with default value
+    op.execute("UPDATE tag SET exclusive = FALSE WHERE exclusive IS NULL")
+
+    # Then alter the column to be non-nullable with a default
+    with op.batch_alter_table("tag") as batch_op:
+        batch_op.alter_column(
+            "exclusive",
+            existing_type=sa.Boolean(),
+            nullable=False,
+            server_default=sa.false(),
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    op.drop_column("tag", "exclusive")