Feat: Add plan option to always compare against prod

Author: VaggelisD

sqlmesh/core/config/plan.py

@@ -20,6 +20,7 @@ class PlanConfig(BaseConfig):
         auto_apply: Whether to automatically apply the new plan after creation.
         use_finalized_state: Whether to compare against the latest finalized environment state, or to use
             whatever state the target environment is currently in.
+        always_compare_against_prod: Whether to always compare against production when planning, even if the target environment exists.
     """
 
     forward_only: bool = False
@@ -30,3 +31,4 @@ class PlanConfig(BaseConfig):
     no_prompts: bool = True
     auto_apply: bool = False
     use_finalized_state: bool = False
+    always_compare_against_prod: bool = False

sqlmesh/core/context.py

@@ -1448,7 +1448,7 @@ def plan_builder(
 
         snapshots = self._snapshots(models_override)
         context_diff = self._context_diff(
-            environment or c.PROD,
+            environment=environment,
             snapshots=snapshots,
             create_from=create_from,
             force_no_diff=restate_models is not None
@@ -2532,6 +2532,25 @@ def _snapshots(
 
         return {name: stored_snapshots.get(s.snapshot_id, s) for name, s in snapshots.items()}
 
+    def _get_target_environment(self, environment: t.Optional[str] = None) -> t.Tuple[str, str]:
+        environment = environment or self.config.default_target_environment
+        environment = Environment.sanitize_name(environment)
+
+        initial_environment = environment
+
+        if self.config.plan.always_compare_against_prod:
+            prod = self.state_reader.get_environment(c.PROD)
+            if prod:
+                logger.warning(
+                    f"Comparing against production environment instead of {environment}. Note that this may lead to "
+                    "additional backfills as accumulated changes are still pushed to the target environment."
+                )
+                environment = c.PROD
+        else:
+            environment = environment or c.PROD
+
+        return environment.lower(), initial_environment.lower()
+
     def _context_diff(
         self,
         environment: str,
@@ -2541,12 +2560,13 @@ def _context_diff(
         ensure_finalized_snapshots: bool = False,
         diff_rendered: bool = False,
     ) -> ContextDiff:
-        environment = Environment.sanitize_name(environment)
+        target_environment, initial_environment = self._get_target_environment(environment)
+
         if force_no_diff:
             return ContextDiff.create_no_diff(environment, self.state_reader)
 
         return ContextDiff.create(
-            environment,
+            environment=target_environment,
             snapshots=snapshots or self.snapshots,
             create_from=create_from or c.PROD,
             state_reader=self.state_reader,
@@ -2556,6 +2576,7 @@ def _context_diff(
             diff_rendered=diff_rendered,
             environment_statements=self._environment_statements,
             gateway_managed_virtual_layer=self.gateway_managed_virtual_layer,
+            initial_environment=initial_environment,
         )
 
     def _destroy(self) -> None:

sqlmesh/core/context_diff.py

@@ -88,6 +88,8 @@ 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(
@@ -102,6 +104,7 @@ def create(
         diff_rendered: bool = False,
         environment_statements: t.Optional[t.List[EnvironmentStatements]] = [],
         gateway_managed_virtual_layer: bool = False,
+        initial_environment: t.Optional[str] = None,
     ) -> ContextDiff:
         """Create a ContextDiff object.
 
@@ -123,6 +126,13 @@ def create(
         environment = environment.lower()
         env = state_reader.get_environment(environment)
 
+        initial_environment = initial_environment or environment
+        initial_env = (
+            env
+            if initial_environment == environment
+            else 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())
@@ -214,6 +224,7 @@ 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),
@@ -224,7 +235,9 @@ def create(
             modified_snapshots=modified_snapshots,
             snapshots=merged_snapshots,
             new_snapshots=new_snapshots,
-            previous_plan_id=env.plan_id if env and not is_new_environment else None,
+            previous_plan_id=initial_env.plan_id
+            if initial_env and not is_new_environment
+            else None,
             previously_promoted_snapshot_ids=previously_promoted_snapshot_ids,
             previous_finalized_snapshots=env.previous_finalized_snapshots if env else None,
             previous_requirements=env.requirements if env else {},

sqlmesh/core/plan/builder.py

@@ -155,7 +155,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.environment,
+            name=self._context_diff.initial_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

@@ -6,7 +6,7 @@
 from datetime import timedelta
 from unittest import mock
 from unittest.mock import patch
-
+import logging
 import os
 import numpy as np
 import pandas as pd
@@ -36,6 +36,7 @@
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.context import Context
 from sqlmesh.core.config.categorizer import CategorizerConfig
+from sqlmesh.core.config.plan import PlanConfig
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.macros import macro
@@ -6129,3 +6130,107 @@ def setup_senario(model_before: str, model_after: str):
         'Binder Error: Referenced column "this_col_does_not_exist" not found in \nFROM clause!'
         in output.stdout
     )
+
+
+@use_terminal_console
+def test_plan_always_compare_against_prod(mocker: MockerFixture, tmp_path: Path):
+    def plan_with_output(ctx: Context, environment: str):
+        with patch.object(logger, "info") as mock_logger:
+            with capture_output() as output:
+                ctx.load()
+                ctx.plan(environment, no_prompts=True, auto_apply=True)
+
+            # Facade logs info "Promoting environment {environment}"
+            assert mock_logger.call_args[0][1] == environment
+
+        return output
+
+    models_dir = tmp_path / "models"
+
+    logger = logging.getLogger("sqlmesh.core.state_sync.db.facade")
+
+    create_temp_file(
+        tmp_path, models_dir / "a.sql", "MODEL (name test.a, kind FULL); SELECT 1 AS col"
+    )
+
+    config = Config(plan=PlanConfig(always_compare_against_prod=True))
+    ctx = Context(paths=[tmp_path], config=config)
+
+    # Case 1: Neither prod nor dev exists, so dev is initialized
+    output = plan_with_output(ctx, "dev")
+
+    assert """`dev` environment will be initialized""" in output.stdout
+
+    # 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 "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
+    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 (
+        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 4: Dev is updated with a metadata change, but comparison against prod shows both the previous and the current changes
+    # so it's still classified as a breaking change
+    create_temp_file(
+        tmp_path,
+        models_dir / "a.sql",
+        "MODEL (name test.a, kind FULL, owner 'test'); 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 (
+        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 (
+        """MODEL (                                                                        
+   name test.a,                                                                 
++  owner test,                                                                  
+   kind FULL                                                                    
+ )                                                                              
+ SELECT                                                                         
+-  5 AS col                                                                     
++  10 AS col"""
+        in output.stdout
+    )