Repository: TobikoData/sqlmesh
Pull request #2292 (2024-03-18 02:17:17)

Fix: Fix partitioning property for Trino iceberg tables
Currently, using iceberg column transforms as the `partitioned_by` model property, eg:

```
MODEL (
     name foo,
     kind INCREMENTAL_BY_TIME_RANGE (
        time_column transacttime,
    ),
    partitioned_by day(transacttime),
  ....
)
```

throw a SQLMesh error:

`Error: partitioned_by field ''day(transacttime)'' does not contain a column`

A closer look at [the docs](https://trino.io/docs/current/connector/iceberg.html#partitioned-tables) reveal that, for an Iceberg table, the partition columns need to be defined as `ARRAY[varchar]`.

So this PR addresses that by wrapping the partition columns in `exp.Array` instead of `exp.Schema` and rendering them as varchar literals.

Note that Trino works with both the following syntax:

`partitioning = ARRAY['day("cola")', '"colb"']`
or
`partitioning = ARRAY['day(cola)','colb']`

(that is, if the column names are quoted within the string literals, it still works as expected)


Patch:
diff --git a/sqlmesh/core/engine_adapter/mixins.py b/sqlmesh/core/engine_adapter/mixins.py
index 8cd56c28..6aebd6a5 100644
--- a/sqlmesh/core/engine_adapter/mixins.py
+++ b/sqlmesh/core/engine_adapter/mixins.py
@@ -146,28 +146,37 @@ class HiveMetastoreTablePropertiesMixin(EngineAdapter):
             properties.append(exp.FileFormatProperty(this=exp.Var(this=storage_format)))
 
         if partitioned_by:
-            for expr in partitioned_by:
-                if not isinstance(expr, exp.Column):
-                    raise SQLMeshError(
-                        f"PARTITIONED BY contains non-column value '{expr.sql(dialect='spark')}'."
-                    )
-
-            property: exp.Property = exp.PartitionedByProperty(
-                this=exp.Schema(expressions=partitioned_by),
-            )
-
             if (
                 self.dialect == "trino"
                 and self.get_catalog_type(catalog_name or self.get_current_catalog()) == "iceberg"
             ):
                 # On the Trino Iceberg catalog, the table property is called "partitioning" - not "partitioned_by"
                 # In addition, partition column transform expressions like `day(col)` or `bucket(col, 5)` are allowed
+                # Also, column names and transforms need to be strings and supplied as an ARRAY[varchar]
                 # ref: https://trino.io/docs/current/connector/iceberg.html#table-properties
-                property = exp.Property(
-                    this=exp.var("PARTITIONING"), value=exp.Schema(expressions=partitioned_by)
+                properties.append(
+                    exp.Property(
+                        this=exp.var("PARTITIONING"),
+                        value=exp.array(
+                            *(
+                                exp.Literal.string(e.sql(dialect=self.dialect))
+                                for e in partitioned_by
+                            )
+                        ),
+                    )
+                )
+            else:
+                for expr in partitioned_by:
+                    if not isinstance(expr, exp.Column):
+                        raise SQLMeshError(
+                            f"PARTITIONED BY contains non-column value '{expr.sql(dialect=self.dialect)}'."
+                        )
+
+                properties.append(
+                    exp.PartitionedByProperty(
+                        this=exp.Schema(expressions=partitioned_by),
+                    )
                 )
-
-            properties.append(property)
 
         if table_description:
             properties.append(
diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py
index 092b1ebf..7c06080b 100644
--- a/sqlmesh/core/table_diff.py
+++ b/sqlmesh/core/table_diff.py
@@ -5,7 +5,8 @@ import typing as t
 
 import pandas as pd
 from sqlglot import exp
-from sqlglot.dialects.dialect import DialectType
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+from sqlglot.optimizer.qualify_columns import quote_identifiers
 
 from sqlmesh.utils.pydantic import PydanticModel
 
@@ -106,7 +107,6 @@ class TableDiff:
         target: TableName,
         on: t.List[str] | exp.Condition,
         where: t.Optional[str | exp.Condition] = None,
-        dialect: DialectType = None,
         limit: int = 20,
         source_alias: t.Optional[str] = None,
         target_alias: t.Optional[str] = None,
@@ -115,24 +115,34 @@ class TableDiff:
         self.adapter = adapter
         self.source = source
         self.target = target
-        self.where = exp.condition(where, dialect=dialect) if where else None
+        self.dialect = adapter.dialect
+        self.where = exp.condition(where, dialect=self.dialect) if where else None
         self.limit = limit
+        self.model_name = model_name
+
         # Support environment aliases for diff output improvement in certain cases
         self.source_alias = source_alias
         self.target_alias = target_alias
-        self.model_name = model_name
 
         if isinstance(on, (list, tuple)):
+            s_table = exp.to_identifier("s", quoted=True)
+            t_table = exp.to_identifier("t", quoted=True)
+
             self.on: exp.Condition = exp.and_(
                 *(
-                    exp.column(c, "s").eq(exp.column(c, "t"))
-                    | (exp.column(c, "s").is_(exp.null()) & exp.column(c, "t").is_(exp.null()))
+                    exp.column(c, s_table).eq(exp.column(c, t_table))
+                    | (
+                        exp.column(c, s_table).is_(exp.null())
+                        & exp.column(c, t_table).is_(exp.null())
+                    )
                     for c in on
                 )
             )
         else:
             self.on = on
 
+        normalize_identifiers(self.on, dialect=self.dialect)
+
         self._source_schema: t.Optional[t.Dict[str, exp.DataType]] = None
         self._target_schema: t.Optional[t.Dict[str, exp.DataType]] = None
         self._row_diff: t.Optional[RowDiff] = None
@@ -234,14 +244,18 @@ class TableDiff:
                 .where(self.where)
             )
 
-            with self.adapter.temp_table(query, name="sqlmesh_temp.diff") as table:
+            query = quote_identifiers(query, dialect=self.dialect)
+            temp_table = exp.table_("diff", db="sqlmesh_temp", quoted=True)
+
+            with self.adapter.temp_table(query, name=temp_table) as table:
                 summary_query = exp.select(
                     exp.func("SUM", "s_exists").as_("s_count"),
                     exp.func("SUM", "t_exists").as_("t_count"),
                     exp.func("SUM", "rows_joined").as_("join_count"),
                     *(exp.func("SUM", name(c)).as_(c.alias) for c in comparisons),
                 ).from_(table)
-                stats_df = self.adapter.fetchdf(summary_query)
+
+                stats_df = self.adapter.fetchdf(summary_query, quote_identifiers=True)
                 stats_df["s_only_count"] = stats_df["s_count"] - stats_df["join_count"]
                 stats_df["t_only_count"] = stats_df["t_count"] - stats_df["join_count"]
                 stats = stats_df.iloc[0].to_dict()
@@ -261,7 +275,7 @@ class TableDiff:
                     .where(exp.column("rows_joined").eq(exp.Literal.number(1)))
                 )
                 column_stats = (
-                    self.adapter.fetchdf(column_stats_query)
+                    self.adapter.fetchdf(column_stats_query, quote_identifiers=True)
                     .T.rename(
                         columns={0: "pct_match"},
                         index=lambda x: str(x).replace("_matches", "") if x else "",
@@ -284,7 +298,7 @@ class TableDiff:
                     )
                     .limit(self.limit)
                 )
-                sample = self.adapter.fetchdf(sample_query)
+                sample = self.adapter.fetchdf(sample_query, quote_identifiers=True)
 
                 joined_sample_cols = [f"s__{c}" for c in index_cols]
                 comparison_cols = [


Test patch:
diff --git a/tests/core/engine_adapter/test_trino.py b/tests/core/engine_adapter/test_trino.py
index 51b4f16d..925c581f 100644
--- a/tests/core/engine_adapter/test_trino.py
+++ b/tests/core/engine_adapter/test_trino.py
@@ -5,7 +5,10 @@ import pytest
 from pytest_mock.plugin import MockerFixture
 from sqlglot import exp, parse_one
 
+import sqlmesh.core.dialect as d
 from sqlmesh.core.engine_adapter import TrinoEngineAdapter
+from sqlmesh.core.model import load_sql_based_model
+from sqlmesh.core.model.definition import SqlModel
 from tests.core.engine_adapter import to_sql_calls
 
 pytestmark = [pytest.mark.engine, pytest.mark.trino]
@@ -148,26 +151,36 @@ def test_partitioned_by_iceberg_transforms(
         return_value="datalake_iceberg",
     )
 
-    columns_to_types = {
-        "cola": exp.DataType.build("INT"),
-        "colb": exp.DataType.build("TEXT"),
-    }
+    expressions = d.parse(
+        f"""
+        MODEL (
+            name test_table,
+            partitioned_by (day(cola), truncate(colb, 8), colc),
+            kind INCREMENTAL_BY_TIME_RANGE(
+                time_column cola,
+            ),
+        );
+
+        SELECT 1::timestamp AS cola, 2::varchar as colb, 'foo' as colc;
+    """
+    )
+    model: SqlModel = t.cast(SqlModel, load_sql_based_model(expressions))
 
     adapter.create_table(
-        "test_table",
-        columns_to_types,
-        partitioned_by=[exp.to_column("day(cola)"), exp.to_column("truncate(colb, 8)")],
+        table_name=model.view_name,
+        columns_to_types=model.columns_to_types_or_raise,
+        partitioned_by=model.partitioned_by,
     )
 
     adapter.ctas(
-        "test_table",
-        parse_one("select 1"),  # type: ignore
-        partitioned_by=[exp.to_column("day(cola)"), exp.to_column("truncate(colb, 8)")],
+        table_name=model.view_name,
+        query_or_df=t.cast(exp.Query, model.query),
+        partitioned_by=model.partitioned_by,
     )
 
     assert to_sql_calls(adapter) == [
-        """CREATE TABLE IF NOT EXISTS "test_table" ("cola" INTEGER, "colb" VARCHAR) WITH (PARTITIONING=ARRAY['day(cola)', 'truncate(colb, 8)'])""",
-        """CREATE TABLE IF NOT EXISTS "test_table" WITH (PARTITIONING=ARRAY['day(cola)', 'truncate(colb, 8)']) AS SELECT 1""",
+        """CREATE TABLE IF NOT EXISTS "test_table" ("cola" TIMESTAMP, "colb" VARCHAR, "colc" VARCHAR) WITH (PARTITIONING=ARRAY['DAY("cola")', 'TRUNCATE("colb", 8)', '"colc"'])""",
+        """CREATE TABLE IF NOT EXISTS "test_table" WITH (PARTITIONING=ARRAY['DAY("cola")', 'TRUNCATE("colb", 8)', '"colc"']) AS SELECT CAST(1 AS TIMESTAMP) AS "cola", CAST(2 AS VARCHAR) AS "colb", \'foo\' AS "colc\"""",
     ]
 
 
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index ffa5eaaa..9d91956e 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -245,17 +245,20 @@ def test_model_qualification():
 
 
 @pytest.mark.parametrize(
-    "partition_by_input, partition_by_output, expected_exception",
+    "partition_by_input, partition_by_output, output_dialect, expected_exception",
     [
-        ("a", ["`a`"], None),
-        ("(a, b)", ["`a`", "`b`"], None),
-        ("TIMESTAMP_TRUNC(`a`, DAY)", ["TIMESTAMP_TRUNC(`a`, DAY)"], None),
-        ("e", "", ConfigError),
-        ("(a, e)", "", ConfigError),
-        ("(a, a)", "", ConfigError),
+        ("a", ["`a`"], "bigquery", None),
+        ("(a, b)", ["`a`", "`b`"], "bigquery", None),
+        ("TIMESTAMP_TRUNC(`a`, DAY)", ["TIMESTAMP_TRUNC(`a`, DAY)"], "bigquery", None),
+        ("e", "", "bigquery", ConfigError),
+        ("(a, e)", "", "bigquery", ConfigError),
+        ("(a, a)", "", "bigquery", ConfigError),
+        ("(day(a),b)", ['DAY("a")', '"b"'], "trino", None),
     ],
 )
-def test_partitioned_by(partition_by_input, partition_by_output, expected_exception):
+def test_partitioned_by(
+    partition_by_input, partition_by_output, output_dialect, expected_exception
+):
     expressions = d.parse(
         f"""
         MODEL (
@@ -280,7 +283,9 @@ def test_partitioned_by(partition_by_input, partition_by_output, expected_except
             model.validate_definition()
     else:
         model.validate_definition()
-        assert [col.sql(dialect="bigquery") for col in model.partitioned_by] == partition_by_output
+        assert [
+            col.sql(dialect=output_dialect) for col in model.partitioned_by
+        ] == partition_by_output
 
 
 def test_no_model_statement():