Repository: TobikoData/sqlmesh
Pull request #4247 (2025-04-24 22:26:16)

feat: add context to linter rules


Patch:
diff --git a/examples/sushi/linter/user.py b/examples/sushi/linter/user.py
index a4a33c0e..3f83c76e 100644
--- a/examples/sushi/linter/user.py
+++ b/examples/sushi/linter/user.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import typing as t
 
+from sqlmesh.core.context import Context
 from sqlmesh.core.linter.rule import Rule, RuleViolation
 from sqlmesh.core.model import Model
 
@@ -10,4 +11,6 @@ class NoMissingOwner(Rule):
     """All models should have an owner specified."""
 
     def check_model(self, model: Model) -> t.Optional[RuleViolation]:
+        assert isinstance(self.context, Context)
+        assert len(self.context.models) > 10
         return self.violation() if not model.owner else None
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 634db636..4a8f99ad 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -2537,7 +2537,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             # Linter may be `None` if the context is not loaded yet
             if linter := self._linters.get(model.project):
                 lint_violation, violations = (
-                    linter.lint_model(model, console=self.console) or found_error
+                    linter.lint_model(model, self, console=self.console) or found_error
                 )
                 if lint_violation:
                     found_error = True
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 4bd8e865..bd58f94d 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -610,6 +610,12 @@ def _create_parser(expression_type: t.Type[exp.Expression], table_keys: t.List[s
                     value = self.expression(ModelKind, this=kind.value, expressions=props)
             elif key == "expression":
                 value = self._parse_conjunction()
+            elif key == "partitioned_by":
+                partitioned_by = self._parse_partitioned_by()
+                if isinstance(partitioned_by.this, exp.Schema):
+                    value = exp.tuple_(*partitioned_by.this.expressions)
+                else:
+                    value = partitioned_by.this
             else:
                 value = self._parse_bracket(self._parse_field(any_token=True))
 
diff --git a/sqlmesh/core/linter/definition.py b/sqlmesh/core/linter/definition.py
index f3bc3c89..a8ebdec0 100644
--- a/sqlmesh/core/linter/definition.py
+++ b/sqlmesh/core/linter/definition.py
@@ -11,6 +11,9 @@ from sqlmesh.core.model import Model
 from sqlmesh.core.linter.rule import Rule, RuleViolation
 from sqlmesh.core.console import LinterConsole, get_console
 
+if t.TYPE_CHECKING:
+    from sqlmesh.core.context import GenericContext
+
 
 def select_rules(all_rules: RuleSet, rule_names: t.Set[str]) -> RuleSet:
     if "all" in rule_names:
@@ -52,7 +55,7 @@ class Linter:
         return Linter(config.enabled, all_rules, rules, warn_rules)
 
     def lint_model(
-        self, model: Model, console: LinterConsole = get_console()
+        self, model: Model, context: GenericContext, console: LinterConsole = get_console()
     ) -> t.Tuple[bool, t.List[AnnotatedRuleViolation]]:
         if not self.enabled:
             return False, []
@@ -62,8 +65,8 @@ class Linter:
         rules = self.rules.difference(ignored_rules)
         warn_rules = self.warn_rules.difference(ignored_rules)
 
-        error_violations = rules.check_model(model)
-        warn_violations = warn_rules.check_model(model)
+        error_violations = rules.check_model(model, context)
+        warn_violations = warn_rules.check_model(model, context)
 
         all_violations: t.List[AnnotatedRuleViolation] = [
             AnnotatedRuleViolation(
@@ -96,11 +99,11 @@ class RuleSet(Mapping[str, type[Rule]]):
     def __init__(self, rules: Iterable[type[Rule]] = ()) -> None:
         self._underlying = {rule.name: rule for rule in rules}
 
-    def check_model(self, model: Model) -> t.List[RuleViolation]:
+    def check_model(self, model: Model, context: GenericContext) -> t.List[RuleViolation]:
         violations = []
 
         for rule in self._underlying.values():
-            violation = rule().check_model(model)
+            violation = rule(context).check_model(model)
 
             if violation:
                 violations.append(violation)
diff --git a/sqlmesh/core/linter/rule.py b/sqlmesh/core/linter/rule.py
index ceee9966..4105e27d 100644
--- a/sqlmesh/core/linter/rule.py
+++ b/sqlmesh/core/linter/rule.py
@@ -9,6 +9,10 @@ from typing import Type
 import typing as t
 
 
+if t.TYPE_CHECKING:
+    from sqlmesh.core.context import GenericContext
+
+
 class _Rule(abc.ABCMeta):
     def __new__(cls: Type[_Rule], clsname: str, bases: t.Tuple, attrs: t.Dict) -> _Rule:
         attrs["name"] = clsname.lower()
@@ -20,6 +24,9 @@ class Rule(abc.ABC, metaclass=_Rule):
 
     name = "rule"
 
+    def __init__(self, context: GenericContext):
+        self.context = context
+
     @abc.abstractmethod
     def check_model(self, model: Model) -> t.Optional[RuleViolation]:
         """The evaluation function that'll check for a violation of this rule."""
diff --git a/sqlmesh/core/model/meta.py b/sqlmesh/core/model/meta.py
index 514a83cb..29c82bc3 100644
--- a/sqlmesh/core/model/meta.py
+++ b/sqlmesh/core/model/meta.py
@@ -5,7 +5,7 @@ from functools import cached_property
 from typing_extensions import Self
 
 from pydantic import Field
-from sqlglot import Dialect, exp
+from sqlglot import Dialect, exp, parse_one
 from sqlglot.helper import ensure_collection, ensure_list
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
@@ -39,6 +39,7 @@ from sqlmesh.utils.pydantic import (
     field_validator,
     list_of_fields_validator,
     model_validator,
+    get_dialect,
 )
 
 if t.TYPE_CHECKING:
@@ -182,6 +183,22 @@ class ModelMeta(_Node):
     def _partition_and_cluster_validator(
         cls, v: t.Any, info: ValidationInfo
     ) -> t.List[exp.Expression]:
+        if (
+            isinstance(v, list)
+            and all(isinstance(i, str) for i in v)
+            and info.field_name == "partitioned_by_"
+        ):
+            # this branch gets hit when we are deserializing from json because `partitioned_by` is stored as a List[str]
+            # however, we should only invoke this if the list contains strings because this validator is also
+            # called by Python models which might pass a List[exp.Expression]
+            string_to_parse = (
+                f"({','.join(v)})"  # recreate the (a, b, c) part of "partitioned_by (a, b, c)"
+            )
+            parsed = parse_one(
+                string_to_parse, into=exp.PartitionedByProperty, dialect=get_dialect(info)
+            )
+            v = parsed.this.expressions if isinstance(parsed.this, exp.Schema) else v
+
         expressions = list_of_fields_validator(v, info.data)
 
         for expression in expressions:
diff --git a/sqlmesh/migrations/v0081_update_partitioned_by.py b/sqlmesh/migrations/v0081_update_partitioned_by.py
new file mode 100644
index 00000000..a88f8c38
--- /dev/null
+++ b/sqlmesh/migrations/v0081_update_partitioned_by.py
@@ -0,0 +1,91 @@
+"""Remove superfluous exp.Paren references from partitioned_by"""
+
+import json
+
+import pandas as pd
+from sqlglot import exp
+
+from sqlmesh.utils.migration import index_text_type
+from sqlmesh.utils.migration import blob_text_type
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    schema = state_sync.schema
+    snapshots_table = "_snapshots"
+    index_type = index_text_type(engine_adapter.dialect)
+    if schema:
+        snapshots_table = f"{schema}.{snapshots_table}"
+
+    new_snapshots = []
+    updated = False
+
+    for (
+        name,
+        identifier,
+        version,
+        snapshot,
+        kind_name,
+        updated_ts,
+        unpaused_ts,
+        ttl_ms,
+        unrestorable,
+    ) in engine_adapter.fetchall(
+        exp.select(
+            "name",
+            "identifier",
+            "version",
+            "snapshot",
+            "kind_name",
+            "updated_ts",
+            "unpaused_ts",
+            "ttl_ms",
+            "unrestorable",
+        ).from_(snapshots_table),
+        quote_identifiers=True,
+    ):
+        parsed_snapshot = json.loads(snapshot)
+
+        if partitioned_by := parsed_snapshot["node"].get("partitioned_by"):
+            new_partitioned_by = []
+            for item in partitioned_by:
+                # rewrite '(foo)' to 'foo'
+                if item.startswith("(") and item.endswith(")"):
+                    item = item[1:-1]
+                    updated = True
+                new_partitioned_by.append(item)
+            parsed_snapshot["node"]["partitioned_by"] = new_partitioned_by
+
+        new_snapshots.append(
+            {
+                "name": name,
+                "identifier": identifier,
+                "version": version,
+                "snapshot": json.dumps(parsed_snapshot),
+                "kind_name": kind_name,
+                "updated_ts": updated_ts,
+                "unpaused_ts": unpaused_ts,
+                "ttl_ms": ttl_ms,
+                "unrestorable": unrestorable,
+            }
+        )
+
+    if new_snapshots and updated:
+        engine_adapter.delete_from(snapshots_table, "TRUE")
+        blob_type = blob_text_type(engine_adapter.dialect)
+
+        engine_adapter.insert_append(
+            snapshots_table,
+            pd.DataFrame(new_snapshots),
+            columns_to_types={
+                "name": exp.DataType.build(index_type),
+                "identifier": exp.DataType.build(index_type),
+                "version": exp.DataType.build(index_type),
+                "snapshot": exp.DataType.build(blob_type),
+                "kind_name": exp.DataType.build(index_type),
+                "updated_ts": exp.DataType.build("bigint"),
+                "unpaused_ts": exp.DataType.build("bigint"),
+                "ttl_ms": exp.DataType.build("bigint"),
+                "unrestorable": exp.DataType.build("boolean"),
+            },
+        )


Test patch:
diff --git a/tests/core/engine_adapter/test_athena.py b/tests/core/engine_adapter/test_athena.py
index 46aabc74..f96f5656 100644
--- a/tests/core/engine_adapter/test_athena.py
+++ b/tests/core/engine_adapter/test_athena.py
@@ -435,3 +435,51 @@ def test_drop_partitions_from_metastore_uses_batches(
     # third call 50-62
     assert calls[2][1]["PartitionsToDelete"][0]["Values"][0] == "50"
     assert calls[2][1]["PartitionsToDelete"][-1]["Values"][0] == "62"
+
+
+def test_iceberg_partition_transforms(adapter: AthenaEngineAdapter):
+    expressions = d.parse(
+        """
+        MODEL (
+            name test_table,
+            kind FULL,
+            table_format iceberg,
+            partitioned_by (month(business_date), bucket(4, colb), colc)
+        );
+
+        SELECT 1::timestamp AS business_date, 2::varchar as colb, 'foo' as colc;
+    """
+    )
+    model: SqlModel = t.cast(SqlModel, load_sql_based_model(expressions))
+
+    assert model.partitioned_by == [
+        exp.Month(this=exp.column("business_date", quoted=True)),
+        exp.PartitionedByBucket(
+            this=exp.column("colb", quoted=True), expression=exp.Literal.number(4)
+        ),
+        exp.column("colc", quoted=True),
+    ]
+
+    adapter.s3_warehouse_location = "s3://bucket/prefix/"
+
+    adapter.create_table(
+        table_name=model.name,
+        columns_to_types=model.columns_to_types_or_raise,
+        partitioned_by=model.partitioned_by,
+        table_format=model.table_format,
+    )
+
+    adapter.ctas(
+        table_name=model.name,
+        columns_to_types=model.columns_to_types_or_raise,
+        partitioned_by=model.partitioned_by,
+        query_or_df=model.ctas_query(),
+        table_format=model.table_format,
+    )
+
+    assert to_sql_calls(adapter) == [
+        # Hive syntax - create table
+        """CREATE TABLE IF NOT EXISTS `test_table` (`business_date` TIMESTAMP, `colb` STRING, `colc` STRING) PARTITIONED BY (MONTH(`business_date`), BUCKET(4, `colb`), `colc`) LOCATION 's3://bucket/prefix/test_table/' TBLPROPERTIES ('table_type'='iceberg')""",
+        # Trino syntax - CTAS
+        """CREATE TABLE IF NOT EXISTS "test_table" WITH (table_type='iceberg', partitioning=ARRAY['MONTH(business_date)', 'BUCKET(colb, 4)', 'colc'], location='s3://bucket/prefix/test_table/', is_external=false) AS SELECT CAST("business_date" AS TIMESTAMP) AS "business_date", CAST("colb" AS VARCHAR) AS "colb", CAST("colc" AS VARCHAR) AS "colc" FROM (SELECT CAST(1 AS TIMESTAMP) AS "business_date", CAST(2 AS VARCHAR) AS "colb", 'foo' AS "colc" LIMIT 0) AS "_subquery\"""",
+    ]
diff --git a/tests/core/test_context.py b/tests/core/test_context.py
index eb68ba8d..cde653b1 100644
--- a/tests/core/test_context.py
+++ b/tests/core/test_context.py
@@ -1978,19 +1978,22 @@ def test_plan_audit_intervals(tmp_path: pathlib.Path, capsys, caplog):
         )
     )
 
-    ctx.plan(
+    plan = ctx.plan(
         environment="dev", auto_apply=True, no_prompts=True, start="2025-02-01", end="2025-02-01"
     )
 
+    date_snapshot = next(s for s in plan.new_snapshots if "date_example" in s.name)
+    timestamp_snapshot = next(s for s in plan.new_snapshots if "timestamp_example" in s.name)
+
     # Case 1: The timestamp audit should be in the inclusive range ['2025-02-01 00:00:00', '2025-02-01 23:59:59.999999']
     assert (
-        """SELECT COUNT(*) FROM (SELECT ("timestamp_id") AS "timestamp_id" FROM (SELECT * FROM "sqlmesh__sqlmesh_audit"."sqlmesh_audit__timestamp_example__2797548448" AS "sqlmesh_audit__timestamp_example__2797548448" WHERE "timestamp_id" BETWEEN CAST('2025-02-01 00:00:00' AS TIMESTAMP) AND CAST('2025-02-01 23:59:59.999999' AS TIMESTAMP)) AS "_q_0" WHERE TRUE GROUP BY ("timestamp_id") HAVING COUNT(*) > 1) AS "audit\""""
+        f"""SELECT COUNT(*) FROM (SELECT ("timestamp_id") AS "timestamp_id" FROM (SELECT * FROM "sqlmesh__sqlmesh_audit"."sqlmesh_audit__timestamp_example__{timestamp_snapshot.version}" AS "sqlmesh_audit__timestamp_example__{timestamp_snapshot.version}" WHERE "timestamp_id" BETWEEN CAST('2025-02-01 00:00:00' AS TIMESTAMP) AND CAST('2025-02-01 23:59:59.999999' AS TIMESTAMP)) AS "_q_0" WHERE TRUE GROUP BY ("timestamp_id") HAVING COUNT(*) > 1) AS "audit\""""
         in caplog.text
     )
 
     # Case 2: The date audit should be in the inclusive range ['2025-02-01', '2025-02-01']
     assert (
-        """SELECT COUNT(*) FROM (SELECT ("date_id") AS "date_id" FROM (SELECT * FROM "sqlmesh__sqlmesh_audit"."sqlmesh_audit__date_example__4100277424" AS "sqlmesh_audit__date_example__4100277424" WHERE "date_id" BETWEEN CAST('2025-02-01' AS DATE) AND CAST('2025-02-01' AS DATE)) AS "_q_0" WHERE TRUE GROUP BY ("date_id") HAVING COUNT(*) > 1) AS "audit\""""
+        f"""SELECT COUNT(*) FROM (SELECT ("date_id") AS "date_id" FROM (SELECT * FROM "sqlmesh__sqlmesh_audit"."sqlmesh_audit__date_example__{date_snapshot.version}" AS "sqlmesh_audit__date_example__{date_snapshot.version}" WHERE "date_id" BETWEEN CAST('2025-02-01' AS DATE) AND CAST('2025-02-01' AS DATE)) AS "_q_0" WHERE TRUE GROUP BY ("date_id") HAVING COUNT(*) > 1) AS "audit\""""
         in caplog.text
     )
 
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 825b39ea..fe9e4787 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -1514,6 +1514,134 @@ def test_render_definition_with_defaults():
     ) == d.format_model_expressions(expected_expressions)
 
 
+def test_render_definition_partitioned_by():
+    # no parenthesis in definition, no parenthesis when rendered
+    model = load_sql_based_model(
+        d.parse(
+            f"""
+        MODEL (
+            name db.table,
+            kind FULL,
+            partitioned_by a
+        );
+
+        select 1 as a;
+        """
+        )
+    )
+
+    assert model.partitioned_by == [exp.column("a", quoted=True)]
+    assert (
+        model.render_definition()[0].sql(pretty=True)
+        == """MODEL (
+  name db.table,
+  kind FULL,
+  partitioned_by "a"
+)"""
+    )
+
+    # single column wrapped in parenthesis in defintion, no parenthesis in rendered
+    model = load_sql_based_model(
+        d.parse(
+            f"""
+        MODEL (
+            name db.table,
+            kind FULL,
+            partitioned_by (a)
+        );
+
+        select 1 as a;
+        """
+        )
+    )
+
+    assert model.partitioned_by == [exp.column("a", quoted=True)]
+    assert (
+        model.render_definition()[0].sql(pretty=True)
+        == """MODEL (
+  name db.table,
+  kind FULL,
+  partitioned_by "a"
+)"""
+    )
+
+    # multiple columns wrapped in parenthesis in definition, parenthesis in rendered
+    model = load_sql_based_model(
+        d.parse(
+            f"""
+        MODEL (
+            name db.table,
+            kind FULL,
+            partitioned_by (a, b)
+        );
+
+        select 1 as a, 2 as b;
+        """
+        )
+    )
+
+    assert model.partitioned_by == [exp.column("a", quoted=True), exp.column("b", quoted=True)]
+    assert (
+        model.render_definition()[0].sql(pretty=True)
+        == """MODEL (
+  name db.table,
+  kind FULL,
+  partitioned_by ("a", "b")
+)"""
+    )
+
+    # multiple columns not wrapped in parenthesis in the definition is an error
+    with pytest.raises(ParseError, match=r"keyword: 'value' missing"):
+        load_sql_based_model(
+            d.parse(
+                f"""
+            MODEL (
+                name db.table,
+                kind FULL,
+                partitioned_by a, b
+            );
+
+            select 1 as a, 2 as b;
+            """
+            )
+        )
+
+    # Iceberg transforms / functions
+    model = load_sql_based_model(
+        d.parse(
+            f"""
+        MODEL (
+            name db.table,
+            kind FULL,
+            partitioned_by (day(a), truncate(b, 4), bucket(c, 3))
+        );
+
+        select 1 as a, 2 as b, 3 as c;
+        """
+        ),
+        dialect="trino",
+    )
+
+    assert model.partitioned_by == [
+        exp.Day(this=exp.column("a", quoted=True)),
+        exp.PartitionByTruncate(
+            this=exp.column("b", quoted=True), expression=exp.Literal.number(4)
+        ),
+        exp.PartitionedByBucket(
+            this=exp.column("c", quoted=True), expression=exp.Literal.number(3)
+        ),
+    ]
+    assert (
+        model.render_definition()[0].sql(pretty=True)
+        == """MODEL (
+  name db.table,
+  dialect trino,
+  kind FULL,
+  partitioned_by (DAY("a"), TRUNCATE("b", 4), BUCKET("c", 3))
+)"""
+    )
+
+
 def test_cron():
     daily = _Node(name="x", cron="@daily")
     assert to_datetime(daily.cron_prev("2020-01-01")) == to_datetime("2019-12-31")
diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py
index f27ad8f3..387c799e 100644
--- a/tests/core/test_snapshot.py
+++ b/tests/core/test_snapshot.py
@@ -2938,3 +2938,37 @@ def test_render_signal(make_snapshot, mocker):
     )
     snapshot_a = make_snapshot(sql_model)
     assert snapshot_a.check_ready_intervals([(0, 1)], mocker.Mock()) == [(0, 1)]
+
+
+def test_partitioned_by_roundtrip(make_snapshot: t.Callable):
+    sql_model = load_sql_based_model(
+        parse("""
+        MODEL (
+            name test_schema.test_model,
+            kind full,
+            partitioned_by (a, bucket(4, b), truncate(3, c), month(d))
+        );
+        SELECT a, b, c, d FROM tbl;
+        """)
+    )
+    snapshot = make_snapshot(sql_model)
+    assert isinstance(snapshot, Snapshot)
+    assert isinstance(snapshot.node, SqlModel)
+
+    assert snapshot.node.partitioned_by == [
+        exp.column("a", quoted=True),
+        exp.PartitionedByBucket(
+            this=exp.column("b", quoted=True), expression=exp.Literal.number(4)
+        ),
+        exp.PartitionByTruncate(
+            this=exp.column("c", quoted=True), expression=exp.Literal.number(3)
+        ),
+        exp.Month(this=exp.column("d", quoted=True)),
+    ]
+
+    # roundtrip through json and ensure we get correct AST nodes on the other end
+    serialized = snapshot.json()
+    deserialized = snapshot.parse_raw(serialized)
+
+    assert isinstance(deserialized.node, SqlModel)
+    assert deserialized.node.partitioned_by == snapshot.node.partitioned_by