apache-airflow-providers-openlineage 1.11.0rc1__tar.gz → 1.11.0rc2__tar.gz

{apache_airflow_providers_openlineage-1.11.0rc1 → apache_airflow_providers_openlineage-1.11.0rc2}/PKG-INFO

@@ -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 → apache_airflow_providers_openlineage-1.11.0rc2}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-openlineage``
 
-Release: ``1.11.0.rc1``
+Release: ``1.11.0.rc2``
 
 
 `OpenLineage <https://openlineage.io/>`__
@@ -69,15 +69,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 → apache_airflow_providers_openlineage-1.11.0rc2}/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

{apache_airflow_providers_openlineage-1.11.0rc1 → apache_airflow_providers_openlineage-1.11.0rc2}/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",

{apache_airflow_providers_openlineage-1.11.0rc1 → apache_airflow_providers_openlineage-1.11.0rc2}/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 = []

{apache_airflow_providers_openlineage-1.11.0rc1 → apache_airflow_providers_openlineage-1.11.0rc2}/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 → apache_airflow_providers_openlineage-1.11.0rc2}/pyproject.toml

@@ -28,7 +28,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-openlineage"
-version = "1.11.0.rc1"
+version = "1.11.0.rc2"
 description = "Provider package apache-airflow-providers-openlineage for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -56,6 +56,7 @@ classifiers = [
 ]
 requires-python = "~=3.8"
 dependencies = [
+    "apache-airflow-providers-common-compat>=1.2.0rc0",
     "apache-airflow-providers-common-sql>=1.6.0rc0",
     "apache-airflow>=2.8.0rc0",
     "attrs>=22.2",