WIP: Fix setting Snowflake query tags #2187
Draft
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR attempts to run the pre- and post-queries as part of the same Snowflake session to address an issue the customer reported when trying to set query tags.
Problem
A customer reported that they could not set query tags in Snowflake when using an Astro SDK 1.7 alpha release ( https://astronomer.zendesk.com/agent/tickets/62628).
They didn't get any errors while running the operator with
session_modifier.pre_queries = ["ALTER SESSION SET QUERY_TAG=MyQueryTag"]
, but they couldn't see the desired queries being annotated with the expected query tags:Context
The way the feature was implemented was to attempt to leverage query tags per session, as described in:
https://www.chaosgenius.io/blog/snowflake-query-tags/
The relevant code in Astro SDK (added in #1898 #1962) is:
astro-sdk/python-sdk/src/astro/databases/base.py
Lines 162 to 168 in 33ca675
https://github.com/astronomer/astro-sdk/blob/33ca6758f8d4052faba21e4579f358cda232dc98/python-sdk/src/astro/databases/base.py#L106C9-L128
As can be seen, the ALTER query statement is being run each time with a different connection/session:
astro-sdk/python-sdk/src/astro/databases/base.py
Lines 96 to 99 in 33ca675
The base method is being used since Snowflake does not override the property connection:
https://github.com/astronomer/astro-sdk/blob/main/python-sdk/src/astro/databases/snowflake.py
What is missing in this PR
Still need to confirm this fix works in practice.
Create an integration test: