Feat: Add additional fields to environment summary (#4267)

Author: crericha

sqlmesh/core/console.py

@@ -28,7 +28,7 @@
 from rich.tree import Tree
 from sqlglot import exp
 
-from sqlmesh.core.environment import EnvironmentNamingInfo
+from sqlmesh.core.environment import EnvironmentNamingInfo, EnvironmentSummary
 from sqlmesh.core.linter.rule import RuleViolation
 from sqlmesh.core.model import Model
 from sqlmesh.core.snapshot import (
@@ -188,7 +188,7 @@ class EnvironmentsConsole(abc.ABC):
     """Console for displaying environments"""
 
     @abc.abstractmethod
-    def print_environments(self, environments_summary: t.Dict[str, int]) -> None:
+    def print_environments(self, environments_summary: t.List[EnvironmentSummary]) -> None:
         """Prints all environment names along with expiry datetime."""
 
     @abc.abstractmethod
@@ -659,7 +659,7 @@ def show_row_diff(
     ) -> None:
         pass
 
-    def print_environments(self, environments_summary: t.Dict[str, int]) -> None:
+    def print_environments(self, environments_summary: t.List[EnvironmentSummary]) -> None:
         pass
 
     def show_intervals(self, snapshot_intervals: t.Dict[Snapshot, SnapshotIntervals]) -> None:
@@ -2089,11 +2089,13 @@ def show_row_diff(
                 self.console.print(f"\n[b][green]{target_name} ONLY[/green] sample rows:[/b]")
                 self.console.print(row_diff.t_sample.to_string(index=False), end="\n\n")
 
-    def print_environments(self, environments_summary: t.Dict[str, int]) -> None:
+    def print_environments(self, environments_summary: t.List[EnvironmentSummary]) -> None:
         """Prints all environment names along with expiry datetime."""
         output = [
-            f"{name} - {time_like_to_str(ts)}" if ts else f"{name} - No Expiry"
-            for name, ts in environments_summary.items()
+            f"{summary.name} - {time_like_to_str(summary.expiration_ts)}"
+            if summary.expiration_ts
+            else f"{summary.name} - No Expiry"
+            for summary in environments_summary
         ]
         output_str = "\n".join([str(len(output)), *output])
         self.log_status_update(f"Number of SQLMesh environments are: {output_str}")

sqlmesh/core/environment.py

@@ -95,32 +95,46 @@ def from_environment_catalog_mapping(
         return cls(**construction_kwargs)
 
 
-class Environment(EnvironmentNamingInfo):
-    """Represents an isolated environment.
-
-    Environments are isolated workspaces that hold pointers to physical tables.
+class EnvironmentSummary(PydanticModel):
+    """Represents summary information of an isolated environment.
 
     Args:
-        snapshots: The snapshots that are part of this environment.
+        name: The name of the environment.
         start_at: The start time of the environment.
         end_at: The end time of the environment.
         plan_id: The ID of the plan that last updated this environment.
         previous_plan_id: The ID of the previous plan that updated this environment.
         expiration_ts: The timestamp when this environment will expire.
         finalized_ts: The timestamp when this environment was finalized.
-        promoted_snapshot_ids: The IDs of the snapshots that are promoted in this environment
-            (i.e. for which the views are created). If not specified, all snapshots are promoted.
-        previous_finalized_snapshots: Snapshots that were part of this environment last time it was finalized.
-        requirements: A mapping of library versions for all the snapshots in this environment.
     """
 
-    snapshots_: t.List[t.Any] = Field(alias="snapshots")
+    name: str
     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
+
+    @property
+    def expired(self) -> bool:
+        return self.expiration_ts is not None and self.expiration_ts <= now_timestamp()
+
+
+class Environment(EnvironmentNamingInfo, EnvironmentSummary):
+    """Represents an isolated environment.
+
+    Environments are isolated workspaces that hold pointers to physical tables.
+
+    Args:
+        snapshots: The snapshots that are part of this environment.
+        promoted_snapshot_ids: The IDs of the snapshots that are promoted in this environment
+            (i.e. for which the views are created). If not specified, all snapshots are promoted.
+        previous_finalized_snapshots: Snapshots that were part of this environment last time it was finalized.
+        requirements: A mapping of library versions for all the snapshots in this environment.
+    """
+
+    snapshots_: t.List[t.Any] = Field(alias="snapshots")
     promoted_snapshot_ids_: t.Optional[t.List[t.Any]] = Field(
         default=None, alias="promoted_snapshot_ids"
     )
@@ -203,8 +217,16 @@ def naming_info(self) -> EnvironmentNamingInfo:
         )
 
     @property
-    def expired(self) -> bool:
-        return self.expiration_ts is not None and self.expiration_ts <= now_timestamp()
+    def summary(self) -> EnvironmentSummary:
+        return EnvironmentSummary(
+            name=self.name,
+            start_at=self.start_at,
+            end_at=self.end_at,
+            plan_id=self.plan_id,
+            previous_plan_id=self.previous_plan_id,
+            expiration_ts=self.expiration_ts,
+            finalized_ts=self.finalized_ts,
+        )
 
     def _convert_list_to_models_and_store(
         self, field: str, type_: t.Type[PydanticType]

sqlmesh/core/state_sync/base.py

@@ -9,7 +9,12 @@
 from sqlglot import __version__ as SQLGLOT_VERSION
 
 from sqlmesh import migrations
-from sqlmesh.core.environment import Environment, EnvironmentNamingInfo, EnvironmentStatements
+from sqlmesh.core.environment import (
+    Environment,
+    EnvironmentNamingInfo,
+    EnvironmentStatements,
+    EnvironmentSummary,
+)
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotId,
@@ -137,11 +142,11 @@ def get_environments(self) -> t.List[Environment]:
         """
 
     @abc.abstractmethod
-    def get_environments_summary(self) -> t.Dict[str, int]:
+    def get_environments_summary(self) -> t.List[EnvironmentSummary]:
         """Fetches all environment names along with expiry datetime.
 
         Returns:
-            A dict of all environment names along with expiry datetime.
+            A list of all environment summaries.
         """
 
     @abc.abstractmethod

sqlmesh/core/state_sync/db/environment.py

@@ -12,7 +12,7 @@
     fetchall,
     fetchone,
 )
-from sqlmesh.core.environment import Environment, EnvironmentStatements
+from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentSummary
 from sqlmesh.utils.migration import index_text_type, blob_text_type
 from sqlmesh.utils.date import now_timestamp, time_like_to_str
 from sqlmesh.utils.errors import SQLMeshError
@@ -211,18 +211,19 @@ def get_environments(self) -> t.List[Environment]:
             for row in fetchall(self.engine_adapter, self._environments_query())
         ]
 
-    def get_environments_summary(self) -> t.Dict[str, int]:
-        """Fetches all environment names along with expiry datetime.
+    def get_environments_summary(self) -> t.List[EnvironmentSummary]:
+        """Fetches summaries for all environments.
 
         Returns:
-            A dict of all environment names along with expiry datetime.
+            A list of all environment summaries.
         """
-        return dict(
-            fetchall(
+        return [
+            self._environment_summmary_from_row(row)
+            for row in fetchall(
                 self.engine_adapter,
-                self._environments_query(required_fields=["name", "expiration_ts"]),
-            ),
-        )
+                self._environments_query(required_fields=list(EnvironmentSummary.all_fields())),
+            )
+        ]
 
     def get_environment(
         self, environment: str, lock_for_update: bool = False
@@ -286,6 +287,11 @@ def get_environment_statements(self, environment: str) -> t.List[EnvironmentStat
     def _environment_from_row(self, row: t.Tuple[str, ...]) -> Environment:
         return Environment(**{field: row[i] for i, field in enumerate(Environment.all_fields())})
 
+    def _environment_summmary_from_row(self, row: t.Tuple[str, ...]) -> EnvironmentSummary:
+        return EnvironmentSummary(
+            **{field: row[i] for i, field in enumerate(EnvironmentSummary.all_fields())}
+        )
+
     def _environments_query(
         self,
         where: t.Optional[str | exp.Expression] = None,

sqlmesh/core/state_sync/db/facade.py

@@ -26,7 +26,7 @@
 
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.engine_adapter import EngineAdapter
-from sqlmesh.core.environment import Environment, EnvironmentStatements
+from sqlmesh.core.environment import Environment, EnvironmentStatements, EnvironmentSummary
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotId,
@@ -330,11 +330,11 @@ def get_environments(self) -> t.List[Environment]:
         """
         return self.environment_state.get_environments()
 
-    def get_environments_summary(self) -> t.Dict[str, int]:
+    def get_environments_summary(self) -> t.List[EnvironmentSummary]:
         """Fetches all environment names along with expiry datetime.
 
         Returns:
-            A dict of all environment names along with expiry datetime.
+            A list of all environment summaries.
         """
         return self.environment_state.get_environments_summary()
 

tests/core/state_sync/test_export_import.py

@@ -437,7 +437,11 @@ def test_import_partial(
     import_state(state_sync, output_file, clear=False)
 
     # StateSync should have "prod", "dev" and "dev2".
-    assert sorted(list(state_sync.get_environments_summary().keys())) == ["dev", "dev2", "prod"]
+    assert sorted(list(env.name for env in state_sync.get_environments_summary())) == [
+        "dev",
+        "dev2",
+        "prod",
+    ]
 
     assert not context.plan(environment="dev", skip_tests=True).has_changes
     assert not context.plan(environment="dev2", skip_tests=True).has_changes

tests/core/state_sync/test_state_sync.py

@@ -3464,8 +3464,8 @@ def test_get_environments_summary(
     )
     state_sync.promote(prod)
 
-    actual = state_sync.get_environments_summary()
-    expected = {"prod": None, "test_environment_a": env_a_ttl, "test_environment_b": env_b_ttl}
+    actual = set(state_sync.get_environments_summary())
+    expected = {prod.summary, env_a.summary, env_b.summary}
     assert actual == expected
 
 
@@ -3493,12 +3493,12 @@ def test_get_environments_summary_only_prod(
     )
     state_sync.promote(prod)
     actual = state_sync.get_environments_summary()
-    expected = {"prod": None}
+    expected = [prod.summary]
     assert actual == expected
 
 
 def test_get_environments_summary_no_env(state_sync: EngineAdapterStateSync) -> None:
-    assert state_sync.get_environments_summary() == {}
+    assert state_sync.get_environments_summary() == []
 
 
 @time_machine.travel("2020-01-05 00:00:00 UTC")