Add set and unset query tags and update base materializations

dbt-labs · Oct 4, 2024 · 67344da · 67344da
1 parent 10e6d59
commit 67344da
Show file tree

Hide file tree

Showing 9 changed files with 43 additions and 0 deletions.
diff --git a/dbt/include/global_project/macros/materializations/configs.sql b/dbt/include/global_project/macros/materializations/configs.sql
@@ -19,3 +19,19 @@
   {% endif %}
   {% do return(config_store_failures) %}
 {% endmacro %}
+
+{% macro set_query_tag() %}
+  {{ return(adapter.dispatch('set_query_tag', 'dbt')()) }}
+{% endmacro %}
+
+{% macro unset_query_tag(original_query_tag) %}
+  {{ return(adapter.dispatch('unset_query_tag', 'dbt')(original_query_tag)) }}
+{% endmacro %}
+
+{% macro default__set_query_tag() %}
+  {{ return("") }}
+{% endmacro%}
+
+{% macro default__unset_query_tag() %}
+  {{ return("") }}
+{% endmacro%}
diff --git a/dbt/include/global_project/macros/materializations/models/clone/clone.sql b/dbt/include/global_project/macros/materializations/models/clone/clone.sql
@@ -1,5 +1,7 @@
 {%- materialization clone, default -%}
 
+  {% set original_query_tag = set_query_tag() %}
+
   {%- set relations = {'relations': []} -%}
 
   {%- if not defer_relation -%}
@@ -64,4 +66,6 @@
 
   {%- endif -%}
 
+  {% do unset_query_tag(original_query_tag) %}
+
 {%- endmaterialization -%}
diff --git a/dbt/include/global_project/macros/materializations/models/incremental/incremental.sql b/dbt/include/global_project/macros/materializations/models/incremental/incremental.sql
@@ -1,6 +1,7 @@
 
 {% materialization incremental, default -%}
 
+  {% set original_query_tag = set_query_tag() %}
   -- relations
   {%- set existing_relation = load_cached_relation(this) -%}
   {%- set target_relation = this.incorporate(type='table') -%}
@@ -90,6 +91,8 @@
 
   {{ run_hooks(post_hooks, inside_transaction=False) }}
 
+  {% do unset_query_tag(original_query_tag) %}
+
   {{ return({'relations': [target_relation]}) }}
 
 {%- endmaterialization %}
diff --git a/dbt/include/global_project/macros/materializations/models/materialized_view.sql b/dbt/include/global_project/macros/materializations/models/materialized_view.sql
@@ -1,4 +1,5 @@
 {% materialization materialized_view, default %}
+    {% set original_query_tag = set_query_tag() %}
     {% set existing_relation = load_cached_relation(this) %}
     {% set target_relation = this.incorporate(type=this.MaterializedView) %}
     {% set intermediate_relation = make_intermediate_relation(target_relation) %}
@@ -17,6 +18,8 @@
 
     {{ materialized_view_teardown(backup_relation, intermediate_relation, post_hooks) }}
 
+    {% do unset_query_tag(original_query_tag) %}
+
     {{ return({'relations': [target_relation]}) }}
 
 {% endmaterialization %}

diff --git a/dbt/include/global_project/macros/materializations/models/table.sql b/dbt/include/global_project/macros/materializations/models/table.sql
@@ -1,5 +1,6 @@
 {% materialization table, default %}
 
+  {% set original_query_tag = set_query_tag() %}
   {%- set existing_relation = load_cached_relation(this) -%}
   {%- set target_relation = this.incorporate(type='table') %}
   {%- set intermediate_relation =  make_intermediate_relation(target_relation) -%}
@@ -60,5 +61,7 @@
 
   {{ run_hooks(post_hooks, inside_transaction=False) }}
 
+  {% do unset_query_tag(original_query_tag) %}
+
   {{ return({'relations': [target_relation]}) }}
 {% endmaterialization %}
diff --git a/dbt/include/global_project/macros/materializations/models/view.sql b/dbt/include/global_project/macros/materializations/models/view.sql
@@ -1,5 +1,6 @@
 {%- materialization view, default -%}
 
+  {% set original_query_tag = set_query_tag() %}
   {%- set existing_relation = load_cached_relation(this) -%}
   {%- set target_relation = this.incorporate(type='view') -%}
   {%- set intermediate_relation =  make_intermediate_relation(target_relation) -%}
@@ -67,6 +68,8 @@
 
   {{ run_hooks(post_hooks, inside_transaction=False) }}
 
+  {% do unset_query_tag(original_query_tag) %}
+
   {{ return({'relations': [target_relation]}) }}
 
 {%- endmaterialization -%}
diff --git a/dbt/include/global_project/macros/materializations/seeds/seed.sql b/dbt/include/global_project/macros/materializations/seeds/seed.sql
@@ -1,5 +1,6 @@
 {% materialization seed, default %}
 
+  {% set original_query_tag = set_query_tag() %}
   {%- set identifier = model['alias'] -%}
   {%- set full_refresh_mode = (should_full_refresh()) -%}
 
@@ -55,6 +56,8 @@
 
   {{ run_hooks(post_hooks, inside_transaction=False) }}
 
+  {% do unset_query_tag(original_query_tag) %}
+
   {{ return({'relations': [target_relation]}) }}
 
 {% endmaterialization %}
diff --git a/dbt/include/global_project/macros/materializations/snapshots/snapshot.sql b/dbt/include/global_project/macros/materializations/snapshots/snapshot.sql
@@ -1,5 +1,7 @@
 {% materialization snapshot, default %}
 
+  {% set original_query_tag = set_query_tag() %}
+
   {%- set target_table = model.get('alias', model.get('name')) -%}
 
   {%- set strategy_name = config.get('strategy') -%}
@@ -102,6 +104,8 @@
 
   {{ run_hooks(post_hooks, inside_transaction=False) }}
 
+  {% do unset_query_tag(original_query_tag) %}
+
   {{ return({'relations': [target_relation]}) }}
 
 {% endmaterialization %}
diff --git a/dbt/include/global_project/macros/materializations/tests/test.sql b/dbt/include/global_project/macros/materializations/tests/test.sql
@@ -1,5 +1,7 @@
 {%- materialization test, default -%}
 
+  {% set original_query_tag = set_query_tag() %}
+
   {% set relations = [] %}
 
   {% if should_store_failures() %}
@@ -55,6 +57,8 @@
 
   {%- endcall %}
 
+  {% do unset_query_tag(original_query_tag) %}
+
   {{ return({'relations': relations}) }}
 
 {%- endmaterialization -%}