apache-airflow-providers-openlineage 1.11.0rc1__py3-none-any.whl → 1.11.0rc2__py3-none-any.whl

airflow/providers/openlineage/extractors/manager.py

@@ -25,6 +25,7 @@ from airflow.providers.openlineage.extractors.bash import BashExtractor
 from airflow.providers.openlineage.extractors.python import PythonExtractor
 from airflow.providers.openlineage.utils.utils import (
     get_unknown_source_attribute_run_facet,
+    translate_airflow_dataset,
     try_import_from_string,
 )
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -90,7 +91,6 @@ class ExtractorManager(LoggingMixin):
             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.
@@ -105,14 +105,22 @@ class ExtractorManager(LoggingMixin):
                 task_metadata = self.validate_task_metadata(task_metadata)
                 if task_metadata:
                     if (not task_metadata.inputs) and (not task_metadata.outputs):
-                        self.extract_inlets_and_outlets(task_metadata, task.inlets, task.outlets)
-
+                        if (hook_lineage := self.get_hook_lineage()) is not None:
+                            inputs, outputs = hook_lineage
+                            task_metadata.inputs = inputs
+                            task_metadata.outputs = outputs
+                        else:
+                            self.extract_inlets_and_outlets(task_metadata, task.inlets, task.outlets)
                     return task_metadata
 
             except Exception as e:
                 self.log.warning(
                     "Failed to extract metadata using found extractor %s - %s %s", extractor, e, task_info
                 )
+        elif (hook_lineage := self.get_hook_lineage()) is not None:
+            inputs, outputs = hook_lineage
+            task_metadata = OperatorLineage(inputs=inputs, outputs=outputs)
+            return task_metadata
         else:
             self.log.debug("Unable to find an extractor %s", task_info)
 
@@ -168,6 +176,30 @@ class ExtractorManager(LoggingMixin):
             if d:
                 task_metadata.outputs.append(d)
 
+    def get_hook_lineage(self) -> tuple[list[Dataset], list[Dataset]] | None:
+        try:
+            from airflow.lineage.hook import get_hook_lineage_collector
+        except ImportError:
+            return None
+
+        if not get_hook_lineage_collector().has_collected:
+            return None
+
+        return (
+            [
+                dataset
+                for dataset_info in get_hook_lineage_collector().collected_datasets.inputs
+                if (dataset := translate_airflow_dataset(dataset_info.dataset, dataset_info.context))
+                is not None
+            ],
+            [
+                dataset
+                for dataset_info in get_hook_lineage_collector().collected_datasets.outputs
+                if (dataset := translate_airflow_dataset(dataset_info.dataset, dataset_info.context))
+                is not None
+            ],
+        )
+
     @staticmethod
     def convert_to_ol_dataset_from_object_storage_uri(uri: str) -> Dataset | None:
         from urllib.parse import urlparse

airflow/providers/openlineage/get_provider_info.py

