Fix: Support insert overwrite with dynamic partitions in the BigQuery adapter (#1188)

Author: izeigerman

setup.py

@@ -45,7 +45,7 @@
         "requests",
         "rich",
         "ruamel.yaml",
-        "sqlglot~=17.4.0",
+        "sqlglot~=17.6.0",
         "fsspec",
     ],
     extras_require={

sqlmesh/core/engine_adapter/base.py

@@ -592,10 +592,11 @@ def _insert_append_pandas_df(
                 ):
                     self.execute(exp.insert(expression, table_name, columns=column_names))
 
-    def insert_overwrite(
+    def insert_overwrite_by_partition(
         self,
         table_name: TableName,
         query_or_df: QueryOrDF,
+        partitioned_by: t.List[exp.Expression],
         columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
     ) -> None:
         self._insert_overwrite_by_condition(
@@ -905,7 +906,7 @@ def execute(
             self.cursor.execute(sql, **kwargs)
 
     @contextlib.contextmanager
-    def temp_table(self, query_or_df: QueryOrDF, name: str = "diff") -> t.Iterator[exp.Table]:
+    def temp_table(self, query_or_df: QueryOrDF, name: TableName = "diff") -> t.Iterator[exp.Table]:
         """A context manager for working a temp table.
 
         The table will be created with a random guid and cleaned up after the block.

sqlmesh/core/engine_adapter/bigquery.py

@@ -354,6 +354,46 @@ def merge(
         ) as source_table:
             return super().merge(target_table, source_table, columns_to_types, unique_key)
 
+    def insert_overwrite_by_partition(
+        self,
+        table_name: TableName,
+        query_or_df: QueryOrDF,
+        partitioned_by: t.List[exp.Expression],
+        columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None,
+    ) -> None:
+        if len(partitioned_by) != 1:
+            raise SQLMeshError(
+                f"Bigquery only supports partitioning by one column, {len(partitioned_by)} were provided."
+            )
+
+        partition_exp = partitioned_by[0]
+        partition_sql = partition_exp.sql(dialect=self.dialect)
+        partition_column = partition_exp.find(exp.Column)
+
+        if not partition_column:
+            raise SQLMeshError(
+                f"The partition expression '{partition_sql}' doesn't contain a column."
+            )
+
+        with self.session(), self.temp_table(query_or_df, name=table_name) as temp_table_name:
+            if columns_to_types is None:
+                columns_to_types = self.columns(temp_table_name)
+
+            partition_type_sql = columns_to_types[partition_column.name].sql(dialect=self.dialect)
+            temp_table_name_sql = temp_table_name.sql(dialect=self.dialect)
+            self.execute(
+                f"DECLARE target_partitions ARRAY<{partition_type_sql}> DEFAULT (SELECT ARRAY_AGG(DISTINCT {partition_sql}) FROM {temp_table_name_sql});"
+            )
+
+            where = t.cast(exp.Condition, partition_exp).isin(unnest="target_partitions")
+
+            self._insert_overwrite_by_condition(
+                table_name,
+                exp.select("*").from_(temp_table_name),
+                where=where,
+                columns_to_types=columns_to_types,
+            )
+
     def _insert_overwrite_by_condition(
         self,
         table_name: TableName,

sqlmesh/core/snapshot/evaluator.py

@@ -800,8 +800,8 @@ def insert(
         **kwargs: t.Any,
     ) -> None:
         if isinstance(model.kind, IncrementalUnmanagedKind) and model.kind.insert_overwrite:
-            self.adapter.insert_overwrite(
-                name, query_or_df, columns_to_types=model.columns_to_types
+            self.adapter.insert_overwrite_by_partition(
+                name, query_or_df, model.partitioned_by, columns_to_types=model.columns_to_types
             )
         else:
             self.append(model, name, query_or_df, snapshots, is_dev, **kwargs)

tests/core/engine_adapter/test_bigquery.py

@@ -1,5 +1,7 @@
 # type: ignore
 import sys
+import typing as t
+import uuid
 
 import pandas as pd
 import pytest
@@ -35,18 +37,48 @@ def test_insert_overwrite_by_time_partition_query(mocker: MockerFixture):
             "ds": exp.DataType.build("string"),
         },
     )
-    sql_calls = [
-        # Python 3.7 support
-        call[0][0].sql(dialect="bigquery", identify=True)
-        if isinstance(call[0], tuple)
-        else call[0].sql(dialect="bigquery", identify=True)
-        for call in execute_mock.call_args_list
-    ]
+    sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == [
         "MERGE INTO `test_table` AS `__MERGE_TARGET__` USING (SELECT * FROM (SELECT `a`, `ds` FROM `tbl`) AS `_subquery` WHERE `ds` BETWEEN '2022-01-01' AND '2022-01-05') AS `__MERGE_SOURCE__` ON FALSE WHEN NOT MATCHED BY SOURCE AND `ds` BETWEEN '2022-01-01' AND '2022-01-05' THEN DELETE WHEN NOT MATCHED THEN INSERT (`a`, `ds`) VALUES (`a`, `ds`)"
     ]
 
 
+def test_insert_overwrite_by_partition_query(mocker: MockerFixture):
+    connection_mock = mocker.NonCallableMock()
+    cursor_mock = mocker.Mock()
+    connection_mock.cursor.return_value = cursor_mock
+
+    adapter = BigQueryEngineAdapter(lambda: connection_mock)
+    execute_mock = mocker.patch(
+        "sqlmesh.core.engine_adapter.bigquery.BigQueryEngineAdapter.execute"
+    )
+
+    temp_table_uuid = uuid.uuid4()
+    uuid4_mock = mocker.patch("uuid.uuid4")
+    uuid4_mock.return_value = temp_table_uuid
+
+    adapter.insert_overwrite_by_partition(
+        "test_schema.test_table",
+        parse_one("SELECT a, ds FROM tbl"),
+        partitioned_by=[
+            d.parse_one("DATETIME_TRUNC(ds, MONTH)"),
+        ],
+        columns_to_types={
+            "a": exp.DataType.build("int"),
+            "ds": exp.DataType.build("DATETIME"),
+        },
+    )
+
+    sql_calls = _to_sql_calls(execute_mock)
+    assert sql_calls == [
+        "CREATE SCHEMA IF NOT EXISTS `test_schema`",
+        f"CREATE TABLE IF NOT EXISTS `test_schema`.`__temp_test_table_{temp_table_uuid.hex}` AS SELECT `a`, `ds` FROM `tbl`",
+        f"DECLARE target_partitions ARRAY<DATETIME> DEFAULT (SELECT ARRAY_AGG(DISTINCT DATETIME_TRUNC(ds, MONTH)) FROM test_schema.__temp_test_table_{temp_table_uuid.hex});",
+        f"MERGE INTO `test_schema`.`test_table` AS `__MERGE_TARGET__` USING (SELECT * FROM (SELECT * FROM `test_schema`.`__temp_test_table_{temp_table_uuid.hex}`) AS `_subquery` WHERE DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`target_partitions`)) AS `__MERGE_SOURCE__` ON FALSE WHEN NOT MATCHED BY SOURCE AND DATETIME_TRUNC(`ds`, MONTH) IN UNNEST(`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_uuid.hex}`",
+    ]
+
+
 def test_insert_overwrite_by_time_partition_pandas(mocker: MockerFixture):
     connection_mock = mocker.NonCallableMock()
     cursor_mock = mocker.Mock()
@@ -130,13 +162,7 @@ def test_replace_query(mocker: MockerFixture):
     )
     adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
 
-    sql_calls = [
-        # Python 3.7 support
-        call[0][0].sql(dialect="bigquery", identify=True)
-        if isinstance(call[0], tuple)
-        else call[0].sql(dialect="bigquery", identify=True)
-        for call in execute_mock.call_args_list
-    ]
+    sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == ["CREATE OR REPLACE TABLE `test_table` AS SELECT `a` FROM `tbl`"]
 
 
@@ -221,13 +247,7 @@ def test_create_table_date_partition(
         clustered_by=["b"],
     )
 
-    sql_calls = [
-        # Python 3.7 support
-        call[0][0].sql(dialect="bigquery", identify=True)
-        if isinstance(call[0], tuple)
-        else call[0].sql(dialect="bigquery", identify=True)
-        for call in execute_mock.call_args_list
-    ]
+    sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == [
         f"CREATE TABLE IF NOT EXISTS `test_table` (`a` int, `b` int) PARTITION BY {partition_by_statement} CLUSTER BY `b`"
     ]
@@ -261,13 +281,7 @@ def test_create_table_time_partition(
         partition_interval_unit=IntervalUnit.HOUR,
     )
 
-    sql_calls = [
-        # Python 3.7 support
-        call[0][0].sql(dialect="bigquery", identify=True)
-        if isinstance(call[0], tuple)
-        else call[0].sql(dialect="bigquery", identify=True)
-        for call in execute_mock.call_args_list
-    ]
+    sql_calls = _to_sql_calls(execute_mock)
     assert sql_calls == [
         f"CREATE TABLE IF NOT EXISTS `test_table` (`a` int, `b` int) PARTITION BY {partition_by_statement}"
     ]
@@ -292,13 +306,7 @@ def test_merge(mocker: MockerFixture):
         },
         unique_key=["id"],
     )
-    sql_calls = [
-        # Python 3.7 support
-        call[0][0].sql(dialect="bigquery")
-        if isinstance(call[0], tuple)
-        else call[0].sql(dialect="bigquery")
-        for call in execute_mock.call_args_list
-    ]
+    sql_calls = _to_sql_calls(execute_mock, identify=False)
     assert sql_calls == [
         "MERGE INTO target AS __MERGE_TARGET__ USING (SELECT id, ts, val FROM source) AS __MERGE_SOURCE__ ON __MERGE_TARGET__.id = __MERGE_SOURCE__.id "
         "WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.id = __MERGE_SOURCE__.id, __MERGE_TARGET__.ts = __MERGE_SOURCE__.ts, __MERGE_TARGET__.val = __MERGE_SOURCE__.val "
@@ -335,13 +343,7 @@ def test_merge(mocker: MockerFixture):
         unique_key=["id"],
     )
 
-    sql_calls = [
-        # Python 3.7 support
-        call[0][0].sql(dialect="bigquery")
-        if isinstance(call[0], tuple)
-        else call[0].sql(dialect="bigquery")
-        for call in execute_mock.call_args_list
-    ]
+    sql_calls = _to_sql_calls(execute_mock, identify=False)
     assert sql_calls == [
         "MERGE INTO target AS __MERGE_TARGET__ USING (SELECT id, ts, val FROM project.dataset.temp_table) AS __MERGE_SOURCE__ ON __MERGE_TARGET__.id = __MERGE_SOURCE__.id "
         "WHEN MATCHED THEN UPDATE SET __MERGE_TARGET__.id = __MERGE_SOURCE__.id, __MERGE_TARGET__.ts = __MERGE_SOURCE__.ts, __MERGE_TARGET__.val = __MERGE_SOURCE__.val "
@@ -399,3 +401,17 @@ def test_begin_end_session(mocker: MockerFixture):
     execute_b_call = connection_mock._client.query.call_args_list[2]
     assert execute_b_call[1]["query"] == "SELECT 3;"
     assert not execute_b_call[1]["job_config"].connection_properties
+
+
+def _to_sql_calls(execute_mock: t.Any, identify: bool = True) -> t.List[str]:
+    output = []
+    for call in execute_mock.call_args_list:
+        # Python 3.7 support
+        value = call[0][0] if isinstance(call[0], tuple) else call[0]
+        sql = (
+            value.sql(dialect="bigquery", identify=identify)
+            if isinstance(value, exp.Expression)
+            else str(value)
+        )
+        output.append(sql)
+    return output

tests/core/test_snapshot_evaluator.py

@@ -277,6 +277,7 @@ def test_evaluate_incremental_unmanaged(
         name="test_schema.test_model",
         query=parse_one("SELECT 1, ds FROM tbl_a"),
         kind=IncrementalUnmanagedKind(insert_overwrite=insert_overwrite),
+        partitioned_by=["ds"],
     )
     snapshot = make_snapshot(model)
     snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
@@ -290,14 +291,19 @@ def test_evaluate_incremental_unmanaged(
         snapshots={},
     )
 
-    expected_call = (
-        adapter_mock.insert_overwrite if insert_overwrite else adapter_mock.insert_append
-    )
-    expected_call.assert_called_once_with(
-        snapshot.table_name(),
-        model.render_query(),
-        columns_to_types=model.columns_to_types,
-    )
+    if insert_overwrite:
+        adapter_mock.insert_overwrite_by_partition.assert_called_once_with(
+            snapshot.table_name(),
+            model.render_query(),
+            [exp.to_column("ds")],
+            columns_to_types=model.columns_to_types,
+        )
+    else:
+        adapter_mock.insert_append.assert_called_once_with(
+            snapshot.table_name(),
+            model.render_query(),
+            columns_to_types=model.columns_to_types,
+        )
 
 
 def test_create_materialized_view(mocker: MockerFixture, adapter_mock, make_snapshot):