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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to