PR Feedback 2, various fixes and improvements

Author: VaggelisD

sqlmesh/core/console.py

@@ -221,6 +221,7 @@ def show_environment_difference_summary(
         self,
         context_diff: ContextDiff,
         no_diff: bool = True,
+        environment: t.Optional[str] = None,
     ) -> None:
         """Displays a summary of differences for the environment."""
 
@@ -647,6 +648,7 @@ def show_environment_difference_summary(
         self,
         context_diff: ContextDiff,
         no_diff: bool = True,
+        environment: t.Optional[str] = None,
     ) -> None:
         pass
 
@@ -1526,18 +1528,21 @@ def show_environment_difference_summary(
         self,
         context_diff: ContextDiff,
         no_diff: bool = True,
+        environment: t.Optional[str] = None,
     ) -> None:
         """Shows a summary of the environment differences.
 
         Args:
             context_diff: The context diff to use to print the summary
             no_diff: Hide the actual environment statement differences.
+            environment: The initial target environment
         """
         if context_diff.is_new_environment:
+            new_environment = environment or context_diff.environment
             msg = (
-                f"\n`{context_diff.environment}` environment will be initialized"
+                f"\n`{new_environment}` environment will be initialized"
                 if not context_diff.create_from_env_exists
-                else f"\nNew environment `{context_diff.environment}` will be created from `{context_diff.create_from}`"
+                else f"\nNew environment `{new_environment}` will be created from `{context_diff.create_from}`"
             )
             self._print(Tree(f"[bold]{msg}\n"))
             if not context_diff.has_snapshot_changes:
