Alter version table to add pre_check_version

Author: georgesittas

sqlmesh/core/state_sync/base.py

@@ -40,6 +40,7 @@ class Versions(PydanticModel):
     schema_version: int = 0
     sqlglot_version: str = "0.0.0"
     sqlmesh_version: str = "0.0.0"
+    pre_check_version: int = 0
 
     @property
     def minor_sqlglot_version(self) -> t.Tuple[int, int]:
@@ -54,9 +55,9 @@ def minor_sqlmesh_version(self) -> t.Tuple[int, int]:
     def _package_version_validator(cls, v: t.Any) -> str:
         return "0.0.0" if v is None else str(v)
 
-    @field_validator("schema_version", mode="before")
+    @field_validator("schema_version", "pre_check_version", mode="before")
     @classmethod
-    def _schema_version_validator(cls, v: t.Any) -> int:
+    def _int_version_validator(cls, v: t.Any) -> int:
         return 0 if v is None else int(v)
 
 
@@ -66,7 +67,11 @@ def _schema_version_validator(cls, v: t.Any) -> int:
 ]
 SCHEMA_VERSION: int = len(MIGRATIONS)
 PRE_CHECK_VERSION: int = (
-    max(idx for idx, migration in enumerate(MIGRATIONS) if hasattr(migration, "pre_check")) + 1
+    max(
+        [idx for idx, migration in enumerate(MIGRATIONS) if hasattr(migration, "pre_check")],
+        default=-1,
+    )
+    + 1
 )
 
 

sqlmesh/core/state_sync/db/version.py

@@ -13,6 +13,7 @@
     SQLMESH_VERSION,
 )
 from sqlmesh.core.state_sync.base import (
+    PRE_CHECK_VERSION,
     SCHEMA_VERSION,
     Versions,
 )
@@ -31,13 +32,15 @@ def __init__(self, engine_adapter: EngineAdapter, schema: t.Optional[str] = None
             "schema_version": exp.DataType.build("int"),
             "sqlglot_version": exp.DataType.build(index_type),
             "sqlmesh_version": exp.DataType.build(index_type),
+            "pre_check_version": exp.DataType.build("int"),
         }
 
     def update_versions(
         self,
         schema_version: int = SCHEMA_VERSION,
         sqlglot_version: str = SQLGLOT_VERSION,
         sqlmesh_version: str = SQLMESH_VERSION,
+        pre_check_version: int = PRE_CHECK_VERSION,
     ) -> None:
         import pandas as pd
 
@@ -51,6 +54,7 @@ def update_versions(
                         "schema_version": schema_version,
                         "sqlglot_version": sqlglot_version,
                         "sqlmesh_version": sqlmesh_version,
+                        "pre_check_version": pre_check_version,
                     }
                 ]
             ),
@@ -69,5 +73,8 @@ def get_versions(self) -> Versions:
             return no_version
 
         return Versions(
-            schema_version=row[0], sqlglot_version=row[1], sqlmesh_version=seq_get(row, 2)
+            schema_version=row[0],
+            sqlglot_version=row[1],
+            sqlmesh_version=seq_get(row, 2),
+            pre_check_version=seq_get(row, 3),
         )

sqlmesh/migrations/v0087_add_pre_check_version.py

@@ -0,0 +1,24 @@
+"""Add new 'pre_check_version' column to the version state table."""
+
+from sqlglot import exp
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    versions_table = "_versions"
+    if schema:
+        versions_table = f"{schema}.{versions_table}"
+
+    alter_table_exp = exp.Alter(
+        this=exp.to_table(versions_table),
+        kind="TABLE",
+        actions=[
+            exp.ColumnDef(
+                this=exp.to_column("pre_check_version"),
+                kind=exp.DataType.build("int"),
+            )
+        ],
+    )
+
+    engine_adapter.execute(alter_table_exp)