Repository: TobikoData/sqlmesh
Pull request #3038 (2024-08-23 06:34:25)

Feat!: Improve environment loading performance by lazy evaluating snapshot and snapshot id fields


Patch:
diff --git a/sqlmesh/core/environment.py b/sqlmesh/core/environment.py
index f89e8a1c..4e32e1c5 100644
--- a/sqlmesh/core/environment.py
+++ b/sqlmesh/core/environment.py
@@ -14,6 +14,7 @@ from sqlmesh.utils.date import TimeLike, now_timestamp
 from sqlmesh.utils.pydantic import PydanticModel, field_validator
 
 T = t.TypeVar("T", bound="EnvironmentNamingInfo")
+PydanticType = t.TypeVar("PydanticType", bound="PydanticModel")
 
 
 class EnvironmentNamingInfo(PydanticModel):
@@ -102,32 +103,52 @@ class Environment(EnvironmentNamingInfo):
         previous_finalized_snapshots: Snapshots that were part of this environment last time it was finalized.
     """
 
-    snapshots: t.List[SnapshotTableInfo]
+    snapshots_: t.List[t.Any] = Field(alias="snapshots")
     start_at: TimeLike
     end_at: t.Optional[TimeLike] = None
     plan_id: str
     previous_plan_id: t.Optional[str] = None
     expiration_ts: t.Optional[int] = None
     finalized_ts: t.Optional[int] = None
-    promoted_snapshot_ids: t.Optional[t.List[SnapshotId]] = None
-    previous_finalized_snapshots: t.Optional[t.List[SnapshotTableInfo]] = None
-
-    @field_validator("snapshots", "previous_finalized_snapshots", mode="before")
+    promoted_snapshot_ids_: t.Optional[t.List[t.Any]] = Field(
+        default=None, alias="promoted_snapshot_ids"
+    )
+    previous_finalized_snapshots_: t.Optional[t.List[t.Any]] = Field(
+        default=None, alias="previous_finalized_snapshots"
+    )
+
+    @field_validator("snapshots_", "previous_finalized_snapshots_", mode="before")
     @classmethod
-    def _convert_snapshots(
-        cls, v: str | t.List[SnapshotTableInfo] | None
-    ) -> t.List[SnapshotTableInfo] | None:
+    def _load_snapshots(cls, v: str | t.List[t.Any] | None) -> t.List[t.Any] | None:
         if isinstance(v, str):
-            return [SnapshotTableInfo.parse_obj(obj) for obj in json.loads(v)]
+            return json.loads(v)
+        if v and not isinstance(next(iter(v)), (dict, SnapshotTableInfo)):
+            raise ValueError("Must be a list of SnapshotTableInfo dicts or objects")
         return v
 
-    @field_validator("promoted_snapshot_ids", mode="before")
+    @field_validator("promoted_snapshot_ids_", mode="before")
     @classmethod
-    def _convert_snapshot_ids(cls, v: str | t.List[SnapshotId]) -> t.List[SnapshotId]:
+    def _load_snapshot_ids(cls, v: str | t.List[t.Any] | None) -> t.List[t.Any] | None:
         if isinstance(v, str):
-            return [SnapshotId.parse_obj(obj) for obj in json.loads(v)]
+            return json.loads(v)
+        if v and not isinstance(next(iter(v)), (dict, SnapshotId)):
+            raise ValueError("Must be a list of SnapshotId dicts or objects")
         return v
 
+    @property
+    def snapshots(self) -> t.List[SnapshotTableInfo]:
+        return self._convert_list_to_models_and_store("snapshots_", SnapshotTableInfo)
+
+    def snapshot_dicts(self) -> t.List[dict]:
+        return self._convert_list_to_dicts(self.snapshots_)
+
+    @property
+    def promoted_snapshot_ids(self) -> t.List[SnapshotId]:
+        return self._convert_list_to_models_and_store("promoted_snapshot_ids_", SnapshotId)
+
+    def promoted_snapshot_id_dicts(self) -> t.List[dict]:
+        return self._convert_list_to_dicts(self.promoted_snapshot_ids_)
+
     @property
     def promoted_snapshots(self) -> t.List[SnapshotTableInfo]:
         if self.promoted_snapshot_ids is None:
@@ -136,6 +157,15 @@ class Environment(EnvironmentNamingInfo):
         promoted_snapshot_ids = set(self.promoted_snapshot_ids)
         return [s for s in self.snapshots if s.snapshot_id in promoted_snapshot_ids]
 
+    @property
+    def previous_finalized_snapshots(self) -> t.List[SnapshotTableInfo]:
+        return self._convert_list_to_models_and_store(
+            "previous_finalized_snapshots_", SnapshotTableInfo
+        )
+
+    def previous_finalized_snapshot_dicts(self) -> t.List[dict]:
+        return self._convert_list_to_dicts(self.previous_finalized_snapshots_)
+
     @property
     def finalized_or_current_snapshots(self) -> t.List[SnapshotTableInfo]:
         return (
@@ -156,3 +186,17 @@ class Environment(EnvironmentNamingInfo):
     @property
     def expired(self) -> bool:
         return self.expiration_ts is not None and self.expiration_ts <= now_timestamp()
+
+    def _convert_list_to_models_and_store(
+        self, field: str, type_: t.Type[PydanticType]
+    ) -> t.List[PydanticType]:
+        value = getattr(self, field)
+        if value and not isinstance(value[0], type_):
+            value = [type_.parse_obj(obj) for obj in value]
+            setattr(self, field, value)
+        return t.cast(t.List[PydanticType], value)
+
+    def _convert_list_to_dicts(self, value: t.Optional[t.List[t.Any]]) -> t.List[dict]:
+        if not value:
+            return []
+        return value if isinstance(value[0], dict) else [v.dict() for v in value]
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index bb81869e..943a2279 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -1495,7 +1495,7 @@ class EngineAdapterStateSync(StateSync):
             ]
 
             if snapshots != environment.snapshots:
-                environment.snapshots = snapshots
+                environment.snapshots_ = snapshots
                 updated_environments.append(environment)
                 if environment.name == c.PROD:
                     updated_prod_environment = environment
@@ -1689,9 +1689,7 @@ def _environment_to_df(environment: Environment) -> pd.DataFrame:
         [
             {
                 "name": environment.name,
-                "snapshots": json.dumps(
-                    [snapshot.dict(mode="json") for snapshot in environment.snapshots]
-                ),
+                "snapshots": json.dumps(environment.snapshot_dicts()),
                 "start_at": time_like_to_str(environment.start_at),
                 "end_at": time_like_to_str(environment.end_at) if environment.end_at else None,
                 "plan_id": environment.plan_id,
@@ -1699,19 +1697,14 @@ def _environment_to_df(environment: Environment) -> pd.DataFrame:
                 "expiration_ts": environment.expiration_ts,
                 "finalized_ts": environment.finalized_ts,
                 "promoted_snapshot_ids": (
-                    json.dumps([s.dict() for s in environment.promoted_snapshot_ids])
+                    json.dumps(environment.promoted_snapshot_id_dicts())
                     if environment.promoted_snapshot_ids is not None
                     else None
                 ),
                 "suffix_target": environment.suffix_target.value,
                 "catalog_name_override": environment.catalog_name_override,
                 "previous_finalized_snapshots": (
-                    json.dumps(
-                        [
-                            snapshot.dict(mode="json")
-                            for snapshot in environment.previous_finalized_snapshots
-                        ]
-                    )
+                    json.dumps(environment.previous_finalized_snapshot_dicts())
                     if environment.previous_finalized_snapshots is not None
                     else None
                 ),


Test patch:
diff --git a/tests/core/test_environment.py b/tests/core/test_environment.py
index 42b75b0c..163bba47 100644
--- a/tests/core/test_environment.py
+++ b/tests/core/test_environment.py
@@ -1,6 +1,8 @@
 import pytest
 
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo
+from sqlmesh.core.snapshot import SnapshotId, SnapshotTableInfo
+from sqlmesh.core.state_sync.engine_adapter import _environment_to_df
 
 
 def test_sanitize_name():
@@ -42,3 +44,38 @@ def test_from_environment_catalog_mapping(mapping, name, expected):
         ).catalog_name_override
         == expected
     )
+
+
+def test_lazy_loading(sushi_context):
+    snapshot_ids = [s.snapshot_id for s in sushi_context.snapshots.values()]
+    snapshot_table_infos = [s.table_info for s in sushi_context.snapshots.values()]
+    env = Environment(
+        name="test",
+        start_at="now",
+        plan_id="plan_1",
+        snapshots=snapshot_table_infos,
+        promoted_snapshot_ids=snapshot_ids,
+        previous_finalized_snapshots=snapshot_table_infos,
+    )
+
+    df = _environment_to_df(env)
+    row = df.to_dict(orient="records")[0]
+    env = Environment(**{field: row[field] for field in Environment.all_fields()})
+
+    assert all(isinstance(snapshot, dict) for snapshot in env.snapshots_)
+    assert all(isinstance(snapshot, SnapshotTableInfo) for snapshot in env.snapshots)
+    assert all(isinstance(s_id, dict) for s_id in env.promoted_snapshot_ids_)
+    assert all(isinstance(s_id, SnapshotId) for s_id in env.promoted_snapshot_ids)
+    assert all(isinstance(snapshot, dict) for snapshot in env.previous_finalized_snapshots_)
+    assert all(
+        isinstance(snapshot, SnapshotTableInfo) for snapshot in env.previous_finalized_snapshots
+    )
+
+    with pytest.raises(ValueError, match="Must be a list of SnapshotTableInfo dicts or objects"):
+        Environment(**{**env.dict(), **{"snapshots": [1, 2, 3]}})
+
+    with pytest.raises(ValueError, match="Must be a list of SnapshotId dicts or objects"):
+        Environment(**{**env.dict(), **{"promoted_snapshot_ids": [1, 2, 3]}})
+
+    with pytest.raises(ValueError, match="Must be a list of SnapshotTableInfo dicts or objects"):
+        Environment(**{**env.dict(), **{"previous_finalized_snapshots": [1, 2, 3]}})
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index 81ad38b8..0efdf043 100644
--- a/tests/core/test_state_sync.py
+++ b/tests/core/test_state_sync.py
@@ -1065,6 +1065,9 @@ def test_delete_expired_environments(state_sync: EngineAdapterStateSync, make_sn
     env_b = env_a.copy(update={"name": "test_environment_b", "expiration_ts": now_ts + 1000})
     state_sync.promote(env_b)
 
+    env_a = Environment(**json.loads(env_a.json()))
+    env_b = Environment(**json.loads(env_b.json()))
+
     assert state_sync.get_environment(env_a.name) == env_a
     assert state_sync.get_environment(env_b.name) == env_b
 
@@ -1215,7 +1218,7 @@ def test_delete_expired_snapshots_promoted(
     assert not state_sync.delete_expired_snapshots()
     assert set(state_sync.get_snapshots(all_snapshots)) == {snapshot.snapshot_id}
 
-    env.snapshots = []
+    env.snapshots_ = []
     state_sync.promote(env)
 
     now_timestamp_mock = mocker.patch("sqlmesh.core.state_sync.engine_adapter.now_timestamp")
diff --git a/tests/schedulers/airflow/test_integration.py b/tests/schedulers/airflow/test_integration.py
index ceeadab3..aac7ffaa 100644
--- a/tests/schedulers/airflow/test_integration.py
+++ b/tests/schedulers/airflow/test_integration.py
@@ -64,7 +64,7 @@ def test_apply_plan_create_backfill_promote(
         airflow_client.apply_plan([snapshot], environment, random_name())
 
     # Verify full environment demotion.
-    environment.snapshots = []
+    environment.snapshots_ = []
     environment.previous_plan_id = environment.plan_id
     environment.plan_id = "new_plan_id"
     _apply_plan_and_block(airflow_client, [], environment)
diff --git a/tests/schedulers/airflow/test_plan.py b/tests/schedulers/airflow/test_plan.py
index 51770325..6d382506 100644
--- a/tests/schedulers/airflow/test_plan.py
+++ b/tests/schedulers/airflow/test_plan.py
@@ -550,13 +550,19 @@ def test_plan_dag_state(snapshot: Snapshot, sushi_context: Context, random_name)
 
     plan_dag_state = PlanDagState.from_state_sync(sushi_context.state_sync)
 
+    def get_hydrated_dag_specs():
+        state_plan_dag_specs = plan_dag_state.get_dag_specs()
+        for state_plan_dag_spec in state_plan_dag_specs:
+            state_plan_dag_spec.environment.snapshots
+        return state_plan_dag_specs
+
     assert not plan_dag_state.get_dag_specs()
 
     plan_dag_state.add_dag_spec(plan_dag_spec)
-    assert plan_dag_state.get_dag_specs() == [plan_dag_spec]
+    assert get_hydrated_dag_specs() == [plan_dag_spec]
 
     plan_dag_state.delete_dag_specs([])
-    assert plan_dag_state.get_dag_specs() == [plan_dag_spec]
+    assert get_hydrated_dag_specs() == [plan_dag_spec]
 
     plan_dag_state.delete_dag_specs(
         [common.plan_application_dag_id(environment_name, "test_request_id")]