Fix!: Change how partitioned_by is parsed so that partition expressions with specialized AST nodes are captured

Author: erindru

sqlmesh/core/dialect.py

@@ -610,6 +610,12 @@ def parse(self: Parser) -> t.Optional[exp.Expression]:
                     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))
 

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\"""",
+    ]

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
     )