Add DB migration

src/zenml/zen_stores/migrations/versions/76b6b9101a0c_add_model_version_producer_run_unique_.py

@@ -0,0 +1,62 @@
+"""Add model version producer run unique constraint [76b6b9101a0c].
+
+Revision ID: 76b6b9101a0c
+Revises: 0.71.0
+Create Date: 2024-12-09 11:18:59.515385
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "76b6b9101a0c"
+down_revision = "0.71.0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("model_version", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "is_numeric",
+                sa.Boolean(),
+                sa.Computed(
+                    "name == number",
+                ),
+                nullable=True,
+            )
+        )
+        batch_op.add_column(
+            sa.Column(
+                "producer_run_id", sqlmodel.sql.sqltypes.GUID(), nullable=True
+            )
+        )
+        batch_op.create_index(
+            "unique_numeric_version_for_pipeline_run",
+            [
+                "model_id",
+                "is_numeric",
+                sa.text(
+                    "CASE WHEN producer_run_id IS NOT NULL THEN producer_run_id ELSE id END"
+                ),
+            ],
+            unique=True,
+        )
+
+    # ### 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_index("unique_numeric_version_for_pipeline_run")
+        batch_op.drop_column("producer_run_id")
+        batch_op.drop_column("is_numeric")
+
+    # ### end Alembic commands ###