Repository: dbt-labs/dbt-core

Issue:
Add query_tag to the model hierarchy
## Feature

### Feature description

After starting to analyze our query history in some depth, I had the realization that it would be useful to be able to specify QUERY_TAG (snowflake concept) anywhere in the hierarchy, like the `materialization` configuration.

This way, you can query snowflake query_history by the query tag. Maybe you set one query tag for your whole dbt project. Maybe you set it per model folder. Maybe you set it on specific models.

https://docs.snowflake.net/manuals/sql-reference/parameters.html#query-tag
https://stackoverflow.com/questions/50965558/snowflake-cli-snowsql-query-tagging

QUERY_TAG is a session parameter, so the solution might be to set it in snowflake's table and view materializations, if the query tag is set?

### Who will this benefit?

Anyone using snowflake who would like more granular information in her query history.



Discussion:
I can see this being useful in a few different situations.
- cost/usage tracking: all queries run by DBT against snowflake being tagged and thus traceable in snowflake query history
- snowflake optimization/debugging: tags at the model level for trending performance and anomaly detection. 

I am considering starting work to resolve this issue and the above use cases make me think there may be at least 2 distinct places one would want to add a tag: 
1. at the connection level
2. at the model level

The second can easily be achieved with model hooks (this is not to say some centralized supported method should not be developed).

I am however not sure how to go about setting query tag immediately after opening a connection so that all queries run regardless of their source are tagged. Any ideas on how to achieve that or other thoughts on the general problem before I dig in would be appreciated.
hey @DrMcTaco - I like the way you're thinking about this feature! I agree - I think we should support query tags both at the model level and at the connection level. We'd very much welcome a PR for this change.

Some rough implementation notes:
- We can add a parameter to the [SnowflakeCredentials dataclass](https://github.com/fishtown-analytics/dbt/blob/eea51bee30d79054637b2adfffc8d3fb0cf99ce3/plugins/snowflake/dbt/adapters/snowflake/connections.py#L29-L41) to capture a `query_tag` string in a profiles.yml target.
- We'll want to alter the session parameter after opening a connection. Check out an example of where we do something similar [in the postgres plugin](https://github.com/fishtown-analytics/dbt/blob/eea51bee30d79054637b2adfffc8d3fb0cf99ce3/plugins/postgres/dbt/adapters/postgres/connections.py#L97-L107).
- Supporting this session parameter in model configs is going to be a tiny bit more involved than setting this field for all connections opened by dbt. We'd need to update each of [dbt's Snowflake materializations](https://github.com/fishtown-analytics/dbt/tree/dev/marian-anderson/plugins/snowflake/dbt/include/snowflake/macros/materializations) to check the config and execute a SQL query to set the session parameter before running a query. There might be a better way to do this though.... let me think about that some more.

Beyond these things, we also need to think about what happens if a query_tag is provided for both the target and a model. The easiest approach would just be to run two `alter session` queries. I think/hope Snowflake will do the right thing there, but it's worth verifying the behavior!

Let me know what you think about all that :)
> The easiest approach would just be to run two alter session queries. I think/hope Snowflake will do the right thing there, but it's worth verifying the behavior!

Looks like you can just run one after another and the second will overwrite the first
<img width="551" alt="alter_tag_cmd" src="https://user-images.githubusercontent.com/25394508/84782687-56007280-afb6-11ea-96cf-5f0005781afd.png">
<img width="1250" alt="alter_tag_history" src="https://user-images.githubusercontent.com/25394508/84782796-74666e00-afb6-11ea-9318-8ee68fb73c18.png">

In other news I have the profile based query_tags working. I am getting a bit hung up finding the right place and way to do it in the model macros. I am not the most familiar with Jinja. 

Can you link to a macro where something similar is done?
Nice!

For the model-level configs, you'll want to add the `query_tag` config here to register the config with the Snowflake adapter plugin:

https://github.com/fishtown-analytics/dbt/blob/74c78ef9279638d14949d4e93ae76166d4fd8022/plugins/snowflake/dbt/adapters/snowflake/impl.py#L20-L27

As far as the jinja code is concerned, I'd grab this config value from the jinja context with `config.get('query_tag')`. If the config is defined, then I'd issue the `alter session` statement at the top of the Snowflake materialization macros:
 - [incremental](https://github.com/fishtown-analytics/dbt/blob/74c78ef9279638d14949d4e93ae76166d4fd8022/plugins/snowflake/dbt/include/snowflake/macros/materializations/incremental.sql#L27-L37)
 - [table](https://github.com/fishtown-analytics/dbt/blob/74c78ef9279638d14949d4e93ae76166d4fd8022/plugins/snowflake/dbt/include/snowflake/macros/materializations/table.sql#L1-L9)
 - [view](https://github.com/fishtown-analytics/dbt/blob/74c78ef9279638d14949d4e93ae76166d4fd8022/plugins/snowflake/dbt/include/snowflake/macros/materializations/view.sql#L1-L4)

I'd try to run this query as early as possible in the materialization code as possible so that all of the SQL (eg. hooks) are executed with the correct query tag. It might also be a good idea to make a helper macro ([like this one](https://github.com/fishtown-analytics/dbt/blob/74c78ef9279638d14949d4e93ae76166d4fd8022/plugins/snowflake/dbt/include/snowflake/macros/adapters.sql#L157-L162)) that actually does the heavy lifting to run this `alter session` statement if a config is defined.

One last thing to consider: dbt uses a connection pool to execute queries. Do we need to "reset" the `query_tag` after running a given model? I want to make sure that the session parameter for a previous model doesn't leak into a subsequent model executed from the same thread.

Resolving patch:
diff --git a/plugins/snowflake/dbt/adapters/snowflake/connections.py b/plugins/snowflake/dbt/adapters/snowflake/connections.py
--- 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:
diff --git a/plugins/snowflake/dbt/adapters/snowflake/impl.py b/plugins/snowflake/dbt/adapters/snowflake/impl.py
--- 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):


Test patch: