Snowflake query tag

CHANGELOG.md

@@ -1,3 +1,6 @@
+## 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)
 
 ## dbt 0.18.0b1 (June 08, 2020)
@@ -13,6 +16,7 @@
 - Added support for BigQuery authorized views ([#1718](https://github.com/fishtown-analytics/dbt/issues/1718), [#2517](https://github.com/fishtown-analytics/dbt/pull/2517))
 - Added support for altering BigQuery column types ([#2546](https://github.com/fishtown-analytics/dbt/issues/2546), [#2547](https://github.com/fishtown-analytics/dbt/pull/2547))
 
+
 ### 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))
@@ -22,8 +26,8 @@ Contributors:
  - [@alf-mindshift](https://github.com/alf-mindshift) ([#2431](https://github.com/fishtown-analytics/dbt/pull/2431))
  - [@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/pull/2517), ([#2521](https://github.com/fishtown-analytics/dbt/pull/2521)), [#2547](https://github.com/fishtown-analytics/dbt/pull/2547))
-
+ - [@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)
 
@@ -34,6 +38,7 @@ Contributors:
 - dbt again respects config aliases in config() calls ([#2557](https://github.com/fishtown-analytics/dbt/issues/2557), [#2559](https://github.com/fishtown-analytics/dbt/pull/2559))
 
 
+
 ## 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') %}
+  {% 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') %}
+  {% 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 %}