Repository: TobikoData/sqlmesh
Pull request #2523 (2024-04-26 17:22:24)

feat: adds plan id to apply notifications
Makes changes to how apply notifications are called in context.py to include plan_id. Also modifies the Airflow plan_request_id to be the same as plan_id so that apply notifications can include links to the Airflow job. All of this will help our users find logs when they want to see what SQL queries SQL Mesh runs or track down an error.
Currently when a SQL Mesh apply fails with Airflow acting as scheduler, the stack trace sent to the notification integration doesn't include any information about why the apply failed. Getting the actual stack trace out of airflow and into the notification would be very complicated, so a good work around is just to include a link to the DAG so users can quickly explore the logs for their run.

isaiahgm: Looks like the test cases for notify() and notify_user() need to be updated or the logic needs to be changed to still be backwards compatible.

Patch:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index ddcd20a1..686d67c4 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1111,18 +1111,25 @@ class GenericContext(BaseContext, t.Generic[C]):
         if plan.uncategorized:
             raise UncategorizedPlanError("Can't apply a plan with uncategorized changes.")
         self.notification_target_manager.notify(
-            NotificationEvent.APPLY_START, environment=plan.environment_naming_info.name
+            NotificationEvent.APPLY_START,
+            environment=plan.environment_naming_info.name,
+            plan_id=plan.plan_id,
         )
         try:
             self._apply(plan, circuit_breaker)
         except Exception as e:
             self.notification_target_manager.notify(
-                NotificationEvent.APPLY_FAILURE, traceback.format_exc()
+                NotificationEvent.APPLY_FAILURE,
+                environment=plan.environment_naming_info.name,
+                plan_id=plan.plan_id,
+                exc=traceback.format_exc(),
             )
             logger.error(f"Apply Failure: {traceback.format_exc()}")
             raise e
         self.notification_target_manager.notify(
-            NotificationEvent.APPLY_END, environment=plan.environment_naming_info.name
+            NotificationEvent.APPLY_END,
+            environment=plan.environment_naming_info.name,
+            plan_id=plan.plan_id,
         )
 
     def invalidate_environment(self, name: str, sync: bool = False) -> None:
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index 0d0ee8b5..844250ab 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -402,14 +402,12 @@ class MacroEvaluator:
             self.columns_to_types_called = True
             return {"__schema_unavailable_at_load__": exp.DataType.build("unknown")}
 
-        if isinstance(model_name, exp.Column):
-            model_name = exp.table_(
-                model_name.this,
-                db=model_name.args.get("table"),
-                catalog=model_name.args.get("db"),
-            )
-
-        columns_to_types = self._schema.find(exp.to_table(model_name))
+        normalized_model_name = normalize_model_name(
+            model_name,
+            default_catalog=self.default_catalog,
+            dialect=self.dialect,
+        )
+        columns_to_types = self._schema.find(exp.to_table(normalized_model_name))
         if columns_to_types is None:
             raise SQLMeshError(f"Schema for model '{model_name}' can't be statically determined.")
 
diff --git a/sqlmesh/core/model/meta.py b/sqlmesh/core/model/meta.py
index e04ee4f5..5d69ab16 100644
--- a/sqlmesh/core/model/meta.py
+++ b/sqlmesh/core/model/meta.py
@@ -144,10 +144,17 @@ class ModelMeta(_Node):
 
         return v
 
-    @field_validator("dialect", "storage_format", mode="before")
-    def _string_validator(cls, v: t.Any) -> t.Optional[str]:
+    @field_validator("storage_format", mode="before")
+    def _storage_format_validator(cls, v: t.Any) -> t.Optional[str]:
         return str_or_exp_to_str(v)
 
