From 331b961819b324c92ea13d6a12ec154117f2d513 Mon Sep 17 00:00:00 2001
From: Dan Rubin <dmrubin3@gmail.com>
Date: Tue, 16 Jun 2020 18:42:56 -0400
Subject: [PATCH 1/7] add query tag to snoflake connection dataclass and set
 tag on new connection if query_tag value is present

---
 plugins/snowflake/dbt/adapters/snowflake/connections.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/plugins/snowflake/dbt/adapters/snowflake/connections.py b/plugins/snowflake/dbt/adapters/snowflake/connections.py
index 4a64f3c3ee8..3f250398413 100644
--- a/plugins/snowflake/dbt/adapters/snowflake/connections.py
+++ b/plugins/snowflake/dbt/adapters/snowflake/connections.py
@@ -38,6 +38,7 @@ class SnowflakeCredentials(Credentials):
     token: Optional[str]
     oauth_client_id: Optional[str]
     oauth_client_secret: Optional[str]
+    query_tag: Optional[str]
     client_session_keep_alive: bool = False
 
     def __post_init__(self):
@@ -211,6 +212,11 @@ def open(cls, connection):
                 **creds.auth_args()
             )
 
+            if creds.query_tag:
+                handle.cursor().execute(
+                    ("alter session set query_tag = '{}'")
+                    .format(creds.query_tag))
+
             connection.handle = handle
             connection.state = 'open'
         except snowflake.connector.errors.Error as e:

From d8b31c02e2823e8ec1522d61c391d2949b043914 Mon Sep 17 00:00:00 2001
From: Dan Rubin <dmrubin3@gmail.com>
Date: Tue, 16 Jun 2020 18:43:13 -0400
Subject: [PATCH 2/7] add query tag to the snowflake adapter config

---
 plugins/snowflake/dbt/adapters/snowflake/impl.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/plugins/snowflake/dbt/adapters/snowflake/impl.py b/plugins/snowflake/dbt/adapters/snowflake/impl.py
index e156998a770..94a559502ff 100644
--- a/plugins/snowflake/dbt/adapters/snowflake/impl.py
+++ b/plugins/snowflake/dbt/adapters/snowflake/impl.py
@@ -25,6 +25,7 @@ class SnowflakeConfig(AdapterConfig):
     secure: Optional[bool] = None
     copy_grants: Optional[bool] = None
     snowflake_warehouse: Optional[str] = None
+    query_tag: Optional[str] = None
 
 
 class SnowflakeAdapter(SQLAdapter):

From 38975ecee0a579808484945510e4bab228168581 Mon Sep 17 00:00:00 2001
From: Dan Rubin <dmrubin3@gmail.com>
Date: Tue, 16 Jun 2020 18:43:26 -0400
Subject: [PATCH 3/7] add snowflake macros to support fetching current session
 query tag, updating the session query tag, and reverting to the old query tag

---
 .../dbt/include/snowflake/macros/adapters.sql | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/plugins/snowflake/dbt/include/snowflake/macros/adapters.sql b/plugins/snowflake/dbt/include/snowflake/macros/adapters.sql
index 1d80180e746..2fdb6ba09ce 100644
--- a/plugins/snowflake/dbt/include/snowflake/macros/adapters.sql
+++ b/plugins/snowflake/dbt/include/snowflake/macros/adapters.sql
@@ -160,3 +160,33 @@
         {{ column_name }} COMMENT $${{ column_dict[column_name]['description'] | replace('$', '[$]') }}$$ {{ ',' if not loop.last else ';' }}
     {% endfor %}
 {% endmacro %}
+
+
+{% macro get_current_query_tag() -%}
+  {{ return(run_query("show parameters like 'query_tag' in session").rows[0]['value']) }}
+{% endmacro %}
+
+
+{% macro set_query_tag() -%}
+  {% set new_query_tag = config.get('query_tag', none) %}
+  {% if new_query_tag %}
+    {% set original_query_tag = get_current_query_tag() %}
+    {{ log("Setting query_tag to '" ~ new_query_tag ~ "'. Will reset to '" ~ original_query_tag ~ "' after materialization.") }}
+    {% do run_query("alter session set query_tag = '{}'".format(new_query_tag)) %}
+    {{ return(original_query_tag)}}
+  {% endif %}
+  {{ return(none)}}
+{% endmacro %}
+
+{% macro unset_query_tag(original_query_tag) -%}
+  {% set new_query_tag = config.get('query_tag', none) %}
+  {% if new_query_tag %}
+    {% if original_query_tag %}
+      {{ log("Resetting query_tag to '" ~ original_query_tag ~ "'.") }}
+      {% do run_query("alter session set query_tag = '{}'".format(original_query_tag)) %}
+    {% else %}
+      {{ log("No original query_tag, unsetting parameter.") }}
+      {% do run_query("alter session unset query_tag") %}
+    {% endif %}
+  {% endif %}
+{% endmacro %} 

