From 8825ccb8bf2c1343f91a5f3c6ea0c7a05376f1ee Mon Sep 17 00:00:00 2001
From: Mike Alfare <13974384+mikealfare@users.noreply.github.com>
Date: Thu, 17 Oct 2024 15:16:35 -0400
Subject: [PATCH] Add query tag back to incremental models (#1212)

* add query tag back to incremental models

* fix query tag test
---
 .../macros/materializations/incremental.sql   |  1 +
 tests/functional/query_tag/test_query_tags.py | 34 +++++++------------
 2 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/dbt/include/snowflake/macros/materializations/incremental.sql b/dbt/include/snowflake/macros/materializations/incremental.sql
index 12645cd11..d73525d6d 100644
--- a/dbt/include/snowflake/macros/materializations/incremental.sql
+++ b/dbt/include/snowflake/macros/materializations/incremental.sql
@@ -58,6 +58,7 @@
 
 {% materialization incremental, adapter='snowflake', supported_languages=['sql', 'python'] -%}
 
+  {% set original_query_tag = set_query_tag() %}
 
   {#-- Set vars --#}
   {%- set full_refresh_mode = (should_full_refresh()) -%}
diff --git a/tests/functional/query_tag/test_query_tags.py b/tests/functional/query_tag/test_query_tags.py
index 421aae1b6..4ddafabb2 100644
--- a/tests/functional/query_tag/test_query_tags.py
+++ b/tests/functional/query_tag/test_query_tags.py
@@ -1,6 +1,7 @@
 import pytest
 from dbt.tests.util import run_dbt
 
+
 snapshots__snapshot_query_tag_sql = """
 {% snapshot snapshot_query_tag %}
     {{
@@ -14,16 +15,15 @@
     }}
     select 1 as id, 'blue' as color
 {% endsnapshot %}
-
 """
 
+
 models__table_model_query_tag_sql = """
 {{ config(materialized = 'table') }}
-
 select 1 as id
-
 """
 
+
 models__models_config_yml = """
 version: 2
 
@@ -33,23 +33,21 @@
       - name: id
         data_tests:
           - unique
-
 """
 
+
 models__view_model_query_tag_sql = """
 {{ config(materialized = 'view') }}
-
 select 1 as id
-
 """
 
+
 models__incremental_model_query_tag_sql = """
 {{ config(materialized = 'incremental', unique_key = 'id') }}
-
 select 1 as id
-
 """
 
+
 macros__check_tag_sql = """
 {% macro check_query_tag() %}
 
@@ -61,12 +59,12 @@
   {% endif %}
 
 {% endmacro %}
-
 """
 
+
 seeds__seed_query_tag_csv = """id
 1
-"""
+""".strip()
 
 
 class TestQueryTag:
@@ -95,20 +93,14 @@ def seeds(self):
     def project_config_update(self, prefix):
         return {
             "config-version": 2,
-            "models": {
-                "tests": {"query_tag": prefix, "post-hook": "{{ check_tag() }}"},
-            },
-            "seeds": {
-                "tests": {"query_tag": prefix, "post-hook": "{{ check_tag() }}"},
-            },
-            "snapshots": {
-                "tests": {"query_tag": prefix, "post-hook": "{{ check_tag() }}"},
-            },
+            "models": {"query_tag": prefix, "post-hook": "{{ check_query_tag() }}"},
+            "seeds": {"query_tag": prefix, "post-hook": "{{ check_query_tag() }}"},
+            "snapshots": {"query_tag": prefix, "post-hook": "{{ check_query_tag() }}"},
             "tests": {"test": {"query_tag": prefix, "post-hook": "{{ check_query_tag() }}"}},
         }
 
     def build_all_with_query_tags(self, project, prefix):
-        run_dbt(["build", "--vars", '{{"check_tag": "{}"}}'.format(prefix)])
+        run_dbt(["build", "--vars", '{{"query_tag": "{}"}}'.format(prefix)])
 
     def test_snowflake_query_tag(self, project, prefix):
         self.build_all_with_query_tags(project, prefix)
@@ -130,7 +122,7 @@ def profiles_config_update(self, prefix):
         return {"query_tag": prefix}
 
     def build_all_with_query_tags(self, project, prefix):
-        run_dbt(["build", "--vars", '{{"check_tag": "{}"}}'.format(prefix)])
+        run_dbt(["build", "--vars", '{{"query_tag": "{}"}}'.format(prefix)])
 
     def test_snowflake_query_tag(self, project, prefix):
         self.build_all_with_query_tags(project, prefix)