mobuchowski commented on code in PR #65727:
URL: https://github.com/apache/airflow/pull/65727#discussion_r3160570038
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -862,10 +878,81 @@ class TaskInstanceInfo(InfoJsonEncodable):
class AssetInfo(InfoJsonEncodable):
- """Defines encoding Airflow Asset object to JSON."""
+ """
+ Defines encoding Airflow Asset object to JSON.
+
+ Output shape: ``{"uri", "extra", "type"[, "source_alias"]}``.
+ ``type`` distinguishes how the asset showed up on the task:
+
+ * ``"asset"`` - statically declared on ``task.outlets`` / ``task.inlets``.
+ * ``"asset_event"`` — emitted dynamically via
``outlet_events[asset].extra``.
+ * ``"asset_event_from_alias"`` — resolved at runtime from an ``AssetAlias``
+ via ``outlet_events[alias].add(asset, ...)``; ``source_alias`` then
carries the alias name.
+
+ ``AssetAlias`` are intentionally NOT included - they are just helpers with
no URI of their own.
+ """
includes = ["uri", "extra"]
+ def __init__(self, obj, asset_type: str = "asset", source_alias: str |
None = None):
+ self.type = asset_type
+ self.source_alias = source_alias
+ super().__init__(obj)
+
+ def _extend_fields(self) -> None:
+ if self.type:
+ self._fields.append("type")
+ if self.source_alias:
+ self._fields.append("source_alias")
+
+
+def get_runtime_outlet_assets(task_instance) -> list[tuple[Asset, AssetAlias |
None]]:
+ """
+ Return ``(asset, alias)`` pairs for every runtime-attached outlet.
+
+ * ``outlet_events[AssetAlias].add(asset, ...)`` → ``(asset, alias)``.
+ * ``outlet_events[Asset].extra = {...}`` → ``(asset, None)``.
+
+ Empty list on AF2 or when ``task_instance`` is not an AF3
``RuntimeTaskInstance``
+ (scheduler / API-server path has no usable worker template context).
+ """
+ if not AIRFLOW_V_3_0_PLUS:
+ return []
+ try:
+ from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+ except ImportError:
+ return []
+ if not isinstance(task_instance, RuntimeTaskInstance):
+ return []
+ if not hasattr(task_instance, "get_template_context"):
+ return []
+ outlet_events = task_instance.get_template_context().get("outlet_events")
+ if outlet_events is None:
+ return []
+ results: list[tuple[Asset, AssetAlias | None]] = []
+ for key in outlet_events:
+ accessor = outlet_events[key]
+ if isinstance(key, AssetAlias):
+ for event in getattr(accessor, "asset_alias_events", []) or []:
+ # ``dest_asset_extra`` was introduced in Airflow 3.1.4 ; on
+ # earlier versions the alias event carries only ``extra``.
+ dest_asset_extra = getattr(event, "dest_asset_extra", None) or
{}
+ results.append(
+ (
+ Asset(
+ name=event.dest_asset_key.name,
+ uri=event.dest_asset_key.uri,
+ # Per-event extra wins over asset-level extra.
Review Comment:
That comment introduces only confusion for me :)
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1733,6 +1844,33 @@ def should_use_external_connection(hook) -> bool:
return True
+def _get_providers_manager_instance():
+ """
+ Return the ProvidersManager instance to use for OpenLineage converter
lookup.
+
+ Picks the task-runtime-safe entry point on Airflow 3.2+ and falls back to
the legacy
+ ``airflow.providers_manager.ProvidersManager`` on 3.0 / 3.1 and 2.x.
Returns ``None``
+ if neither import is available (very old Airflow).
Review Comment:
>neither import is available (very old Airflow).
That's not a realistic scenario IMO, we should support 2.11 at this moment
at the lowest.
##########
providers/openlineage/tests/unit/openlineage/extractors/test_manager.py:
##########
@@ -523,3 +523,206 @@ def
test_get_hook_lineage_passes_failed_state(hook_lineage_collector):
sql_extras = mock_sql_fn.call_args.kwargs["sql_extras"]
assert len(sql_extras) == 1
assert
sql_extras[0].value[SqlJobHookLineageExtra.VALUE__SQL_STATEMENT.value] ==
"SELECT 1"
+
+
+def test_extract_inlets_and_outlets_converts_asset_inlet_outlet():
+ """An Airflow ``Asset`` entry goes through ``translate_airflow_asset``
into an OL Dataset."""
+ extractor_manager = ExtractorManager()
+ task = PythonOperator(
+ task_id="task_id",
+ python_callable=lambda x: x,
+ inlets=[Asset(uri="s3://bucket/in.txt", extra={})],
+ outlets=[Asset(uri="s3://bucket/out.txt", extra={})],
+ )
+ lineage = OperatorLineage()
+
+ translated = [
+ OpenLineageDataset(namespace="s3://bucket", name="in.txt"),
+ OpenLineageDataset(namespace="s3://bucket", name="out.txt"),
+ ]
+ with patch(
+
"airflow.providers.openlineage.extractors.manager.translate_airflow_asset",
+ side_effect=translated,
+ ):
+ extractor_manager.extract_inlets_and_outlets(lineage, task)
+
+ assert lineage.inputs == [OpenLineageDataset(namespace="s3://bucket",
name="in.txt")]
+ assert lineage.outputs == [OpenLineageDataset(namespace="s3://bucket",
name="out.txt")]
+
+
+def test_extract_inlets_and_outlets_appends_runtime_outlets():
+ """Runtime outlets (alias resolutions + dynamic events) land on outputs."""
+ extractor_manager = ExtractorManager()
+ task = PythonOperator(task_id="t", python_callable=lambda: 1, outlets=[])
+ lineage = OperatorLineage()
+
+ runtime_assets = [(Asset(uri="s3://bucket/runtime.txt", extra={}), None)]
+ with (
+ patch(
+
"airflow.providers.openlineage.extractors.manager.get_runtime_outlet_assets",
+ return_value=runtime_assets,
+ ),
+ patch(
+
"airflow.providers.openlineage.extractors.manager.translate_airflow_asset",
+ return_value=OpenLineageDataset(namespace="s3://bucket",
name="runtime.txt"),
+ ),
+ ):
+ extractor_manager.extract_inlets_and_outlets(lineage, task,
task_instance=MagicMock())
Review Comment:
Maybe use the `_make_runtime_ti` more widely?
##########
providers/openlineage/src/airflow/providers/openlineage/extractors/manager.py:
##########
@@ -174,18 +175,30 @@ def _get_extractor(self, task: BaseOperator) ->
BaseExtractor | None:
return extractor(task)
return None
- def extract_inlets_and_outlets(self, task_metadata: OperatorLineage, task)
-> None:
+ 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:
- d = self.convert_to_ol_dataset(i)
- if d:
+ if d := self.convert_to_ol_dataset(i):
task_metadata.inputs.append(d)
for o in task.outlets:
- d = self.convert_to_ol_dataset(o)
- if d:
+ 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)
Review Comment:
Just to think - over the whole `get_runtime_outlet_assets ->
translate_airflow_asset -> convert_to_ol_dataset` there's multiple places we
can early return.
Maybe we should have some aggregation of "expected" vs "actual" amount of
assets in the event, so we can have single aggregated info/warn level log that
would inform user about the potential metadata loss.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
