Repository: TobikoData/sqlmesh
Pull request #3056 (2024-08-27 15:57:39)

Fix(bigquery): do not use execution_project as default catalog in partitions query
Biquery "projects" are associated with both object location (catalog in 3-level names) and with computation (computations are executed by a specific project).

Our BQ connection configuration provides an [`execution_project` field](#2181) that allows users to execute the computations with a different project than the location of the objects.

Bigquery's default catalog is the user-specified connection field `project`, if present. Otherwise, it is the project to which the BQ client is connected (which is `execution_project` when present).

In the `incremental_by_partition` model kind, we manually construct a query of the BQ partitions metadata view. Because that query is manually constructed it does not pass through the `@set_catalog()` decorator and was incorrectly using the client connection `execution_project` instead of the user-specified `project` in the object name.


Patch:
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 85f71c09..eb4b8880 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -27,6 +27,7 @@ from sqlmesh.utils.pydantic import (
     field_validator,
     model_validator,
     model_validator_v1_args,
+    field_validator_v1_args,
 )
 
 if sys.version_info >= (3, 9):
@@ -740,6 +741,19 @@ class BigQueryConnectionConfig(ConnectionConfig):
 
     type_: Literal["bigquery"] = Field(alias="type", default="bigquery")
 
+    @field_validator("execution_project")
+    @field_validator_v1_args
+    def validate_execution_project(
+        cls,
+        v: t.Optional[str],
+        values: t.Dict[str, t.Any],
+    ) -> t.Optional[str]:
+        if v and not values.get("project"):
+            raise ConfigError(
+                "If the `execution_project` field is specified, you must also specify the `project` field to provide a default object location."
+            )
+        return v
+
     @property
     def _connection_kwargs_keys(self) -> t.Set[str]:
         return set()
diff --git a/sqlmesh/core/engine_adapter/bigquery.py b/sqlmesh/core/engine_adapter/bigquery.py
index ca893a27..fdae3a4e 100644
--- a/sqlmesh/core/engine_adapter/bigquery.py
+++ b/sqlmesh/core/engine_adapter/bigquery.py
@@ -331,7 +331,7 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
         table_ = exp.to_table(table).copy()
 
         if not table_.catalog:
-            table_.set("catalog", exp.to_identifier(self.client.project))
+            table_.set("catalog", exp.to_identifier(self.default_catalog))
 
         return bigquery.Table(
             table_ref=self._table_name(table_),
@@ -389,6 +389,7 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
                 partition_type_sql,
                 granularity=granularity,
                 agg_func="ARRAY_AGG",
+                catalog=temp_table_name.catalog or self.default_catalog,
             )
 
             self.execute(
@@ -675,7 +676,7 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
         # The BigQuery Client's list_tables method does not support filtering by table name, so we have to
         # resort to using SQL instead.
         schema = to_schema(schema_name)
-        catalog = schema.catalog or self.get_current_catalog()
+        catalog = schema.catalog or self.default_catalog
         query = exp.select(
             exp.column("table_catalog").as_("catalog"),
             exp.column("table_name").as_("name"),
@@ -795,7 +796,7 @@ def select_partitions_expr(
     data_type: t.Union[str, exp.DataType],
     granularity: t.Optional[str] = None,
     agg_func: str = "MAX",
-    database: t.Optional[str] = None,
+    catalog: t.Optional[str] = None,
 ) -> str:
     """Generates a SQL expression that aggregates partition values for a table.
 
@@ -805,14 +806,14 @@ def select_partitions_expr(
         data_type: The data type of the partition column.
         granularity: The granularity of the partition. Supported values are: 'day', 'month', 'year' and 'hour'.
         agg_func: The aggregation function to use.
-        database: The database (BigQuery project ID) of the table.
+        catalog: The catalog (BigQuery project ID) of the table.
 
     Returns:
         A SELECT statement that aggregates partition values for a table.
     """
     partitions_table_name = f"`{schema}`.INFORMATION_SCHEMA.PARTITIONS"
-    if database:
-        partitions_table_name = f"`{database}`.{partitions_table_name}"
+    if catalog:
+        partitions_table_name = f"`{catalog}`.{partitions_table_name}"
 
     if isinstance(data_type, exp.DataType):
         data_type = data_type.sql(dialect="bigquery")
diff --git a/sqlmesh/dbt/model.py b/sqlmesh/dbt/model.py
index bc9b50fb..2cf0c412 100644
--- a/sqlmesh/dbt/model.py
+++ b/sqlmesh/dbt/model.py
@@ -448,7 +448,7 @@ class ModelConfig(BaseModelConfig):
             "{{ adapter.resolve_identifier(this) }}",
             data_type,
             granularity=self.partition_by.get("granularity"),
-            database="{{ target.database }}",
+            catalog="{{ target.database }}",
         )
 
         data_type = data_type.upper()


Test patch:
diff --git a/tests/core/engine_adapter/test_bigquery.py b/tests/core/engine_adapter/test_bigquery.py
index caf86e51..ad6e6fbe 100644
--- a/tests/core/engine_adapter/test_bigquery.py
+++ b/tests/core/engine_adapter/test_bigquery.py
@@ -51,7 +51,7 @@ def test_insert_overwrite_by_partition_query(
     execute_mock = mocker.patch(
         "sqlmesh.core.engine_adapter.bigquery.BigQueryEngineAdapter.execute"
     )
-
+    adapter._default_catalog = "test_project"
     temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
     table_name = "test_schema.test_table"
     temp_table_id = "abcdefgh"
@@ -73,7 +73,7 @@ def test_insert_overwrite_by_partition_query(
     assert sql_calls == [
         "CREATE SCHEMA IF NOT EXISTS `test_schema`",
         f"CREATE TABLE IF NOT EXISTS `test_schema`.`__temp_test_table_{temp_table_id}` PARTITION BY DATETIME_TRUNC(`ds`, MONTH) AS SELECT `a`, `ds` FROM `tbl`",
-        f"DECLARE _sqlmesh_target_partitions_ ARRAY<DATETIME> DEFAULT (SELECT ARRAY_AGG(PARSE_DATETIME('%Y%m', partition_id)) FROM `test_schema`.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '__temp_test_table_{temp_table_id}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__');",
+        f"DECLARE _sqlmesh_target_partitions_ ARRAY<DATETIME> DEFAULT (SELECT ARRAY_AGG(PARSE_DATETIME('%Y%m', partition_id)) FROM `test_project`.`test_schema`.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '__temp_test_table_{temp_table_id}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__');",
         f"MERGE INTO `test_schema`.`test_table` AS `__MERGE_TARGET__` USING (SELECT `a`, `ds` FROM (SELECT * FROM `test_schema`.`__temp_test_table_{temp_table_id}`) AS `_subquery` WHERE DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`)) AS `__MERGE_SOURCE__` ON FALSE WHEN NOT MATCHED BY SOURCE AND DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`) THEN DELETE WHEN NOT MATCHED THEN INSERT (`a`, `ds`) VALUES (`a`, `ds`)",
         f"DROP TABLE IF EXISTS `test_schema`.`__temp_test_table_{temp_table_id}`",
     ]
@@ -94,7 +94,7 @@ def test_insert_overwrite_by_partition_query_unknown_column_types(
         "a": exp.DataType.build("int"),
         "ds": exp.DataType.build("DATETIME"),
     }
-
+    adapter._default_catalog = "test_project"
     temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
     table_name = "test_schema.test_table"
     temp_table_id = "abcdefgh"
@@ -120,7 +120,7 @@ def test_insert_overwrite_by_partition_query_unknown_column_types(
     assert sql_calls == [
         "CREATE SCHEMA IF NOT EXISTS `test_schema`",
         f"CREATE TABLE IF NOT EXISTS `test_schema`.`__temp_test_table_{temp_table_id}` PARTITION BY DATETIME_TRUNC(`ds`, MONTH) AS SELECT `a`, `ds` FROM `tbl`",
-        f"DECLARE _sqlmesh_target_partitions_ ARRAY<DATETIME> DEFAULT (SELECT ARRAY_AGG(PARSE_DATETIME('%Y%m', partition_id)) FROM `test_schema`.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '__temp_test_table_{temp_table_id}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__');",
+        f"DECLARE _sqlmesh_target_partitions_ ARRAY<DATETIME> DEFAULT (SELECT ARRAY_AGG(PARSE_DATETIME('%Y%m', partition_id)) FROM `test_project`.`test_schema`.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '__temp_test_table_{temp_table_id}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__');",
         f"MERGE INTO `test_schema`.`test_table` AS `__MERGE_TARGET__` USING (SELECT `a`, `ds` FROM (SELECT * FROM `test_schema`.`__temp_test_table_{temp_table_id}`) AS `_subquery` WHERE DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`)) AS `__MERGE_SOURCE__` ON FALSE WHEN NOT MATCHED BY SOURCE AND DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`_sqlmesh_target_partitions_`) THEN DELETE WHEN NOT MATCHED THEN INSERT (`a`, `ds`) VALUES (`a`, `ds`)",
         f"DROP TABLE IF EXISTS `test_schema`.`__temp_test_table_{temp_table_id}`",
     ]
@@ -651,7 +651,7 @@ def test_select_partitions_expr():
             "{{ adapter.resolve_identifier(this) }}",
             "date",
             granularity="day",
-            database="{{ target.database }}",
+            catalog="{{ target.database }}",
         )
         == "SELECT MAX(PARSE_DATE('%Y%m%d', partition_id)) FROM `{{ target.database }}`.`{{ adapter.resolve_schema(this) }}`.INFORMATION_SCHEMA.PARTITIONS WHERE table_name = '{{ adapter.resolve_identifier(this) }}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__'"
     )
diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py
index 6fb158ce..249eefc1 100644
--- a/tests/core/test_connection_config.py
+++ b/tests/core/test_connection_config.py
@@ -584,6 +584,9 @@ def test_bigquery(make_config):
     assert config.get_catalog() == "project"
     assert config.is_recommended_for_state_sync is False
 
+    with pytest.raises(ConfigError, match="you must also specify the `project` field"):
+        make_config(type="bigquery", execution_project="execution_project")
+
 
 def test_postgres(make_config):
     config = make_config(