You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "kacpermuda (via GitHub)" <gi...@apache.org> on 2025/11/07 16:37:07 UTC
[PR] feat: Adjust HookLevelLineage with add_extra [airflow]
kacpermuda opened a new pull request, #58057:
URL: https://github.com/apache/airflow/pull/58057
<!--
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.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
Followup to #57620, adjusting common.compat provider to work with new Hook Level Lineage feature.
<!-- Please keep an empty line above the dashes. -->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Adjust common.compat HookLevelLineage collector for new add_extra method [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on PR #58057:
URL: https://github.com/apache/airflow/pull/58057#issuecomment-3551891929
@ashb How do you feel about the above? Can we proceed with the polyfill approach?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Adjust common.compat HookLevelLineage collector for new add_extra method [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on PR #58057:
URL: https://github.com/apache/airflow/pull/58057#issuecomment-3532980231
> It's definitely not beautiful... but can't think of better way of making this work on older versions.
Could not agree more, and I think it can bring a lot of value for older Airflow versions.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Adjust common.compat HookLevelLineage collector for new add_extra method [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #58057:
URL: https://github.com/apache/airflow/pull/58057#discussion_r2557089799
##########
providers/common/compat/src/airflow/providers/common/compat/lineage/hook.py:
##########
@@ -72,17 +202,34 @@ def collected_assets_compat(collector) -> HookLineage:
setattr(
collector.__class__,
"collected_assets",
- property(lambda collector: collected_assets_compat(collector)),
+ property(lambda c: _compat_collected_assets(c)),
)
+ # Add `add_extra` polyfill for AF 2.x, needs to be called after setting `collected_assets`
+ collector = _add_extra_polyfill(collector)
+
return collector
def get_hook_lineage_collector():
- # Dataset has been renamed as Asset in 3.0
- if AIRFLOW_V_3_0_PLUS:
- from airflow.lineage.hook import get_hook_lineage_collector
+ """
+ Get hook lineage collector with appropriate compatibility layers.
+
+ - AF 2.x: Apply both dataset->asset rename and then `add_extra` polyfill
+ - AF 3.0-3.1: Apply only `add_extra` polyfill, no renaming needed
+ - AF 3.2+: Use native implementation (no renaming or polyfill needed)
+ """
+ from airflow.lineage.hook import get_hook_lineage_collector as get_global_collector
+
+ global_collector = get_global_collector()
+
+ # AF 2.x: needs dataset->asset rename + `add_extra` polyfill
+ if not AIRFLOW_V_3_0_PLUS:
+ return _get_af2_asset_compat_hook_lineage_collector(global_collector)
- return get_hook_lineage_collector()
+ # AF 3.0-3.1: needs only polyfill for `add_extra`
+ if not AIRFLOW_V_3_2_PLUS:
+ return _add_extra_polyfill(global_collector)
Review Comment:
Agree, thanks for the comment. Added functions that check method/attr presence on collector and we now make decisions based on that. Adjusted the docstring to match that logic. Let me know if that's enough.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Adjust common.compat HookLevelLineage collector for new add_extra method [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb merged PR #58057:
URL: https://github.com/apache/airflow/pull/58057
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Adjust common.compat HookLevelLineage collector for new add_extra method [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on PR #58057:
URL: https://github.com/apache/airflow/pull/58057#issuecomment-3570451144
Looking now.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Adjust common.compat HookLevelLineage collector for new add_extra method [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb commented on code in PR #58057:
URL: https://github.com/apache/airflow/pull/58057#discussion_r2556041699
##########
providers/common/compat/src/airflow/providers/common/compat/lineage/hook.py:
##########
@@ -16,24 +16,153 @@
# under the License.
from __future__ import annotations
-from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS
+from typing import TYPE_CHECKING
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS
-def _get_asset_compat_hook_lineage_collector():
- from airflow.lineage.hook import get_hook_lineage_collector
+if TYPE_CHECKING:
+ from typing import Any
- collector = get_hook_lineage_collector()
+ from airflow.lineage.hook import LineageContext
+
+def _add_extra_polyfill(collector):
+ """
+ Add support for extra lineage information on Airflow versions < 3.2.
+
+ This polyfill adds the `add_extra` method and modifies `collected_assets` property to include
+ extra lineage information. Should be called after renaming from dataset to asset. It's rewriting
+ `collected_assets` method and not `collected_datasets` method.
+ """
+ # We already added it, skip
+ if hasattr(collector, "add_extra"):
+ return collector
+
+ import hashlib
+ import json
+ import types
+ from collections import defaultdict
+
+ import attr
+
+ from airflow.lineage.hook import HookLineage as _BaseHookLineage
+
+ # Add `extra` to HookLineage returned by `collected_assets` property
+ @attr.define
+ class ExtraLineageInfo:
+ """
+ Holds lineage information for arbitrary non-asset metadata.
+
+ This class represents additional lineage context captured during a hook execution that is not
+ associated with a specific asset. It includes the metadata payload itself, the count of
+ how many times it has been encountered, and the context in which it was encountered.
+ """
+
+ key: str
+ value: Any
+ count: int
+ context: LineageContext
+
+ @attr.define
+ class HookLineage(_BaseHookLineage):
+ # mypy is not happy, as base class is using other ExtraLineageInfo, but this code will never
+ # run on AF3.2, where this other one is used, so this is fine - we can ignore.
+ extra: list[ExtraLineageInfo] = attr.field(factory=list) # type: ignore[assignment]
+
+ # Initialize extra tracking attributes
+ collector._extra = {}
+ collector._extra_counts = defaultdict(int)
+
+ # Save the original `collected_assets` getter
+ _original_collected_assets = collector.__class__.collected_assets
+
+ def _compat_collected_assets(self) -> HookLineage:
+ """Get the collected hook lineage information."""
+ # call the original `collected_assets` getter
+ lineage = _original_collected_assets.fget(self)
+ extra_list = [
+ ExtraLineageInfo(
+ key=key,
+ value=value,
+ count=self._extra_counts[count_key],
+ context=context,
+ )
+ for count_key, (key, value, context) in self._extra.items()
+ ]
+ return HookLineage(
+ inputs=lineage.inputs,
+ outputs=lineage.outputs,
+ extra=extra_list,
+ )
+
+ setattr(
+ collector.__class__,
+ "collected_assets",
+ property(lambda c: _compat_collected_assets(c)),
+ )
Review Comment:
Does this need to be setattr?
I think this will work just as well
```suggestion
type(collector).collected_assets = property(_compat_collected_assets)
```
##########
providers/common/compat/src/airflow/providers/common/compat/lineage/hook.py:
##########
@@ -72,17 +202,34 @@ def collected_assets_compat(collector) -> HookLineage:
setattr(
collector.__class__,
"collected_assets",
- property(lambda collector: collected_assets_compat(collector)),
+ property(lambda c: _compat_collected_assets(c)),
)
+ # Add `add_extra` polyfill for AF 2.x, needs to be called after setting `collected_assets`
+ collector = _add_extra_polyfill(collector)
+
return collector
def get_hook_lineage_collector():
- # Dataset has been renamed as Asset in 3.0
- if AIRFLOW_V_3_0_PLUS:
- from airflow.lineage.hook import get_hook_lineage_collector
+ """
+ Get hook lineage collector with appropriate compatibility layers.
+
+ - AF 2.x: Apply both dataset->asset rename and then `add_extra` polyfill
+ - AF 3.0-3.1: Apply only `add_extra` polyfill, no renaming needed
+ - AF 3.2+: Use native implementation (no renaming or polyfill needed)
+ """
+ from airflow.lineage.hook import get_hook_lineage_collector as get_global_collector
+
+ global_collector = get_global_collector()
+
+ # AF 2.x: needs dataset->asset rename + `add_extra` polyfill
+ if not AIRFLOW_V_3_0_PLUS:
+ return _get_af2_asset_compat_hook_lineage_collector(global_collector)
- return get_hook_lineage_collector()
+ # AF 3.0-3.1: needs only polyfill for `add_extra`
+ if not AIRFLOW_V_3_2_PLUS:
+ return _add_extra_polyfill(global_collector)
Review Comment:
I feel uneasy about this version-based approach, espeically as it isn't the Airflow version we need to be targetting, but the apache-airflow-task-sdk version, as in Airflow core 3.2 timeframe there is a decent chance we will have removed the dependency, meaning TaskSDK could run without apache-airflow-core installed at all.
Elsewhere we have:
```
if all(
getattr(collector, asset_method_name, None)
for asset_method_name in ("add_input_asset", "add_output_asset", "collected_assets", "add_extra")
):
```
I think we need to do something similar here
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Adjust common.compat HookLevelLineage collector for new add_extra method [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on PR #58057:
URL: https://github.com/apache/airflow/pull/58057#issuecomment-3540715206
I considered that approach as well, but ultimately decided to go with the polyfill implementation. The code isn’t very complex in the end, and since it’s only executed once for a singleton collector instance, the overhead is minimal. I understand that maintaining it could be more challenging compared to a no-op, and that ideally, users should upgrade Airflow to benefit from the latest features. However, in my view, the value this solution provides outweighs the effort required to implement and maintain it.
That said, if this is a hard no or a blocker for you, I’m happy to adjust the code and switch to a no-op. But if you think there’s room to proceed with the polyfill approach, I’d be glad to make any changes needed to align it better.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat: Adjust common.compat HookLevelLineage collector for new add_extra method [airflow]
Posted by "kacpermuda (via GitHub)" <gi...@apache.org>.
kacpermuda commented on code in PR #58057:
URL: https://github.com/apache/airflow/pull/58057#discussion_r2557077290
##########
providers/common/compat/src/airflow/providers/common/compat/lineage/hook.py:
##########
@@ -16,24 +16,153 @@
# under the License.
from __future__ import annotations
-from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS
+from typing import TYPE_CHECKING
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS
-def _get_asset_compat_hook_lineage_collector():
- from airflow.lineage.hook import get_hook_lineage_collector
+if TYPE_CHECKING:
+ from typing import Any
- collector = get_hook_lineage_collector()
+ from airflow.lineage.hook import LineageContext
+
+def _add_extra_polyfill(collector):
+ """
+ Add support for extra lineage information on Airflow versions < 3.2.
+
+ This polyfill adds the `add_extra` method and modifies `collected_assets` property to include
+ extra lineage information. Should be called after renaming from dataset to asset. It's rewriting
+ `collected_assets` method and not `collected_datasets` method.
+ """
+ # We already added it, skip
+ if hasattr(collector, "add_extra"):
+ return collector
+
+ import hashlib
+ import json
+ import types
+ from collections import defaultdict
+
+ import attr
+
+ from airflow.lineage.hook import HookLineage as _BaseHookLineage
+
+ # Add `extra` to HookLineage returned by `collected_assets` property
+ @attr.define
+ class ExtraLineageInfo:
+ """
+ Holds lineage information for arbitrary non-asset metadata.
+
+ This class represents additional lineage context captured during a hook execution that is not
+ associated with a specific asset. It includes the metadata payload itself, the count of
+ how many times it has been encountered, and the context in which it was encountered.
+ """
+
+ key: str
+ value: Any
+ count: int
+ context: LineageContext
+
+ @attr.define
+ class HookLineage(_BaseHookLineage):
+ # mypy is not happy, as base class is using other ExtraLineageInfo, but this code will never
+ # run on AF3.2, where this other one is used, so this is fine - we can ignore.
+ extra: list[ExtraLineageInfo] = attr.field(factory=list) # type: ignore[assignment]
+
+ # Initialize extra tracking attributes
+ collector._extra = {}
+ collector._extra_counts = defaultdict(int)
+
+ # Save the original `collected_assets` getter
+ _original_collected_assets = collector.__class__.collected_assets
+
+ def _compat_collected_assets(self) -> HookLineage:
+ """Get the collected hook lineage information."""
+ # call the original `collected_assets` getter
+ lineage = _original_collected_assets.fget(self)
+ extra_list = [
+ ExtraLineageInfo(
+ key=key,
+ value=value,
+ count=self._extra_counts[count_key],
+ context=context,
+ )
+ for count_key, (key, value, context) in self._extra.items()
+ ]
+ return HookLineage(
+ inputs=lineage.inputs,
+ outputs=lineage.outputs,
+ extra=extra_list,
+ )
+
+ setattr(
+ collector.__class__,
+ "collected_assets",
+ property(lambda c: _compat_collected_assets(c)),
+ )
Review Comment:
Correct, adjusted.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org