mobuchowski commented on code in PR #73655:
URL: https://github.com/apache/airflow/pull/73655#discussion_r4258169003
##########
providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py:
##########
@@ -189,12 +138,20 @@ def _process_data_from_api(data: list[dict[str, Any]]) ->
list[dict[str, Any]]:
def _get_queries_details_from_snowflake(
- hook: SnowflakeHook | SnowflakeSqlApiHook, query_ids: list[str]
+ hook: SnowflakeHook | SnowflakeSqlApiHook,
+ query_ids: list[str],
+ task_start_time: datetime.datetime | None = None,
) -> dict[str, dict[str, Any]]:
"""Retrieve execution details for specific queries from Snowflake's query
history."""
if not query_ids:
return {}
query_condition = f"IN {tuple(query_ids)}" if len(query_ids) > 1 else f"=
'{query_ids[0]}'"
+ # QUERY_HISTORY() defaults to scanning the last 7 days
(`END_TIME_RANGE_START`/`END_TIME_RANGE_END`
+ # default to now / now - 10080 minutes), which is a needlessly expensive
account-wide scan when we
+ # already know roughly when these queries ran. Bound it to the task's own
execution window instead,
+ # with a small buffer on both sides for clock skew between Airflow and
Snowflake.
+ range_start = (task_start_time or timezone.utcnow()) -
datetime.timedelta(minutes=5)
+ range_end = timezone.utcnow() + datetime.timedelta(minutes=5)
Review Comment:
Why add `+ datetime.timedelta(minutes=5)`
I don't think it makes sense to add the buffer there.
https://docs.snowflake.com/en/sql-reference/functions/query_history#arguments
```
If the time range does not fall within the last 7 days, an error is returned.
```
If you want to make sure you process it till the end, you might as well not
pass the query end there.
##########
providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py:
##########
@@ -230,44 +191,15 @@ def _get_queries_details_from_snowflake(
return {row["QUERY_ID"]: row for row in result} if result else {}
-def _create_snowflake_event_pair(
- job_namespace: str,
- job_name: str,
- start_time: datetime.datetime,
- end_time: datetime.datetime,
- is_successful: bool,
- run_facets: dict | None = None,
- job_facets: dict | None = None,
-) -> tuple[RunEvent, RunEvent]:
- """Create a pair of OpenLineage RunEvents representing the start and end
of a Snowflake job execution."""
- from openlineage.client.event_v2 import Job, Run, RunEvent, RunState
- from openlineage.client.uuid import generate_new_uuid
-
- run = Run(runId=str(generate_new_uuid()), facets=run_facets or {})
- job = Job(namespace=job_namespace, name=job_name, facets=job_facets or {})
-
- start = RunEvent(
- eventType=RunState.START,
- eventTime=start_time.isoformat(),
- run=run,
- job=job,
- )
- end = RunEvent(
- eventType=RunState.COMPLETE if is_successful else RunState.FAIL,
- eventTime=end_time.isoformat(),
- run=run,
- job=job,
- )
- return start, end
-
-
-@require_openlineage_version(provider_min_version="2.5.0")
+@require_openlineage_version(provider_min_version="2.16.0")
def emit_openlineage_events_for_snowflake_queries(
task_instance,
hook: SnowflakeHook | SnowflakeSqlApiHook | None = None,
query_ids: list[str] | None = None,
query_source_namespace: str | None = None,
query_for_extra_metadata: bool = False,
+ default_database: str | None = None,
+ default_schema: str | None = None,
Review Comment:
Add new arguments at the end, to preserve compatibility with previous
positional arguments
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