@@ -1788,6 +1793,7 @@ def _prompt_categorize(
             self.show_environment_difference_summary(
                 plan.context_diff,
                 no_diff=no_diff,
+                environment=plan_builder.environment_naming_info.name,
             )
 
         if plan.context_diff.has_changes:
@@ -2898,18 +2904,21 @@ def show_environment_difference_summary(
         self,
         context_diff: ContextDiff,
         no_diff: bool = True,
+        environment: t.Optional[str] = None,
     ) -> None:
         """Shows a summary of the environment differences.
 
         Args:
             context_diff: The context diff to use to print the summary.
             no_diff: Hide the actual environment statements differences.
+            environment: The initial target environment
         """
         if context_diff.is_new_environment:
+            new_environment = environment or context_diff.environment
             msg = (
-                f"\n**`{context_diff.environment}` environment will be initialized**"
+                f"\n**`{new_environment}` environment will be initialized**"
                 if not context_diff.create_from_env_exists
-                else f"\n**New environment `{context_diff.environment}` will be created from `{context_diff.create_from}`**"
+                else f"\n**New environment `{new_environment}` will be created from `{context_diff.create_from}`**"
             )
             self._print(msg)
             if not context_diff.has_snapshot_changes:
@@ -3495,6 +3504,7 @@ def show_environment_difference_summary(
         self,
         context_diff: ContextDiff,
         no_diff: bool = True,
+        environment: t.Optional[str] = None,
     ) -> None:
         self._write("Environment Difference Summary:")
 

sqlmesh/core/context.py

@@ -1477,13 +1477,14 @@ def plan_builder(
 
         snapshots = self._snapshots(models_override)
         context_diff = self._context_diff(
-            environment=environment,
+            environment or c.PROD,
             snapshots=snapshots,
             create_from=create_from,
             force_no_diff=restate_models is not None
             or (backfill_models is not None and not backfill_models),
             ensure_finalized_snapshots=self.config.plan.use_finalized_state,
             diff_rendered=diff_rendered,
+            always_compare_against_prod=self.config.plan.always_compare_against_prod,
         )
         modified_model_names = {
             *context_diff.modified_snapshots,
@@ -1517,6 +1518,7 @@ def plan_builder(
 
         return self.PLAN_BUILDER_TYPE(
             context_diff=context_diff,
+            environment=environment or c.PROD,
             start=start,
             end=end,
             execution_time=execution_time,
@@ -1628,6 +1630,7 @@ def diff(self, environment: t.Optional[str] = None, detailed: bool = False) -> b
         self.console.show_environment_difference_summary(
             context_diff,
             no_diff=not detailed,
+            environment=environment,
         )
         if context_diff.has_changes:
             self.console.show_model_difference_summary(
@@ -2607,14 +2610,14 @@ def _context_diff(
         force_no_diff: bool = False,
         ensure_finalized_snapshots: bool = False,
         diff_rendered: bool = False,
+        always_compare_against_prod: bool = False,
     ) -> ContextDiff:
         environment = Environment.sanitize_name(environment)
-
         if force_no_diff:
             return ContextDiff.create_no_diff(environment, self.state_reader)
 
         return ContextDiff.create(
-            environment=environment or c.PROD,
+            environment,
             snapshots=snapshots or self.snapshots,
             create_from=create_from or c.PROD,
             state_reader=self.state_reader,
@@ -2625,7 +2628,7 @@ def _context_diff(
             environment_statements=self._environment_statements,
             gateway_managed_virtual_layer=self.config.gateway_managed_virtual_layer,
             infer_python_dependencies=self.config.infer_python_dependencies,
-            always_compare_against_prod=self.config.plan.always_compare_against_prod,
+            always_compare_against_prod=always_compare_against_prod,
         )
 
     def _destroy(self) -> None:

sqlmesh/core/context_diff.py

@@ -89,8 +89,6 @@ class ContextDiff(PydanticModel):
     """Environment statements."""
     diff_rendered: bool = False
     """Whether the diff should compare raw vs rendered models"""
-    initial_environment: str = ""
-    """The initial target environment (e.g 'dev'), if the plan option `always_compare_to_prod` is set"""
 
     @classmethod
     def create(
@@ -106,7 +104,6 @@ def create(
         environment_statements: t.Optional[t.List[EnvironmentStatements]] = [],
         gateway_managed_virtual_layer: bool = False,
         infer_python_dependencies: bool = True,
-        initial_environment: t.Optional[str] = None,
         always_compare_against_prod: bool = False,
     ) -> ContextDiff:
         """Create a ContextDiff object.
@@ -133,33 +130,34 @@ def create(
             The ContextDiff object.
         """
         initial_environment = environment.lower()
-
-        environment = _get_target_environment(
-            environment, state_reader, always_compare_against_prod
-        )
-
-        env = state_reader.get_environment(environment)
-        initial_env = (
-            env
-            if initial_environment == environment
-            else state_reader.get_environment(initial_environment)
-        )
+        initial_env = state_reader.get_environment(initial_environment)
 
         create_from_env_exists = False
-        if env is None or env.expired:
-            env = state_reader.get_environment(create_from.lower())
+        if initial_env is None or initial_env.expired:
+            initial_env = state_reader.get_environment(create_from.lower())
 
-            if not env and create_from != c.PROD:
+            if not initial_env and create_from != c.PROD:
                 get_console().log_warning(
                     f"The environment name '{create_from}' was passed to the `plan` command's `--create-from` argument, but '{create_from}' does not exist. Initializing new environment '{environment}' from scratch."
                 )
 
             is_new_environment = True
-            create_from_env_exists = env is not None
+            create_from_env_exists = initial_env is not None
             previously_promoted_snapshot_ids = set()
         else:
             is_new_environment = False
-            previously_promoted_snapshot_ids = {s.snapshot_id for s in env.promoted_snapshots}
+            previously_promoted_snapshot_ids = {
+                s.snapshot_id for s in initial_env.promoted_snapshots
+            }
+
+        # Find the proper environment to diff against, this might be different than the "initial" (i.e user provided) environment
+        # e.g it will default to prod if the plan option `always_compare_against_prod` is set.
+        environment = _get_diff_environment(environment, state_reader, always_compare_against_prod)
+        env = (
+            initial_env
+            if (initial_environment == environment)
+            else state_reader.get_environment(environment)
+        )
 
         environment_snapshot_infos = []
         if env:
@@ -237,7 +235,6 @@ def create(
 
         return ContextDiff(
             environment=environment,
-            initial_environment=initial_environment,
             is_new_environment=is_new_environment,
             is_unfinalized_environment=bool(env and not env.finalized_ts),
             normalize_environment_name=is_new_environment or bool(env and env.normalize_name),
@@ -279,9 +276,8 @@ def create_no_diff(cls, environment: str, state_reader: StateReader) -> ContextD
 
         snapshots = state_reader.get_snapshots(env.snapshots)
 
-        environment = env.name
         return ContextDiff(
-            environment=environment,
+            environment=env.name,
             is_new_environment=False,
             is_unfinalized_environment=False,
             normalize_environment_name=env.normalize_name,
@@ -300,7 +296,6 @@ def create_no_diff(cls, environment: str, state_reader: StateReader) -> ContextD
             previous_environment_statements=[],
             previous_gateway_managed_virtual_layer=env.gateway_managed,
             gateway_managed_virtual_layer=env.gateway_managed,
-            initial_environment=environment,
         )
 
     @property
@@ -499,7 +494,7 @@ def text_diff(self, name: str) -> str:
             return ""
 
 
-def _get_target_environment(
+def _get_diff_environment(
     environment: str, state_reader: StateReader, always_compare_against_prod: bool = False
 ) -> str:
     if always_compare_against_prod:

sqlmesh/core/plan/builder.py

@@ -6,7 +6,7 @@
 from collections import defaultdict
 from functools import cached_property
 
-
+from sqlmesh.core import constants as c
 from sqlmesh.core.console import PlanBuilderConsole, get_console
 from sqlmesh.core.config import (
     AutoCategorizationMode,
@@ -115,6 +115,7 @@ def __init__(
         interval_end_per_model: t.Optional[t.Dict[str, int]] = None,
         console: t.Optional[PlanBuilderConsole] = None,
         user_provided_flags: t.Optional[t.Dict[str, UserProvidedFlags]] = None,
+        environment: str = c.PROD,
     ):
         self._context_diff = context_diff
         self._no_gaps = no_gaps
@@ -156,7 +157,7 @@ def __init__(
         self.override_end = end is not None
         self.environment_naming_info = EnvironmentNamingInfo.from_environment_catalog_mapping(
             environment_catalog_mapping or {},
-            name=self._context_diff.initial_environment,
+            name=environment,
             suffix_target=environment_suffix_target,
             normalize_name=self._context_diff.normalize_environment_name,
             gateway_managed=self._context_diff.gateway_managed_virtual_layer,

tests/core/test_integration.py

@@ -6224,6 +6224,11 @@ def plan_with_output(ctx: Context, environment: str):
 
         return output
 
+    def assert_environments(ctx: Context, input_env: str, promote_env: str, diff_env: str):
+        plan_builder = ctx.plan_builder(input_env)
+        assert plan_builder.environment_naming_info.name == promote_env
+        assert plan_builder.build().context_diff.environment == diff_env
+
     models_dir = tmp_path / "models"
 
     logger = logging.getLogger("sqlmesh.core.state_sync.db.facade")
@@ -6239,44 +6244,40 @@ def plan_with_output(ctx: Context, environment: str):
     output = plan_with_output(ctx, "dev")
 
     assert """`dev` environment will be initialized""" in output.stdout
+    assert_environments(ctx, input_env="dev", promote_env="dev", diff_env="dev")
 
     # Case 2: Prod does not exist, so dev is updated
     create_temp_file(
         tmp_path, models_dir / "a.sql", "MODEL (name test.a, kind FULL); SELECT 5 AS col"
     )
 
-    plan = ctx.plan_builder("dev").build()
-
-    assert plan.context_diff.initial_environment == "dev"
-    assert plan.context_diff.environment == "dev"
-
     output = plan_with_output(ctx, "dev")
 
+    assert_environments(ctx, input_env="dev", promote_env="dev", diff_env="dev")
     assert "Differences from the `dev` environment" in output.stdout
 
     # Case 3: Prod is initialized, so plan comparisons moving forward should be against prod
     output = plan_with_output(ctx, "prod")
-
     assert "`prod` environment will be initialized" in output.stdout
 
-    # Case 4: Dev is updated with a breaking change, so plan comparisons moving forward should be against prod
+    assert_environments(ctx, input_env="prod", promote_env="prod", diff_env="prod")
+
+    # Case 4: Dev is updated with a breaking change. Prod exists now so plan comparisons moving forward should be against prod
     create_temp_file(
         tmp_path, models_dir / "a.sql", "MODEL (name test.a, kind FULL); SELECT 10 AS col"
     )
     ctx.load()
 
     plan = ctx.plan_builder("dev").build()
 
-    assert plan.context_diff.initial_environment == "dev"
-    assert plan.context_diff.environment == "prod"
+    assert_environments(ctx, input_env="dev", promote_env="dev", diff_env="prod")
 
     assert (
         next(iter(plan.context_diff.snapshots.values())).change_category
         == SnapshotChangeCategory.BREAKING
     )
 
     output = plan_with_output(ctx, "dev")
-
     assert "Differences from the `prod` environment" in output.stdout
 
     # Case 5: Dev is updated with a metadata change, but comparison against prod shows both the previous and the current changes
@@ -6290,16 +6291,14 @@ def plan_with_output(ctx: Context, environment: str):
 
     plan = ctx.plan_builder("dev").build()
 
-    assert plan.context_diff.initial_environment == "dev"
-    assert plan.context_diff.environment == "prod"
+    assert_environments(ctx, input_env="dev", promote_env="dev", diff_env="prod")
 
     assert (
         next(iter(plan.context_diff.snapshots.values())).change_category
         == SnapshotChangeCategory.BREAKING
     )
 
     output = plan_with_output(ctx, "dev")
-
     assert "Differences from the `prod` environment" in output.stdout
 
     assert (
@@ -6313,3 +6312,8 @@ def plan_with_output(ctx: Context, environment: str):
 +  10 AS col"""
         in output.stdout
     )
+
+    # Case 6: Check that we can still run Context::diff() against any environment
+    for environment in ["dev", "prod"]:
+        context_diff = ctx._context_diff(environment)
+        assert context_diff.environment == environment