Alembic migration fix for databases with scheduled pipelines with 2+ runs

src/zenml/models/v2/core/pipeline_deployment.py

@@ -139,10 +139,10 @@ class PipelineDeploymentResponseMetadata(WorkspaceScopedResponseMetadata):
     client_environment: Dict[str, str] = Field(
         default={}, title="The client environment for this deployment."
     )
-    client_version: str = Field(
+    client_version: Optional[str] = Field(
         title="The version of the ZenML installation on the client side."
     )
-    server_version: str = Field(
+    server_version: Optional[str] = Field(
         title="The version of the ZenML installation on the server side."
     )
     pipeline: Optional[PipelineResponse] = Field(
@@ -219,7 +219,7 @@ def client_environment(self) -> Dict[str, str]:
         return self.get_metadata().client_environment
 
     @property
-    def client_version(self) -> str:
+    def client_version(self) -> Optional[str]:
         """The `client_version` property.
 
         Returns:
@@ -228,7 +228,7 @@ def client_version(self) -> str:
         return self.get_metadata().client_version
 
     @property
-    def server_version(self) -> str:
+    def server_version(self) -> Optional[str]:
         """The `server_version` property.
 
         Returns:

src/zenml/zen_stores/migrations/versions/37835ce041d2_optimizing_database.py

@@ -9,6 +9,7 @@
 import sqlmodel
 from alembic import op
 from sqlalchemy.dialects import mysql
+from sqlalchemy.engine.base import Connection
 from sqlalchemy.sql import text
 
 # revision identifiers, used by Alembic.
@@ -21,21 +22,41 @@
 def upgrade() -> None:
     """Upgrade database schema and/or data, creating a new revision."""
     # ### commands auto generated by Alembic - please adjust! ###
+    bind: Connection = op.get_bind()
+    column_info = sa.inspect(bind).get_columns("pipeline_deployment")
+    columns = [c["name"] for c in column_info]
+
     with op.batch_alter_table("pipeline_deployment", schema=None) as batch_op:
-        batch_op.add_column(
-            sa.Column(
+        if "client_version" not in columns:
+            batch_op.add_column(
+                sa.Column(
+                    "client_version",
+                    sqlmodel.sql.sqltypes.AutoString(),
+                    nullable=True,
+                )
+            )
+        else:
+            batch_op.alter_column(
                 "client_version",
-                sqlmodel.sql.sqltypes.AutoString(),
+                existing_type=sqlmodel.sql.sqltypes.AutoString(),
                 nullable=True,
             )
-        )
-        batch_op.add_column(
-            sa.Column(
+
+        if "server_version" not in columns:
+            batch_op.add_column(
+                sa.Column(
+                    "server_version",
+                    sqlmodel.sql.sqltypes.AutoString(),
+                    nullable=True,
+                )
+            )
+        else:
+            batch_op.alter_column(
                 "server_version",
-                sqlmodel.sql.sqltypes.AutoString(),
+                existing_type=sqlmodel.sql.sqltypes.AutoString(),
                 nullable=True,
             )
-        )
+
         batch_op.alter_column(
             "pipeline_configuration",
             existing_type=sa.TEXT(),
@@ -52,7 +73,7 @@ def upgrade() -> None:
         """
         UPDATE pipeline_deployment
         SET client_version = (
-            SELECT pipeline_run.client_version
+            SELECT max(pipeline_run.client_version)
             FROM pipeline_run
             WHERE pipeline_run.deployment_id = pipeline_deployment.id
         )
@@ -69,7 +90,7 @@ def upgrade() -> None:
         """
         UPDATE pipeline_deployment
         SET server_version = (
-            SELECT pipeline_run.server_version
+            SELECT max(pipeline_run.server_version)
             FROM pipeline_run
             WHERE pipeline_run.deployment_id = pipeline_deployment.id
         )