Snowflake query tag

CHANGELOG.md

@@ -9,8 +9,10 @@
 - 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))
 - Added support for BigQuery authorized views ([#1718](https://github.com/fishtown-analytics/dbt/issues/1718), [#2517](https://github.com/fishtown-analytics/dbt/issues/2517))
 
+
 ### 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))
 - Remove misleading "Opening a new connection" log message in set_connection_name. ([#2511](https://github.com/fishtown-analytics/dbt/issues/2511))
@@ -21,7 +23,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) ([#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)
 
@@ -31,6 +34,7 @@ Contributors:
 - `packages.yml` revision/version fields can be float-like again (`revision: '1.0'` is valid). ([#2518](https://github.com/fishtown-analytics/dbt/issues/2518), [#2535](https://github.com/fishtown-analytics/dbt/pull/2535))
 
 
+
 ## dbt 0.17.0 (June 08, 2020)
 
 ### Fixes

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:

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):

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 %} 

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 %}

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 %}

plugins/snowflake/dbt/include/snowflake/macros/materializations/view.sql

@@ -1,9 +1,13 @@
 {% 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') %}
     {% do persist_docs(target_relation, model, for_columns=false) %}
 
     {% do return(to_return) %}
 
+    {% do unset_query_tag(original_query_tag) %}
+
 {%- endmaterialization %}