PR feedback

Author: erindru

docs/concepts/state.md

@@ -74,7 +74,14 @@ The state file is a simple `json` file that looks like:
     /* object for every Virtual Data Environment in the project. key = environment name, value = environment details */
     "environments": {
         "prod": {
-            "..."
+            /* information about the environment itself */
+            "environment": {
+                "..."
+            },
+            /* information about any before_all / after_all statements for this environment */
+            "statements": [
+                "..."
+            ]
         }
     }
 }

sqlmesh/core/console.py

@@ -1129,6 +1129,8 @@ def stop_state_export(self, success: bool, output_file: Path) -> None:
             self.state_export_progress.stop()
             self.state_export_progress = None
 
+            self.log_status_update("")
+
             if success:
                 self.log_success(f"State exported successfully to '{output_file.as_posix()}'")
             else:
@@ -1252,6 +1254,8 @@ def stop_state_import(self, success: bool, input_file: Path) -> None:
             self.state_import_progress.stop()
             self.state_import_progress = None
 
+            self.log_status_update("")
+
             if success:
                 self.log_success(f"State imported successfully from '{input_file.as_posix()}'")
             else:

sqlmesh/core/state_sync/common.py

@@ -8,7 +8,8 @@
 
 from sqlmesh.core.console import Console
 from sqlmesh.core.dialect import schema_
-from sqlmesh.core.environment import Environment
+from sqlmesh.utils.pydantic import PydanticModel
+from sqlmesh.core.environment import Environment, EnvironmentStatements
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.core.snapshot import Snapshot
 
