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

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_WORKSPACE_NAME,
+    ENV_ZENML_DEFAULT_WORKSPACE_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_WORKSPACE_NAME, DEFAULT_WORKSPACE_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")

zenml/zen_stores/migrations/versions/b557b2871693_update_step_run_input_types.py

@@ -17,17 +17,21 @@ depends_on = None
 
 def upgrade() -> None:
     """Upgrade database schema and/or data, creating a new revision."""
-    op.execute("""
+    op.execute(
+        """
         UPDATE step_run_input_artifact
         SET type = 'step_output'
         WHERE type = 'default'
-    """)
+    """
+    )
 
 
 def downgrade() -> None:
     """Downgrade database schema and/or data back to the previous revision."""
-    op.execute("""
+    op.execute(
+        """
         UPDATE step_run_input_artifact
         SET type = 'default'
         WHERE type = 'step_output'
-    """)
+    """
+    )

zenml/zen_stores/migrations/versions/cc269488e5a9_separate_run_metadata.py

@@ -41,10 +41,12 @@ def upgrade() -> None:
     connection = op.get_bind()
 
     run_metadata_data = connection.execute(
-        sa.text("""
+        sa.text(
+            """
             SELECT id, resource_id, resource_type
             FROM run_metadata
-        """)
+        """
+        )
     ).fetchall()
 
     # Prepare data with new UUIDs for bulk insert
@@ -107,20 +109,24 @@ def downgrade() -> None:
 
     # Fetch data from `run_metadata_resource`
     run_metadata_resource_data = connection.execute(
-        sa.text("""
+        sa.text(
+            """
             SELECT resource_id, resource_type, run_metadata_id
             FROM run_metadata_resource
-        """)
+        """
+        )
     ).fetchall()
 
     # Update `run_metadata` with the data from `run_metadata_resource`
     for row in run_metadata_resource_data:
         connection.execute(
-            sa.text("""
+            sa.text(
+                """
                 UPDATE run_metadata
                 SET resource_id = :resource_id, resource_type = :resource_type
                 WHERE id = :run_metadata_id
-            """),
+            """
+            ),
             {
                 "resource_id": row.resource_id,
                 "resource_type": row.resource_type,

zenml/zen_stores/migrations/versions/f1d723fd723b_add_secret_private_attr.py

@@ -0,0 +1,61 @@
+"""add secret private attr [f1d723fd723b].
+
+Revision ID: f1d723fd723b
+Revises: 41b28cae31ce
+Create Date: 2025-02-21 19:40:14.596681
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "f1d723fd723b"
+down_revision = "41b28cae31ce"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # Step 1: Add private column as nullable
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("private", sa.Boolean(), nullable=True))
+
+    # Step 2: Update existing records based on scope
+    connection = op.get_bind()
+    connection.execute(
+        sa.text("UPDATE secret SET private = TRUE WHERE scope = 'user';")
+    )
+    connection.execute(
+        sa.text("UPDATE secret SET private = FALSE WHERE scope != 'user';")
+    )
+
+    # Step 3: Make private column non-nullable and drop scope
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.alter_column(
+            "private", existing_type=sa.Boolean(), nullable=False
+        )
+        batch_op.drop_column("scope")
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("scope", sa.VARCHAR(), nullable=True))
+
+    # Restore scope values based on private flag
+    connection = op.get_bind()
+    connection.execute(
+        sa.text("UPDATE secret SET scope = 'user' WHERE private = TRUE;")
+    )
+    connection.execute(
+        sa.text("UPDATE secret SET scope = 'workspace' WHERE private = FALSE;")
+    )
+
+    # Make scope non-nullable and drop private
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.alter_column(
+            "scope", existing_type=sa.VARCHAR(), nullable=False
+        )
+        batch_op.drop_column("private")

zenml/zen_stores/migrations/versions/f76a368a25a5_add_stack_description.py

@@ -0,0 +1,35 @@
+"""add stack description [f76a368a25a5].
+
+Revision ID: f76a368a25a5
+Revises: f1d723fd723b
+Create Date: 2025-02-22 19:03:32.568076
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "f76a368a25a5"
+down_revision = "f1d723fd723b"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "description",
+                sqlmodel.sql.sqltypes.AutoString(),
+                nullable=True,
+            )
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.drop_column("description")