Merge pull request #2555 from DrMcTaco/snowflake-query-tag

Snowflake query tag

CHANGELOG.md

@@ -1,10 +1,12 @@
 ## dbt 0.18.0 (Release TBD)
 
 ### Features
+- Added support for Snowflake query tags at the connection and model level ([#1030](https:/fishtown-analytics/dbt/issues/1030), [#2555](https:/fishtown-analytics/dbt/pull/2555/))
 - Added option to specify profile when connecting to Redshift via IAM ([#2437](https:/fishtown-analytics/dbt/issues/2437), [#2581](https:/fishtown-analytics/dbt/pull/2581))
 
 Contributors:
 - [@brunomurino](https:/brunomurino) ([#2437](https:/fishtown-analytics/dbt/pull/2581))
+- [@DrMcTaco](https:/DrMcTaco) ([#1030](https:/fishtown-analytics/dbt/issues/1030)),[#2555](https:/fishtown-analytics/dbt/pull/2555/))
 
 
 ## dbt 0.18.0b1 (June 08, 2020)
@@ -21,6 +23,7 @@ Contributors:
 - Added support for altering BigQuery column types ([#2546](https:/fishtown-analytics/dbt/issues/2546), [#2547](https:/fishtown-analytics/dbt/pull/2547))
 - Include row counts and bytes processed in log output for all BigQuery statement types ([#2526](https:/fishtown-analytics/dbt/issues/2526))
 
+
 ### Fixes
 - Fixed an error in create_adapter_plugins.py script when -dependency arg not passed ([#2507](https:/fishtown-analytics/dbt/issues/2507), [#2508](https:/fishtown-analytics/dbt/pull/2508))
 - Remove misleading "Opening a new connection" log message in set_connection_name. ([#2511](https:/fishtown-analytics/dbt/issues/2511))
@@ -34,7 +37,6 @@ Contributors:
  - [@azhard](https:/azhard) ([#2517](https:/fishtown-analytics/dbt/pull/2517), ([#2521](https:/fishtown-analytics/dbt/pull/2521)), [#2547](https:/fishtown-analytics/dbt/pull/2547))
  - [@alepuccetti](https:/alepuccetti) ([#2526](https:/fishtown-analytics/dbt/issues/2526))
 
-
 ## dbt 0.17.1 (Release TBD)
 
 ### Fixes
@@ -44,6 +46,7 @@ Contributors:
 - dbt again respects config aliases in config() calls ([#2557](https:/fishtown-analytics/dbt/issues/2557), [#2559](https:/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 %}