From 67344daa0d4e124a139149230e72d67d50a6f2d5 Mon Sep 17 00:00:00 2001
From: trouze <tyler@tylerrouze.com>
Date: Fri, 4 Oct 2024 17:28:46 -0500
Subject: [PATCH] Add set and unset query tags and update base materializations

---
 .../macros/materializations/configs.sql          | 16 ++++++++++++++++
 .../materializations/models/clone/clone.sql      |  4 ++++
 .../models/incremental/incremental.sql           |  3 +++
 .../models/materialized_view.sql                 |  3 +++
 .../macros/materializations/models/table.sql     |  3 +++
 .../macros/materializations/models/view.sql      |  3 +++
 .../macros/materializations/seeds/seed.sql       |  3 +++
 .../materializations/snapshots/snapshot.sql      |  4 ++++
 .../macros/materializations/tests/test.sql       |  4 ++++
 9 files changed, 43 insertions(+)

diff --git a/dbt/include/global_project/macros/materializations/configs.sql b/dbt/include/global_project/macros/materializations/configs.sql
index d15ccb8e..11f1a009 100644
--- 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
index 56d80082..c321d496 100644
--- 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
index f932751a..b047e0cb 100644
--- 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
index a39f8aa2..b846aedd 100644
--- 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
index 3d1122ef..c7c1de8e 100644
--- 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
index 59ac6c4b..1a76d080 100644
--- 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
index 4ee4fb80..c672b5e7 100644
--- 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
index 5daead4c..f9a090a4 100644
--- 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
index ba205a9b..469ac4fa 100644
--- 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 -%}