Feat: allow a single macro function call in model definitions (#3032)

Author: georgesittas

sqlmesh/core/model/definition.py

@@ -1673,8 +1673,10 @@ def load_sql_based_model(
     if query_or_seed_insert is not None and (
         isinstance(query_or_seed_insert, (exp.Query, d.JinjaQuery))
         or (
+            # Macro functions are allowed in place of model queries only when there are no
+            # other statements in the model definition, otherwise they would be ambiguous
             isinstance(query_or_seed_insert, d.MacroFunc)
-            and query_or_seed_insert.this.name.lower() == "union"
+            and (query_or_seed_insert.this.name.lower() == "union" or len(expressions) == 2)
         )
     ):
         return create_sql_model(
@@ -1745,9 +1747,6 @@ def create_sql_model(
         used_variables: The set of variable names used by this model.
     """
     if not isinstance(query, (exp.Query, d.JinjaQuery, d.MacroFunc)):
-        # Users are not expected to pass in a single MacroFunc instance for a model's query;
-        # this is an implementation detail which allows us to create python models that return
-        # SQL, either in the form of SQLGlot expressions or just plain strings.
         raise_config_error(
             "A query is required and must be a SELECT statement, a UNION statement, or a JINJA_QUERY block",
             path,
@@ -2081,7 +2080,10 @@ def _split_sql_model_statements(
             if (
                 isinstance(expr, (exp.Query, d.JinjaQuery))
                 or expr == INSERT_SEED_MACRO_CALL
-                or (isinstance(expr, d.MacroFunc) and expr.this.name.lower() == "union")
+                or (
+                    isinstance(expr, d.MacroFunc)
+                    and (expr.this.name.lower() == "union" or length == 1)
+                )
             ):
                 query_positions.append((expr, idx))
             sql_statements.append(expr)

sqlmesh/core/renderer.py

@@ -431,7 +431,9 @@ def render(
             if not query:
                 return None
             if not isinstance(query, exp.Query):
-                raise_config_error(f"Query needs to be a SELECT or a UNION {query}.", self._path)
+                raise_config_error(
+                    f"Model query needs to be a SELECT or a UNION, got {query}.", self._path
+                )
                 raise
 
             if optimize:

tests/core/test_model.py

@@ -376,12 +376,40 @@ def test_no_query():
         );
 
         @DEF(x, 1)
-    """
+        """
     )
 
     with pytest.raises(ConfigError) as ex:
-        load_sql_based_model(expressions, path=Path("test_location"))
-    assert "have a SELECT" in str(ex.value)
+        model = load_sql_based_model(expressions, path=Path("test_location"))
+        model.validate_definition()
+
+    assert "Model query needs to be a SELECT or a UNION, got @DEF(x, 1)." in str(ex.value)
+
+
+def test_single_macro_as_query(assert_exp_eq):
+    @macro()
+    def select_query(evaluator, *projections):
+        return exp.select(*[f'{p} AS "{p}"' for p in projections])
+
+    expressions = d.parse(
+        """
+        MODEL (
+            name test
+        );
+
+        @SELECT_QUERY(1, 2, 3)
+        """
+    )
+    model = load_sql_based_model(expressions)
+    assert_exp_eq(
+        model.render_query(),
+        """
+        SELECT
+          1 AS "1",
+          2 AS "2",
+          3 AS "3"
+        """,
+    )
 
 
 def test_partition_key_is_missing_in_query():