Feat: Improve the cli before all after all diff and include python env (#4116)

Author: themisvaltinos

sqlmesh/core/audit/definition.py

@@ -16,6 +16,8 @@
     default_catalog_validator,
     depends_on_validator,
     expression_validator,
+    sort_python_env,
+    sorted_python_env_payloads,
 )
 from sqlmesh.core.model.common import make_python_env, single_value_or_tuple
 from sqlmesh.core.node import _Node
@@ -236,7 +238,7 @@ def depends_on(self) -> t.Set[str]:
     @property
     def sorted_python_env(self) -> t.List[t.Tuple[str, Executable]]:
         """Returns the python env sorted by executable kind and then var name."""
-        return sorted(self.python_env.items(), key=lambda x: (x[1].kind, x[0]))
+        return sort_python_env(self.python_env)
 
     @property
     def data_hash(self) -> str:
@@ -337,12 +339,7 @@ def render_definition(
         jinja_expressions = []
         python_expressions = []
         if include_python:
-            python_env = d.PythonCode(
-                expressions=[
-                    v.payload if v.is_import or v.is_definition else f"{k} = {v.payload}"
-                    for k, v in self.sorted_python_env
-                ]
-            )
+            python_env = d.PythonCode(expressions=sorted_python_env_payloads(self.python_env))
             if python_env.expressions:
                 python_expressions.append(python_env)
 

sqlmesh/core/console.py

@@ -1311,9 +1311,9 @@ def show_model_difference_summary(
             self._print(f"[bold]Requirements:\n{context_diff.requirements_diff()}")
 
         if context_diff.has_environment_statements_changes:
-            self._print(
-                f"[bold]Environment statements:\n{context_diff.environment_statements_diff()}"
-            )
+            self._print("[bold]Environment statements:\n")
+            for type, diff in context_diff.environment_statements_diff():
+                self._print(Syntax(diff, type, line_numbers=False))
 
         self._show_summary_tree_for(
             context_diff,
@@ -2463,7 +2463,9 @@ def show_model_difference_summary(
             self._print(f"Requirements:\n{context_diff.requirements_diff()}")
 
         if context_diff.has_environment_statements_changes:
-            self._print(f"Environment statements:\n{context_diff.environment_statements_diff()}")
+            self._print("[bold]Environment statements:\n")
+            for _, diff in context_diff.environment_statements_diff():
+                self._print(diff)
 
         added_snapshots = {context_diff.snapshots[s_id] for s_id in context_diff.added}
         added_snapshot_models = {s for s in added_snapshots if s.is_model}
@@ -2976,7 +2978,9 @@ def show_model_difference_summary(
             self._write(f"Requirements:\n{context_diff.requirements_diff()}")
 
         if context_diff.has_environment_statements_changes:
-            self._write(f"Environment statements:\n{context_diff.environment_statements_diff()}")
+            self._write("Environment statements:\n")
+            for _, diff in context_diff.environment_statements_diff():
+                self._write(diff)
 
         for added in context_diff.new_snapshots:
             self._write(f"  Added: {added}")

sqlmesh/core/context_diff.py

@@ -19,6 +19,7 @@
 from sqlmesh.core import constants as c
 from sqlmesh.core.console import get_console
 from sqlmesh.core.macros import RuntimeStage
+from sqlmesh.core.model.common import sorted_python_env_payloads
 from sqlmesh.core.snapshot import Snapshot, SnapshotId, SnapshotTableInfo
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.pydantic import PydanticModel
@@ -340,22 +341,43 @@ def requirements_diff(self) -> str:
             )
         )
 
-    def environment_statements_diff(self) -> str:
+    def environment_statements_diff(self) -> t.List[t.Tuple[str, str]]:
         def extract_statements(statements: t.List[EnvironmentStatements], attr: str) -> t.List[str]:
-            return [str(stmt) for statement in statements for stmt in getattr(statement, attr)]
-
-        def format_diff(runtime_stage: str) -> str:
-            previous = extract_statements(self.previous_environment_statements, runtime_stage)
-            current = extract_statements(self.environment_statements, runtime_stage)
-            return (
-                f"  {runtime_stage}:\n    " + "\n    ".join(ndiff(previous, current)) + "\n"
-                if previous or current
-                else ""
-            )
+            return [
+                expr
+                for statement in statements
+                for expr in (
+                    sorted_python_env_payloads(statement.python_env)
+                    if attr == "python_env"
+                    else getattr(statement, attr)
+                )
+            ]
 
-        return format_diff(RuntimeStage.BEFORE_ALL.value) + format_diff(
-            RuntimeStage.AFTER_ALL.value
-        )
+        def compute_diff(attribute: str) -> t.Optional[t.Tuple[str, str]]:
+            previous = extract_statements(self.previous_environment_statements, attribute)
+            current = extract_statements(self.environment_statements, attribute)
+
+            if previous == current:
+                return None
+
+            diff_lines = list(ndiff(previous, current))
+            diff_text = attribute if not attribute == "python_env" else "dependencies"
+            diff_text += ":\n"
+
+            if any(line.startswith(("-", "+")) for line in diff_lines):
+                diff_text += "  " + "\n  ".join(diff_lines) + "\n"
+
+            return "python" if attribute == "python_env" else "sql", diff_text
+
+        return [
+            diff
+            for attribute in [
+                RuntimeStage.BEFORE_ALL.value,
+                RuntimeStage.AFTER_ALL.value,
+                "python_env",
+            ]
+            if (diff := compute_diff(attribute)) is not None
+        ]
 
     @property
     def environment_snapshots(self) -> t.List[SnapshotTableInfo]:

sqlmesh/core/model/common.py

@@ -352,6 +352,19 @@ def depends_on(cls: t.Type, v: t.Any, info: ValidationInfo) -> t.Optional[t.Set[
     return v
 
 
+def sort_python_env(python_env: t.Dict[str, Executable]) -> t.List[t.Tuple[str, Executable]]:
+    """Returns the python env sorted."""
+    return sorted(python_env.items(), key=lambda x: (x[1].kind, x[0]))
+
+
+def sorted_python_env_payloads(python_env: t.Dict[str, Executable]) -> t.List[str]:
+    """Returns the payloads of the sorted python env."""
+    return [
+        v.payload if v.is_import or v.is_definition else f"{k} = {v.payload}"
+        for k, v in sort_python_env(python_env)
+    ]
+
+
 expression_validator: t.Callable = field_validator(
     "query",
     "expressions_",

sqlmesh/core/model/definition.py

@@ -30,6 +30,7 @@
     make_python_env,
     parse_dependencies,
     single_value_or_tuple,
+    sorted_python_env_payloads,
 )
 from sqlmesh.core.model.meta import ModelMeta, FunctionCall
 from sqlmesh.core.model.kind import (
@@ -257,12 +258,7 @@ def render_definition(
         jinja_expressions = []
         python_expressions = []
         if include_python:
-            python_env = d.PythonCode(
-                expressions=[
-                    v.payload if v.is_import or v.is_definition else f"{k} = {v.payload}"
-                    for k, v in self.sorted_python_env
-                ]
-            )
+            python_env = d.PythonCode(expressions=sorted_python_env_payloads(self.python_env))
             if python_env.expressions:
                 python_expressions.append(python_env)
 

tests/core/test_plan.py

@@ -4,6 +4,8 @@
 from unittest.mock import patch
 
 import pytest
+from sqlmesh.utils.metaprogramming import Executable
+from tests.core.test_table_diff import create_test_console, strip_ansi_codes
 import time_machine
 from pytest_mock.plugin import MockerFixture
 from sqlglot import parse_one
@@ -3000,9 +3002,13 @@ def test_plan_environment_statements_diff(make_snapshot):
         previous_finalized_snapshots=None,
         environment_statements=[
             EnvironmentStatements(
-                before_all=["CREATE OR REPLACE TABLE table_1 AS SELECT 1"],
+                before_all=["CREATE OR REPLACE TABLE table_1 AS SELECT 1", "@test_macro()"],
                 after_all=["CREATE OR REPLACE TABLE table_2 AS SELECT 2"],
-                python_env={},
+                python_env={
+                    "test_macro": Executable(
+                        payload="def test_macro(evaluator):\n    return 'one'"
+                    ),
+                },
             )
         ],
         previous_gateway_managed_virtual_layer=False,
@@ -3011,11 +3017,21 @@ def test_plan_environment_statements_diff(make_snapshot):
 
     assert context_diff.has_changes
     assert context_diff.has_environment_statements_changes
-    assert (
-        context_diff.environment_statements_diff()
-        == """  before_all:
-    + CREATE OR REPLACE TABLE table_1 AS SELECT 1
-  after_all:
-    + CREATE OR REPLACE TABLE table_2 AS SELECT 2
-"""
-    )
+
+    console_output, terminal_console = create_test_console()
+    for _, diff in context_diff.environment_statements_diff():
+        terminal_console._print(diff)
+    output = console_output.getvalue()
+    stripped = strip_ansi_codes(output)
+    expected_output = (
+        "before_all:\n"
+        "  + CREATE OR REPLACE TABLE table_1 AS SELECT 1\n"
+        "  + @test_macro()\n\n"
+        "after_all:\n"
+        "  + CREATE OR REPLACE TABLE table_2 AS SELECT 2\n\n"
+        "dependencies:\n"
+        "  + def test_macro(evaluator):\n"
+        "    return 'one'"
+    )
+    assert stripped == expected_output
+    console_output.close()