Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

WIP: Fix setting Snowflake query tags #2187

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Conversation

tatiana
Copy link
Collaborator

@tatiana tatiana commented Sep 20, 2024

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:

It is assumed that the param query_modifier is used in the decorator and specifies the alter session query tag as pre_queries to send that query tag before any query is launched by the function. Still, it happens to create two different sessions, and the queries launched by the function don't have the tag.

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:

for sql_query in query_modifier.pre_queries:
_ = self.run_single_sql_query(sql_query, parameters)
result = self.run_single_sql_query(sql, parameters)
for sql_query in query_modifier.post_queries:
_ = self.run_single_sql_query(sql_query, parameters)

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:

@property
def connection(self) -> sqlalchemy.engine.base.Connection:
"""Return a Sqlalchemy connection object for the given database."""
return self.sqlalchemy_engine.connect()

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:

  • An operator creating the query tag session_modifier.pre_queries = ["ALTER SESSION SET QUERY_TAG=MyQueryTag"]
  • A check, before executing the operator, that there was no QueryTag.
  • A check, after running the operator, that the query tag was available in the database after running the operator, e.g.:
USE ROLE ACCOUNTADMIN;

SELECT user_name, role_name, query_tag
FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY
WHERE query_tag = 'MyQueryTag'

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant