mobuchowski commented on code in PR #52161:
URL: https://github.com/apache/airflow/pull/52161#discussion_r2174705185
##########
providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py:
##########
@@ -204,9 +205,29 @@ def _run_single_query_with_hook(hook: SnowflakeHook, sql:
str) -> list[dict]:
return result
+def _run_single_query_with_api_hook(hook: SnowflakeSqlApiHook, sql: str) ->
list[dict[str, Any]]:
+ """Execute a query against Snowflake API without adding extra logging or
instrumentation."""
+ # `hook.execute_query` resets the query_ids, so we need to save them and
re-assign after we're done
+ query_ids_before_execution = list(hook.query_ids)
+ try:
+ _query_ids = hook.execute_query(sql=sql, statement_count=0)
+ hook.wait_for_query(query_id=_query_ids[0], raise_error=True,
poll_interval=1, timeout=3)
+ return
hook.get_result_from_successful_sql_api_query(query_id=_query_ids[0])
+ finally:
+ hook.query_ids = query_ids_before_execution
Review Comment:
Wouldn't it be better to move the raw `execute_query` logic into something
like internal `_execute_query` and keep the `query_ids` behavior, but just call
the `_` version when you don't want to store the query id?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