From be61e6724a4cb863fabbb77149b82252f0820b62 Mon Sep 17 00:00:00 2001
From: Dan Rubin <dmrubin3@gmail.com>
Date: Tue, 16 Jun 2020 18:43:44 -0400
Subject: [PATCH 4/7] add query tag update macros to each snowflake
 materialization

---
 .../snowflake/macros/materializations/incremental.sql        | 4 ++++
 .../dbt/include/snowflake/macros/materializations/table.sql  | 5 +++++
 .../dbt/include/snowflake/macros/materializations/view.sql   | 4 ++++
 3 files changed, 13 insertions(+)

diff --git a/plugins/snowflake/dbt/include/snowflake/macros/materializations/incremental.sql b/plugins/snowflake/dbt/include/snowflake/macros/materializations/incremental.sql
index 1c8bfbde3c1..188795b1537 100644
--- a/plugins/snowflake/dbt/include/snowflake/macros/materializations/incremental.sql
+++ b/plugins/snowflake/dbt/include/snowflake/macros/materializations/incremental.sql
@@ -26,6 +26,8 @@
 
 {% materialization incremental, adapter='snowflake' -%}
 
+  {% set original_query_tag = set_query_tag() %}
+
   {%- set unique_key = config.get('unique_key') -%}
   {%- set full_refresh_mode = (should_full_refresh()) -%}
 
@@ -74,6 +76,8 @@
   {% set target_relation = target_relation.incorporate(type='table') %}
   {% do persist_docs(target_relation, model) %}
 
+  {% do unset_query_tag(original_query_tag) %}
+
   {{ return({'relations': [target_relation]}) }}
 
 {%- endmaterialization %}
diff --git a/plugins/snowflake/dbt/include/snowflake/macros/materializations/table.sql b/plugins/snowflake/dbt/include/snowflake/macros/materializations/table.sql
index 84d78781864..2724f60b29c 100644
--- a/plugins/snowflake/dbt/include/snowflake/macros/materializations/table.sql
+++ b/plugins/snowflake/dbt/include/snowflake/macros/materializations/table.sql
@@ -1,4 +1,7 @@
 {% materialization table, adapter='snowflake' %}
+
+  {% set original_query_tag = set_query_tag() %}
+
   {%- set identifier = model['alias'] -%}
 
   {%- set old_relation = adapter.get_relation(database=database, schema=schema, identifier=identifier) -%}
@@ -32,6 +35,8 @@
 
   {% do persist_docs(target_relation, model) %}
 
+  {% do unset_query_tag(original_query_tag) %}
+
   {{ return({'relations': [target_relation]}) }}
 
 {% endmaterialization %}
diff --git a/plugins/snowflake/dbt/include/snowflake/macros/materializations/view.sql b/plugins/snowflake/dbt/include/snowflake/macros/materializations/view.sql
index 700bb3807bf..b3858c67478 100644
--- a/plugins/snowflake/dbt/include/snowflake/macros/materializations/view.sql
+++ b/plugins/snowflake/dbt/include/snowflake/macros/materializations/view.sql
@@ -1,4 +1,6 @@
 {% materialization view, adapter='snowflake' -%}
+
+    {% set original_query_tag = set_query_tag() %}
     {% set to_return = create_or_replace_view() %}
 
     {% set target_relation = this.incorporate(type='view') %}
@@ -6,4 +8,6 @@
 
     {% do return(to_return) %}
 
+    {% do unset_query_tag(original_query_tag) %}
+
 {%- endmaterialization %}

From 5c1840914ba99c002abb8467af1b9bfd0ac383ec Mon Sep 17 00:00:00 2001
From: Dan Rubin <dmrubin3@gmail.com>
Date: Tue, 16 Jun 2020 18:43:58 -0400
Subject: [PATCH 5/7] update CHANGELOG.md

