Fix!: bigquery create table requires correct trunc method for column type (#1048)

* partitioned_by config field removed for dbt. Use partition_by field instead.

Author: crericha

sqlmesh/dbt/model.py

@@ -44,8 +44,6 @@ class ModelConfig(BaseModelConfig):
     Args:
         sql: The model sql
         time_column: The name of the time column
-        partitioned_by: List of columns to partition by. time_column will automatically be
-            included, if specified.
         cron: A cron string specifying how often the model should be refreshed, leveraging the
             [croniter](https://github.com/kiorky/croniter) library.
         dialect: The SQL dialect that the model's query is written in. By default,
@@ -60,14 +58,12 @@ class ModelConfig(BaseModelConfig):
         materialized: How the model will be materialized in the database
         sql_header: SQL statement to inject above create table/view as
         unique_key: List of columns that define row uniqueness for the model
-        partition_by: Dictionary of bigquery partition by parameters ([dbt bigquery config](https://docs.getdbt.com/reference/resource-configs/bigquery-configs)).
-            If partitioned_by is set, this field will be ignored.
+        partition_by: List of partition columns or dictionary of bigquery partition by parameters ([dbt bigquery config](https://docs.getdbt.com/reference/resource-configs/bigquery-configs)).
     """
 
     # sqlmesh fields
     sql: SqlStr = SqlStr("")
     time_column: t.Optional[str] = None
-    partitioned_by: t.Optional[t.List[str]] = None
     cron: t.Optional[str] = None
     dialect: t.Optional[str] = None
     batch_size: t.Optional[int] = None
@@ -92,7 +88,6 @@ class ModelConfig(BaseModelConfig):
     @validator(
         "unique_key",
         "cluster_by",
-        "partitioned_by",
         pre=True,
     )
     def _validate_list(cls, v: t.Union[str, t.List[str]]) -> t.List[str]:
@@ -111,9 +106,7 @@ def _validate_partition_by(cls, v: t.Any) -> t.Union[t.List[str], t.Dict[str, t.
         if isinstance(v, dict):
             if not v.get("field"):
                 raise ConfigError("'field' key required for partition_by.")
-            if not v.get("granularity"):
-                v["granularity"] = "day"
-            return v
+            return {"data_type": "date", "granularity": "day", **v}
         raise ConfigError(f"Invalid format for partition_by '{v}'")
 
     _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
@@ -204,28 +197,37 @@ def _extract_sql_config(self) -> None:
         self._sql_no_config = SqlStr(no_config)
         self._sql_embedded_config = SqlStr(embedded_config)
 
+    @property
+    def _big_query_partition_by_expr(self) -> exp.Expression:
+        assert isinstance(self.partition_by, dict)
+        data_type = self.partition_by["data_type"].lower()
+        if data_type == "int64" or (
+            data_type == "date" and self.partition_by["granularity"].lower() == "day"
+        ):
+            return exp.to_column(self.partition_by["field"])
+
+        return TIME_TYPE_TO_TRUNC_EXPR[data_type](
+            this=exp.to_column(self.partition_by["field"]),
+            unit=exp.var(self.partition_by["granularity"].upper()),
+        )
+
     def to_sqlmesh(self, context: DbtContext) -> Model:
         """Converts the dbt model into a SQLMesh model."""
         dialect = self.model_dialect or context.dialect
         query = d.jinja_query(self.sql_no_config)
 
         optional_kwargs: t.Dict[str, t.Any] = {}
 
-        if self.partitioned_by:
-            optional_kwargs["partitioned_by"] = [
-                d.parse_one(val, dialect=dialect) for val in self.partitioned_by
-            ]
-        elif self.partition_by and isinstance(self.partition_by, list):
-            optional_kwargs["partitioned_by"] = [exp.to_column(val) for val in self.partition_by]
-        elif self.partition_by and isinstance(self.partition_by, dict):
-            optional_kwargs["partitioned_by"] = [
-                exp.TimestampTrunc(
-                    this=exp.to_column(self.partition_by["field"]),
-                    unit=exp.var(self.partition_by["granularity"]),
-                )
-            ]
+        if self.partition_by:
+            optional_kwargs["partitioned_by"] = (
+                [exp.to_column(val) for val in self.partition_by]
+                if isinstance(self.partition_by, list)
+                else self._big_query_partition_by_expr
+            )
+
         if self.cluster_by:
             optional_kwargs["clustered_by"] = self.cluster_by
+
         for field in ["cron"]:
             field_val = getattr(self, field, None) or self.meta.get(field, None)
             if field_val:
@@ -243,3 +245,10 @@ def to_sqlmesh(self, context: DbtContext) -> Model:
             **optional_kwargs,
             **self.sqlmesh_model_kwargs(context),
         )
+
+
+TIME_TYPE_TO_TRUNC_EXPR = {
+    "date": exp.DateTrunc,
+    "datetime": exp.DatetimeTrunc,
+    "timestamp": exp.TimestampTrunc,
+}

tests/dbt/test_config.py

@@ -72,7 +72,7 @@ def test_model_to_sqlmesh_fields(sushi_test_project: Project):
         description="test model",
         sql="SELECT 1 AS a FROM foo",
         start="Jan 1 2023",
-        partitioned_by=["a"],
+        partition_by=["a"],
         cluster_by=["a"],
         cron="@hourly",
         batch_size=5,

tests/dbt/test_transformation.py

@@ -454,55 +454,49 @@ def test_parsetime_adapter_call(
 def test_partition_by(sushi_test_project: Project):
     context = sushi_test_project.context
     model_config = ModelConfig(
-        dialect="bigquery",
+        dialect="spark",
         name="model",
         schema="test",
         package_name="package",
         materialized="table",
         unique_key="ds",
-        partitioned_by="ds",
+        partition_by="ds",
         sql="""SELECT 1 AS one, ds, ts FROM foo""",
     )
     assert model_config.to_sqlmesh(context).partitioned_by == [exp.to_column("ds")]
 
-    model_config.partitioned_by = "DATE_TRUNC(ds, MONTH)"  # type: ignore
-    assert model_config.to_sqlmesh(context).partitioned_by == [
-        parse_one(model_config.partitioned_by[0], read="bigquery")  # type: ignore
-    ]
+    assert model_config.partition_by == ["ds"]
+    assert model_config.to_sqlmesh(context).partitioned_by == [exp.to_column("ds")]
 
-    model_config.partitioned_by = ["ds", "ts"]
+    model_config.partition_by = ["ds", "ts"]
     assert model_config.to_sqlmesh(context).partitioned_by == [
         exp.to_column("ds"),
         exp.to_column("ts"),
     ]
 
     model_config = ModelConfig(
-        dialect="spark",
+        dialect="bigquery",
         name="model",
         schema="test",
         package_name="package",
         materialized="table",
         unique_key="ds",
-        partition_by="ds",
+        partition_by={"field": "ds", "granularity": "month"},
         sql="""SELECT 1 AS one, ds FROM foo""",
     )
-    assert model_config.partition_by == ["ds"]
-    assert model_config.to_sqlmesh(context).partitioned_by == [exp.to_column("ds")]
+    assert (
+        model_config.to_sqlmesh(context).partitioned_by[0].sql(dialect="bigquery")
+        == "DATE_TRUNC(ds, MONTH)"
+    )
 
-    model_config.partition_by = ["ds"]
-    assert model_config.partition_by == ["ds"]
+    model_config.partition_by = {"field": "ds", "data_type": "timestamp", "granularity": "day"}
+    assert (
+        model_config.to_sqlmesh(context).partitioned_by[0].sql(dialect="bigquery")
+        == "TIMESTAMP_TRUNC(ds, DAY)"
+    )
+
+    model_config.partition_by = {"field": "ds", "data_type": "int64", "granularity": "day"}
     assert model_config.to_sqlmesh(context).partitioned_by == [exp.to_column("ds")]
 
-    model_config = ModelConfig(
-        dialect="bigquery",
-        name="model",
-        schema="test",
-        package_name="package",
-        materialized="table",
-        unique_key="ds",
-        partition_by={"field": "ds", "granularity": "month"},
-        sql="""SELECT 1 AS one, ds FROM foo""",
-    )
-    assert model_config.to_sqlmesh(context).partitioned_by == [
-        parse_one("TIMESTAMP_TRUNC(ds, MONTH)", read="bigquery")
-    ]
+    model_config.partition_by = {"field": "ds", "data_type": "date", "granularity": "day"}
+    assert model_config.to_sqlmesh(context).partitioned_by == [exp.to_column("ds")]