Move get_target_environment to ContextDiff

Author: VaggelisD

sqlmesh/core/context.py

@@ -2599,25 +2599,6 @@ 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,
@@ -2627,13 +2608,13 @@ def _context_diff(
         ensure_finalized_snapshots: bool = False,
         diff_rendered: bool = False,
     ) -> ContextDiff:
-        target_environment, initial_environment = self._get_target_environment(environment)
+        environment = Environment.sanitize_name(environment)
 
         if force_no_diff:
             return ContextDiff.create_no_diff(environment, self.state_reader)
 
         return ContextDiff.create(
-            environment=target_environment,
+            environment=environment,
             snapshots=snapshots or self.snapshots,
             create_from=create_from or c.PROD,
             state_reader=self.state_reader,
@@ -2644,7 +2625,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,
-            initial_environment=initial_environment,
+            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.
@@ -106,6 +110,7 @@ def create(
         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.
 
@@ -130,10 +135,12 @@ def create(
         Returns:
             The ContextDiff object.
         """
-        environment = environment.lower()
-        env = state_reader.get_environment(environment)
+        initial_environment = environment
+        environment = _get_target_environment(
+            environment, state_reader, always_compare_against_prod
+        )
 
-        initial_environment = initial_environment or environment
+        env = state_reader.get_environment(environment)
         initial_env = (
             env
             if initial_environment == environment
@@ -492,6 +499,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],