+    @field_validator("dialect", mode="before")
+    def _dialect_validator(cls, v: t.Any) -> t.Optional[str]:
+        # dialects are parsed as identifiers and may get normalized as uppercase,
+        # so this ensures they'll be stored as lowercase
+        dialect = str_or_exp_to_str(v)
+        return dialect and dialect.lower()
+
     @field_validator("partitioned_by_", mode="before")
     @field_validator_v1_args
     def _partition_by_validator(
diff --git a/sqlmesh/core/notification_target.py b/sqlmesh/core/notification_target.py
index 7da4ac51..80236885 100644
--- a/sqlmesh/core/notification_target.py
+++ b/sqlmesh/core/notification_target.py
@@ -114,23 +114,29 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
         """
 
     @notify(NotificationEvent.APPLY_START)
-    def notify_apply_start(self, environment: str) -> None:
+    def notify_apply_start(self, environment: str, plan_id: str) -> None:
         """Notify when an apply starts.
 
         Args:
             environment: The target environment of the plan.
+            plan_id: plan_id that is being applied.
         """
-        self.send(NotificationStatus.INFO, f"Plan apply started for environment `{environment}`.")
+        self.send(
+            NotificationStatus.INFO,
+            f"Plan {plan_id} apply started for environment `{environment}`.",
+        )
 
     @notify(NotificationEvent.APPLY_END)
-    def notify_apply_end(self, environment: str) -> None:
+    def notify_apply_end(self, environment: str, plan_id: str) -> None:
         """Notify when an apply ends.
 
         Args:
             environment: The target environment of the plan.
+            plan_id: plan_id that was applied.
         """
         self.send(
-            NotificationStatus.SUCCESS, f"Plan apply finished for environment `{environment}`."
+            NotificationStatus.SUCCESS,
+            f"Plan {plan_id} apply finished for environment `{environment}`.",
         )
 
     @notify(NotificationEvent.RUN_START)
@@ -164,13 +170,19 @@ class BaseNotificationTarget(PydanticModel, frozen=True):
         self.send(NotificationStatus.SUCCESS, "SQLMesh migration finished.")
 
     @notify(NotificationEvent.APPLY_FAILURE)
-    def notify_apply_failure(self, exc: str) -> None:
+    def notify_apply_failure(self, environment: str, plan_id: str, exc: str) -> None:
         """Notify in the case of an apply failure.
 
         Args:
+            environment: The target environment of the run.
+            plan_id: The plan id of the failed apply
             exc: The exception stack trace.
         """
-        self.send(NotificationStatus.FAILURE, "Plan apply failed.", exc=exc)
+        self.send(
+            NotificationStatus.FAILURE,
+            f"Plan {plan_id} in environment `{environment}` apply failed.",
+            exc=exc,
+        )
 
     @notify(NotificationEvent.RUN_FAILURE)
     def notify_run_failure(self, exc: str) -> None:
diff --git a/sqlmesh/core/plan/evaluator.py b/sqlmesh/core/plan/evaluator.py
index 3458f3ac..c3e3b80f 100644
--- a/sqlmesh/core/plan/evaluator.py
+++ b/sqlmesh/core/plan/evaluator.py
@@ -32,7 +32,6 @@ from sqlmesh.core.user import User
 from sqlmesh.schedulers.airflow import common as airflow_common
 from sqlmesh.schedulers.airflow.client import AirflowClient, BaseAirflowClient
 from sqlmesh.schedulers.airflow.mwaa_client import MWAAClient
-from sqlmesh.utils import random_id
 from sqlmesh.utils.errors import SQLMeshError
 
 logger = logging.getLogger(__name__)
@@ -302,7 +301,7 @@ class BaseAirflowPlanEvaluator(PlanEvaluator):
     def evaluate(
         self, plan: Plan, circuit_breaker: t.Optional[t.Callable[[], bool]] = None
     ) -> None:
-        plan_request_id = random_id()
+        plan_request_id = plan.plan_id
         self._apply_plan(plan, plan_request_id)
 
         if self.blocking:
diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py
index 77decd1d..cebc10cd 100644
--- a/sqlmesh/core/renderer.py
+++ b/sqlmesh/core/renderer.py
@@ -100,10 +100,17 @@ class BaseExpressionRenderer:
             return self._cache
 
         if self._model_fqn and "this_model" not in kwargs:
-            kwargs["this_model"] = self._to_table_mapping(
-                [snapshots[self._model_fqn]] if snapshots and self._model_fqn in snapshots else [],
-                deployability_index,
-            ).get(self._model_fqn, self._model_fqn)
+            kwargs["this_model"] = exp.to_table(
+                self._to_table_mapping(
+                    (
+                        [snapshots[self._model_fqn]]
+                        if snapshots and self._model_fqn in snapshots
+                        else []
+                    ),
+                    deployability_index,
+                ).get(self._model_fqn, self._model_fqn),
+                dialect=self._dialect,
+            ).sql(dialect=self._dialect, identify=True)
 
         expressions = [self._expression]
 


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index bf936a56..70729d83 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -1802,6 +1802,12 @@ def test_case_sensitivity(assert_exp_eq):
         """,
     )
 
