Adap 540/add copy partitions with static partitions (#1100)

Co-authored-by: Colin Rogers <[email protected]>

Author: VersusFacit

dbt-bigquery/.changes/unreleased/Fixes-20250514-185407.yaml

@@ -0,0 +1,6 @@
+kind: Fixes
+body: Add support for copy_partitions when used with static partitions.
+time: 2025-05-14T18:54:07.360083-07:00
+custom:
+    Author: versusfacit Kayrnt
+    Issue: "540"

dbt-bigquery/src/dbt/include/bigquery/macros/materializations/incremental_strategy/insert_overwrite.sql

@@ -40,62 +40,144 @@
 {% macro bq_insert_overwrite_sql(
     tmp_relation, target_relation, sql, unique_key, partition_by, partitions, dest_columns, tmp_relation_exists, copy_partitions
 ) %}
-  {% if partitions is not none and partitions != [] %} {# static #}
-      {{ bq_static_insert_overwrite_sql(tmp_relation, target_relation, sql, partition_by, partitions, dest_columns, tmp_relation_exists, copy_partitions) }}
-  {% else %} {# dynamic #}
-      {{ bq_dynamic_insert_overwrite_sql(tmp_relation, target_relation, sql, unique_key, partition_by, dest_columns, tmp_relation_exists, copy_partitions) }}
-  {% endif %}
+
+    {# static #}
+    {% if partitions is not none and partitions != [] %}
+        {{ bq_static_insert_overwrite_sql(tmp_relation, target_relation, sql, partition_by, partitions, dest_columns, tmp_relation_exists, copy_partitions) }}
+
+    {# dynamic #}
+    {% else %}
+        {{ bq_dynamic_insert_overwrite_sql(tmp_relation, target_relation, sql, unique_key, partition_by, dest_columns, tmp_relation_exists, copy_partitions) }}
+    {% endif %}
+
 {% endmacro %}
 
-{% macro bq_static_insert_overwrite_sql(
-    tmp_relation, target_relation, sql, partition_by, partitions, dest_columns, tmp_relation_exists, copy_partitions
+{#
+    -- static partitions refer to a fixed set of partition values that are
+    -- known ahead of time and do not depend on source data or runtime
+    -- conditions. these values are typically hardcoded or configured
+    -- externally (e.g., via dbt variables or macros).
+    --
+    -- with `insert_overwrite`, dbt uses a predefined partition filter
+    -- expression to overwrite specific partitions
+    -- (e.g., where partition_column = '2024-05-01').
+    -- it excels where using a fixed minimal partition list, no need to
+    -- query source data for partitions, and insert statements can be
+    -- batched. this also keeps things deterministic.
+#}
+
+{% macro bq_static_select_insert_overwrite_sql(tmp_relation, sql, partition_by, tmp_relation_exists) %}
+  {%- set source_sql -%}
+  (
+    {% if tmp_relation_exists %}
+      select
+        {% if partition_by.time_ingestion_partitioning %}
+          {{ partition_by.insertable_time_partitioning_field() }},
+        {% endif %}
+        * from {{ tmp_relation }}
+    {%- elif partition_by.time_ingestion_partitioning -%}
+      {{ wrap_with_time_ingestion_partitioning_sql(partition_by, sql, true) }}
+    {%- else -%}
+      {{ sql }}
+    {%- endif %}
+  )
+  {%- endset -%}
+  {{ return(source_sql) }}
+{% endmacro %}
+
+
+{#
+  -- Static-partition + `copy_partitions=true` should inline the literals / foldable constants
+  -- supplied via the `partitions=` config.
+
+  -- The inline method will still copy (truncate) a partition even
+  -- if the incremental run produced no rows for that date, because the user
+  -- explicitly listed it in `partitions=`.
+#}
+{% macro bq_static_copy_partitions_insert_overwrite_sql(
+  tmp_relation, target_relation, sql, partition_by, partitions, tmp_relation_exists
 ) %}
 
-      {% set predicate -%}
-          {{ partition_by.render_wrapped(alias='DBT_INTERNAL_DEST') }} in (
-              {{ partitions | join (', ') }}
-          )
-      {%- endset %}
+    {%- if tmp_relation_exists is false -%}
+        {%- set source_sql = bq_static_select_insert_overwrite_sql(tmp_relation, sql, partition_by, tmp_relation_exists) %}
 
-      {%- set source_sql -%}
-        (
-          {% if partition_by.time_ingestion_partitioning and tmp_relation_exists -%}
-          select
-            {{ partition_by.insertable_time_partitioning_field() }},
-            * from {{ tmp_relation }}
-          {% elif tmp_relation_exists -%}
-            select
-            * from {{ tmp_relation }}
-          {%- elif partition_by.time_ingestion_partitioning -%}
-            {{ wrap_with_time_ingestion_partitioning_sql(partition_by, sql, True) }}
-          {%- else -%}
-            {{sql}}
-          {%- endif %}
-
-        )
-      {%- endset -%}
+        {# -- we run temp table creation in a separate script to move to partitions copy if it doesn't already exist #}
+        {%- call statement('create_tmp_relation_for_copy', language='sql') -%}
+          {{ bq_create_table_as(partition_by, true, tmp_relation, source_sql, 'sql')
+        }}
+        {%- endcall %}
+    {%- endif -%}
 
-      {% if copy_partitions %}
-          {% do bq_copy_partitions(tmp_relation, target_relation, partitions, partition_by) %}
-      {% else %}
+    {%- set partitions_sql -%}
+	select
+	    cast(partition_literal as timestamp) as partition_ts
+	from unnest([
+	    {{ partitions | join(', ') }}
+	]) as partition_literal
+    {%- endset -%}
 
-      {#-- In case we're putting the model SQL _directly_ into the MERGE statement,
-         we need to prepend the MERGE statement with the user-configured sql_header,
-         which may be needed to resolve that model SQL (e.g. referencing a variable or UDF in the header)
-         in the "temporary table exists" case, we save the model SQL result as a temp table first, wherein the
-         sql_header is included by the create_table_as macro.
-      #}
-      -- 1. run the merge statement
-      {{ get_insert_overwrite_merge_sql(target_relation, source_sql, dest_columns, [predicate], include_sql_header = not tmp_relation_exists) }};
+    {%- set resolved_partitions = run_query(partitions_sql).columns[0].values() -%}
 
-      {%- if tmp_relation_exists -%}
-      -- 2. clean up the temp table
-      drop table if exists {{ tmp_relation }};
-      {%- endif -%}
+    {% do bq_copy_partitions(tmp_relation, target_relation, resolved_partitions, partition_by) %}
 
-  {% endif %}
+    {# clean up temp table #}
+    drop table if exists {{ tmp_relation }}
 {% endmacro %}
 
+
+{% macro bq_static_insert_overwrite_sql(
+    tmp_relation, target_relation, sql, partition_by, partitions, dest_columns, tmp_relation_exists, copy_partitions
+) %}
+
+    {%- if copy_partitions %}
+        {{ bq_static_copy_partitions_insert_overwrite_sql(tmp_relation, target_relation, sql, partition_by, partitions, tmp_relation_exists) }}
+    {% else -%}
+        {% set predicate -%}
+            {{ partition_by.render_wrapped(alias='dbt_internal_dest') }} in (
+                {{ partitions | join (', ') }}
+            )
+        {%- endset %}
+
+        {%- set source_sql = bq_static_select_insert_overwrite_sql(tmp_relation, sql, partition_by, tmp_relation_exists) %}
+	{#
+	  -- when the model sql is inserted directly into the merge statement,
+	  -- we need to prepend it with the user-defined `sql_header`. this is
+	  -- important when the model sql references elements like variables or udfs
+	  -- defined in the header.
+
+	  -- in the case where a temporary table is created first (i.e., the
+	  -- "temp table exists" path), the `sql_header` is already included via
+	  -- the `create_table_as` macro, so no additional handling is needed.
+	#}
+
+        -- 1. run the merge statement
+        {{ get_insert_overwrite_merge_sql(target_relation, source_sql, dest_columns, [predicate], include_sql_header = not tmp_relation_exists) }};
+
+        {%- if tmp_relation_exists -%}
+        -- 2. clean up the temp table
+        drop table if exists {{ tmp_relation }};
+        {%- endif -%}
+
+    {%- endif -%}
+{% endmacro %}
+
+
+{#
+    -- dynamic partitions refer to the set partition values that are
+    -- determined at runtime, based on either the contents of source
+    -- data (e.g. values in an updated_date column) or external runtime
+    -- parameters (e.g. time of day, system params).
+    --
+    -- with `insert_overwrite`, this allows dbt to compute which
+    -- partitions to overwrite dynamically using a partition filter
+    -- expression (e.g., where partition_column >= current_date()).
+    --
+    -- this enables targeted incremental updates by overwriting only the
+    -- affected partitions, rather than replacing the entire table.
+    -- this reduces latency and cost by limiting the scope of writes.
+#}
+
+
 {% macro bq_dynamic_copy_partitions_insert_overwrite_sql(
   tmp_relation, target_relation, sql, unique_key, partition_by, dest_columns, tmp_relation_exists, copy_partitions
   ) %}

dbt-bigquery/tests/functional/adapter/incremental/incremental_strategy_fixtures.py

@@ -336,6 +336,52 @@
 }}
 
 
+with data as (
+
+    {% if not is_incremental() %}
+
+        select 1 as id, cast('2020-01-01' as date) as date_day union all
+        select 2 as id, cast('2020-01-01' as date) as date_day union all
+        select 3 as id, cast('2020-01-01' as date) as date_day union all
+        select 4 as id, cast('2020-01-01' as date) as date_day
+
+    {% else %}
+
+        -- we want to overwrite the 4 records in the 2020-01-01 partition
+        -- with the 2 records below, but add two more in the 2020-01-02 partition
+        select 10 as id, cast('2020-01-01' as date) as date_day union all
+        select 20 as id, cast('2020-01-01' as date) as date_day union all
+        select 30 as id, cast('2020-01-02' as date) as date_day union all
+        select 40 as id, cast('2020-01-02' as date) as date_day
+
+    {% endif %}
+
+)
+
+select * from data
+
+{% if is_incremental() %}
+where date_day in ({{ config.get("partitions") | join(",") }})
+{% endif %}
+-- Test comment to prevent recurrence of https://github.com/dbt-labs/dbt-bigquery/issues/896
+""".lstrip()
+
+overwrite_copy_partitions_with_partitions_sql = """
+{{
+    config(
+        materialized="incremental",
+        incremental_strategy='insert_overwrite',
+        cluster_by="id",
+        partitions=["CAST('2020-01-01' AS DATE)","'2020-01-02'"],
+        partition_by={
+            "field": "date_day",
+            "data_type": "date",
+            "copy_partitions": true
+        }
+    )
+}}
+
+
 with data as (
 
     {% if not is_incremental() %}

dbt-bigquery/tests/functional/adapter/incremental/test_incremental_strategies.py

@@ -22,6 +22,7 @@
     overwrite_day_sql,
     overwrite_day_with_copy_partitions_sql,
     overwrite_partitions_sql,
+    overwrite_copy_partitions_with_partitions_sql,
     overwrite_range_sql,
     overwrite_time_sql,
     overwrite_day_with_time_ingestion_sql,
@@ -45,6 +46,7 @@ def models(self):
             "incremental_overwrite_day.sql": overwrite_day_sql,
             "incremental_overwrite_day_with_copy_partitions.sql": overwrite_day_with_copy_partitions_sql,
             "incremental_overwrite_partitions.sql": overwrite_partitions_sql,
+            "incremental_overwrite_copy_partitions_with_partitions.sql": overwrite_copy_partitions_with_partitions_sql,
             "incremental_overwrite_range.sql": overwrite_range_sql,
             "incremental_overwrite_time.sql": overwrite_time_sql,
             "incremental_overwrite_day_with_time_partition.sql": overwrite_day_with_time_ingestion_sql,
@@ -67,17 +69,22 @@ def seeds(self):
     def test__bigquery_assert_incremental_configurations_apply_the_right_strategy(self, project):
         run_dbt(["seed"])
         results = run_dbt()
-        assert len(results) == 12
+        assert len(results) == 13
 
         results = run_dbt()
-        assert len(results) == 12
+        assert len(results) == 13
+
         incremental_strategies = [
             ("incremental_merge_range", "merge_expected"),
             ("incremental_merge_time", "merge_expected"),
             ("incremental_merge_time_with_require_partition_view", "merge_expected"),
             ("incremental_overwrite_time", "incremental_overwrite_time_expected"),
             ("incremental_overwrite_date", "incremental_overwrite_date_expected"),
             ("incremental_overwrite_partitions", "incremental_overwrite_date_expected"),
+            (
+                "incremental_overwrite_copy_partitions_with_partitions",
+                "incremental_overwrite_date_expected",
+            ),
             ("incremental_overwrite_day", "incremental_overwrite_day_expected"),
             ("incremental_overwrite_range", "incremental_overwrite_range_expected"),
             (