add migration framework (#654)

* add migration framework

* move init to migration and refactor for pr

* make postgres engine adapter so it can inherit table_exists

* add logging

Author: tobymao

sqlmesh/core/context.py

@@ -322,7 +322,9 @@ def state_sync(self) -> StateSync:
                 raise ConfigError(
                     "The operation is not supported when using a read-only state sync"
                 )
-            self._state_sync.init_schema()
+
+            if self._state_sync.get_versions(validate=False).schema_version == 0:
+                self._state_sync.migrate()
         return self._state_sync
 
     @property
@@ -807,6 +809,13 @@ def audit(
             self.console.show_sql(f"{error.query}")
         self.console.log_status_update("Done.")
 
+    def migrate(self) -> None:
+        """Migrates SQLMesh to the current running version.
+
+        Please contact your SQLMesh administrator before doing this.
+        """
+        self.state_sync.migrate()
+
     def close(self) -> None:
         """Releases all resources allocated by this context."""
         self.snapshot_evaluator.close()

sqlmesh/core/engine_adapter/__init__.py

@@ -9,6 +9,7 @@
 from sqlmesh.core.engine_adapter.databricks import DatabricksSparkSessionEngineAdapter
 from sqlmesh.core.engine_adapter.databricks_api import DatabricksSQLEngineAdapter
 from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter
+from sqlmesh.core.engine_adapter.postgres import PostgresEngineAdapter
 from sqlmesh.core.engine_adapter.redshift import RedshiftEngineAdapter
 from sqlmesh.core.engine_adapter.shared import TransactionType
 from sqlmesh.core.engine_adapter.snowflake import SnowflakeEngineAdapter
@@ -21,7 +22,7 @@
     "snowflake": SnowflakeEngineAdapter,
     "databricks": DatabricksSparkSessionEngineAdapter,
     "redshift": RedshiftEngineAdapter,
-    "postgres": EngineAdapterWithIndexSupport,
+    "postgres": PostgresEngineAdapter,
     "mysql": EngineAdapterWithIndexSupport,
     "mssql": EngineAdapterWithIndexSupport,
 }

sqlmesh/core/engine_adapter/postgres.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+
+from sqlmesh.core.engine_adapter.base import (
+    EngineAdapter,
+    EngineAdapterWithIndexSupport,
+)
+
+if t.TYPE_CHECKING:
+    from sqlmesh.core._typing import TableName
+
+
+class PostgresBaseEngineAdapter(EngineAdapter):
+    def table_exists(self, table_name: TableName) -> bool:
+        """
+        Redshift/Postgres doesn't support describe so I'm using what the redshift cursor does to check if a table
+        exists. We don't use this directly because we still want all execution to go through our execute method
+
+        Reference: https://github.com/aws/amazon-redshift-python-driver/blob/master/redshift_connector/cursor.py#L528-L553
+        """
+        table = exp.to_table(table_name)
+
+        # Redshift doesn't support catalog
+        if table.args.get("catalog"):
+            return False
+
+        query = exp.select("1").from_("information_schema.tables")
+        where = exp.condition(f"table_name = '{table.alias_or_name}'")
+
+        schema = table.text("db")
+        if schema:
+            where = where.and_(f"table_schema = '{schema}'")
+
+        self.execute(query.where(where))
+
+        result = self.cursor.fetchone()
+
+        return result[0] == 1 if result is not None else False
+
+
+class PostgresEngineAdapter(PostgresBaseEngineAdapter, EngineAdapterWithIndexSupport):
+    DIALECT = "postgres"

sqlmesh/core/engine_adapter/redshift.py

@@ -8,15 +8,15 @@
 
 from sqlmesh.core.dialect import pandas_to_sql
 from sqlmesh.core.engine_adapter._typing import DF_TYPES, Query
-from sqlmesh.core.engine_adapter.base import EngineAdapter
+from sqlmesh.core.engine_adapter.postgres import PostgresBaseEngineAdapter
 from sqlmesh.core.engine_adapter.shared import DataObject
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
     from sqlmesh.core.engine_adapter._typing import QueryOrDF
 
 
-class RedshiftEngineAdapter(EngineAdapter):
+class RedshiftEngineAdapter(PostgresBaseEngineAdapter):
     DIALECT = "redshift"
     DEFAULT_BATCH_SIZE = 1000
 
@@ -145,32 +145,6 @@ def replace_query(
     def _short_hash(self) -> str:
         return uuid.uuid4().hex[:8]
 
-    def table_exists(self, table_name: TableName) -> bool:
-        """
-        Redshift doesn't support describe so I'm using what the redshift cursor does to check if a table
-        exists. We don't use this directly because we still want all execution to go through our execute method
-
-        Reference: https://github.com/aws/amazon-redshift-python-driver/blob/master/redshift_connector/cursor.py#L528-L553
-        """
-        table = exp.to_table(table_name)
-
-        # Redshift doesn't support catalog
-        if table.args.get("catalog"):
-            return False
-
-        q: str = (
-            f"SELECT 1 FROM information_schema.tables WHERE table_name = '{table.alias_or_name}'"
-        )
-        database_name = table.args.get("db")
-        if database_name:
-            q += f" AND table_schema = '{database_name}'"
-
-        self.execute(q)
-
-        result = self.cursor.fetchone()
-
-        return result[0] == 1 if result is not None else False
-
     def _get_data_objects(
         self, schema_name: str, catalog_name: t.Optional[str] = None
     ) -> t.List[DataObject]:

sqlmesh/core/state_sync/__init__.py

@@ -13,6 +13,6 @@
 The provided `sqlmesh.core.state_sync.EngineAdapterStateSync` leverages an existing engine
 adapter to read and write state to the underlying data store.
 """
-from sqlmesh.core.state_sync.base import StateReader, StateSync
+from sqlmesh.core.state_sync.base import StateReader, StateSync, Versions
 from sqlmesh.core.state_sync.common import CommonStateSyncMixin
 from sqlmesh.core.state_sync.engine_adapter import EngineAdapterStateSync

sqlmesh/core/state_sync/base.py

@@ -1,8 +1,14 @@
 from __future__ import annotations
 
 import abc
+import importlib
+import logging
+import pkgutil
 import typing as t
 
+from sqlglot import __version__ as SQLGLOT_VERSION
+
+from sqlmesh import migrations
 from sqlmesh.core import scheduler
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.snapshot import (
@@ -14,8 +20,30 @@
     SnapshotNameVersionLike,
     SnapshotTableInfo,
 )
+from sqlmesh.utils import major_minor
 from sqlmesh.utils.date import TimeLike, now, to_datetime
 from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.utils.pydantic import PydanticModel
+
+logger = logging.getLogger(__name__)
+
+
+class Versions(PydanticModel):
+    """Represents the various versions of dependencies in the state sync."""
+
+    schema_version: int
+    sqlglot_version: str
+
+    @property
+    def minor_sqlglot_version(self) -> t.Tuple[int, int]:
+        return major_minor(self.sqlglot_version)
+
+
+MIGRATIONS = [
+    importlib.import_module(f"sqlmesh.migrations.{migration}")
+    for migration in sorted(info.name for info in pkgutil.iter_modules(migrations.__path__))
+]
+SCHEMA_VERSION: int = len(MIGRATIONS)
 
 
 class StateReader(abc.ABC):
@@ -168,13 +196,55 @@ def missing_intervals(
                 missing[snapshot] = intervals
         return missing
 
+    def get_versions(self, validate: bool = True) -> Versions:
+        """Get the current versions of the SQLMesh schema and libraries.
 
-class StateSync(StateReader, abc.ABC):
-    """Abstract base class for snapshot and environment state management."""
+        Args:
+            validate: Whether or not to raise error if the running version is ahead of state.
+
+        Returns:
+            The versions object.
+        """
+        versions = self._get_versions()
+
+        def raise_error(lib: str, local: str | int, remote: str | int, ahead: bool = False) -> None:
+            if ahead:
+                raise SQLMeshError(
+                    f"{lib} (local) is using version '{local}' which is ahead of '{remote}' (remote). Please run a migration."
+                )
+            raise SQLMeshError(
+                f"{lib} (local) is using version '{local}' which is behind '{remote}' (remote). Please upgrade {lib}."
+            )
+
+        if SCHEMA_VERSION < versions.schema_version:
+            raise_error("SQLMesh", SCHEMA_VERSION, versions.schema_version)
+
+        if major_minor(SQLGLOT_VERSION) < major_minor(versions.sqlglot_version):
+            raise_error("SQLGlot", SQLGLOT_VERSION, versions.sqlglot_version)
+
+        if validate:
+            if SCHEMA_VERSION > versions.schema_version:
+                raise_error("SQLMesh", SCHEMA_VERSION, versions.schema_version, ahead=True)
+
+            if major_minor(SQLGLOT_VERSION) > major_minor(versions.sqlglot_version):
+                raise_error("SQLGlot", SQLGLOT_VERSION, versions.sqlglot_version, ahead=True)
+
+        return versions
 
     @abc.abstractmethod
-    def init_schema(self) -> None:
-        """Optional initialization of the sync."""
+    def _get_versions(self, lock_for_update: bool = False) -> Versions:
+        """Queries the store to get the current versions of SQLMesh and deps.
+
+        Args:
+            lock_for_update: Whether or not the usage of this method plans to update the row.
+
+        Returns:
+            The versions object.
+        """
+
+
+class StateSync(StateReader, abc.ABC):
+    """Abstract base class for snapshot and environment state management."""
 
     @abc.abstractmethod
     def push_snapshots(self, snapshots: t.Iterable[Snapshot]) -> None:
@@ -291,3 +361,30 @@ def unpause_snapshots(
             unpaused_dt: The datetime object which indicates when target snapshots
                 were unpaused.
         """
+
+    def migrate(self) -> None:
+        """Migrate the state sync to the latest SQLMesh / SQLGlot version."""
+        versions = self.get_versions(validate=False)
+        migrations = MIGRATIONS[versions.schema_version :]
+
+        if not migrations and major_minor(SQLGLOT_VERSION) == versions.minor_sqlglot_version:
+            return
+
+        for migration in migrations:
+            logger.info(f"Applying migration {migration}")
+            migration.migrate(self)
+
+        self._migrate_rows()
+        self._update_versions()
+
+    @abc.abstractmethod
+    def _migrate_rows(self) -> None:
+        """Migrate all rows in the state sync, including snapshots and environments."""
+
+    @abc.abstractmethod
+    def _update_versions(
+        self,
+        schema_version: int = SCHEMA_VERSION,
+        sqlglot_version: str = SQLGLOT_VERSION,
+    ) -> None:
+        """Update the schema versions to the latest running versions."""