This is an automated email from the ASF dual-hosted git repository.
mobuchowski pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 22f86cb5514 Snowflake, Databricks: emit query lineage via OpenLineage
public API (#73655)
22f86cb5514 is described below
commit 22f86cb55147222149cfe3fe14ad155e14f762aa
Author: Kacper Muda <[email protected]>
AuthorDate: Wed Oct 14 17:23:51 2026 +0200
Snowflake, Databricks: emit query lineage via OpenLineage public API
(#73655)
* Snowflake, Databricks: emit query lineage via OpenLineage public API
Snowflake and Databricks each duplicated their own logic for building
and emitting per-query OpenLineage lineage events, instead of using
the OpenLineage provider's own public API for that. This replaces
both with calls to the provider's `emit_query_lineage` API.
This also fixes a bug where the START and COMPLETE events for a query
shared a single mutable Run/Job object, letting a facet added while
emitting one event leak into, or get corrupted alongside, the other.
* Fix expected-events fixture for the SQL hook-lineage job rename
The hook-lineage query sub-job name changed from "...query.N" to
"...hll_query.N" to avoid colliding with the "...query.N" naming used
by the Snowflake/Databricks query-lineage emitters added in this PR.
The policy/hook-lineage e2e fixture still expected the old name.
* Bump openlineage version in pyproject.toml
* pre commit
# Conflicts:
# uv.lock
* address review comments
---
providers/databricks/docs/index.rst | 2 +-
providers/databricks/pyproject.toml | 2 +-
.../providers/databricks/hooks/databricks_sql.py | 5 +-
.../providers/databricks/utils/openlineage.py | 123 +--
.../unit/databricks/hooks/test_databricks_sql.py | 16 +-
.../unit/databricks/utils/test_openlineage.py | 947 +++++----------------
.../openlineage/utils/sql_hook_lineage.py | 28 +-
.../openlineage_policy_hook_lineage_dag.json | 16 +-
.../openlineage/utils/test_sql_hook_lineage.py | 6 +-
providers/snowflake/docs/index.rst | 2 +-
providers/snowflake/pyproject.toml | 2 +-
.../airflow/providers/snowflake/hooks/snowflake.py | 5 +-
.../providers/snowflake/utils/openlineage.py | 174 +---
.../tests/unit/snowflake/hooks/test_snowflake.py | 17 +-
.../tests/unit/snowflake/utils/test_openlineage.py | 886 +++++--------------
15 files changed, 552 insertions(+), 1679 deletions(-)
diff --git a/providers/databricks/docs/index.rst
b/providers/databricks/docs/index.rst
index 7aed991add9..2ccab4f4743 100644
--- a/providers/databricks/docs/index.rst
+++ b/providers/databricks/docs/index.rst
@@ -158,7 +158,7 @@ Extra Dependencies
``google`` ``apache-airflow-providers-google>=10.24.0``
``sdk`` ``databricks-sdk==0.10.0``
``standard`` ``apache-airflow-providers-standard``
-``openlineage`` ``apache-airflow-providers-openlineage>=2.3.0``
+``openlineage`` ``apache-airflow-providers-openlineage>=2.16.0``
``sqlalchemy`` ``databricks-sqlalchemy>=1.0.2``
==================
================================================================================================================================================================
diff --git a/providers/databricks/pyproject.toml
b/providers/databricks/pyproject.toml
index 28910f667f5..12903d9b8ee 100644
--- a/providers/databricks/pyproject.toml
+++ b/providers/databricks/pyproject.toml
@@ -101,7 +101,7 @@ dependencies = [
"apache-airflow-providers-standard"
]
"openlineage" = [
- "apache-airflow-providers-openlineage>=2.3.0"
+ "apache-airflow-providers-openlineage>=2.16.0"
]
"sqlalchemy" = [
"databricks-sqlalchemy>=1.0.2",
diff --git
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py
index 24be7acc927..8475188d502 100644
---
a/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py
+++
b/providers/databricks/src/airflow/providers/databricks/hooks/databricks_sql.py
@@ -502,7 +502,8 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
self.log.debug("openlineage: getting connection to get database info")
connection = self.get_connection(self.get_conn_id())
- namespace =
SQLParser.create_namespace(self.get_openlineage_database_info(connection))
+ database_info = self.get_openlineage_database_info(connection)
+ namespace = SQLParser.create_namespace(database_info)
self.log.info("Separate OpenLineage events will be emitted for each
Databricks query_id.")
emit_openlineage_events_for_databricks_queries(
@@ -511,6 +512,8 @@ class DatabricksSqlHook(BaseDatabricksHook, DbApiHook):
query_ids=self.query_ids,
query_for_extra_metadata=True,
query_source_namespace=namespace,
+ default_database=database_info.database,
+ default_schema=self.get_openlineage_default_schema(),
)
if len(self.query_ids) == 1:
diff --git
a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py
b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py
index d3ae8f9cc36..e21758dd506 100644
--- a/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py
+++ b/providers/databricks/src/airflow/providers/databricks/utils/openlineage.py
@@ -33,9 +33,6 @@ from airflow.providers.common.compat.openlineage.utils.spark
import (
from airflow.providers.common.compat.sdk import timezone
if TYPE_CHECKING:
- from openlineage.client.event_v2 import RunEvent
- from openlineage.client.facet_v2 import JobFacet
-
from airflow.providers.common.compat.sdk import Context
from airflow.providers.databricks.hooks.databricks import DatabricksHook
from airflow.providers.databricks.hooks.databricks_sql import
DatabricksSqlHook
@@ -63,6 +60,7 @@ def _get_parent_run_facet(task_instance):
lineage_job_name,
lineage_job_namespace,
lineage_root_job_name,
+ lineage_root_job_namespace,
lineage_root_run_id,
lineage_run_id,
)
@@ -72,14 +70,8 @@ def _get_parent_run_facet(task_instance):
parent_job_namespace = lineage_job_namespace()
root_parent_run_id = lineage_root_run_id(task_instance)
- rot_parent_job_name = lineage_root_job_name(task_instance)
-
- try: # Added in OL provider 2.9.0, try to use it if possible
- from airflow.providers.openlineage.plugins.macros import
lineage_root_job_namespace
-
- root_parent_job_namespace = lineage_root_job_namespace(task_instance)
- except ImportError:
- root_parent_job_namespace = lineage_job_namespace()
+ root_parent_job_name = lineage_root_job_name(task_instance)
+ root_parent_job_namespace = lineage_root_job_namespace(task_instance)
return parent_run.ParentRunFacet(
run=parent_run.Run(runId=parent_run_id),
@@ -90,7 +82,7 @@ def _get_parent_run_facet(task_instance):
root=parent_run.Root(
run=parent_run.RootRun(runId=root_parent_run_id),
job=parent_run.RootJob(
- name=rot_parent_job_name,
+ name=root_parent_job_name,
namespace=root_parent_job_namespace,
),
),
@@ -153,38 +145,7 @@ def _get_queries_details_from_databricks(
return query_details
-def _create_ol_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 query 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_databricks_queries(
task_instance,
hook: DatabricksSqlHook | DatabricksHook | None = None,
@@ -193,6 +154,8 @@ def emit_openlineage_events_for_databricks_queries(
query_for_extra_metadata: bool = False,
additional_run_facets: dict | None = None,
additional_job_facets: dict | None = None,
+ default_database: str | None = None,
+ default_schema: str | None = None,
) -> None:
"""
Emit OpenLineage events for executed Databricks queries.
@@ -219,17 +182,15 @@ def emit_openlineage_events_for_databricks_queries(
Must be `False` if `hook` is not provided.
additional_run_facets: Additional run facets to include in OpenLineage
events.
additional_job_facets: Additional job facets to include in OpenLineage
events.
+ default_database: Default database used to qualify table references
parsed out of each query's
+ text that don't already carry their own database qualifier. Callers
with a hook typically pass
+ `hook.get_openlineage_database_info(connection).database`.
+ default_schema: Default schema used to qualify table references parsed
out of each query's text
+ that don't already carry their own schema qualifier. Callers with a
hook typically pass
+ `hook.get_openlineage_default_schema()`.
"""
- from openlineage.client.facet_v2 import job_type_job
-
- from airflow.providers.common.compat.openlineage.facet import (
- ErrorMessageRunFacet,
- ExternalQueryRunFacet,
- RunFacet,
- SQLJobFacet,
- )
- from airflow.providers.openlineage.conf import namespace
- from airflow.providers.openlineage.plugins.listener import
get_openlineage_listener
+ from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
+ from airflow.providers.openlineage.api import emit_query_lineage
log.info("OpenLineage will emit events for Databricks queries.")
@@ -278,19 +239,6 @@ def emit_openlineage_events_for_databricks_queries(
)
default_state = "finished" if default_state in ("running", "success") else
default_state
- log.debug("Generating OpenLineage facets")
- common_run_facets = {"parent": _get_parent_run_facet(task_instance)}
- common_job_facets: dict[str, JobFacet] = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- integration="DATABRICKS",
- processingType="BATCH",
- )
- }
- additional_run_facets = additional_run_facets or {}
- additional_job_facets = additional_job_facets or {}
-
- events: list[RunEvent] = []
for counter, query_id in enumerate(query_ids, 1):
query_metadata = databricks_metadata.get(query_id, {})
log.debug(
@@ -300,36 +248,29 @@ def emit_openlineage_events_for_databricks_queries(
query_metadata if query_metadata else "not found",
)
- query_specific_run_facets: dict[str, RunFacet] = {
- "externalQuery": ExternalQueryRunFacet(externalQueryId=query_id,
source=query_source_namespace)
- }
- if query_metadata.get("error_message"):
- query_specific_run_facets["error"] = ErrorMessageRunFacet(
- message=query_metadata["error_message"],
- programmingLanguage="SQL",
- )
-
- query_specific_job_facets = {}
+ query_job_facets = dict(additional_job_facets or {})
if query_metadata.get("query_text"):
- query_specific_job_facets["sql"] =
SQLJobFacet(query=query_metadata["query_text"])
-
- log.debug("Creating OpenLineage event pair for query ID: %s", query_id)
- event_batch = _create_ol_event_pair(
- job_namespace=namespace(),
-
job_name=f"{task_instance.dag_id}.{task_instance.task_id}.query.{counter}",
+ query_job_facets["sql"] =
SQLJobFacet(query=query_metadata["query_text"])
+
+ log.debug("Emitting OpenLineage event pair for query ID: %s", query_id)
+ emit_query_lineage(
+ query_id=query_id,
+ query_source_namespace=query_source_namespace,
+ # Not forwarding query_text: that would additionally parse it into
input/output datasets,
+ # which is a bigger behavior change than attaching the raw SQL
text as a job facet above.
+ query_text=None,
start_time=query_metadata.get("start_time") or default_event_time,
end_time=query_metadata.get("end_time") or default_event_time,
# Only finished status means it completed without failures
is_successful=(query_metadata.get("status") or
default_state).lower() == "finished",
- run_facets={**query_specific_run_facets, **common_run_facets,
**additional_run_facets},
- job_facets={**query_specific_job_facets, **common_job_facets,
**additional_job_facets},
+ error_message=query_metadata.get("error_message"),
+ default_database=default_database,
+ default_schema=default_schema,
+
job_name=f"{task_instance.dag_id}.{task_instance.task_id}.query.{counter}",
+ task_instance=task_instance,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=query_job_facets,
)
- events.extend(event_batch)
-
- log.debug("Generated %s OpenLineage events; emitting now.", len(events))
- adapter = get_openlineage_listener().adapter
- for event in events:
- adapter.emit(event)
log.info("OpenLineage has successfully finished processing information
about Databricks queries.")
return
diff --git
a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
index 4d736014efd..c8d24bd4c9d 100644
--- a/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
+++ b/providers/databricks/tests/unit/databricks/hooks/test_databricks_sql.py
@@ -662,7 +662,10 @@ def
test_get_openlineage_database_specific_lineage_with_single_query_id(mock_emi
hook = DatabricksSqlHook()
hook.query_ids = ["query1"]
hook.get_connection = mock.MagicMock()
- hook.get_openlineage_database_info = lambda x:
mock.MagicMock(authority="auth", scheme="scheme")
+ hook.get_openlineage_database_info = lambda x: mock.MagicMock(
+ authority="auth", scheme="scheme", database="mydb"
+ )
+ hook.get_openlineage_default_schema = lambda: "myschema"
ti = mock.MagicMock()
@@ -674,6 +677,8 @@ def
test_get_openlineage_database_specific_lineage_with_single_query_id(mock_emi
"query_source_namespace": "scheme://auth",
"task_instance": ti,
"query_for_extra_metadata": True,
+ "default_database": "mydb",
+ "default_schema": "myschema",
}
)
assert result == OperatorLineage(
@@ -686,7 +691,10 @@ def
test_get_openlineage_database_specific_lineage_with_multiple_query_ids(mock_
hook = DatabricksSqlHook()
hook.query_ids = ["query1", "query2"]
hook.get_connection = mock.MagicMock()
- hook.get_openlineage_database_info = lambda x:
mock.MagicMock(authority="auth", scheme="scheme")
+ hook.get_openlineage_database_info = lambda x: mock.MagicMock(
+ authority="auth", scheme="scheme", database="mydb"
+ )
+ hook.get_openlineage_default_schema = lambda: "myschema"
ti = mock.MagicMock()
@@ -698,6 +706,8 @@ def
test_get_openlineage_database_specific_lineage_with_multiple_query_ids(mock_
"query_source_namespace": "scheme://auth",
"task_instance": ti,
"query_for_extra_metadata": True,
+ "default_database": "mydb",
+ "default_schema": "myschema",
}
)
assert result is None
@@ -711,7 +721,7 @@ def
test_get_openlineage_database_specific_lineage_with_old_openlineage_provider
hook.get_openlineage_database_info = lambda x:
mock.MagicMock(authority="auth", scheme="scheme")
expected_err = (
- "OpenLineage provider version `1.99.0` is lower than required `2.5.0`,
"
+ "OpenLineage provider version `1.99.0` is lower than required
`2.16.0`, "
"skipping function `emit_openlineage_events_for_databricks_queries`
execution"
)
with pytest.raises(AirflowOptionalProviderFeatureException,
match=expected_err):
diff --git
a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
index e88e50febaa..84c01e15f53 100644
--- a/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
+++ b/providers/databricks/tests/unit/databricks/utils/test_openlineage.py
@@ -23,19 +23,13 @@ from types import SimpleNamespace
from unittest import mock
import pytest
-from openlineage.client.event_v2 import Job, Run, RunEvent, RunState
-from openlineage.client.facet_v2 import job_type_job, parent_run
+import time_machine
-from airflow.providers.common.compat.openlineage.facet import (
- ErrorMessageRunFacet,
- ExternalQueryRunFacet,
- SQLJobFacet,
-)
+from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, timezone
from airflow.providers.databricks.hooks.databricks import DatabricksHook
from airflow.providers.databricks.hooks.databricks_sql import DatabricksSqlHook
from airflow.providers.databricks.utils.openlineage import (
- _create_ol_event_pair,
_extract_new_clusters_from_databricks_job,
_get_parent_run_facet,
_get_queries_details_from_databricks,
@@ -48,6 +42,8 @@ from airflow.providers.databricks.utils.openlineage import (
from airflow.providers.openlineage.conf import namespace
from airflow.utils.state import TaskInstanceState
+FROZEN_NOW = timezone.datetime(2025, 1, 1, 12, 0, 0)
+
def test_get_parent_run_facet():
logical_date = timezone.datetime(2025, 1, 1)
@@ -212,68 +208,17 @@ def
test_get_queries_details_from_databricks_no_data_found(mock_api_call):
assert details == {}
[email protected]("is_successful", [True, False])
[email protected]("openlineage.client.uuid.generate_new_uuid")
-def test_create_ol_event_pair_success(mock_generate_uuid, is_successful):
- fake_uuid = "01941f29-7c00-7087-8906-40e512c257bd"
- mock_generate_uuid.return_value = fake_uuid
-
- job_namespace = "test_namespace"
- job_name = "test_job"
- start_time = timezone.datetime(2021, 1, 1, 10, 0, 0)
- end_time = timezone.datetime(2021, 1, 1, 10, 30, 0)
- run_facets = {"run_key": "run_value"}
- job_facets = {"job_key": "job_value"}
-
- start_event, end_event = _create_ol_event_pair(
- job_namespace,
- job_name,
- start_time,
- end_time,
- is_successful=is_successful,
- run_facets=run_facets,
- job_facets=job_facets,
- )
-
- assert start_event.eventType == RunState.START
- assert start_event.eventTime == start_time.isoformat()
- assert end_event.eventType == RunState.COMPLETE if is_successful else
RunState.FAIL
- assert end_event.eventTime == end_time.isoformat()
-
- assert start_event.run.runId == fake_uuid
- assert start_event.run.facets == run_facets
-
- assert start_event.job.namespace == job_namespace
- assert start_event.job.name == job_name
- assert start_event.job.facets == job_facets
-
- assert start_event.run is end_event.run
- assert start_event.job == end_event.job
-
-
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
-def test_emit_openlineage_events_for_databricks_queries(mock_generate_uuid,
mock_version, time_machine):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
+def
test_emit_openlineage_events_for_databricks_queries(mock_emit_query_lineage,
mock_version):
query_ids = ["query1", "query2", "query3"]
original_query_ids = copy.deepcopy(query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
- mock_dagrun = mock.MagicMock(logical_date=logical_date, clear_number=0)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state=TaskInstanceState.FAILED, # This will be query default state if
no metadata found
- dag_run=mock_dagrun,
)
- mock_ti.get_template_context.return_value = {"dag_run": mock_dagrun}
fake_metadata = {
"query1": {
@@ -295,21 +240,12 @@ def
test_emit_openlineage_events_for_databricks_queries(mock_generate_uuid, mock
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
+ default_database = "my_catalog"
+ default_schema = "my_schema"
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with (
- mock.patch(
-
"airflow.providers.databricks.utils.openlineage._get_queries_details_from_databricks",
- return_value=fake_metadata,
- ),
- mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ),
+ with mock.patch(
+
"airflow.providers.databricks.utils.openlineage._get_queries_details_from_databricks",
+ return_value=fake_metadata,
):
emit_openlineage_events_for_databricks_queries(
query_ids=query_ids,
@@ -317,772 +253,301 @@ def
test_emit_openlineage_events_for_databricks_queries(mock_generate_uuid, mock
task_instance=mock_ti,
hook=mock.MagicMock(),
query_for_extra_metadata=True,
+ default_database=default_database,
+ default_schema=default_schema,
additional_run_facets=additional_run_facets,
additional_job_facets=additional_job_facets,
)
- assert query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 6 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="DATABRICKS",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
-
- expected_calls = [
- mock.call( # Query1: START event
- RunEvent(
- eventTime="2020-07-21T18:44:46.200000+00:00",
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets={
- "sql": SQLJobFacet(query="SELECT * FROM table1"),
- **expected_common_job_facets,
- },
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event
- RunEvent(
- eventTime="2020-07-21T18:44:47.200000+00:00",
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets={
- "sql": SQLJobFacet(query="SELECT * FROM table1"),
- **expected_common_job_facets,
- },
- ),
- )
- ),
- mock.call( # Query2: START event
- RunEvent(
- eventTime="2020-07-21T18:44:48.200000+00:00",
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query2",
source="databricks_ns"
- ),
- "error": ErrorMessageRunFacet(
- message="Error occurred",
programmingLanguage="SQL"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.2",
- facets={
- "sql": SQLJobFacet(query="SELECT * FROM table2"),
- **expected_common_job_facets,
- },
- ),
- )
- ),
- mock.call( # Query2: FAIL event
- RunEvent(
- eventTime="2020-07-21T18:44:49.200000+00:00",
- eventType=RunState.FAIL,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query2",
source="databricks_ns"
- ),
- "error": ErrorMessageRunFacet(
- message="Error occurred",
programmingLanguage="SQL"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.2",
- facets={
- "sql": SQLJobFacet(query="SELECT * FROM table2"),
- **expected_common_job_facets,
- },
- ),
- )
- ),
- mock.call( # Query3: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(), # no metadata
for query3
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query3",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.3",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query3: FAIL event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(), # no metadata
for query3
- eventType=RunState.FAIL,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query3",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.3",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
+ assert query_ids == original_query_ids # Verify that the input query_ids
list is unchanged.
- assert fake_adapter.emit.call_args_list == expected_calls
+ expected_calls = [
+ mock.call(
+ query_id="query1",
+ query_source_namespace="databricks_ns",
+ query_text=None,
+ default_database=default_database,
+ default_schema=default_schema,
+ start_time=fake_metadata["query1"]["start_time"],
+ end_time=fake_metadata["query1"]["end_time"],
+ is_successful=True,
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets={
+ "custom_job": "value_job",
+ "sql": SQLJobFacet(query="SELECT * FROM table1"),
+ },
+ ),
+ mock.call(
+ query_id="query2",
+ query_source_namespace="databricks_ns",
+ query_text=None,
+ default_database=default_database,
+ default_schema=default_schema,
+ start_time=fake_metadata["query2"]["start_time"],
+ end_time=fake_metadata["query2"]["end_time"],
+ is_successful=False,
+ error_message="Error occurred",
+ job_name="dag_id.task_id.query.2",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets={
+ "custom_job": "value_job",
+ "sql": SQLJobFacet(query="SELECT * FROM table2"),
+ },
+ ),
+ mock.call(
+ query_id="query3",
+ query_source_namespace="databricks_ns",
+ query_text=None,
+ default_database=default_database,
+ default_schema=default_schema,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=False, # no metadata for query3, default state
("failed") is used
+ error_message=None,
+ job_name="dag_id.task_id.query.3",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ ),
+ ]
+ assert mock_emit_query_lineage.call_args_list == expected_calls
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def test_emit_openlineage_events_for_databricks_queries_without_metadata(
- mock_generate_uuid, mock_version, time_machine
+ mock_emit_query_lineage, mock_version
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
query_ids = ["query1"]
original_query_ids = copy.deepcopy(query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state=TaskInstanceState.SUCCESS, # This will be query default state
if no metadata found
- dag_run=mock.MagicMock(logical_date=logical_date, clear_number=0),
)
- mock_ti.get_template_context.return_value = {
- "dag_run": mock.MagicMock(logical_date=logical_date, clear_number=0)
- }
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ):
- emit_openlineage_events_for_databricks_queries(
- query_ids=query_ids,
- query_source_namespace="databricks_ns",
- task_instance=mock_ti,
- hook=mock.MagicMock(),
- # query_for_extra_metadata=False, # False by default
- additional_run_facets=additional_run_facets,
- additional_job_facets=additional_job_facets,
- )
-
- assert query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 2 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="DATABRICKS",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
-
- expected_calls = [
- mock.call( # Query1: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
+ emit_openlineage_events_for_databricks_queries(
+ query_ids=query_ids,
+ query_source_namespace="databricks_ns",
+ task_instance=mock_ti,
+ hook=mock.MagicMock(),
+ # query_for_extra_metadata=False, # False by default
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
- assert fake_adapter.emit.call_args_list == expected_calls
+ assert query_ids == original_query_ids # Verify that the input query_ids
list is unchanged.
+ mock_emit_query_lineage.assert_called_once_with(
+ query_id="query1",
+ query_source_namespace="databricks_ns",
+ query_text=None,
+ default_database=None,
+ default_schema=None,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=True, # no metadata, default state ("success") is used
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets={"custom_job": "value_job"},
+ )
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def
test_emit_openlineage_events_for_databricks_queries_without_explicit_query_ids(
- mock_generate_uuid, mock_version, time_machine
+ mock_emit_query_lineage, mock_version
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
query_ids = ["query1"]
hook = mock.MagicMock()
hook.query_ids = query_ids
original_query_ids = copy.deepcopy(query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state=TaskInstanceState.RUNNING, # This will be query default state
if no metadata found
- dag_run=mock.MagicMock(logical_date=logical_date, clear_number=0),
)
- mock_ti.get_template_context.return_value = {
- "dag_run": mock.MagicMock(logical_date=logical_date, clear_number=0)
- }
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ):
- emit_openlineage_events_for_databricks_queries(
- query_source_namespace="databricks_ns",
- task_instance=mock_ti,
- hook=hook,
- # query_for_extra_metadata=False, # False by default
- additional_run_facets=additional_run_facets,
- additional_job_facets=additional_job_facets,
- )
-
- assert query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 2 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="DATABRICKS",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
-
- expected_calls = [
- mock.call( # Query1: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
+ emit_openlineage_events_for_databricks_queries(
+ query_source_namespace="databricks_ns",
+ task_instance=mock_ti,
+ hook=hook,
+ # query_for_extra_metadata=False, # False by default
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
- assert fake_adapter.emit.call_args_list == expected_calls
+ assert query_ids == original_query_ids # Verify that the input query_ids
list is unchanged.
+ mock_emit_query_lineage.assert_called_once_with(
+ query_id="query1",
+ query_source_namespace="databricks_ns",
+ query_text=None,
+ default_database=None,
+ default_schema=None,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=True, # no metadata, "running" default state is treated
as "finished"
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets={"custom_job": "value_job"},
+ )
@mock.patch(
"airflow.providers.openlineage.sqlparser.SQLParser.create_namespace",
return_value="databricks_ns"
)
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def
test_emit_openlineage_events_for_databricks_queries_without_explicit_query_ids_and_namespace(
- mock_generate_uuid, mock_version, mock_parser, time_machine
+ mock_emit_query_lineage, mock_version, mock_parser
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
query_ids = ["query1"]
hook = mock.MagicMock()
hook.query_ids = query_ids
original_query_ids = copy.deepcopy(query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state=TaskInstanceState.RUNNING, # This will be query default state
if no metadata found
- dag_run=mock.MagicMock(logical_date=logical_date, clear_number=0),
)
- mock_ti.get_template_context.return_value = {
- "dag_run": mock.MagicMock(logical_date=logical_date, clear_number=0)
- }
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ):
- emit_openlineage_events_for_databricks_queries(
- task_instance=mock_ti,
- hook=hook,
- # query_for_extra_metadata=False, # False by default
- additional_run_facets=additional_run_facets,
- additional_job_facets=additional_job_facets,
- )
-
- assert query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 2 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="DATABRICKS",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
-
- expected_calls = [
- mock.call( # Query1: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
+ emit_openlineage_events_for_databricks_queries(
+ task_instance=mock_ti,
+ hook=hook,
+ # query_for_extra_metadata=False, # False by default
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
- assert fake_adapter.emit.call_args_list == expected_calls
+ assert query_ids == original_query_ids # Verify that the input query_ids
list is unchanged.
+ mock_emit_query_lineage.assert_called_once_with(
+ query_id="query1",
+ query_source_namespace="databricks_ns", # resolved from the hook via
the patched SQLParser
+ query_text=None,
+ default_database=None,
+ default_schema=None,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=True, # no metadata, "running" default state is treated
as "finished"
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets={"custom_job": "value_job"},
+ )
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def
test_emit_openlineage_events_for_databricks_queries_without_explicit_query_ids_and_namespace_raw_ns(
- mock_generate_uuid, mock_version, time_machine
+ mock_emit_query_lineage, mock_version
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
query_ids = ["query1"]
hook = DatabricksHook()
hook.query_ids = query_ids
hook.host = "some_host"
original_query_ids = copy.deepcopy(query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state=TaskInstanceState.RUNNING, # This will be query default state
if no metadata found
- dag_run=mock.MagicMock(logical_date=logical_date, clear_number=0),
)
- mock_ti.get_template_context.return_value = {
- "dag_run": mock.MagicMock(logical_date=logical_date, clear_number=0)
- }
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ):
- emit_openlineage_events_for_databricks_queries(
- task_instance=mock_ti,
- hook=hook,
- # query_for_extra_metadata=False, # False by default
- additional_run_facets=additional_run_facets,
- additional_job_facets=additional_job_facets,
- )
-
- assert query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 2 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="DATABRICKS",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
-
- expected_calls = [
- mock.call( # Query1: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks://some_host"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks://some_host"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
+ emit_openlineage_events_for_databricks_queries(
+ task_instance=mock_ti,
+ hook=hook,
+ # query_for_extra_metadata=False, # False by default
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
- assert fake_adapter.emit.call_args_list == expected_calls
+ assert query_ids == original_query_ids # Verify that the input query_ids
list is unchanged.
+ mock_emit_query_lineage.assert_called_once_with(
+ query_id="query1",
+ query_source_namespace="databricks://some_host", # DatabricksHook
lacks OL database info methods
+ query_text=None,
+ default_database=None,
+ default_schema=None,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=True, # no metadata, "running" default state is treated
as "finished"
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets={"custom_job": "value_job"},
+ )
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def
test_emit_openlineage_events_for_databricks_queries_ith_query_ids_and_hook_query_ids(
- mock_generate_uuid, mock_version, time_machine
+ mock_emit_query_lineage, mock_version
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
hook = DatabricksSqlHook()
hook.query_ids = ["query2", "query3"]
query_ids = ["query1"]
original_query_ids = copy.deepcopy(query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state=TaskInstanceState.SUCCESS, # This will be query default state
if no metadata found
- dag_run=mock.MagicMock(logical_date=logical_date, clear_number=0),
)
- mock_ti.get_template_context.return_value = {
- "dag_run": mock.MagicMock(logical_date=logical_date, clear_number=0)
- }
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ):
- emit_openlineage_events_for_databricks_queries(
- query_ids=query_ids,
- query_source_namespace="databricks_ns",
- task_instance=mock_ti,
- hook=hook,
- # query_for_extra_metadata=False, # False by default
- additional_run_facets=additional_run_facets,
- additional_job_facets=additional_job_facets,
- )
-
- assert query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 2 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="DATABRICKS",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
-
- expected_calls = [
- mock.call( # Query1: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1",
source="databricks_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
+ emit_openlineage_events_for_databricks_queries(
+ query_ids=query_ids,
+ query_source_namespace="databricks_ns",
+ task_instance=mock_ti,
+ hook=hook,
+ # query_for_extra_metadata=False, # False by default
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
- assert fake_adapter.emit.call_args_list == expected_calls
+ # The explicitly passed `query_ids=["query1"]` takes precedence over
`hook.query_ids`.
+ assert query_ids == original_query_ids # Verify that the input query_ids
list is unchanged.
+ mock_emit_query_lineage.assert_called_once_with(
+ query_id="query1",
+ query_source_namespace="databricks_ns",
+ query_text=None,
+ default_database=None,
+ default_schema=None,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=True, # no metadata, default state ("success") is used
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets={"custom_job": "value_job"},
+ )
@mock.patch("importlib.metadata.version", return_value="3.0.0")
@@ -1181,7 +646,7 @@ def
test_emit_openlineage_events_with_old_openlineage_provider(mock_version):
return_value=fake_listener,
):
expected_err = (
- "OpenLineage provider version `1.99.0` is lower than required
`2.5.0`, "
+ "OpenLineage provider version `1.99.0` is lower than required
`2.16.0`, "
"skipping function
`emit_openlineage_events_for_databricks_queries` execution"
)
diff --git
a/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
b/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
index 065fea24779..f673971ac1a 100644
---
a/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
+++
b/providers/openlineage/src/airflow/providers/openlineage/utils/sql_hook_lineage.py
@@ -194,7 +194,7 @@ def emit_lineage_from_sql_extras(task_instance, sql_extras:
list, is_successful:
events.extend(
_create_ol_event_pair(
task_instance=task_instance,
-
job_name=f"{task_instance.dag_id}.{task_instance.task_id}.query.{query_count}",
+
job_name=f"{task_instance.dag_id}.{task_instance.task_id}.hll_query.{query_count}",
is_successful=is_successful,
inputs=query_lineage.inputs,
outputs=query_lineage.outputs,
@@ -248,10 +248,12 @@ def _create_ol_event_pair(
Handles parent-run facet generation, run-ID creation and event timestamps
so callers only need to supply the query-specific facets and datasets.
"""
+ job_namespace = lineage_job_namespace()
+
parent_facet = _get_parent_run_facet(
parent_run_id=lineage_run_id(task_instance),
parent_job_name=lineage_job_name(task_instance),
- parent_job_namespace=lineage_job_namespace(),
+ parent_job_namespace=job_namespace,
root_parent_run_id=lineage_root_run_id(task_instance),
root_parent_job_name=lineage_root_job_name(task_instance),
root_parent_job_namespace=lineage_root_job_namespace(task_instance),
@@ -262,28 +264,28 @@ def _create_ol_event_pair(
)
}
- run = Run(
- runId=str(generate_new_uuid(instant=_get_logical_date(task_instance))),
- facets={**(run_facets or {}), **parent_facet},
- )
- job = Job(
- namespace=lineage_job_namespace(), name=job_name,
facets={**(job_facets or {}), **job_type_facet}
- )
+ run_id = str(generate_new_uuid(instant=_get_logical_date(task_instance)))
+ base_run_facets = {**(run_facets or {}), **parent_facet}
+
+ base_job_facets = {**(job_facets or {}), **job_type_facet}
now = timezone.utcnow()
+ # Each event gets its own `Run`/`Job` `facets` dict (same
`runId`/namespace/name) so that facets the
+ # OpenLineage client adds during `emit()` of one event (e.g.
environmentVariables, tags,
+ # sourceCodeLocation) can't leak into, or be seen as already-present by,
the other event in the pair.
start = RunEvent(
eventType=RunState.START,
eventTime=(start_event_time or now).isoformat(),
- run=run,
- job=job,
+ run=Run(runId=run_id, facets=dict(base_run_facets)),
+ job=Job(namespace=job_namespace, name=job_name,
facets=dict(base_job_facets)),
inputs=inputs or [],
outputs=outputs or [],
)
end = RunEvent(
eventType=RunState.COMPLETE if is_successful else RunState.FAIL,
eventTime=(end_event_time or now).isoformat(),
- run=run,
- job=job,
+ run=Run(runId=run_id, facets=dict(base_run_facets)),
+ job=Job(namespace=job_namespace, name=job_name,
facets=dict(base_job_facets)),
inputs=inputs or [],
outputs=outputs or [],
)
diff --git
a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_policy_hook_lineage_dag.json
b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_policy_hook_lineage_dag.json
index af3d3e5ea95..4daa8b78b21 100644
---
a/providers/openlineage/tests/system/openlineage/expected_events/openlineage_policy_hook_lineage_dag.json
+++
b/providers/openlineage/tests/system/openlineage/expected_events/openlineage_policy_hook_lineage_dag.json
@@ -70,7 +70,7 @@
{
"eventType": "START",
"job": {
- "name":
"openlineage_policy_hook_lineage_dag.t_with_sql_hook.query.1",
+ "name":
"openlineage_policy_hook_lineage_dag.t_with_sql_hook.hll_query.1",
"facets": {
"sql": {
"query": "SELECT 1 AS test_col"
@@ -81,7 +81,7 @@
{
"eventType": "COMPLETE",
"job": {
- "name":
"openlineage_policy_hook_lineage_dag.t_with_sql_hook.query.1",
+ "name":
"openlineage_policy_hook_lineage_dag.t_with_sql_hook.hll_query.1",
"facets": {
"sql": {
"query": "SELECT 1 AS test_col"
@@ -105,21 +105,21 @@
"$absent": "hook_lineage=False suppresses SQL child query events",
"eventType": "START",
"job": {
- "name": "openlineage_policy_hook_lineage_dag.t_no_sql_hook.query.1"
+ "name":
"openlineage_policy_hook_lineage_dag.t_no_sql_hook.hll_query.1"
}
},
{
"$absent": "hook_lineage=False suppresses SQL child query events",
"eventType": "COMPLETE",
"job": {
- "name": "openlineage_policy_hook_lineage_dag.t_no_sql_hook.query.1"
+ "name":
"openlineage_policy_hook_lineage_dag.t_no_sql_hook.hll_query.1"
}
},
{
"$absent": "hook_lineage=False suppresses SQL child query events",
"eventType": "FAIL",
"job": {
- "name": "openlineage_policy_hook_lineage_dag.t_no_sql_hook.query.1"
+ "name":
"openlineage_policy_hook_lineage_dag.t_no_sql_hook.hll_query.1"
}
},
{
@@ -154,21 +154,21 @@
"$absent": "extract_operator_metadata=False skips all extraction incl.
SQL child events",
"eventType": "START",
"job": {
- "name":
"openlineage_policy_hook_lineage_dag.t_sql_extract_false.query.1"
+ "name":
"openlineage_policy_hook_lineage_dag.t_sql_extract_false.hll_query.1"
}
},
{
"$absent": "extract_operator_metadata=False skips all extraction incl.
SQL child events",
"eventType": "COMPLETE",
"job": {
- "name":
"openlineage_policy_hook_lineage_dag.t_sql_extract_false.query.1"
+ "name":
"openlineage_policy_hook_lineage_dag.t_sql_extract_false.hll_query.1"
}
},
{
"$absent": "extract_operator_metadata=False skips all extraction incl.
SQL child events",
"eventType": "FAIL",
"job": {
- "name":
"openlineage_policy_hook_lineage_dag.t_sql_extract_false.query.1"
+ "name":
"openlineage_policy_hook_lineage_dag.t_sql_extract_false.hll_query.1"
}
}
]
diff --git
a/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
b/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
index f2e97c568ca..b2a716c97e7 100644
---
a/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
+++
b/providers/openlineage/tests/unit/openlineage/utils/test_sql_hook_lineage.py
@@ -128,7 +128,7 @@ class TestEmitLineageFromSqlExtras:
self.mock_build.assert_called_once()
call = self.mock_build.call_args
assert call.kwargs["task_instance"] is mock_ti
- assert call.kwargs["job_name"] == "dag_id.task_id.query.1"
+ assert call.kwargs["job_name"] == "dag_id.task_id.hll_query.1"
assert call.kwargs["is_successful"] is True
assert call.kwargs["inputs"] == parsed_inputs
assert call.kwargs["outputs"] == parsed_outputs
@@ -160,8 +160,8 @@ class TestEmitLineageFromSqlExtras:
assert self.mock_build.call_count == 2
first, second = self.mock_build.call_args_list
- assert first.kwargs["job_name"] == "dag_id.task_id.query.1"
- assert second.kwargs["job_name"] == "dag_id.task_id.query.2"
+ assert first.kwargs["job_name"] == "dag_id.task_id.hll_query.1"
+ assert second.kwargs["job_name"] == "dag_id.task_id.hll_query.2"
adapter = self.mock_listener.return_value.adapter
assert adapter.emit.call_args_list == [
diff --git a/providers/snowflake/docs/index.rst
b/providers/snowflake/docs/index.rst
index 93c23845fcc..c307c9f68c7 100644
--- a/providers/snowflake/docs/index.rst
+++ b/providers/snowflake/docs/index.rst
@@ -158,7 +158,7 @@ Extra Dependencies
=================== ======================================================
``common.ai`` ``apache-airflow-providers-common-ai[openai]>=0.10.0``
``microsoft.azure`` ``apache-airflow-providers-microsoft-azure>=12.8.0``
-``openlineage`` ``apache-airflow-providers-openlineage>=2.3.0``
+``openlineage`` ``apache-airflow-providers-openlineage>=2.16.0``
=================== ======================================================
Downloading official packages
diff --git a/providers/snowflake/pyproject.toml
b/providers/snowflake/pyproject.toml
index 84de0a45667..7c7f7e580e2 100644
--- a/providers/snowflake/pyproject.toml
+++ b/providers/snowflake/pyproject.toml
@@ -94,7 +94,7 @@ dependencies = [
"apache-airflow-providers-microsoft-azure>=12.8.0"
]
"openlineage" = [
- "apache-airflow-providers-openlineage>=2.3.0"
+ "apache-airflow-providers-openlineage>=2.16.0"
]
[dependency-groups]
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
index 943c144d1ab..6ac118b0db8 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/hooks/snowflake.py
@@ -1119,7 +1119,8 @@ class SnowflakeHook(DbApiHook):
self.log.debug("openlineage: getting connection to get database info")
connection = self.get_connection(self.get_conn_id())
- namespace =
SQLParser.create_namespace(self.get_openlineage_database_info(connection))
+ database_info = self.get_openlineage_database_info(connection)
+ namespace = SQLParser.create_namespace(database_info)
self.log.info("Separate OpenLineage events will be emitted for each
query_id.")
emit_openlineage_events_for_snowflake_queries(
@@ -1128,6 +1129,8 @@ class SnowflakeHook(DbApiHook):
query_ids=self.query_ids,
query_for_extra_metadata=True,
query_source_namespace=namespace,
+ default_database=database_info.database,
+ default_schema=self.get_openlineage_default_schema(),
)
if len(self.query_ids) == 1:
diff --git
a/providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py
b/providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py
index 0dd70c7493b..f001afec78a 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/utils/openlineage.py
@@ -26,9 +26,6 @@ from airflow.providers.common.compat.openlineage.check import
require_openlineag
from airflow.providers.common.compat.sdk import timezone
if TYPE_CHECKING:
- from openlineage.client.event_v2 import RunEvent
- from openlineage.client.facet_v2 import JobFacet
-
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.providers.snowflake.hooks.snowflake_sql_api import
SnowflakeSqlApiHook
@@ -108,54 +105,6 @@ def fix_snowflake_sqlalchemy_uri(uri: str) -> str:
return urlunparse((parts.scheme, hostname, parts.path, parts.params,
parts.query, parts.fragment))
-def _get_parent_run_facet(task_instance):
- """
- Retrieve the ParentRunFacet associated with a specific Airflow task
instance.
-
- This facet helps link OpenLineage events of child jobs - such as queries
executed within
- external systems (e.g., Snowflake) by the Airflow task - to the original
Airflow task execution.
- Establishing this connection enables better lineage tracking and
observability.
- """
- from openlineage.client.facet_v2 import parent_run
-
- from airflow.providers.openlineage.plugins.macros import (
- lineage_job_name,
- lineage_job_namespace,
- lineage_root_job_name,
- lineage_root_run_id,
- lineage_run_id,
- )
-
- parent_run_id = lineage_run_id(task_instance)
- parent_job_name = lineage_job_name(task_instance)
- parent_job_namespace = lineage_job_namespace()
-
- root_parent_run_id = lineage_root_run_id(task_instance)
- rot_parent_job_name = lineage_root_job_name(task_instance)
-
- try: # Added in OL provider 2.9.0, try to use it if possible
- from airflow.providers.openlineage.plugins.macros import
lineage_root_job_namespace
-
- root_parent_job_namespace = lineage_root_job_namespace(task_instance)
- except ImportError:
- root_parent_job_namespace = lineage_job_namespace()
-
- return parent_run.ParentRunFacet(
- run=parent_run.Run(runId=parent_run_id),
- job=parent_run.Job(
- namespace=parent_job_namespace,
- name=parent_job_name,
- ),
- root=parent_run.Root(
- run=parent_run.RootRun(runId=root_parent_run_id),
- job=parent_run.RootJob(
- name=rot_parent_job_name,
- namespace=root_parent_job_namespace,
- ),
- ),
- )
-
-
def _run_single_query_with_hook(hook: SnowflakeHook, sql: str) -> list[dict]:
"""Execute a query against Snowflake without adding extra logging or
instrumentation."""
with closing(hook.get_conn()) as conn:
@@ -189,12 +138,16 @@ 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]}'"
+ # Bound the default 7-day QUERY_HISTORY() scan to the task's execution
window, with a skew buffer.
+ range_start = (task_start_time or timezone.utcnow()) -
datetime.timedelta(minutes=5)
#
https://docs.snowflake.com/en/sql-reference/account-usage#differences-between-account-usage-and-information-schema
# INFORMATION_SCHEMA.QUERY_HISTORY has no latency, so it's better than
ACCOUNT_USAGE.QUERY_HISTORY
# https://docs.snowflake.com/en/sql-reference/functions/query_history
@@ -204,7 +157,10 @@ def _get_queries_details_from_snowflake(
"SELECT "
"QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
"FROM "
- "table(snowflake.information_schema.query_history()) "
+ "table(snowflake.information_schema.query_history("
+ f"end_time_range_start=>to_timestamp_tz('{range_start.isoformat()}'), "
+ "result_limit=>10000"
+ ")) "
f"WHERE "
f"QUERY_ID {query_condition}"
f";"
@@ -230,38 +186,7 @@ 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,
@@ -270,6 +195,8 @@ def emit_openlineage_events_for_snowflake_queries(
query_for_extra_metadata: bool = False,
additional_run_facets: dict | None = None,
additional_job_facets: dict | None = None,
+ default_database: str | None = None,
+ default_schema: str | None = None,
) -> None:
"""
Emit OpenLineage events for executed Snowflake queries.
@@ -296,17 +223,15 @@ def emit_openlineage_events_for_snowflake_queries(
Must be `False` if `hook` is not provided.
additional_run_facets: Additional run facets to include in OpenLineage
events.
additional_job_facets: Additional job facets to include in OpenLineage
events.
+ default_database: Default database used to qualify table references
parsed out of each query's
+ text that don't already carry their own database qualifier. Callers
with a hook typically pass
+ `hook.get_openlineage_database_info(connection).database`.
+ default_schema: Default schema used to qualify table references parsed
out of each query's text
+ that don't already carry their own schema qualifier. Callers with a
hook typically pass
+ `hook.get_openlineage_default_schema()`.
"""
- from openlineage.client.facet_v2 import job_type_job
-
- from airflow.providers.common.compat.openlineage.facet import (
- ErrorMessageRunFacet,
- ExternalQueryRunFacet,
- RunFacet,
- SQLJobFacet,
- )
- from airflow.providers.openlineage.conf import namespace
- from airflow.providers.openlineage.plugins.listener import
get_openlineage_listener
+ from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
+ from airflow.providers.openlineage.api import emit_query_lineage
log.info("OpenLineage will emit events for Snowflake queries.")
@@ -337,7 +262,9 @@ def emit_openlineage_events_for_snowflake_queries(
if query_for_extra_metadata and hook:
log.debug("Retrieving metadata for %s queries from Snowflake.",
len(query_ids))
- snowflake_metadata = _get_queries_details_from_snowflake(hook,
query_ids)
+ snowflake_metadata = _get_queries_details_from_snowflake(
+ hook, query_ids, task_start_time=getattr(task_instance,
"start_date", None)
+ )
else:
log.debug("`query_for_extra_metadata` is False. No extra metadata fill
be fetched from Snowflake.")
snowflake_metadata = {}
@@ -352,18 +279,6 @@ def emit_openlineage_events_for_snowflake_queries(
)
default_state = "success" if default_state == "running" else default_state
- common_run_facets = {"parent": _get_parent_run_facet(task_instance)}
- common_job_facets: dict[str, JobFacet] = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- integration="SNOWFLAKE",
- processingType="BATCH",
- )
- }
- additional_run_facets = additional_run_facets or {}
- additional_job_facets = additional_job_facets or {}
-
- events: list[RunEvent] = []
for counter, query_id in enumerate(query_ids, 1):
query_metadata = snowflake_metadata.get(query_id, {})
log.debug(
@@ -373,36 +288,33 @@ def emit_openlineage_events_for_snowflake_queries(
query_metadata if query_metadata else "not found",
)
- query_specific_run_facets: dict[str, RunFacet] = {
- "externalQuery": ExternalQueryRunFacet(externalQueryId=query_id,
source=query_source_namespace)
- }
- if query_metadata.get("ERROR_MESSAGE"):
- query_specific_run_facets["error"] = ErrorMessageRunFacet(
- message=f"{query_metadata.get('ERROR_CODE')} :
{query_metadata['ERROR_MESSAGE']}",
- programmingLanguage="SQL",
- )
-
- query_specific_job_facets = {}
+ query_job_facets = dict(additional_job_facets or {})
if query_metadata.get("QUERY_TEXT"):
- query_specific_job_facets["sql"] =
SQLJobFacet(query=query_metadata["QUERY_TEXT"])
+ query_job_facets["sql"] =
SQLJobFacet(query=query_metadata["QUERY_TEXT"])
- log.debug("Creating OpenLineage event pair for query ID: %s", query_id)
- event_batch = _create_snowflake_event_pair(
- job_namespace=namespace(),
-
job_name=f"{task_instance.dag_id}.{task_instance.task_id}.query.{counter}",
+ error_message = None
+ if query_metadata.get("ERROR_MESSAGE"):
+ error_message = f"{query_metadata.get('ERROR_CODE')} :
{query_metadata['ERROR_MESSAGE']}"
+
+ log.debug("Emitting OpenLineage event pair for query ID: %s", query_id)
+ emit_query_lineage(
+ query_id=query_id,
+ query_source_namespace=query_source_namespace,
+ # Not forwarding query_text: that would additionally parse it into
input/output datasets,
+ # which is a bigger behavior change than attaching the raw SQL
text as a job facet below.
+ query_text=None,
start_time=query_metadata.get("START_TIME", default_event_time),
end_time=query_metadata.get("END_TIME", default_event_time),
# `EXECUTION_STATUS` can be `success`, `fail` or `incident`
(Snowflake outage, so still failure)
is_successful=query_metadata.get("EXECUTION_STATUS",
default_state).lower() == "success",
- run_facets={**query_specific_run_facets, **common_run_facets,
**additional_run_facets},
- job_facets={**query_specific_job_facets, **common_job_facets,
**additional_job_facets},
+ error_message=error_message,
+ default_database=default_database,
+ default_schema=default_schema,
+
job_name=f"{task_instance.dag_id}.{task_instance.task_id}.query.{counter}",
+ task_instance=task_instance,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=query_job_facets,
)
- events.extend(event_batch)
-
- log.debug("Generated %s OpenLineage events; emitting now.", len(events))
- adapter = get_openlineage_listener().adapter
- for event in events:
- adapter.emit(event)
log.info("OpenLineage has successfully finished processing information
about Snowflake queries.")
return
diff --git a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py
b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py
index 8b13b8f8912..d30f8dcbf51 100644
--- a/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py
+++ b/providers/snowflake/tests/unit/snowflake/hooks/test_snowflake.py
@@ -1218,7 +1218,10 @@ class TestPytestSnowflakeHook:
hook = SnowflakeHook(snowflake_conn_id="test_conn")
hook.query_ids = ["query1"]
hook.get_connection = mock.MagicMock()
- hook.get_openlineage_database_info = lambda x:
mock.MagicMock(authority="auth", scheme="scheme")
+ hook.get_openlineage_database_info = lambda x: mock.MagicMock(
+ authority="auth", scheme="scheme", database="mydb"
+ )
+ hook.get_openlineage_default_schema = lambda: "myschema"
ti = mock.MagicMock()
@@ -1230,6 +1233,8 @@ class TestPytestSnowflakeHook:
"query_source_namespace": "scheme://auth",
"task_instance": ti,
"query_for_extra_metadata": True,
+ "default_database": "mydb",
+ "default_schema": "myschema",
}
)
assert result == OperatorLineage(
@@ -1243,7 +1248,10 @@ class TestPytestSnowflakeHook:
hook = SnowflakeHook(snowflake_conn_id="test_conn")
hook.query_ids = ["query1", "query2"]
hook.get_connection = mock.MagicMock()
- hook.get_openlineage_database_info = lambda x:
mock.MagicMock(authority="auth", scheme="scheme")
+ hook.get_openlineage_database_info = lambda x: mock.MagicMock(
+ authority="auth", scheme="scheme", database="mydb"
+ )
+ hook.get_openlineage_default_schema = lambda: "myschema"
ti = mock.MagicMock()
@@ -1255,6 +1263,8 @@ class TestPytestSnowflakeHook:
"query_source_namespace": "scheme://auth",
"task_instance": ti,
"query_for_extra_metadata": True,
+ "default_database": "mydb",
+ "default_schema": "myschema",
}
)
assert result is None
@@ -1265,9 +1275,10 @@ class TestPytestSnowflakeHook:
hook.query_ids = ["query1", "query2"]
hook.get_connection = mock.MagicMock()
hook.get_openlineage_database_info = lambda x:
mock.MagicMock(authority="auth", scheme="scheme")
+ hook.get_openlineage_default_schema = lambda: "myschema"
expected_err = (
- "OpenLineage provider version `1.99.0` is lower than required
`2.5.0`, "
+ "OpenLineage provider version `1.99.0` is lower than required
`2.16.0`, "
"skipping function `emit_openlineage_events_for_snowflake_queries`
execution"
)
with pytest.raises(AirflowOptionalProviderFeatureException,
match=expected_err):
diff --git a/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
b/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
index 6b27372e337..06a42d3abc9 100644
--- a/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
+++ b/providers/snowflake/tests/unit/snowflake/utils/test_openlineage.py
@@ -21,21 +21,13 @@ import datetime
from unittest import mock
import pytest
-from openlineage.client.event_v2 import Job, Run, RunEvent, RunState
-from openlineage.client.facet_v2 import job_type_job, parent_run
+import time_machine
-from airflow.providers.common.compat.openlineage.facet import (
- ErrorMessageRunFacet,
- ExternalQueryRunFacet,
- SQLJobFacet,
-)
+from airflow.providers.common.compat.openlineage.facet import SQLJobFacet
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, timezone
-from airflow.providers.openlineage.conf import namespace
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from airflow.providers.snowflake.hooks.snowflake_sql_api import
SnowflakeSqlApiHook
from airflow.providers.snowflake.utils.openlineage import (
- _create_snowflake_event_pair,
- _get_parent_run_facet,
_get_queries_details_from_snowflake,
_process_data_from_api,
_run_single_query_with_api_hook,
@@ -46,6 +38,12 @@ from airflow.providers.snowflake.utils.openlineage import (
)
from airflow.utils.state import TaskInstanceState
+FROZEN_NOW = timezone.datetime(2025, 1, 1, 12, 0, 0)
+_EXPECTED_RANGE_START = (FROZEN_NOW -
datetime.timedelta(minutes=5)).isoformat()
+_EXPECTED_TIME_FILTER = (
+ f"end_time_range_start=>to_timestamp_tz('{_EXPECTED_RANGE_START}'),
result_limit=>10000"
+)
+
@pytest.mark.parametrize(
("source", "target"),
@@ -115,30 +113,6 @@ def test_fix_account_name(name, expected):
)
-def test_get_parent_run_facet():
- logical_date = timezone.datetime(2025, 1, 1)
- dr = mock.MagicMock(logical_date=logical_date, clear_number=0)
- mock_ti = mock.MagicMock(
- dag_id="dag_id",
- task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
- state=TaskInstanceState.SUCCESS,
- dag_run=dr,
- )
- mock_ti.get_template_context.return_value = {"dag_run": dr}
-
- result = _get_parent_run_facet(mock_ti)
-
- assert result.run.runId == "01941f29-7c00-7087-8906-40e512c257bd"
- assert result.job.namespace == namespace()
- assert result.job.name == "dag_id.task_id"
- assert result.root.run.runId == "01941f29-7c00-743e-b109-28b18d0a19c5"
- assert result.root.job.namespace == namespace()
- assert result.root.job.name == "dag_id"
-
-
def test_process_data_from_api():
data = [
{
@@ -255,6 +229,7 @@ def
test_get_queries_details_from_snowflake_empty_query_ids():
assert details == {}
+@time_machine.travel(FROZEN_NOW, tick=False)
@mock.patch("airflow.providers.snowflake.utils.openlineage._run_single_query_with_hook")
def
test_get_queries_details_from_snowflake_single_query(mock_run_single_query):
hook = SnowflakeHook(snowflake_conn_id="test_conn")
@@ -275,13 +250,14 @@ def
test_get_queries_details_from_snowflake_single_query(mock_run_single_query):
details = _get_queries_details_from_snowflake(hook, query_ids)
expected_query = (
"SELECT QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
- "FROM table(snowflake.information_schema.query_history()) "
+ f"FROM
table(snowflake.information_schema.query_history({_EXPECTED_TIME_FILTER})) "
"WHERE QUERY_ID = 'ABC';"
)
mock_run_single_query.assert_called_once_with(hook=hook,
sql=expected_query)
assert details == {"ABC": fake_result[0]}
+@time_machine.travel(FROZEN_NOW, tick=False)
@mock.patch("airflow.providers.snowflake.utils.openlineage._run_single_query_with_api_hook")
def
test_get_queries_details_from_snowflake_single_query_api_hook(mock_run_single_query_api):
hook = SnowflakeSqlApiHook(snowflake_conn_id="test_conn")
@@ -303,7 +279,7 @@ def
test_get_queries_details_from_snowflake_single_query_api_hook(mock_run_singl
expected_query = (
"SELECT QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
- "FROM table(snowflake.information_schema.query_history()) "
+ f"FROM
table(snowflake.information_schema.query_history({_EXPECTED_TIME_FILTER})) "
"WHERE QUERY_ID = 'ABC';"
)
expected_details = {
@@ -319,6 +295,7 @@ def
test_get_queries_details_from_snowflake_single_query_api_hook(mock_run_singl
assert details == {"ABC": expected_details}
+@time_machine.travel(FROZEN_NOW, tick=False)
@mock.patch("airflow.providers.snowflake.utils.openlineage._run_single_query_with_hook")
def
test_get_queries_details_from_snowflake_multiple_queries(mock_run_single_query):
hook = SnowflakeHook(snowflake_conn_id="test_conn")
@@ -350,13 +327,14 @@ def
test_get_queries_details_from_snowflake_multiple_queries(mock_run_single_que
expected_query_condition = f"IN {tuple(query_ids)}"
expected_query = (
"SELECT QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
- "FROM table(snowflake.information_schema.query_history()) "
+ f"FROM
table(snowflake.information_schema.query_history({_EXPECTED_TIME_FILTER})) "
f"WHERE QUERY_ID {expected_query_condition};"
)
mock_run_single_query.assert_called_once_with(hook=hook,
sql=expected_query)
assert details == {row["QUERY_ID"]: row for row in fake_result}
+@time_machine.travel(FROZEN_NOW, tick=False)
@mock.patch("airflow.providers.snowflake.utils.openlineage._run_single_query_with_api_hook")
def
test_get_queries_details_from_snowflake_multiple_queries_api_hook(mock_run_single_query_api):
hook = SnowflakeSqlApiHook(snowflake_conn_id="test_conn")
@@ -388,7 +366,7 @@ def
test_get_queries_details_from_snowflake_multiple_queries_api_hook(mock_run_s
expected_query_condition = f"IN {tuple(query_ids)}"
expected_query = (
"SELECT QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
- "FROM table(snowflake.information_schema.query_history()) "
+ f"FROM
table(snowflake.information_schema.query_history({_EXPECTED_TIME_FILTER})) "
f"WHERE QUERY_ID {expected_query_condition};"
)
expected_details = [
@@ -415,6 +393,7 @@ def
test_get_queries_details_from_snowflake_multiple_queries_api_hook(mock_run_s
assert details == {row["QUERY_ID"]: row for row in expected_details}
+@time_machine.travel(FROZEN_NOW, tick=False)
@mock.patch("airflow.providers.snowflake.utils.openlineage._run_single_query_with_hook")
def
test_get_queries_details_from_snowflake_no_data_found(mock_run_single_query):
hook = SnowflakeHook(snowflake_conn_id="test_conn")
@@ -426,13 +405,14 @@ def
test_get_queries_details_from_snowflake_no_data_found(mock_run_single_query)
expected_query_condition = f"IN {tuple(query_ids)}"
expected_query = (
"SELECT QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
- "FROM table(snowflake.information_schema.query_history()) "
+ f"FROM
table(snowflake.information_schema.query_history({_EXPECTED_TIME_FILTER})) "
f"WHERE QUERY_ID {expected_query_condition};"
)
mock_run_single_query.assert_called_once_with(hook=hook,
sql=expected_query)
assert details == {}
+@time_machine.travel(FROZEN_NOW, tick=False)
@mock.patch("airflow.providers.snowflake.utils.openlineage._run_single_query_with_api_hook")
def
test_get_queries_details_from_snowflake_no_data_found_api_hook(mock_run_single_query_api):
hook = SnowflakeSqlApiHook(snowflake_conn_id="test_conn")
@@ -444,13 +424,14 @@ def
test_get_queries_details_from_snowflake_no_data_found_api_hook(mock_run_sing
expected_query_condition = f"IN {tuple(query_ids)}"
expected_query = (
"SELECT QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
- "FROM table(snowflake.information_schema.query_history()) "
+ f"FROM
table(snowflake.information_schema.query_history({_EXPECTED_TIME_FILTER})) "
f"WHERE QUERY_ID {expected_query_condition};"
)
mock_run_single_query_api.assert_called_once_with(hook=hook,
sql=expected_query)
assert details == {}
+@time_machine.travel(FROZEN_NOW, tick=False)
@mock.patch("airflow.providers.snowflake.utils.openlineage._run_single_query_with_hook")
def test_get_queries_details_from_snowflake_error(mock_run_single_query):
hook = SnowflakeHook(snowflake_conn_id="test_conn")
@@ -462,13 +443,14 @@ def
test_get_queries_details_from_snowflake_error(mock_run_single_query):
expected_query_condition = f"IN {tuple(query_ids)}"
expected_query = (
"SELECT QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
- "FROM table(snowflake.information_schema.query_history()) "
+ f"FROM
table(snowflake.information_schema.query_history({_EXPECTED_TIME_FILTER})) "
f"WHERE QUERY_ID {expected_query_condition};"
)
mock_run_single_query.assert_called_once_with(hook=hook,
sql=expected_query)
assert details == {}
+@time_machine.travel(FROZEN_NOW, tick=False)
@mock.patch("airflow.providers.snowflake.utils.openlineage._run_single_query_with_api_hook")
def
test_get_queries_details_from_snowflake_error_api_hook(mock_run_single_query_api):
hook = SnowflakeSqlApiHook(snowflake_conn_id="test_conn")
@@ -480,13 +462,14 @@ def
test_get_queries_details_from_snowflake_error_api_hook(mock_run_single_query
expected_query_condition = f"IN {tuple(query_ids)}"
expected_query = (
"SELECT QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
- "FROM table(snowflake.information_schema.query_history()) "
+ f"FROM
table(snowflake.information_schema.query_history({_EXPECTED_TIME_FILTER})) "
f"WHERE QUERY_ID {expected_query_condition};"
)
mock_run_single_query_api.assert_called_once_with(hook=hook,
sql=expected_query)
assert details == {}
+@time_machine.travel(FROZEN_NOW, tick=False)
@mock.patch("airflow.providers.snowflake.utils.openlineage._process_data_from_api")
@mock.patch("airflow.providers.snowflake.utils.openlineage._run_single_query_with_api_hook")
def test_get_queries_details_from_snowflake_error_api_hook_process_data(
@@ -502,7 +485,7 @@ def
test_get_queries_details_from_snowflake_error_api_hook_process_data(
expected_query_condition = f"IN {tuple(query_ids)}"
expected_query = (
"SELECT QUERY_ID, EXECUTION_STATUS, START_TIME, END_TIME, QUERY_TEXT,
ERROR_CODE, ERROR_MESSAGE "
- "FROM table(snowflake.information_schema.query_history()) "
+ f"FROM
table(snowflake.information_schema.query_history({_EXPECTED_TIME_FILTER})) "
f"WHERE QUERY_ID {expected_query_condition};"
)
mock_run_single_query_api.assert_called_once_with(hook=hook,
sql=expected_query)
@@ -510,70 +493,19 @@ def
test_get_queries_details_from_snowflake_error_api_hook_process_data(
assert details == {}
[email protected]("is_successful", [True, False])
[email protected]("openlineage.client.uuid.generate_new_uuid")
-def test_create_snowflake_event_pair_success(mock_generate_uuid,
is_successful):
- fake_uuid = "01941f29-7c00-7087-8906-40e512c257bd"
- mock_generate_uuid.return_value = fake_uuid
-
- job_namespace = "test_namespace"
- job_name = "test_job"
- start_time = timezone.datetime(2021, 1, 1, 10, 0, 0)
- end_time = timezone.datetime(2021, 1, 1, 10, 30, 0)
- run_facets = {"run_key": "run_value"}
- job_facets = {"job_key": "job_value"}
-
- start_event, end_event = _create_snowflake_event_pair(
- job_namespace,
- job_name,
- start_time,
- end_time,
- is_successful=is_successful,
- run_facets=run_facets,
- job_facets=job_facets,
- )
-
- assert start_event.eventType == RunState.START
- assert start_event.eventTime == start_time.isoformat()
- assert end_event.eventType == RunState.COMPLETE if is_successful else
RunState.FAIL
- assert end_event.eventTime == end_time.isoformat()
-
- assert start_event.run.runId == fake_uuid
- assert start_event.run.facets == run_facets
-
- assert start_event.job.namespace == job_namespace
- assert start_event.job.name == job_name
- assert start_event.job.facets == job_facets
-
- assert start_event.run is end_event.run
- assert start_event.job == end_event.job
-
-
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def test_emit_openlineage_events_for_snowflake_queries_with_extra_metadata(
- mock_generate_uuid, mock_version, time_machine
+ mock_emit_query_lineage, mock_version
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
query_ids = ["query1", "query2", "query3"]
original_query_ids = copy.deepcopy(query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
- mock_dagrun = mock.MagicMock(logical_date=logical_date, clear_number=0)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state=TaskInstanceState.FAILED, # This will be query default state if
no metadata found
- dag_run=mock_dagrun,
)
- mock_ti.get_template_context.return_value = {"dag_run": mock_dagrun}
fake_metadata = {
"query1": {
@@ -596,21 +528,12 @@ def
test_emit_openlineage_events_for_snowflake_queries_with_extra_metadata(
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
+ default_database = "MY_DB"
+ default_schema = "MY_SCHEMA"
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with (
- mock.patch(
-
"airflow.providers.snowflake.utils.openlineage._get_queries_details_from_snowflake",
- return_value=fake_metadata,
- ),
- mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ),
+ with mock.patch(
+
"airflow.providers.snowflake.utils.openlineage._get_queries_details_from_snowflake",
+ return_value=fake_metadata,
):
emit_openlineage_events_for_snowflake_queries(
query_ids=query_ids,
@@ -618,651 +541,254 @@ def
test_emit_openlineage_events_for_snowflake_queries_with_extra_metadata(
task_instance=mock_ti,
hook=mock.MagicMock(),
query_for_extra_metadata=True,
+ default_database=default_database,
+ default_schema=default_schema,
additional_run_facets=additional_run_facets,
additional_job_facets=additional_job_facets,
)
- assert query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 6 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="SNOWFLAKE",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
+ assert query_ids == original_query_ids # Verify that the input query_ids
list is unchanged.
- expected_calls = [
- mock.call( # Query1: START event
- RunEvent(
-
eventTime=fake_metadata["query1"]["START_TIME"].isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets={
- "sql": SQLJobFacet(query="SELECT * FROM table1"),
- **expected_common_job_facets,
- },
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event
- RunEvent(
- eventTime=fake_metadata["query1"]["END_TIME"].isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets={
- "sql": SQLJobFacet(query="SELECT * FROM table1"),
- **expected_common_job_facets,
- },
- ),
- )
- ),
- mock.call( # Query2: START event
- RunEvent(
-
eventTime=fake_metadata["query2"]["START_TIME"].isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query2", source="snowflake_ns"
- ),
- "error": ErrorMessageRunFacet(
- message="ERR001 : Error occurred",
programmingLanguage="SQL"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.2",
- facets={
- "sql": SQLJobFacet(query="SELECT * FROM table2"),
- **expected_common_job_facets,
- },
- ),
- )
- ),
- mock.call( # Query2: FAIL event
- RunEvent(
- eventTime=fake_metadata["query2"]["END_TIME"].isoformat(),
- eventType=RunState.FAIL,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query2", source="snowflake_ns"
- ),
- "error": ErrorMessageRunFacet(
- message="ERR001 : Error occurred",
programmingLanguage="SQL"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.2",
- facets={
- "sql": SQLJobFacet(query="SELECT * FROM table2"),
- **expected_common_job_facets,
- },
- ),
- )
- ),
- mock.call( # Query3: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(), # no metadata
for query3
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query3", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.3",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query3: FAIL event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(), # no metadata
for query3
- eventType=RunState.FAIL,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query3", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.3",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
-
- assert fake_adapter.emit.call_args_list == expected_calls
+ expected_calls = [
+ mock.call(
+ query_id="query1",
+ query_source_namespace="snowflake_ns",
+ query_text=None,
+ default_database=default_database,
+ default_schema=default_schema,
+ start_time=fake_metadata["query1"]["START_TIME"],
+ end_time=fake_metadata["query1"]["END_TIME"],
+ is_successful=True,
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets={
+ "custom_job": "value_job",
+ "sql": SQLJobFacet(query="SELECT * FROM table1"),
+ },
+ ),
+ mock.call(
+ query_id="query2",
+ query_source_namespace="snowflake_ns",
+ query_text=None,
+ default_database=default_database,
+ default_schema=default_schema,
+ start_time=fake_metadata["query2"]["START_TIME"],
+ end_time=fake_metadata["query2"]["END_TIME"],
+ is_successful=False,
+ error_message="ERR001 : Error occurred",
+ job_name="dag_id.task_id.query.2",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets={
+ "custom_job": "value_job",
+ "sql": SQLJobFacet(query="SELECT * FROM table2"),
+ },
+ ),
+ mock.call(
+ query_id="query3",
+ query_source_namespace="snowflake_ns",
+ query_text=None,
+ default_database=default_database,
+ default_schema=default_schema,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=False, # no metadata for query3, default state
("failed") is used
+ error_message=None,
+ job_name="dag_id.task_id.query.3",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ ),
+ ]
+ assert mock_emit_query_lineage.call_args_list == expected_calls
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def test_emit_openlineage_events_for_snowflake_queries_without_extra_metadata(
- mock_generate_uuid, mock_version, time_machine
+ mock_emit_query_lineage, mock_version
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
query_ids = ["query1"]
original_query_ids = copy.deepcopy(query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state=TaskInstanceState.SUCCESS, # This will be query default state
if no metadata found
- dag_run=mock.MagicMock(logical_date=logical_date, clear_number=0),
)
- mock_ti.get_template_context.return_value = {
- "dag_run": mock.MagicMock(logical_date=logical_date, clear_number=0)
- }
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
+ hook = mock.MagicMock()
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ):
- emit_openlineage_events_for_snowflake_queries(
- query_ids=query_ids,
- query_source_namespace="snowflake_ns",
- task_instance=mock_ti,
- hook=mock.MagicMock(),
- # query_for_extra_metadata=False, # False by default
- additional_run_facets=additional_run_facets,
- additional_job_facets=additional_job_facets,
- )
-
- assert query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 2 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="SNOWFLAKE",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
+ emit_openlineage_events_for_snowflake_queries(
+ query_ids=query_ids,
+ query_source_namespace="snowflake_ns",
+ task_instance=mock_ti,
+ hook=hook,
+ # query_for_extra_metadata=False, # False by default
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
- expected_calls = [
- mock.call( # Query1: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
-
- assert fake_adapter.emit.call_args_list == expected_calls
+ assert query_ids == original_query_ids # Verify that the input query_ids
list is unchanged.
+ mock_emit_query_lineage.assert_called_once_with(
+ query_id="query1",
+ query_source_namespace="snowflake_ns",
+ query_text=None,
+ default_database=None,
+ default_schema=None,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=True, # no metadata, default state ("success") is used
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def test_emit_openlineage_events_for_snowflake_queries_without_query_ids(
- mock_generate_uuid, mock_version, time_machine
+ mock_emit_query_lineage, mock_version
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
hook = mock.MagicMock()
hook.query_ids = ["query1"]
original_query_ids = copy.deepcopy(hook.query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state=TaskInstanceState.RUNNING, # Success will be query default
state if no metadata found
- dag_run=mock.MagicMock(logical_date=logical_date, clear_number=0),
)
- mock_ti.get_template_context.return_value = {
- "dag_run": mock.MagicMock(logical_date=logical_date, clear_number=0)
- }
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ):
- emit_openlineage_events_for_snowflake_queries(
- query_ids=[],
- query_source_namespace="snowflake_ns",
- task_instance=mock_ti,
- hook=hook,
- # query_for_extra_metadata=False, # False by default
- additional_run_facets=additional_run_facets,
- additional_job_facets=additional_job_facets,
- )
-
- assert hook.query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 2 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="SNOWFLAKE",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
+ emit_openlineage_events_for_snowflake_queries(
+ query_ids=[],
+ query_source_namespace="snowflake_ns",
+ task_instance=mock_ti,
+ hook=hook,
+ # query_for_extra_metadata=False, # False by default
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
- expected_calls = [
- mock.call( # Query1: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
-
- assert fake_adapter.emit.call_args_list == expected_calls
+ assert hook.query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
+ mock_emit_query_lineage.assert_called_once_with(
+ query_id="query1",
+ query_source_namespace="snowflake_ns",
+ query_text=None,
+ default_database=None,
+ default_schema=None,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=True, # no metadata, "running" default state is treated
as success
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
@mock.patch("airflow.providers.openlineage.sqlparser.SQLParser.create_namespace",
return_value="snowflake_ns")
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def
test_emit_openlineage_events_for_snowflake_queries_without_query_ids_and_namespace(
- mock_generate_uuid, mock_version, mock_parser, time_machine
+ mock_emit_query_lineage, mock_version, mock_parser
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
hook = mock.MagicMock()
hook.query_ids = ["query1"]
original_query_ids = copy.deepcopy(hook.query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state="running", # Success will be query default state if no metadata
found
- dag_run=mock.MagicMock(logical_date=logical_date, clear_number=0),
)
- mock_ti.get_template_context.return_value = {
- "dag_run": mock.MagicMock(logical_date=logical_date, clear_number=0)
- }
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ):
- emit_openlineage_events_for_snowflake_queries(
- query_ids=[],
- query_source_namespace=None,
- task_instance=mock_ti,
- hook=hook,
- # query_for_extra_metadata=False, # False by default
- additional_run_facets=additional_run_facets,
- additional_job_facets=additional_job_facets,
- )
-
- assert hook.query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 2 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="SNOWFLAKE",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
+ emit_openlineage_events_for_snowflake_queries(
+ query_ids=[],
+ query_source_namespace=None,
+ task_instance=mock_ti,
+ hook=hook,
+ # query_for_extra_metadata=False, # False by default
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
- expected_calls = [
- mock.call( # Query1: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query1", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
-
- assert fake_adapter.emit.call_args_list == expected_calls
+ assert hook.query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
+ mock_emit_query_lineage.assert_called_once_with(
+ query_id="query1",
+ query_source_namespace="snowflake_ns", # resolved from the hook via
the patched SQLParser
+ query_text=None,
+ default_database=None,
+ default_schema=None,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=True, # no metadata, "running" default state is treated
as success
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
@mock.patch("importlib.metadata.version", return_value="3.0.0")
[email protected]("openlineage.client.uuid.generate_new_uuid")
[email protected]("airflow.providers.openlineage.api.emit_query_lineage")
+@time_machine.travel(FROZEN_NOW, tick=False)
def
test_emit_openlineage_events_for_snowflake_queries_with_query_ids_and_hook_query_ids(
- mock_generate_uuid, mock_version, time_machine
+ mock_emit_query_lineage, mock_version
):
- fake_uuid = "01958e68-03a2-79e3-9ae9-26865cc40e2f"
- mock_generate_uuid.return_value = fake_uuid
-
- default_event_time = timezone.datetime(2025, 1, 5, 0, 0, 0)
- time_machine.move_to(default_event_time, tick=False)
-
hook = mock.MagicMock()
hook.query_ids = ["query1"]
original_query_ids = copy.deepcopy(hook.query_ids)
- logical_date = timezone.datetime(2025, 1, 1)
mock_ti = mock.MagicMock(
dag_id="dag_id",
task_id="task_id",
- map_index=1,
- try_number=1,
- logical_date=logical_date,
state="running", # Success will be query default state if no metadata
found
- dag_run=mock.MagicMock(logical_date=logical_date, clear_number=0),
)
- mock_ti.get_template_context.return_value = {
- "dag_run": mock.MagicMock(logical_date=logical_date, clear_number=0)
- }
additional_run_facets = {"custom_run": "value_run"}
additional_job_facets = {"custom_job": "value_job"}
- fake_adapter = mock.MagicMock()
- fake_adapter.emit = mock.MagicMock()
- fake_listener = mock.MagicMock()
- fake_listener.adapter = fake_adapter
-
- with mock.patch(
-
"airflow.providers.openlineage.plugins.listener.get_openlineage_listener",
- return_value=fake_listener,
- ):
- emit_openlineage_events_for_snowflake_queries(
- query_ids=["query2"],
- query_source_namespace="snowflake_ns",
- task_instance=mock_ti,
- hook=hook,
- # query_for_extra_metadata=False, # False by default
- additional_run_facets=additional_run_facets,
- additional_job_facets=additional_job_facets,
- )
-
- assert hook.query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
- assert fake_adapter.emit.call_count == 2 # Expect two events per
query.
-
- expected_common_job_facets = {
- "jobType": job_type_job.JobTypeJobFacet(
- jobType="QUERY",
- processingType="BATCH",
- integration="SNOWFLAKE",
- ),
- "custom_job": "value_job",
- }
- expected_common_run_facets = {
- "parent": parent_run.ParentRunFacet(
-
run=parent_run.Run(runId="01941f29-7c00-7087-8906-40e512c257bd"),
- job=parent_run.Job(namespace=namespace(),
name="dag_id.task_id"),
- root=parent_run.Root(
-
run=parent_run.RootRun(runId="01941f29-7c00-743e-b109-28b18d0a19c5"),
- job=parent_run.RootJob(namespace=namespace(),
name="dag_id"),
- ),
- ),
- "custom_run": "value_run",
- }
+ emit_openlineage_events_for_snowflake_queries(
+ query_ids=["query2"],
+ query_source_namespace="snowflake_ns",
+ task_instance=mock_ti,
+ hook=hook,
+ # query_for_extra_metadata=False, # False by default
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
- expected_calls = [
- mock.call( # Query1: START event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.START,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query2", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- mock.call( # Query1: COMPLETE event (no metadata)
- RunEvent(
- eventTime=default_event_time.isoformat(),
- eventType=RunState.COMPLETE,
- run=Run(
- runId=fake_uuid,
- facets={
- "externalQuery": ExternalQueryRunFacet(
- externalQueryId="query2", source="snowflake_ns"
- ),
- **expected_common_run_facets,
- },
- ),
- job=Job(
- namespace=namespace(),
- name="dag_id.task_id.query.1",
- facets=expected_common_job_facets,
- ),
- )
- ),
- ]
-
- assert fake_adapter.emit.call_args_list == expected_calls
+ # The explicitly passed `query_ids=["query2"]` takes precedence over
`hook.query_ids`.
+ assert hook.query_ids == original_query_ids # Verify that the input
query_ids list is unchanged.
+ mock_emit_query_lineage.assert_called_once_with(
+ query_id="query2",
+ query_source_namespace="snowflake_ns",
+ query_text=None,
+ default_database=None,
+ default_schema=None,
+ start_time=FROZEN_NOW,
+ end_time=FROZEN_NOW,
+ is_successful=True, # no metadata, "running" default state is treated
as success
+ error_message=None,
+ job_name="dag_id.task_id.query.1",
+ task_instance=mock_ti,
+ additional_run_facets=additional_run_facets,
+ additional_job_facets=additional_job_facets,
+ )
@mock.patch("importlib.metadata.version", return_value="3.0.0")
@@ -1354,7 +880,7 @@ def
test_emit_openlineage_events_with_old_openlineage_provider(mock_version):
return_value=fake_listener,
):
expected_err = (
- "OpenLineage provider version `1.99.0` is lower than required
`2.5.0`, "
+ "OpenLineage provider version `1.99.0` is lower than required
`2.16.0`, "
"skipping function `emit_openlineage_events_for_snowflake_queries`
execution"
)
