Source code for airflow.providers.openlineage.extractors.manager
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
import importlib
import types
from collections.abc import Iterator
from typing import TYPE_CHECKING
from airflow.providers.common.compat.openlineage.utils.utils import translate_airflow_asset
from airflow.providers.openlineage import conf
from airflow.providers.openlineage.extractors import BaseExtractor, OperatorLineage
from airflow.providers.openlineage.extractors.base import (
OL_METHOD_NAME_COMPLETE,
OL_METHOD_NAME_START,
DefaultExtractor,
)
from airflow.providers.openlineage.extractors.bash import BashExtractor
from airflow.providers.openlineage.extractors.python import PythonExtractor
from airflow.providers.openlineage.utils.emission_policy import EmissionPolicy
from airflow.providers.openlineage.utils.utils import (
get_runtime_outlet_assets,
get_unknown_source_attribute_run_facet,
try_import_from_string,
)
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.state import TaskInstanceState
if TYPE_CHECKING:
from openlineage.client.event_v2 import Dataset
from airflow.providers.common.compat.lineage.entities import Table
from airflow.providers.common.compat.sdk import BaseOperator
def _is_hook_lineage_collector_created() -> bool:
"""
Return False only if the hook lineage collector was certainly never created in this process.
Hooks report lineage through this process-wide collector, so if it was never created, nothing was
collected. Creating it just to read it back empty imports every provider's asset URI handlers,
which takes seconds. The check reads Airflow internals, so any state it does not recognize, such
as a getter that is not a plain function, counts as created.
"""
try:
from airflow.sdk import lineage
except ImportError:
# Airflow < 3.2 keeps the collector in a module global instead of a cached getter.
try:
hook = importlib.import_module("airflow.lineage.hook")
except ImportError:
return True
if getattr(hook, "_hook_lineage_collector", True) is not None:
return True
return type(getattr(hook, "get_hook_lineage_collector", None)) is not types.FunctionType
try:
return lineage.get_hook_lineage_collector.cache_info().currsize != 0
except AttributeError:
return True
def _iter_extractor_types() -> Iterator[type[BaseExtractor]]:
if PythonExtractor is not None:
yield PythonExtractor
if BashExtractor is not None:
yield BashExtractor
[docs]
class ExtractorManager(LoggingMixin):
"""Class abstracting management of custom extractors."""
def __init__(self):
super().__init__()
# Built-in Extractors like Bash and Python
for extractor in _iter_extractor_types():
for operator_class in extractor.get_operator_classnames():
self.extractors[operator_class] = extractor
for extractor_path in conf.custom_extractors():
extractor: type[BaseExtractor] | None = try_import_from_string(extractor_path)
if not extractor:
self.log.warning(
"OpenLineage is unable to import custom extractor `%s`; will ignore it.",
extractor_path,
)
continue
for operator_class in extractor.get_operator_classnames():
if operator_class in self.extractors:
self.log.warning(
"Duplicate OpenLineage custom extractor found for `%s`. "
"`%s` will be used instead of `%s`",
operator_class,
extractor_path,
self.extractors[operator_class],
)
self.extractors[operator_class] = extractor
self.log.debug(
"Registered custom OpenLineage extractor `%s` for class `%s`",
extractor_path,
operator_class,
)
[docs]
def add_extractor(self, operator_class: str, extractor: type[BaseExtractor]):
self.extractors[operator_class] = extractor
[docs]
def extract_metadata(
self,
dagrun,
task,
task_instance_state: TaskInstanceState,
task_instance,
controls: EmissionPolicy | None = None,
) -> OperatorLineage:
if controls is None:
controls = EmissionPolicy.defaults()
extractor = self._get_extractor(task, source_code_enabled=controls.include_source_code)
task_info = (
f"task_type={task.task_type} "
f"airflow_dag_id={task.dag_id} "
f"task_id={task.task_id} "
f"airflow_run_id={dagrun.run_id} "
)
if extractor:
# Extracting advanced metadata is only possible when extractor for particular operator
# is defined. Without it, we can't extract any input or output data.
try:
self.log.debug(
"Using extractor %s %s",
extractor.__class__.__name__,
str(task_info),
)
if task_instance_state == TaskInstanceState.RUNNING:
task_metadata = extractor.extract()
elif task_instance_state == TaskInstanceState.FAILED:
if callable(getattr(extractor, "extract_on_failure", None)):
task_metadata = extractor.extract_on_failure(task_instance)
else:
task_metadata = extractor.extract_on_complete(task_instance)
else:
task_metadata = extractor.extract_on_complete(task_instance)
self.log.debug(
"Found task metadata for operation %s: %s",
task.task_id,
str(task_metadata),
)
task_metadata = self.validate_task_metadata(task_metadata) or OperatorLineage()
# If no inputs and outputs are present - check Hook Lineage if enabled
if (not task_metadata.inputs) and (not task_metadata.outputs):
if controls.hook_lineage:
hook_lineage = self.get_hook_lineage(task_instance, task_instance_state)
if hook_lineage is not None:
task_metadata = task_metadata.merge(hook_lineage)
else: # Last resort - check manual annotations
self.extract_inlets_and_outlets(task_metadata, task, task_instance)
else:
self.log.info(
"Skipping OpenLineage hook lineage collection for task '%s' due to emission_policy.",
task.task_id,
)
self.extract_inlets_and_outlets(task_metadata, task, task_instance)
return task_metadata
except Exception as e:
self.log.warning(
"Failed to extract OpenLineage metadata using found extractor %s - %s %s",
extractor,
e,
task_info,
)
self.log.debug("OpenLineage extraction failure details:", exc_info=True)
else:
# No extractor found.
if controls.hook_lineage:
# Fall back to hook lineage. This call must be wrapped in try/except: it runs
# emit_lineage_from_sql_extras → _create_ol_event_pair which is not guarded
# internally. An uncaught exception here would propagate up to the listener's
# @print_warning decorator, silently suppressing the task-level event.
try:
hook_lineage = self.get_hook_lineage(task_instance, task_instance_state)
except Exception as e:
self.log.warning(
"Failed to extract OpenLineage hook lineage %s: %s. Task event will be emitted without lineage.",
task_info,
e,
)
self.log.debug("OpenLineage hook lineage failure details:", exc_info=True)
hook_lineage = None
if hook_lineage is not None:
return hook_lineage
else:
self.log.info(
"Skipping OpenLineage hook lineage collection for task '%s' due to emission_policy.",
task.task_id,
)
# No extractor and no hook lineage (or hook lineage disabled) — always emit the
# unknownSourceAttribute facet and fall back to manually-declared inlets/outlets.
self.log.debug("Unable to find an OpenLineage extractor %s", task_info)
task_metadata = OperatorLineage(
run_facets=get_unknown_source_attribute_run_facet(task=task),
)
self.extract_inlets_and_outlets(task_metadata, task, task_instance)
return task_metadata
return OperatorLineage()
[docs]
def get_extractor_class(self, task: BaseOperator) -> type[BaseExtractor] | None:
if task.task_type in self.extractors:
return self.extractors[task.task_type]
def method_exists(method_name):
return callable(getattr(task, method_name, None))
if method_exists(OL_METHOD_NAME_START) or method_exists(OL_METHOD_NAME_COMPLETE):
return self.default_extractor
return None
def _get_extractor(self, task: BaseOperator, source_code_enabled: bool = True) -> BaseExtractor | None:
extractor_cls = self.get_extractor_class(task)
self.log.debug("extractor for %s is %s", task.task_type, extractor_cls)
if extractor_cls is None:
return None
# Set the flag after construction rather than passing it as a constructor kwarg:
# custom extractors may still use the historically-public `__init__(self, operator)`
# signature, which would raise TypeError on an unexpected keyword and drop the whole
# task event. Built-in extractors read `source_code_enabled` at extraction time.
extractor = extractor_cls(task)
extractor.source_code_enabled = source_code_enabled
return extractor
[docs]
def extract_inlets_and_outlets(
self,
task_metadata: OperatorLineage,
task,
task_instance=None,
) -> None:
if task.inlets or task.outlets:
self.log.debug("Manually extracting lineage metadata from inlets and outlets")
for i in task.inlets:
if d := self.convert_to_ol_dataset(i):
task_metadata.inputs.append(d)
for o in task.outlets:
if d := self.convert_to_ol_dataset(o):
task_metadata.outputs.append(d)
# Add runtime-emitted outlets (alias resolutions + dynamic asset events), deduped
# by (namespace, name) against both static outputs and each other.
seen = {(d.namespace, d.name) for d in task_metadata.outputs}
for asset, _alias in get_runtime_outlet_assets(task_instance):
ol = translate_airflow_asset(asset, None)
if ol is not None and (ol.namespace, ol.name) not in seen:
task_metadata.outputs.append(ol)
seen.add((ol.namespace, ol.name))
[docs]
def get_hook_lineage(
self,
task_instance=None,
task_instance_state: TaskInstanceState | None = None,
) -> OperatorLineage | None:
"""
Extract lineage from the Hook Lineage Collector.
Combines two sources into a single :class:`OperatorLineage`:
* **Asset-based** inputs/outputs reported via ``add_input_asset`` / ``add_output_asset``.
* **SQL-based** lineage from ``sql_job`` extras reported via
:func:`~airflow.providers.common.sql.hooks.lineage.send_sql_hook_lineage`.
When ``task_instance`` is provided, each extra is parsed and separate per-query
OpenLineage events are emitted.
Returns ``None`` when nothing was collected.
"""
try:
from airflow.providers.common.compat.lineage.hook import get_hook_lineage_collector
from airflow.providers.common.sql.hooks.lineage import SqlJobHookLineageExtra
except ImportError:
return None
if not _is_hook_lineage_collector_created():
return None
collector = get_hook_lineage_collector()
if not hasattr(collector, "has_collected"):
return None
if not collector.has_collected:
return None
self.log.debug("OpenLineage will extract lineage from Hook Lineage Collector.")
collected = collector.collected_assets
# Asset-based inputs/outputs - keep only assets that can be translated to OL datasets
inputs = [
asset
for asset_info in collected.inputs
if (asset := translate_airflow_asset(asset_info.asset, asset_info.context)) is not None
]
outputs = [
asset
for asset_info in collected.outputs
if (asset := translate_airflow_asset(asset_info.asset, asset_info.context)) is not None
]
# SQL-based lineage - keep only SQL extra with query_text or job_id.
sql_extras = [
info
for info in collected.extra
if info.key == SqlJobHookLineageExtra.KEY.value
and (
info.value.get(SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value)
or info.value.get(SqlJobHookLineageExtra.VALUE__JOB_ID.value)
)
]
if sql_extras:
from airflow.providers.openlineage.utils.sql_hook_lineage import emit_lineage_from_sql_extras
self.log.debug("Found %s sql_job extra(s) in Hook Lineage Collector.", len(sql_extras))
emit_lineage_from_sql_extras(
task_instance=task_instance,
sql_extras=sql_extras,
is_successful=task_instance_state != TaskInstanceState.FAILED,
)
if not inputs and not outputs:
return None
return OperatorLineage(inputs=inputs, outputs=outputs)
@staticmethod
[docs]
def convert_to_ol_dataset_from_object_storage_uri(uri: str) -> Dataset | None:
from urllib.parse import urlparse
from openlineage.client.event_v2 import Dataset
if "/" not in uri:
return None
try:
scheme, netloc, path, params, _, _ = urlparse(uri)
except Exception:
return None
common_schemas = {
"s3": "s3",
"gs": "gs",
"gcs": "gs",
"hdfs": "hdfs",
"file": "file",
}
for found, final in common_schemas.items():
if scheme.startswith(found):
return Dataset(namespace=f"{final}://{netloc}", name=path.lstrip("/"))
return Dataset(namespace=scheme, name=f"{netloc}{path}")
@staticmethod
[docs]
def convert_to_ol_dataset_from_table(table: Table) -> Dataset:
from openlineage.client.event_v2 import Dataset
from openlineage.client.facet_v2 import (
DatasetFacet,
documentation_dataset,
ownership_dataset,
schema_dataset,
)
facets: dict[str, DatasetFacet] = {}
if table.columns:
facets["schema"] = schema_dataset.SchemaDatasetFacet(
fields=[
schema_dataset.SchemaDatasetFacetFields(
name=column.name,
type=column.data_type,
description=column.description,
)
for column in table.columns
]
)
if table.owners:
facets["ownership"] = ownership_dataset.OwnershipDatasetFacet(
owners=[
ownership_dataset.Owner(
# f.e. "user:John Doe <jdoe@company.com>" or just "user:<jdoe@company.com>"
name=f"user:"
f"{user.first_name + ' ' if user.first_name else ''}"
f"{user.last_name + ' ' if user.last_name else ''}"
f"<{user.email}>",
type="",
)
for user in table.owners
]
)
if table.description:
facets["documentation"] = documentation_dataset.DocumentationDatasetFacet(
description=table.description
)
return Dataset(
namespace=f"{table.cluster}",
name=f"{table.database}.{table.name}",
facets=facets,
)
@staticmethod
[docs]
def convert_to_ol_dataset(obj) -> Dataset | None:
from openlineage.client.event_v2 import Dataset
from airflow.providers.common.compat.assets import Asset
from airflow.providers.common.compat.lineage.entities import File, Table
if isinstance(obj, Dataset):
return obj
if isinstance(obj, Asset):
return translate_airflow_asset(obj, None)
if isinstance(obj, Table):
return ExtractorManager.convert_to_ol_dataset_from_table(obj)
if isinstance(obj, File):
return ExtractorManager.convert_to_ol_dataset_from_object_storage_uri(obj.url)
return None
[docs]
def validate_task_metadata(self, task_metadata) -> OperatorLineage | None:
try:
return OperatorLineage(
inputs=task_metadata.inputs,
outputs=task_metadata.outputs,
run_facets=task_metadata.run_facets,
job_facets=task_metadata.job_facets,
)
except AttributeError:
self.log.warning("OpenLineage extractor returns non-valid metadata: `%s`", task_metadata)
return None