Feat: add optional @limit macro (#2275)

* wip: add plumbing for limit structural macro

* feat: update to consume required upstream api change

* chore: add tested doc examples of usage to optional_expression function

Author: z3z1ma

docs/concepts/macros/sqlmesh_macros.md

@@ -526,6 +526,7 @@ SQLMesh's macro system has six operators that correspond to different clauses in
 - `@GROUP_BY`: grouping `GROUP BY` clause
 - `@HAVING`: group by filtering `HAVING` clause
 - `@ORDER_BY`: ordering `ORDER BY` clause
+- `@LIMIT`: limiting `LIMIT` clause
 
 Each of these operators is used to dynamically add the code for its corresponding clause to a model's SQL query.
 
@@ -746,6 +747,22 @@ FROM all_cities
 ORDER BY city_pop
 ```
 
+#### @LIMIT operator
+
+```sql linenums="1"
+SELECT *
+FROM all_cities
+@LIMIT(True) 10
+```
+
+renders to
+
+```sql linenums="1"
+SELECT *
+FROM all_cities
+LIMIT 10
+```
+
 ## User-defined macro functions
 
 User-defined macro functions allow the same macro code to be used in multiple models.

sqlmesh/core/dialect.py

@@ -225,7 +225,7 @@ def _parse_macro(self: Parser, keyword_macro: str = "") -> t.Optional[exp.Expres
     return self.expression(MacroVar, this=field.this)
 
 
-KEYWORD_MACROS = {"WITH", "JOIN", "WHERE", "GROUP_BY", "HAVING", "ORDER_BY"}
+KEYWORD_MACROS = {"WITH", "JOIN", "WHERE", "GROUP_BY", "HAVING", "ORDER_BY", "LIMIT"}
 
 
 def _parse_matching_macro(self: Parser, name: str) -> t.Optional[exp.Expression]:
@@ -251,6 +251,8 @@ def _parse_body_macro(self: Parser) -> t.Tuple[str, t.Optional[exp.Expression]]:
         return ("having", self._parse_having())
     if name == "ORDER_BY":
         return ("order", self._parse_order())
+    if name == "LIMIT":
+        return ("limit", self._parse_limit())
     return ("", None)
 
 
@@ -323,6 +325,20 @@ def _parse_order(
     return macro
 
 
+def _parse_limit(
+    self: Parser,
+    this: t.Optional[exp.Expression] = None,
+    top: bool = False,
+    skip_limit_token: bool = False,
+) -> t.Optional[exp.Expression]:
+    macro = _parse_matching_macro(self, "TOP" if top else "LIMIT")
+    if not macro:
+        return self.__parse_limit(this, top=top, skip_limit_token=skip_limit_token)  # type: ignore
+
+    macro.this.append("expressions", self.__parse_limit(this, top=top, skip_limit_token=True))  # type: ignore
+    return macro
+
+
 def _parse_props(self: Parser) -> t.Optional[exp.Expression]:
     key = self._parse_id_var(any_token=True)
     if not key:
@@ -762,6 +778,7 @@ def extend_sqlglot() -> None:
     _override(Parser, _parse_group)
     _override(Parser, _parse_with)
     _override(Parser, _parse_having)
+    _override(Parser, _parse_limit)
     _override(Parser, _parse_lambda)
     _override(Parser, _parse_types)
     _override(Parser, _parse_if)

sqlmesh/core/macros.py

@@ -539,158 +539,48 @@ def filter_(evaluator: MacroEvaluator, *args: t.Any) -> t.List[t.Any]:
     return list(filter(lambda arg: evaluator.eval_expression(func(arg)), items))
 
 
-@macro("WITH")
-def with_(
+def _optional_expression(
     evaluator: MacroEvaluator,
     condition: exp.Condition,
-    expression: exp.With,
-) -> t.Optional[exp.With]:
-    """Inserts WITH expression when the condition is True
+    expression: exp.Expression,
+) -> t.Optional[exp.Expression]:
+    """Inserts expression when the condition is True
 
-    Example:
+    The following examples express the usage of this function in the context of the macros which wrap it.
+
+    Examples:
         >>> from sqlglot import parse_one
         >>> from sqlmesh.core.macros import MacroEvaluator
         >>> sql = "@WITH(True) all_cities as (select * from city) select all_cities"
         >>> MacroEvaluator().transform(parse_one(sql)).sql()
         'WITH all_cities AS (SELECT * FROM city) SELECT all_cities'
-
-    Args:
-        evaluator: MacroEvaluator that invoked the macro
-        condition: Condition expression
-        expression: With expression
-    Returns:
-        With expression if the conditional is True; otherwise None
-    """
-    return expression if evaluator.eval_expression(condition) else None
-
-
-@macro()
-def join(
-    evaluator: MacroEvaluator,
-    condition: exp.Condition,
-    expression: exp.Join,
-) -> t.Optional[exp.Join]:
-    """Inserts JOIN expression when the condition is True
-
-    Example:
-        >>> from sqlglot import parse_one
-        >>> from sqlmesh.core.macros import MacroEvaluator
-        >>> sql = "select * from city @JOIN(True) country on city.country = country.name"
-        >>> MacroEvaluator().transform(parse_one(sql)).sql()
-        'SELECT * FROM city JOIN country ON city.country = country.name'
-
         >>> sql = "select * from city left outer @JOIN(True) country on city.country = country.name"
         >>> MacroEvaluator().transform(parse_one(sql)).sql()
         'SELECT * FROM city LEFT OUTER JOIN country ON city.country = country.name'
-
-    Args:
-        evaluator: MacroEvaluator that invoked the macro
-        condition: Condition expression
-        expression: Join expression
-    Returns:
-        Join expression if the conditional is True; otherwise None
-    """
-    return expression if evaluator.eval_expression(condition) else None
-
-
-@macro()
-def where(
-    evaluator: MacroEvaluator,
-    condition: exp.Condition,
-    expression: exp.Where,
-) -> t.Optional[exp.Where]:
-    """Inserts WHERE expression when the condition is True
-
-    Example:
-        >>> from sqlglot import parse_one
-        >>> from sqlmesh.core.macros import MacroEvaluator
-        >>> sql = "select * from city @WHERE(True) population > 100 and country = 'Mexico'"
-        >>> MacroEvaluator().transform(parse_one(sql)).sql()
-        "SELECT * FROM city WHERE population > 100 AND country = 'Mexico'"
-
-    Args:
-        evaluator: MacroEvaluator that invoked the macro
-        condition: Condition expression
-        expression: Where expression
-    Returns:
-        Where expression if condition is True; otherwise None
-    """
-    return expression if evaluator.eval_expression(condition) else None
-
-
-@macro()
-def group_by(
-    evaluator: MacroEvaluator,
-    condition: exp.Condition,
-    expression: exp.Group,
-) -> t.Optional[exp.Group]:
-    """Inserts GROUP BY expression when the condition is True
-
-    Example:
-        >>> from sqlglot import parse_one
-        >>> from sqlmesh.core.macros import MacroEvaluator, group_by
         >>> sql = "select * from city @GROUP_BY(True) country, population"
         >>> MacroEvaluator().transform(parse_one(sql)).sql()
         'SELECT * FROM city GROUP BY country, population'
-
-    Args:
-        evaluator: MacroEvaluator that invoked the macro
-        condition: Condition expression
-        expression: Group expression
-    Returns:
-        Group expression if the condition is True; otherwise None
-    """
-    return expression if evaluator.eval_expression(condition) else None
-
-
-@macro()
-def having(
-    evaluator: MacroEvaluator,
-    condition: exp.Condition,
-    expression: exp.Having,
-) -> t.Optional[exp.Having]:
-    """Inserts HAVING expression when the condition is True
-
-    Example:
-        >>> from sqlglot import parse_one
-        >>> from sqlmesh.core.macros import MacroEvaluator
         >>> sql = "select * from city group by country @HAVING(True) population > 100 and country = 'Mexico'"
         >>> MacroEvaluator().transform(parse_one(sql)).sql()
         "SELECT * FROM city GROUP BY country HAVING population > 100 AND country = 'Mexico'"
 
     Args:
         evaluator: MacroEvaluator that invoked the macro
         condition: Condition expression
-        expression: Having expression
+        expression: SQL expression
     Returns:
-        Having expression if the condition is True; otherwise None
+        Expression if the conditional is True; otherwise None
     """
     return expression if evaluator.eval_expression(condition) else None
 
 
-@macro()
-def order_by(
-    evaluator: MacroEvaluator,
-    condition: exp.Condition,
-    expression: exp.Order,
-) -> t.Optional[exp.Order]:
-    """Inserts ORDER BY expression when the condition is True
-
-    Example:
-        >>> from sqlglot import parse_one
-        >>> from sqlmesh.core.macros import MacroEvaluator
-        >>> sql = "select * from city @ORDER_BY(True) population, name DESC"
-        >>> MacroEvaluator().transform(parse_one(sql)).sql()
-        'SELECT * FROM city ORDER BY population, name DESC'
-
-    Args:
-        evaluator: MacroEvaluator that invoked the macro
-        condition: Condition expression
-        expression: Order expression
-    Returns:
-        Order expression if the condition is True; otherwise None
-    """
-    return expression if evaluator.eval_expression(condition) else None
+with_ = macro("WITH")(_optional_expression)
+join = macro("JOIN")(_optional_expression)
+where = macro("WHERE")(_optional_expression)
+group_by = macro("GROUP_BY")(_optional_expression)
+having = macro("HAVING")(_optional_expression)
+order_by = macro("ORDER_BY")(_optional_expression)
+limit = macro("LIMIT")(_optional_expression)
 
 
 @macro("eval")

tests/core/test_macros.py

@@ -272,6 +272,16 @@ def test_ast_correctness(macro_evaluator):
             "SELECT * FROM city",
             {"do_order": False},
         ),
+        (
+            """select * from city @LIMIT(@do_limit) 10""",
+            "SELECT * FROM city LIMIT 10",
+            {"do_limit": True},
+        ),
+        (
+            """select * from city @LIMIT(@do_limit) 10""",
+            "SELECT * FROM city",
+            {"do_limit": False},
+        ),
         (
             """select @if(TRUE, 1, 0)""",
             "SELECT 1",