Repository: TobikoData/sqlmesh
Pull request #2313 (2024-03-20 20:24:29)

Fix: Reset the 'effective_from' attribute for migrated snapshots
If not reset, the next time the migrated snapshot is executed it'd detect that there are no intervals associated with its `identifier` after the `effective_from` timestamp and conclude that additional backfill is necessary in order to meet the configured `effective_from` setting.


Patch:
diff --git a/docs/concepts/macros/sqlmesh_macros.md b/docs/concepts/macros/sqlmesh_macros.md
index 22ea3d73..ebbea6fa 100644
--- a/docs/concepts/macros/sqlmesh_macros.md
+++ b/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.
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 2dc328ed..8ae4d294 100644
--- a/sqlmesh/core/dialect.py
+++ b/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)
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index ed51936a..f8edb22d 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -539,121 +539,27 @@ 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'"
@@ -661,36 +567,20 @@ def having(
     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")
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index f63d5731..3f65ecb2 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -736,7 +736,8 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
         previous_ids = {s.snapshot_id(self.name) for s in self.previous_versions}
         if self.identifier == other.identifier or (
             # Indirect Non-Breaking snapshots share the dev table with its previous version.
-            self.is_indirect_non_breaking
+            # The same applies to migrated snapshots.
+            (self.is_indirect_non_breaking or self.migrated)
             and other.snapshot_id in previous_ids
         ):
             for start, end in other.dev_intervals:
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index b390cbbc..25f8097e 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -1097,6 +1097,8 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
                 logger.exception("Could not compute fingerprint for %s", snapshot.snapshot_id)
                 return
 
+            # Reset the effective_from date for the new snapshot to avoid unexpected backfills.
+            new_snapshot.effective_from = None
             new_snapshot.previous_versions = snapshot.all_versions
             new_snapshot.migrated = True
             if not new_snapshot.temp_version:


Test patch:
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index 073f37b1..c05b67bf 100644
--- a/tests/core/test_macros.py
+++ b/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",
diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py
index caac7a3d..8ad01eb8 100644
--- a/tests/core/test_snapshot.py
+++ b/tests/core/test_snapshot.py
@@ -198,7 +198,7 @@ def test_add_interval(snapshot: Snapshot, make_snapshot):
     ]
 
 
-def test_add_interval_dev(snapshot: Snapshot):
+def test_add_interval_dev(snapshot: Snapshot, make_snapshot):
     snapshot.version = "existing_version"
     snapshot.change_category = SnapshotChangeCategory.FORWARD_ONLY
 
@@ -209,6 +209,18 @@ def test_add_interval_dev(snapshot: Snapshot):
     assert snapshot.intervals == [(to_timestamp("2020-01-01"), to_timestamp("2020-01-02"))]
     assert snapshot.dev_intervals == [(to_timestamp("2020-01-02"), to_timestamp("2020-01-03"))]
 
+    new_snapshot = make_snapshot(snapshot.model)
+    new_snapshot.merge_intervals(snapshot)
+    assert new_snapshot.intervals == [(to_timestamp("2020-01-01"), to_timestamp("2020-01-02"))]
+    assert new_snapshot.dev_intervals == []
+
+    new_snapshot = make_snapshot(snapshot.model)
+    new_snapshot.previous_versions = snapshot.all_versions
+    new_snapshot.migrated = True
+    new_snapshot.merge_intervals(snapshot)
+    assert new_snapshot.intervals == [(to_timestamp("2020-01-01"), to_timestamp("2020-01-02"))]
+    assert new_snapshot.dev_intervals == [(to_timestamp("2020-01-02"), to_timestamp("2020-01-03"))]
+
 
 def test_add_interval_partial(snapshot: Snapshot, make_snapshot):
     snapshot.add_interval("2023-01-01 00:00:00", "2023-01-01 23:59:59")