@@ -96,6 +97,11 @@ def chunk_iterable(iterable: t.Iterable[T], size: int = 10) -> t.Iterable[t.Iter
         yield itertools.chain([first], itertools.islice(iterator, size - 1))
 
 
+class EnvironmentWithStatements(PydanticModel):
+    environment: Environment
+    statements: t.List[EnvironmentStatements] = []
+
+
 class StateStream(abc.ABC):
     """
     Represents a stream of state either going into the StateSync (perhaps loaded from a file)
@@ -114,5 +120,5 @@ def snapshots(self) -> t.Iterable[Snapshot]:
 
     @property
     @abc.abstractmethod
-    def environments(self) -> t.Iterable[Environment]:
-        """A stream of Environment objects"""
+    def environments(self) -> t.Iterable[EnvironmentWithStatements]:
+        """A stream of Environments with any EnvironmentStatements attached"""

sqlmesh/core/state_sync/db/facade.py

@@ -51,6 +51,7 @@
     transactional,
     StateStream,
     chunk_iterable,
+    EnvironmentWithStatements,
 )
 from sqlmesh.core.state_sync.db.interval import IntervalState
 from sqlmesh.core.state_sync.db.environment import EnvironmentState
@@ -482,11 +483,13 @@ def snapshots(self) -> t.Iterable[Snapshot]:
                     yield from state_sync.get_snapshots(chunk).values()
 
             @property
-            def environments(self) -> t.Iterable[Environment]:
-                if environment_names:
-                    yield from selected_environments
-                else:
-                    yield from state_sync.get_environments()
+            def environments(self) -> t.Iterable[EnvironmentWithStatements]:
+                envs = selected_environments if environment_names else state_sync.get_environments()
+
+                for env in envs:
+                    yield EnvironmentWithStatements(
+                        environment=env, statements=state_sync.get_environment_statements(env.name)
+                    )
 
         return _DumpStateStream()
 
@@ -515,7 +518,7 @@ def import_(self, stream: StateStream, clear: bool = True) -> None:
             )
             overwrite_existing_snapshots = (
                 not clear
-            )  # if clear=True, all existing snapshjots were dropped anyway
+            )  # if clear=True, all existing snapshots were dropped anyway
             self.snapshot_state.push_snapshots(
                 snapshot_iterator, overwrite=overwrite_existing_snapshots
             )
@@ -529,12 +532,12 @@ def import_(self, stream: StateStream, clear: bool = True) -> None:
                 }
             )
 
-        existing_environments = set(self.get_environments_summary().keys()) if not clear else set()
-        for environment in stream.environments:
-            if not clear and environment.name in existing_environments:
-                self.environment_state.update_environment(environment)
-            else:
-                self.promote(environment)
+        for environment_with_statements in stream.environments:
+            environment = environment_with_statements.environment
+            self.environment_state.update_environment(environment)
+            self.environment_state.update_environment_statements(
+                environment.name, environment.plan_id, environment_with_statements.statements
+            )
 
         self.update_auto_restatements(auto_restatements)
 

sqlmesh/core/state_sync/export_import.py

@@ -3,10 +3,9 @@
 from sqlmesh.core.state_sync import StateSync
 from sqlmesh.core.snapshot import Snapshot
 from sqlmesh.utils.date import now, to_tstz
-from sqlmesh.core.environment import Environment
 from sqlmesh.utils.pydantic import _expression_encoder
 from sqlmesh.core.state_sync import Versions
-from sqlmesh.core.state_sync.common import StateStream
+from sqlmesh.core.state_sync.common import StateStream, EnvironmentWithStatements
 from sqlmesh.core.console import Console
 from pathlib import Path
 from sqlmesh.core.console import NoopConsole
@@ -45,7 +44,7 @@ def snapshots(self) -> t.Iterable[Snapshot]:
             return iter(snapshots.values())
 
         @property
-        def environments(self) -> t.Iterable[Environment]:
+        def environments(self) -> t.Iterable[EnvironmentWithStatements]:
             return []
 
     return _LocalStateStream()
@@ -71,11 +70,11 @@ def _dump_snapshots(
 
     @streamable_dict
     def _dump_environments(
-        environment_stream: t.Iterable[Environment],
+        environment_stream: t.Iterable[EnvironmentWithStatements],
     ) -> t.Iterator[t.Tuple[str, t.Any]]:
         console.update_state_export_progress(environment_count=0)
         for idx, env in enumerate(environment_stream):
-            yield env.name, _dump_pydantic_model(env)
+            yield env.environment.name, _dump_pydantic_model(env)
             console.update_state_export_progress(environment_count=idx + 1)
 
     @streamable_dict
@@ -129,12 +128,14 @@ def snapshots(self) -> t.Iterable[Snapshot]:
             console.update_state_import_progress(snapshots_complete=True)
 
         @property
-        def environments(self) -> t.Iterable[Environment]:
+        def environments(self) -> t.Iterable[EnvironmentWithStatements]:
             stream = data()["environments"]
 
             console.update_state_import_progress(environment_count=0)
             for idx, (_, raw_environment) in enumerate(stream.items()):
-                environment = Environment.model_validate(to_standard_types(raw_environment))
+                environment = EnvironmentWithStatements.model_validate(
+                    to_standard_types(raw_environment)
+                )
                 yield environment
                 console.update_state_import_progress(environment_count=idx + 1)
 

tests/core/state_sync/test_export_import.py

@@ -1,6 +1,6 @@
 import pytest
 from pathlib import Path
-from sqlmesh.core.state_sync import StateSync, EngineAdapterStateSync
+from sqlmesh.core.state_sync import StateSync, EngineAdapterStateSync, CachingStateSync
 from sqlmesh.core.state_sync.export_import import export_state, import_state
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.core import constants as c
@@ -566,3 +566,52 @@ def test_roundtrip_includes_auto_restatements(
 
     plan = context.plan(skip_tests=True)
     assert not plan.has_changes
+
+
+def test_roundtrip_includes_environment_statements(tmp_path: Path) -> None:
+    config = Config(
+        gateways={
+            "main": GatewayConfig(
+                connection=DuckDBConnectionConfig(database=str(tmp_path / "warehouse.db")),
+                state_connection=DuckDBConnectionConfig(database=str(tmp_path / "state.db")),
+            )
+        },
+        default_gateway="main",
+        model_defaults=ModelDefaultsConfig(
+            dialect="duckdb",
+        ),
+        before_all=["select 1 as before_all"],
+        after_all=["select 2 as after_all"],
+    )
+
+    context = Context(paths=tmp_path, config=config)
+    context.plan(auto_apply=True)
+
+    state_file = tmp_path / "state_dump.json"
+    context.export_state(state_file)
+
+    environments = json.loads(state_file.read_text(encoding="utf8"))["environments"]
+
+    assert environments["prod"]["statements"][0]["before_all"][0] == "select 1 as before_all"
+    assert environments["prod"]["statements"][0]["after_all"][0] == "select 2 as after_all"
+
+    assert not context.plan().has_changes
+
+    state_sync = context.state_sync
+    assert isinstance(state_sync, CachingStateSync)
+    assert isinstance(state_sync.state_sync, EngineAdapterStateSync)
+
+    # show state destroyed
+    state_sync.state_sync.engine_adapter.drop_schema("sqlmesh", cascade=True)  # type: ignore
+    with pytest.raises(SQLMeshError, match=r"Please run a migration"):
+        state_sync.get_versions(validate=True)
+
+    state_sync.migrate(default_catalog=None)
+    import_state(state_sync, state_file)
+
+    assert not context.plan().has_changes
+
+    environment_statements = state_sync.get_environment_statements("prod")
+    assert len(environment_statements) == 1
+    assert environment_statements[0].before_all[0] == "select 1 as before_all"
+    assert environment_statements[0].after_all[0] == "select 2 as after_all"