Fix: correctly extract catalog name for table properties expression builder (#2196)

* Extract and pass catalog name to partition exp builder

* Add tests that pass a table to _build_create_table_exp

* Ensure schema.this is a table expression

Author: treysp

sqlmesh/core/engine_adapter/base.py

@@ -688,11 +688,17 @@ def _build_create_table_exp(
         **kwargs: t.Any,
     ) -> exp.Create:
         exists = False if replace else exists
+        catalog_name = None
         if not isinstance(table_name_or_schema, exp.Schema):
             table_name_or_schema = exp.to_table(table_name_or_schema)
+            catalog_name = table_name_or_schema.catalog
+        else:
+            if isinstance(table_name_or_schema.this, exp.Table):
+                catalog_name = table_name_or_schema.this.catalog
+
         properties = (
             self._build_table_properties_exp(
-                **kwargs, table=table_name_or_schema.this, columns_to_types=columns_to_types
+                **kwargs, catalog_name=catalog_name, columns_to_types=columns_to_types
             )
             if kwargs
             else None
@@ -1864,7 +1870,7 @@ def _table_properties_to_expressions(
 
     def _build_table_properties_exp(
         self,
-        table: exp.Table,
+        catalog_name: t.Optional[str] = None,
         storage_format: t.Optional[str] = None,
         partitioned_by: t.Optional[t.List[exp.Expression]] = None,
         partition_interval_unit: t.Optional[IntervalUnit] = None,

sqlmesh/core/engine_adapter/bigquery.py

@@ -459,7 +459,7 @@ def _build_description_property_exp(self, description: str) -> exp.Property:
 
     def _build_table_properties_exp(
         self,
-        table: exp.Table,
+        catalog_name: t.Optional[str] = None,
         storage_format: t.Optional[str] = None,
         partitioned_by: t.Optional[t.List[exp.Expression]] = None,
         partition_interval_unit: t.Optional[IntervalUnit] = None,

sqlmesh/core/engine_adapter/mixins.py

@@ -128,7 +128,7 @@ def _insert_overwrite_by_condition(
 class HiveMetastoreTablePropertiesMixin(EngineAdapter):
     def _build_table_properties_exp(
         self,
-        table: exp.Table,
+        catalog_name: t.Optional[str] = None,
         storage_format: t.Optional[str] = None,
         partitioned_by: t.Optional[t.List[exp.Expression]] = None,
         partition_interval_unit: t.Optional[IntervalUnit] = None,
@@ -155,7 +155,7 @@ def _build_table_properties_exp(
 
             if (
                 self.dialect == "trino"
-                and self.get_catalog_type(table.catalog or self.get_current_catalog()) == "iceberg"
+                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

tests/core/engine_adapter/test_trino.py

@@ -2,7 +2,7 @@
 
 import pytest
 from pytest_mock.plugin import MockerFixture
-from sqlglot import exp
+from sqlglot import exp, parse_one
 
 from sqlmesh.core.engine_adapter import TrinoEngineAdapter
 from tests.core.engine_adapter import to_sql_calls
@@ -73,8 +73,11 @@ def test_partitioned_by_hive(
 
     adapter.create_table("test_table", columns_to_types, partitioned_by=[exp.to_column("colb")])
 
+    adapter.ctas("test_table", parse_one("select 1"), partitioned_by=[exp.to_column("colb")])  # type: ignore
+
     assert to_sql_calls(adapter) == [
-        """CREATE TABLE IF NOT EXISTS "test_table" ("cola" INTEGER, "colb" VARCHAR) WITH (PARTITIONED_BY=ARRAY['colb'])"""
+        """CREATE TABLE IF NOT EXISTS "test_table" ("cola" INTEGER, "colb" VARCHAR) WITH (PARTITIONED_BY=ARRAY['colb'])""",
+        """CREATE TABLE IF NOT EXISTS "test_table" WITH (PARTITIONED_BY=ARRAY['colb']) AS SELECT 1""",
     ]
 
 
@@ -96,8 +99,11 @@ def test_partitioned_by_iceberg(
 
     adapter.create_table("test_table", columns_to_types, partitioned_by=[exp.to_column("colb")])
 
+    adapter.ctas("test_table", parse_one("select 1"), partitioned_by=[exp.to_column("colb")])  # type: ignore
+
     assert to_sql_calls(adapter) == [
-        """CREATE TABLE IF NOT EXISTS "test_table" ("cola" INTEGER, "colb" VARCHAR) WITH (PARTITIONING=ARRAY['colb'])"""
+        """CREATE TABLE IF NOT EXISTS "test_table" ("cola" INTEGER, "colb" VARCHAR) WITH (PARTITIONING=ARRAY['colb'])""",
+        """CREATE TABLE IF NOT EXISTS "test_table" WITH (PARTITIONING=ARRAY['colb']) AS SELECT 1""",
     ]
 
 
@@ -122,8 +128,15 @@ def test_partitioned_by_iceberg_transforms(
         partitioned_by=[exp.to_column("day(cola)"), exp.to_column("truncate(colb, 8)")],
     )
 
+    adapter.ctas(
+        "test_table",
+        parse_one("select 1"),  # type: ignore
+        partitioned_by=[exp.to_column("day(cola)"), exp.to_column("truncate(colb, 8)")],
+    )
+
     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" ("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""",
     ]
 
 
@@ -141,13 +154,27 @@ def test_partitioned_by_with_multiple_catalogs_same_server(
         "datalake.test_schema.test_table", columns_to_types, partitioned_by=[exp.to_column("colb")]
     )
 
+    adapter.ctas(
+        "datalake.test_schema.test_table",
+        parse_one("select 1"),  # type: ignore
+        partitioned_by=[exp.to_column("colb")],
+    )
+
     adapter.create_table(
         "datalake_iceberg.test_schema.test_table",
         columns_to_types,
         partitioned_by=[exp.to_column("colb")],
     )
 
+    adapter.ctas(
+        "datalake_iceberg.test_schema.test_table",
+        parse_one("select 1"),  # type: ignore
+        partitioned_by=[exp.to_column("colb")],
+    )
+
     assert to_sql_calls(adapter) == [
         """CREATE TABLE IF NOT EXISTS "datalake"."test_schema"."test_table" ("cola" INTEGER, "colb" VARCHAR) WITH (PARTITIONED_BY=ARRAY['colb'])""",
+        """CREATE TABLE IF NOT EXISTS "datalake"."test_schema"."test_table" WITH (PARTITIONED_BY=ARRAY['colb']) AS SELECT 1""",
         """CREATE TABLE IF NOT EXISTS "datalake_iceberg"."test_schema"."test_table" ("cola" INTEGER, "colb" VARCHAR) WITH (PARTITIONING=ARRAY['colb'])""",
+        """CREATE TABLE IF NOT EXISTS "datalake_iceberg"."test_schema"."test_table" WITH (PARTITIONING=ARRAY['colb']) AS SELECT 1""",
     ]