This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch backport-22f86cb-v3-4-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 26431fad3a346e8c766a38b83b4fbd76b37d45a5 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 (cherry picked from commit 22f86cb55147222149cfe3fe14ad155e14f762aa) --- 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 5e0b4914ce6..61814dab216 100644 --- a/providers/snowflake/docs/index.rst +++ b/providers/snowflake/docs/index.rst @@ -159,7 +159,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 c8b362e0acf..b682aabdd7e 100644 --- a/providers/snowflake/pyproject.toml +++ b/providers/snowflake/pyproject.toml @@ -99,7 +99,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" )
