Repository: TobikoData/sqlmesh (python)

Issue:
Unable to parse macro call in UNION ALL
[Public slack context](https://tobiko-data.slack.com/archives/C044BRE5W4S/p1725557364459159)

Repro:

```SQL
@macro()
def union_macro(evaluator):
 return "SELECT 2 as col"

MODEL (...)

SELECT 1 as col UNION ALL @union_macro()
```

In `_parse_set_operations()` we'll attempt to `_parse_select()` for the RHS which doesn't have a path for parsing placeholders yet; This means that we'll return `None` instead of parsing the macro, causing the error in the discussion.



Pull request description:
fix: Allow macros in set operations
Fixes #3102

Patch:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index add5f099..8d24cb3d 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -327,6 +327,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         self._metrics: UniqueKeyDict[str, Metric] = UniqueKeyDict("metrics")
         self._jinja_macros = JinjaMacroRegistry()
         self._default_catalog: t.Optional[str] = None
+        self._loaded: bool = False
 
         self.path, self.config = t.cast(t.Tuple[Path, C], next(iter(self.configs.items())))
 
@@ -550,6 +551,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             project_name=self.config.project,
         )
 
+        self._loaded = True
         return self
 
     @python_api_analytics
@@ -586,6 +588,11 @@ class GenericContext(BaseContext, t.Generic[C]):
             state_sync_type=self.state_sync.state_type(),
         )
 
+        if not self._loaded:
+            # Signals should be loaded to run correctly.
+            with sys_path(*self.configs):
+                self._loader.load_signals(self)
+
         success = False
         try:
             success = self._run(
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 77b569cb..164c87ac 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -314,6 +314,31 @@ def _parse_join(
     return macro
 
 
+def _parse_select(
+    self: Parser,
+    nested: bool = False,
+    table: bool = False,
+    parse_subquery_alias: bool = True,
+    parse_set_operation: bool = True,
+) -> t.Optional[exp.Expression]:
+    select = self.__parse_select(  # type: ignore
+        nested=nested,
+        table=table,
+        parse_subquery_alias=parse_subquery_alias,
+        parse_set_operation=parse_set_operation,
+    )
+
+    if (
+        not select
+        and not parse_set_operation
+        and self._match_pair(TokenType.PARAMETER, TokenType.VAR, advance=False)
+    ):
+        self._advance()
+        return _parse_macro(self)
+
+    return select
+
+
 def _parse_where(self: Parser, skip_where_token: bool = False) -> t.Optional[exp.Expression]:
     macro = _parse_matching_macro(self, "WHERE")
     if not macro:
@@ -846,6 +871,7 @@ def extend_sqlglot() -> None:
                 MacroDef,
             )
 
+    _override(Parser, _parse_select)
     _override(Parser, _parse_statement)
     _override(Parser, _parse_join)
     _override(Parser, _parse_order)
diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py
index da77ff35..5d193306 100644
--- a/sqlmesh/core/engine_adapter/databricks.py
+++ b/sqlmesh/core/engine_adapter/databricks.py
@@ -100,7 +100,7 @@ class DatabricksEngineAdapter(SparkEngineAdapter):
         if not self._use_spark_session:
             return False
         return (
-            RuntimeEnv.get().is_databricks and str_to_bool(os.environ["IS_SERVERLESS"])
+            RuntimeEnv.get().is_databricks and str_to_bool(os.environ.get("IS_SERVERLESS", "False"))
         ) or bool(self._extra_config["databricks_connect_use_serverless"])
 
     @property
diff --git a/sqlmesh/core/loader.py b/sqlmesh/core/loader.py
index 52c121a9..76cdbd8a 100644
--- a/sqlmesh/core/loader.py
+++ b/sqlmesh/core/loader.py
@@ -123,6 +123,11 @@ class Loader(abc.ABC):
         )
         return project
 