---
 CHANGELOG.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b5578fd1b1a..4e92a68b434 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@
 - Added intersection syntax for model selector ([#2167](https://github.com/fishtown-analytics/dbt/issues/2167), [#2417](https://github.com/fishtown-analytics/dbt/pull/2417))
 - Extends model selection syntax with at most n-th parent/children `dbt run --models 3+m1+2` ([#2052](https://github.com/fishtown-analytics/dbt/issues/2052), [#2485](https://github.com/fishtown-analytics/dbt/pull/2485))
 - Added support for renaming BigQuery relations ([#2520](https://github.com/fishtown-analytics/dbt/issues/2520), [#2521](https://github.com/fishtown-analytics/dbt/pull/2521))
+- Added support for Snowflake query tags at the connection and model level ([#1030](https://github.com/fishtown-analytics/dbt/issues/1030))
 
 ### Fixes
 - Fixed an error in create_adapter_plugins.py script when -dependency arg not passed ([#2507](https://github.com/fishtown-analytics/dbt/issues/2507), [#2508](https://github.com/fishtown-analytics/dbt/pull/2508))
@@ -20,6 +21,8 @@ Contributors:
  - [@scarrucciu](https://github.com/scarrucciu) ([#2508](https://github.com/fishtown-analytics/dbt/pull/2508))
  - [@southpolemonkey](https://github.com/southpolemonkey)([#2511](https://github.com/fishtown-analytics/dbt/issues/2511))
  - [@azhard](https://github.com/azhard) ([#2521](https://github.com/fishtown-analytics/dbt/pull/2521)
+ - [@DrMcTaco](https://github.com/DrMcTaco) ([#1030](https://github.com/fishtown-analytics/dbt/issues/1030)))
+
 
 ## dbt 0.17.0 (June 08, 2020)
 

From 2cc53321ae9d3a0618564713598da9b013207f54 Mon Sep 17 00:00:00 2001
From: Dan Rubin <dmrubin3@gmail.com>
Date: Wed, 17 Jun 2020 13:01:16 -0400
Subject: [PATCH 6/7] remove incorrect 2nd arg to config.get

---
 plugins/snowflake/dbt/include/snowflake/macros/adapters.sql | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/plugins/snowflake/dbt/include/snowflake/macros/adapters.sql b/plugins/snowflake/dbt/include/snowflake/macros/adapters.sql
index 2fdb6ba09ce..8c59271090e 100644
--- a/plugins/snowflake/dbt/include/snowflake/macros/adapters.sql
+++ b/plugins/snowflake/dbt/include/snowflake/macros/adapters.sql
@@ -168,7 +168,7 @@
 
 
 {% macro set_query_tag() -%}
-  {% set new_query_tag = config.get('query_tag', none) %}
+  {% set new_query_tag = config.get('query_tag') %}
   {% if new_query_tag %}
     {% set original_query_tag = get_current_query_tag() %}
     {{ log("Setting query_tag to '" ~ new_query_tag ~ "'. Will reset to '" ~ original_query_tag ~ "' after materialization.") }}
@@ -179,7 +179,7 @@
 {% endmacro %}
 
 {% macro unset_query_tag(original_query_tag) -%}
-  {% set new_query_tag = config.get('query_tag', none) %}
+  {% set new_query_tag = config.get('query_tag') %}
   {% if new_query_tag %}
     {% if original_query_tag %}
       {{ log("Resetting query_tag to '" ~ original_query_tag ~ "'.") }}

From c89ea1ad56057c4295aef75bf3751c0613c8634e Mon Sep 17 00:00:00 2001
From: Dan Rubin <dmrubin3@gmail.com>
Date: Wed, 24 Jun 2020 09:15:17 -0400
Subject: [PATCH 7/7] fix CHANGELOG formatting

---
 CHANGELOG.md | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d3ad9d9b9ee..a71b717710e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,8 +1,15 @@
-## Next
-- Added support for Snowflake query tags at the connection and model level ([#1030](https://github.com/fishtown-analytics/dbt/issues/1030))
-
 ## dbt 0.18.0 (Release TBD)
 
+
+### Features
+- Added support for Snowflake query tags at the connection and model level ([#1030](https://github.com/fishtown-analytics/dbt/issues/1030), [#2555](https://github.com/fishtown-analytics/dbt/pull/2555/))
+
+
+### Fixes
+
+Contributors:
+- [@DrMcTaco](https://github.com/DrMcTaco) ([#1030](https://github.com/fishtown-analytics/dbt/issues/1030)),[#2555](https://github.com/fishtown-analytics/dbt/pull/2555/))
+
 ## dbt 0.18.0b1 (June 08, 2020)
 
 
@@ -27,7 +34,6 @@ Contributors:
  - [@scarrucciu](https://github.com/scarrucciu) ([#2508](https://github.com/fishtown-analytics/dbt/pull/2508))
  - [@southpolemonkey](https://github.com/southpolemonkey) ([#2511](https://github.com/fishtown-analytics/dbt/issues/2511))
  - [@azhard](https://github.com/azhard) ([#2517](https://github.com/fishtown-analytics/dbt/issues/2517), ([#2521](https://github.com/fishtown-analytics/dbt/pull/2521)))
- - [@DrMcTaco](https://github.com/DrMcTaco) ([#1030](https://github.com/fishtown-analytics/dbt/issues/1030)))
 
 ## dbt 0.17.1 (Release TBD)