Merge pull request #22 from get-select/allow_extra_metadata

Add support for custom query comment keys and values
get-select · Jan 18, 2024 · fa868f4 · fa868f4
2 parents 56ecce4 + f1eb9a5
commit fa868f4
Show file tree

Hide file tree

Showing 8 changed files with 52 additions and 27 deletions.
diff --git a/.changes/2.3.2.md b/.changes/2.3.2.md
@@ -0,0 +1,10 @@
+## dbt-snowflake-query-tags 2.3.2 - January 18, 2024
+
+### Features
+
+- Support for query tags from profiles.yml and environment variables ([#21](https://github.com/get-select/dbt-snowflake-query-tags/pull/21))
+- Add support for custom query comment keys and values ([#22](https://github.com/get-select/dbt-snowflake-query-tags/pull/22))
+
+### Contributors
+- [@maddoc1](https://github.com/maddoc1) (Features)
+
diff --git a/.changes/unreleased/Features-20240114-205231.yaml b/.changes/unreleased/Features-20240114-205231.yaml
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,17 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html),
 and is generated by [Changie](https://github.com/miniscruff/changie).
 
+## dbt-snowflake-query-tags 2.3.2 - January 18, 2024
+
+### Features
+
+- Support for query tags from profiles.yml and environment variables ([#21](https://github.com/get-select/dbt-snowflake-query-tags/pull/21))
+- Add support for custom query comment keys and values ([#22](https://github.com/get-select/dbt-snowflake-query-tags/pull/22))
+
+### Contributors
+- [@maddoc1](https://github.com/maddoc1) (Features)
+
+
 ## dbt-snowflake-query-tags 2.3.1 - August 18, 2023
 
 ### Features

diff --git a/README.md b/README.md
@@ -6,7 +6,7 @@ An example query comment contains:
 
 ```json
 {
-    "dbt_snowflake_query_tags_version": "2.3.1",
+    "dbt_snowflake_query_tags_version": "2.3.2",
     "app": "dbt",
     "dbt_version": "1.4.0",
     "project_name": "my_project",
@@ -45,7 +45,7 @@ Query tags are used solely for attaching the `is_incremental` flag, as this isn'
 
 ```json
 {
-    "dbt_snowflake_query_tags_version": "2.3.1",
+    "dbt_snowflake_query_tags_version": "2.3.2",
     "app": "dbt",
     "is_incremental": true
 }
@@ -100,7 +100,15 @@ That's it! All dbt-issued queries will now be tagged.
 
 ### Query comments
 
-To extend the information added in the query comments, use [meta](https://docs.getdbt.com/reference/resource-configs/meta) or [tag](https://docs.getdbt.com/reference/resource-configs/tags) configs. These are automatically added to the query comments.
+Both [meta](https://docs.getdbt.com/reference/resource-configs/meta) and [tag](https://docs.getdbt.com/reference/resource-configs/tags) configs are automatically added to the query comments.
+
+To add arbitrary keys and values to the comments, you can use the `extra` kwarg when calling `dbt_snowflake_query_tags.get_query_comment`. For example, to add a `run_started_at` key and value to the comment, we can do:
+
+```yaml
+query-comment:
+  comment: '{{ dbt_snowflake_query_tags.get_query_comment(node, extra={"run_started_at": builtins.run_started_at | string }) }}'
+  append: true # Snowflake removes prefixed comments.
+```
 
 ### Query tags
 
@@ -119,6 +127,11 @@ Model
 select ...
 ```
 
+Results in the following query tag. The additional information is added by this package.
+```
+'{"team": "data", "app": "dbt", "dbt_snowflake_query_tags_version": "2.3.2", "is_incremental": true}'
+```
+
 Note that using a non-mapping type in the `query_tag` config will result in a warning, and the config being ignored.
 
 Model
@@ -135,11 +148,9 @@ Warning:
 dbt-snowflake-query-tags warning: the query_tag config value of 'data team' is not a mapping type, so is being ignored. If you'd like to add additional query tag information, use a mapping type instead, or remove it to avoid this message.
 ```
 
-This results in a final query tag without 'data team' being present.
-
 #### Profiles.yml
 
-Additionally, you can set the `query_tag` in the `profiles.yml`. This must be a valid json object.
+Additionally, you can set the `query_tag` value in the `profiles.yml`. This must be a valid json object.
 
 profiles.yml
 ```yml
@@ -165,7 +176,7 @@ dbt_project.yml:
 
 Results in a final query tag of
 ```
-'{"team": "data", "job_name": "daily", "app": "dbt", "dbt_snowflake_query_tags_version": "2.3.1", "is_incremental": true}'
+'{"team": "data", "job_name": "daily", "app": "dbt", "dbt_snowflake_query_tags_version": "2.3.2", "is_incremental": true}'
 ```
 
 ## Contributing

diff --git a/dbt_project.yml b/dbt_project.yml
@@ -1,3 +1,3 @@
 name: 'dbt_snowflake_query_tags'
-version: '2.3.1'
+version: '2.3.2'
 config-version: 2
diff --git a/integration_test_project/dbt_project.yml b/integration_test_project/dbt_project.yml
@@ -13,5 +13,5 @@ dispatch:
       - dbt
 
 query-comment:
-  comment: '{{ dbt_snowflake_query_tags.get_query_comment(node) }}'
+  comment: '{{ dbt_snowflake_query_tags.get_query_comment(node, extra={"run_started_at": builtins.run_started_at | string }) }}'
   append: true # Snowflake removes prefixed comments.
diff --git a/macros/query_comment.sql b/macros/query_comment.sql
@@ -1,8 +1,8 @@
-{% macro get_query_comment(node) %}
-    {%- set comment_dict = {} -%}
+{% macro get_query_comment(node, extra = {}) %}
+    {%- set comment_dict = extra -%}
     {%- do comment_dict.update(
         app='dbt',
-        dbt_snowflake_query_tags_version='2.3.1',
+        dbt_snowflake_query_tags_version='2.3.2',
         dbt_version=dbt_version,
         project_name=project_name,
         target_name=target.name,

diff --git a/macros/query_tags.sql b/macros/query_tags.sql
@@ -4,14 +4,14 @@
 
 {% macro default__set_query_tag() -%}
     {# Get session level query tag #}
-    {% set session_query_tag = get_current_query_tag() %}
-    {% set session_query_tag_parsed = {} %}
+    {% set original_query_tag = get_current_query_tag() %}
+    {% set original_query_tag_parsed = {} %}
 
-    {% if session_query_tag %}
-        {% if fromjson(session_query_tag) is mapping %}
-            {% set session_query_tag_parsed = fromjson(session_query_tag) %}
+    {% if original_query_tag %}
+        {% if fromjson(original_query_tag) is mapping %}
+            {% set original_query_tag_parsed = fromjson(original_query_tag) %}
         {% else %}
-            {% do log("dbt-snowflake-query-tags warning: the session level query tag value of '{}' is not a mapping type, so is being ignored. If you'd like to add additional query tag information, use a mapping type instead, or remove it to avoid this message.".format(session_query_tag), True) %}
+            {% do log("dbt-snowflake-query-tags warning: the session level query tag value of '{}' is not a mapping type, so is being ignored. If you'd like to add additional query tag information, use a mapping type instead, or remove it to avoid this message.".format(original_query_tag), True) %}
         {% endif %}
     {% endif %}
 
@@ -32,12 +32,12 @@
     {% set query_tag = {} %} {# If the user has set the query tag config as a non mapping type, start fresh #}
     {% endif %}
 
-    {% do query_tag.update(session_query_tag_parsed) %}
+    {% do query_tag.update(original_query_tag_parsed) %}
     {% do query_tag.update(env_var_query_tags) %}
 
     {%- do query_tag.update(
         app='dbt',
-        dbt_snowflake_query_tags_version='2.3.1',
+        dbt_snowflake_query_tags_version='2.3.2',
     ) -%}
 
     {% if thread_id %}
@@ -55,7 +55,6 @@
     {% endif %}
 
     {% set query_tag_json = tojson(query_tag) %}
-    {% set original_query_tag = get_current_query_tag() %}
     {{ log("Setting query_tag to '" ~ query_tag_json ~ "'. Will reset to '" ~ original_query_tag ~ "' after materialization.") }}
     {% do run_query("alter session set query_tag = '{}'".format(query_tag_json)) %}
     {{ return(original_query_tag)}}