@@ -53,6 +53,7 @@ def get_provider_info():
         "dependencies": [
             "apache-airflow>=2.8.0",
             "apache-airflow-providers-common-sql>=1.6.0",
+            "apache-airflow-providers-common-compat>=1.2.0",
             "attrs>=22.2",
             "openlineage-integration-common>=1.16.0",
             "openlineage-python>=1.16.0",

airflow/providers/openlineage/plugins/openlineage.py

@@ -25,6 +25,7 @@ from airflow.providers.openlineage.plugins.macros import (
     lineage_parent_id,
     lineage_run_id,
 )
+from airflow.providers.openlineage.utils.utils import IS_AIRFLOW_2_10_OR_HIGHER
 
 
 class OpenLineageProviderPlugin(AirflowPlugin):
@@ -39,6 +40,10 @@ class OpenLineageProviderPlugin(AirflowPlugin):
     if not conf.is_disabled():
         macros = [lineage_job_namespace, lineage_job_name, lineage_run_id, lineage_parent_id]
         listeners = [get_openlineage_listener()]
+        if IS_AIRFLOW_2_10_OR_HIGHER:
+            from airflow.lineage.hook import HookLineageReader
+
+            hook_lineage_readers = [HookLineageReader]
     else:
         macros = []
         listeners = []

airflow/providers/openlineage/utils/utils.py

@@ -33,7 +33,7 @@ from packaging.version import Version
 from airflow import __version__ as AIRFLOW_VERSION
 from airflow.datasets import Dataset
 from airflow.exceptions import AirflowProviderDeprecationWarning  # TODO: move this maybe to Airflow's logic?
-from airflow.models import DAG, BaseOperator, MappedOperator, Operator
+from airflow.models import DAG, BaseOperator, MappedOperator
 from airflow.providers.openlineage import conf
 from airflow.providers.openlineage.plugins.facets import (
     AirflowDagRunFacet,
@@ -423,7 +423,7 @@ def get_airflow_job_facet(dag_run: DagRun) -> dict[str, AirflowJobFacet]:
         return {}
     return {
         "airflow": AirflowJobFacet(
-            taskTree=_get_parsed_dag_tree(dag_run.dag),
+            taskTree={},  # caused OOM errors, to be removed, see #41587
             taskGroups=_get_task_groups_details(dag_run.dag),
             tasks=_get_tasks_details(dag_run.dag),
         )
@@ -439,43 +439,6 @@ def get_airflow_state_run_facet(dag_run: DagRun) -> dict[str, AirflowStateRunFac
     }
 
 
-def _get_parsed_dag_tree(dag: DAG) -> dict:
-    """
-    Get DAG's tasks hierarchy representation.
-
-    While the task dependencies are defined as following:
-    task >> [task_2, task_4] >> task_7
-    task_3 >> task_5
-    task_6  # has no dependencies, it's a root and a leaf
-
-    The result of this function will look like:
-    {
-        "task": {
-            "task_2": {
-                "task_7": {}
-            },
-            "task_4": {
-                "task_7": {}
-            }
-        },
-        "task_3": {
-            "task_5": {}
-        },
-        "task_6": {}
-    }
-    """
-
-    def get_downstream(task: Operator, current_dict: dict):
-        current_dict[task.task_id] = {}
-        for tmp_task in sorted(task.downstream_list, key=lambda x: x.task_id):
-            get_downstream(tmp_task, current_dict[task.task_id])
-
-    task_dict: dict = {}
-    for t in sorted(dag.roots, key=lambda x: x.task_id):
-        get_downstream(t, task_dict)
-    return task_dict
-
-
 def _get_tasks_details(dag: DAG) -> dict:
     tasks = {
         single_task.task_id: {
@@ -487,8 +450,9 @@ def _get_tasks_details(dag: DAG) -> dict:
             "ui_label": single_task.label,
             "is_setup": single_task.is_setup,
             "is_teardown": single_task.is_teardown,
+            "downstream_task_ids": sorted(single_task.downstream_task_ids),
         }
-        for single_task in dag.tasks
+        for single_task in sorted(dag.tasks, key=lambda x: x.task_id)
     }
 
     return tasks

{apache_airflow_providers_openlineage-1.11.0rc1.dist-info → apache_airflow_providers_openlineage-1.11.0rc2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-openlineage
-Version: 1.11.0rc1
+Version: 1.11.0rc2
 Summary: Provider package apache-airflow-providers-openlineage for Apache Airflow
 Keywords: airflow-provider,openlineage,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -21,6 +21,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
+Requires-Dist: apache-airflow-providers-common-compat>=1.2.0rc0
 Requires-Dist: apache-airflow-providers-common-sql>=1.6.0rc0
 Requires-Dist: apache-airflow>=2.8.0rc0
 Requires-Dist: attrs>=22.2
@@ -80,7 +81,7 @@ Provides-Extra: common.sql
 
 Package ``apache-airflow-providers-openlineage``
 
-Release: ``1.11.0.rc1``
+Release: ``1.11.0.rc2``
 
 
 `OpenLineage <https://openlineage.io/>`__
@@ -107,15 +108,16 @@ The package supports the following python versions: 3.8,3.9,3.10,3.11,3.12
 Requirements
 ------------
 
-=======================================  ==================
-PIP package                              Version required
-=======================================  ==================
-``apache-airflow``                       ``>=2.8.0``
-``apache-airflow-providers-common-sql``  ``>=1.6.0``
-``attrs``                                ``>=22.2``
-``openlineage-integration-common``       ``>=1.16.0``
-``openlineage-python``                   ``>=1.16.0``
-=======================================  ==================
+==========================================  ==================
+PIP package                                 Version required
+==========================================  ==================
+``apache-airflow``                          ``>=2.8.0``
+``apache-airflow-providers-common-sql``     ``>=1.6.0``
+``apache-airflow-providers-common-compat``  ``>=1.2.0``
+``attrs``                                   ``>=22.2``
+``openlineage-integration-common``          ``>=1.16.0``
+``openlineage-python``                      ``>=1.16.0``
+==========================================  ==================
 
 Cross provider package dependencies
 -----------------------------------

{apache_airflow_providers_openlineage-1.11.0rc1.dist-info → apache_airflow_providers_openlineage-1.11.0rc2.dist-info}/RECORD

@@ -1,12 +1,12 @@
 airflow/providers/openlineage/LICENSE,sha256=FFb4jd2AXnOOf7XLP04pQW6jbdhG49TxlGY6fFpCV1Y,13609
 airflow/providers/openlineage/__init__.py,sha256=jUD56ffTq0DAqqtYoTQi2pjOt51XQnVZsPuAJPt0UV0,1499
 airflow/providers/openlineage/conf.py,sha256=paV6AHxPJPdQyVMoTGBRHymT-bpbs4Tovttqy9oicMs,5151
-airflow/providers/openlineage/get_provider_info.py,sha256=I8lTHWr3IccBOsCicCr_sR304Ayij9mVGc0dwI9zgtI,8824
+airflow/providers/openlineage/get_provider_info.py,sha256=0qKj7j0nBUYIadp5B3ZxTj8_bKVwnDaT02Z1_41qWTo,8885
 airflow/providers/openlineage/sqlparser.py,sha256=c7q3VVw41S87ZFozrkrEr2oZK79N12mC3KdDs9V1IuM,15581
 airflow/providers/openlineage/extractors/__init__.py,sha256=I0X4f6zUniclyD9zT0DFHRImpCpJVP4MkPJT3cd7X5I,1081
 airflow/providers/openlineage/extractors/base.py,sha256=olafVPBKxeGjqXerCYM0vj2q78Lm4ErWjex_R6nhjKY,6635
 airflow/providers/openlineage/extractors/bash.py,sha256=3aR0PXs8fzRLibRxXN1R8wMZnGzyCur7mjpy8e5GC4A,2583
-airflow/providers/openlineage/extractors/manager.py,sha256=uuSGLZgJHu0tu3RGzrc41uw6SahdgU3816T26FtnFEY,10539
+airflow/providers/openlineage/extractors/manager.py,sha256=Z5fZN6dt67xfJFtlD-_Xnc-tYpxOp66o-ASm6Sn2Vpk,11958
 airflow/providers/openlineage/extractors/python.py,sha256=hVWOplMlBimrpPKPeW6vm75a8OmAYMU1oJzqMz8Jh90,3171
 airflow/providers/openlineage/facets/AirflowDagRunFacet.json,sha256=ie6c-J3-wGgk80WDTGWePz18o6DbW--TNM7BMF4WfcU,2251
 airflow/providers/openlineage/facets/AirflowDebugRunFacet.json,sha256=_zA5gFqGje5MOH1SmdMeA5ViOHvW_pV4oijEAvkuBbY,768
@@ -19,12 +19,12 @@ airflow/providers/openlineage/plugins/adapter.py,sha256=mdXF8ZgyYLgjyWForHQnfKxn
 airflow/providers/openlineage/plugins/facets.py,sha256=VvyMYR6ONkC95q5FdNmohv0scbA1Ej_B5cQ97as5GvA,4161
 airflow/providers/openlineage/plugins/listener.py,sha256=vPtYWYYzKsNw6cR3MsXM4rfnFGoc-xPDGa_6jfTkuLU,18754
 airflow/providers/openlineage/plugins/macros.py,sha256=hgFA3ZdQibyn4KXIOsKYBm4WRKDLA5q6Asscx5rvNfM,3076
-airflow/providers/openlineage/plugins/openlineage.py,sha256=rsRUW_zpXVAglzsgQRv5T9VWYY7CMQl0qRWm8-3oqDA,1678
+airflow/providers/openlineage/plugins/openlineage.py,sha256=T0L5Yxpyq_wzs2_hltJCMY5NKzgsYp0vuEn8LppV5PU,1915
 airflow/providers/openlineage/utils/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
 airflow/providers/openlineage/utils/selective_enable.py,sha256=dFJ7wK7J_-BFwcOKp9tqFOSrASV3lmLv7HtRkEuMk3Q,3087
 airflow/providers/openlineage/utils/sql.py,sha256=bnuU9WvjVKcWVMN3cUp0jaHtU5_ZRM5I1OP1WhIdztg,9583
-airflow/providers/openlineage/utils/utils.py,sha256=gBaWzD4ifBPLAtp4TuUAEzVy_w1L_bQ4S9yfMp_uO2Q,24366
-apache_airflow_providers_openlineage-1.11.0rc1.dist-info/entry_points.txt,sha256=GAx0_i2OeZzqaiiiYuA-xchICDXiCT5kVqpKSxsOjt4,214
-apache_airflow_providers_openlineage-1.11.0rc1.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-apache_airflow_providers_openlineage-1.11.0rc1.dist-info/METADATA,sha256=QGnBVXjzazlTPQqGJTrqVuoqpAIfK444OzKxf6YvZ-c,6387
-apache_airflow_providers_openlineage-1.11.0rc1.dist-info/RECORD,,
+airflow/providers/openlineage/utils/utils.py,sha256=i8bxV6T3tU2hOsMDlV9lo2Pnf-b9lm-xpYSUUprcd-Q,23509
+apache_airflow_providers_openlineage-1.11.0rc2.dist-info/entry_points.txt,sha256=GAx0_i2OeZzqaiiiYuA-xchICDXiCT5kVqpKSxsOjt4,214
+apache_airflow_providers_openlineage-1.11.0rc2.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+apache_airflow_providers_openlineage-1.11.0rc2.dist-info/METADATA,sha256=zPTUcHhwoSwvaJylG8UmcZwKz9s_cQFnemAnH7VrN9s,6534
+apache_airflow_providers_openlineage-1.11.0rc2.dist-info/RECORD,,