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

@@ -1480,7 +1480,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
@@ -2630,11 +2630,12 @@ def _context_diff(
         diff_rendered: 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=environment,
             snapshots=snapshots or self.snapshots,
             create_from=create_from or c.PROD,
             state_reader=self.state_reader,
@@ -2645,6 +2646,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,
         )
 
     def _destroy(self) -> None:

sqlmesh/core/context_diff.py

@@ -14,6 +14,8 @@
 
 import sys
 import typing as t
+import logging
+
 from difflib import ndiff, unified_diff
 from functools import cached_property
 from sqlmesh.core import constants as c
@@ -38,6 +40,8 @@
 
 IGNORED_PACKAGES = {"sqlmesh", "sqlglot"}
 
+logger = logging.getLogger(__name__)
+
 
 class ContextDiff(PydanticModel):
     """ContextDiff is an object representing the difference between two environments.
@@ -88,6 +92,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(
@@ -103,6 +109,8 @@ 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.
 
@@ -127,8 +135,17 @@ def create(
         Returns:
             The ContextDiff object.
         """
-        environment = environment.lower()
+        initial_environment = environment
+        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)
+        )
 
         create_from_env_exists = False
         if env is None or env.expired:
@@ -222,6 +239,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),
@@ -232,7 +250,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 {},
@@ -261,8 +281,9 @@ def create_no_diff(cls, environment: str, state_reader: StateReader) -> ContextD
 
         snapshots = state_reader.get_snapshots(env.snapshots)
 
+        environment = env.name
         return ContextDiff(
-            environment=env.name,
+            environment=environment,
             is_new_environment=False,
             is_unfinalized_environment=False,
             normalize_environment_name=env.normalize_name,
@@ -281,6 +302,7 @@ 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
@@ -479,6 +501,23 @@ def text_diff(self, name: str) -> str:
             return ""
 
 
+def _get_target_environment(
+    environment: str, state_reader: StateReader, always_compare_against_prod: bool = False
+) -> str:
+    if always_compare_against_prod:
+        prod = 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()
+
+
 def _build_requirements(
     provided_requirements: t.Dict[str, str],
     excluded_requirements: t.Set[str],

sqlmesh/core/plan/builder.py

@@ -159,7 +159,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  # noqa: TID253
 import pandas as pd  # noqa: TID253
@@ -37,6 +37,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
@@ -6252,3 +6253,107 @@ def test_render_path_instead_of_model(tmp_path: Path):
 
     # Case 3: Render the model successfully
     assert ctx.render("test_model").sql() == 'SELECT 1 AS "col"'
+
+
+@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 5: 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
+    )