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)