+    model = load_sql_based_model(
+        d.parse("MODEL (name test, dialect hive); SELECT 1 AS c"),
+        dialect="snowflake",
+    )
+    assert model.dialect == "hive"
+
 
 def test_batch_size_validation():
     expressions = d.parse(
@@ -3814,7 +3820,8 @@ def test_this_model() -> None:
     expressions = d.parse(
         """
         MODEL (
-            name db.table,
+            name `project-1.table`,
+            dialect bigquery,
         );
 
         JINJA_STATEMENT_BEGIN;
@@ -3832,18 +3839,38 @@ def test_this_model() -> None:
     )
     model = load_sql_based_model(expressions)
 
-    assert model.render_query_or_raise().sql() == '''SELECT '"db"."table"' AS "x"'''
+    assert (
+        model.render_query_or_raise().sql(dialect="bigquery")
+        == """SELECT '`project-1`.`table`' AS `x`"""
+    )
 
-    assert model.render_pre_statements()[0].sql() == """COPY "db"."table" TO 'a'"""
-    assert model.render_post_statements()[0].sql() == """COPY "db"."table" TO 'b'"""
+    assert (
+        model.render_pre_statements()[0].sql(dialect="bigquery")
+        == """COPY `project-1`.`table` TO 'a'"""
+    )
+    assert (
+        model.render_post_statements()[0].sql(dialect="bigquery")
+        == """COPY `project-1`.`table` TO 'b'"""
+    )
 
     snapshot = Snapshot.from_node(model, nodes={})
+
+    assert (
+        model.render_query_or_raise(
+            start="2020-01-01",
+            snapshots={snapshot.name: snapshot},
+        ).sql(dialect="bigquery")
+        == """SELECT '`project-1`.`table`' AS `x`"""
+    )
+
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
     assert (
         model.render_query_or_raise(
             start="2020-01-01",
             snapshots={snapshot.name: snapshot},
-        ).sql()
-        == '''SELECT '"db"."table"' AS "x"'''
+        ).sql(dialect="bigquery")
+        == "SELECT '`sqlmesh__project-1`.`project_1__table__1178373199`' AS `x`"
     )
 
 
diff --git a/tests/core/test_notification_target.py b/tests/core/test_notification_target.py
index 027a1d89..9adff6fb 100644
--- a/tests/core/test_notification_target.py
+++ b/tests/core/test_notification_target.py
@@ -34,37 +34,43 @@ def notification_target_manager_with_spy(mocker) -> tuple[NotificationTargetMana
 
 def test_notify(notification_target_manager_with_spy):
     notification_target_manager, spy = notification_target_manager_with_spy
-    notification_target_manager.notify(NotificationEvent.APPLY_START, "prod")
+    notification_target_manager.notify(NotificationEvent.APPLY_START, "prod", "a-plan-id")
     spy.assert_called_once_with(
         mock.ANY,
         NotificationStatus.INFO,
-        "Plan apply started for environment `prod`.",
+        "Plan a-plan-id apply started for environment `prod`.",
     )
 
     spy.reset_mock()
-    notification_target_manager.notify(NotificationEvent.APPLY_END, "prod")
+    notification_target_manager.notify(NotificationEvent.APPLY_END, "prod", "a-plan-id")
     spy.assert_called_once_with(
         mock.ANY,
         NotificationStatus.SUCCESS,
-        "Plan apply finished for environment `prod`.",
+        "Plan a-plan-id apply finished for environment `prod`.",
     )
 
     # No notification target configured for APPLY_FAILURE event
     spy.reset_mock()
-    notification_target_manager.notify(NotificationEvent.APPLY_FAILURE, ValueError())
+    notification_target_manager.notify(
+        NotificationEvent.APPLY_FAILURE, "prod", "a-plan-id", ValueError()
+    )
     spy.assert_not_called()
 
 
 def test_notify_user(notification_target_manager_with_spy):
     notification_target_manager, spy = notification_target_manager_with_spy
-    notification_target_manager.notify_user(NotificationEvent.APPLY_START, "test_user", "prod")
+    notification_target_manager.notify_user(
+        NotificationEvent.APPLY_START, "test_user", "prod", "a-plan-id"
+    )
     spy.assert_called_once_with(
         mock.ANY,
         NotificationStatus.INFO,
-        "Plan apply started for environment `prod`.",
+        "Plan a-plan-id apply started for environment `prod`.",
     )
 
     # No notification target configured for APPLY_END event for test_user
     spy.reset_mock()
-    notification_target_manager.notify_user(NotificationEvent.APPLY_END, "test_user", "prod")
+    notification_target_manager.notify_user(
+        NotificationEvent.APPLY_END, "test_user", "prod", "a-plan-id"
+    )
     spy.assert_not_called()