+    def load_signals(self, context: GenericContext) -> None:
+        """Loads signals for the built-in scheduler."""
+        self._context = context
+        self._load_signals()
+
     def reload_needed(self) -> bool:
         """
         Checks for any modifications to the files the macros and models depend on
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 96f5fbc9..b69eec2a 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -559,16 +559,28 @@ class _Model(ModelMeta, frozen=True):
 
         referenced_audits = []
 
-        for audit_name, _ in self.audits + default_audits:
+        for audit_name, audit_args in self.audits + default_audits:
             if audit_name in self.inline_audits:
                 referenced_audits.append(self.inline_audits[audit_name])
             elif audit_name in audits:
                 referenced_audits.append(audits[audit_name])
-            elif audit_name not in BUILT_IN_AUDITS:
-                raise_config_error(
-                    f"Unknown audit '{audit_name}' referenced in model '{self.name}'",
-                    self._path,
-                )
+            else:
+                audit = BUILT_IN_AUDITS.get(audit_name)
+                if not audit:
+                    raise_config_error(
+                        f"Unknown audit '{audit_name}' referenced in model '{self.name}'",
+                        self._path,
+                    )
+
+                # Builtin audits are generally not included in order to reduce fingerprint size,
+                # but those that override `blocking` need to be included because otherwise doing
+                # `audit.blocking` will always return the builtin audit's default value
+                blocking = audit_args.get("blocking")
+                if blocking:
+                    referenced_audits.append(
+                        audit.copy(update={"blocking": blocking == exp.true()})  # type: ignore
+                    )
+
         return referenced_audits
 
     def text_diff(self, other: Node) -> str:


Test patch:
diff --git a/tests/core/test_context.py b/tests/core/test_context.py
index 3f163ca1..14e31c37 100644
--- a/tests/core/test_context.py
+++ b/tests/core/test_context.py
@@ -292,6 +292,59 @@ def test_plan_execution_time():
     )
 
 
+def test_override_builtin_audit_blocking_mode():
+    context = Context(config=Config())
+    context.upsert_model(
+        load_sql_based_model(
+            parse(
+                """
+                MODEL(
+                    name db.x,
+                    kind FULL,
+                    audits (
+                        not_null(columns := [c], blocking := false),
+                        unique_values(columns := [c]),
+                    )
+                );
+
+                SELECT NULL AS c
+                """
+            )
+        )
+    )
+
+    plan = context.plan(auto_apply=True, no_prompts=True)
+    new_snapshot = next(iter(plan.context_diff.new_snapshots.values()))
+
+    # Even though there are two builtin audits referenced in the above definition, we only
+    # store the one that overrides `blocking` in the snapshot; the other one isn't needed
+    assert len(new_snapshot.audits) == 1
+    assert new_snapshot.audits[0].name == "not_null"
+    assert new_snapshot.audits[0].blocking is False
+
+    context = Context(config=Config())
+    context.upsert_model(
+        load_sql_based_model(
+            parse(
+                """
+                MODEL(
+                    name db.x,
+                    kind FULL,
+                    audits (
+                        not_null_non_blocking(columns := [c], blocking := true)
+                    )
+                );
+
+                SELECT NULL AS c
+                """
+            )
+        )
+    )
+
+    with pytest.raises(SQLMeshError):
+        context.plan(auto_apply=True, no_prompts=True)
+
+
 def test_python_model_empty_df_raises(sushi_context, capsys):
     @model(
         "memory.sushi.test_model",
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index 0c8d21dd..53063461 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -89,6 +89,10 @@ def macro_evaluator() -> MacroEvaluator:
     ):
         return sum([a1, a2])
 
+    @macro()
+    def test_select_macro(evaluator):
+        return "SELECT 1 AS col"
+
     return MacroEvaluator(
         "hive",
         {"test": Executable(name="test", payload="def test(_):\n    return 'test'")},
@@ -907,7 +911,7 @@ def test_date_spine(assert_exp_eq, dialect, date_part):
         FROM (
             SELECT
                 date_{date_part}
-            FROM _generated_dates 
+            FROM _generated_dates
         ) AS _generated_dates
         """
     assert_exp_eq(evaluator.transform(parse_one(date_spine_macro)), expected_sql, dialect=dialect)
@@ -943,3 +947,10 @@ def test_date_spine_error_handling(macro_evaluator):
         str(e.value.__cause__)
         == "Invalid date range - start_date '2024-12-31' is after end_date '2022-01-01'."
     )
+
+
+def test_macro_union(assert_exp_eq, macro_evaluator: MacroEvaluator):
+    sql = "SELECT 1 AS col UNION ALL @TEST_SELECT_MACRO()"
+    expected_sql = "SELECT 1 AS col UNION ALL SELECT 1 AS col"
+
+    assert_exp_eq(macro_evaluator.transform(parse_one(sql)), expected_sql)