fix(airflow): fix tox tests + update docs (#11125)

datahub-project · Aug 12, 2024 · 3d4b3b9 · 3d4b3b9
1 parent 3155914
commit 3d4b3b9
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 15 deletions.
diff --git a/docs/lineage/airflow.md b/docs/lineage/airflow.md
@@ -18,7 +18,7 @@ There's two actively supported implementations of the plugin, with different Air
 | Approach  | Airflow Version | Notes                                                                       |
 | --------- | --------------- | --------------------------------------------------------------------------- |
 | Plugin v2 | 2.3.4+          | Recommended. Requires Python 3.8+                                           |
-| Plugin v1 | 2.1+            | No automatic lineage extraction; may not extract lineage if the task fails. |
+| Plugin v1 | 2.1 - 2.8       | No automatic lineage extraction; may not extract lineage if the task fails. |
 
 If you're using Airflow older than 2.1, it's possible to use the v1 plugin with older versions of `acryl-datahub-airflow-plugin`. See the [compatibility section](#compatibility) for more details.
 
@@ -84,7 +84,7 @@ enabled = True  # default
 
 ### Installation
 
-The v1 plugin requires Airflow 2.1+ and Python 3.8+. If you're on older versions, it's still possible to use an older version of the plugin. See the [compatibility section](#compatibility) for more details.
+The v1 plugin requires Airflow 2.1 - 2.8 and Python 3.8+. If you're on older versions, it's still possible to use an older version of the plugin. See the [compatibility section](#compatibility) for more details.
 
 If you're using Airflow 2.3+, we recommend using the v2 plugin instead. If you need to use the v1 plugin with Airflow 2.3+, you must also set the environment variable `DATAHUB_AIRFLOW_PLUGIN_USE_V1_PLUGIN=true`.
 

diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py
@@ -143,6 +143,20 @@ def wrapper(*args, **kwargs):
     return cast(_F, wrapper)
 
 
+def _render_templates(task_instance: "TaskInstance") -> "TaskInstance":
+    # Render templates in a copy of the task instance.
+    # This is necessary to get the correct operator args in the extractors.
+    try:
+        task_instance_copy = copy.deepcopy(task_instance)
+        task_instance_copy.render_templates()
+        return task_instance_copy
+    except Exception as e:
+        logger.info(
+            f"Error rendering templates in DataHub listener. Jinja-templated variables will not be extracted correctly: {e}"
+        )
+        return task_instance
+
+
 class DataHubListener:
     __name__ = "DataHubListener"
 
@@ -360,15 +374,7 @@ def on_task_instance_running(
             f"DataHub listener got notification about task instance start for {task_instance.task_id}"
         )
 
-        # Render templates in a copy of the task instance.
-        # This is necessary to get the correct operator args in the extractors.
-        try:
-            task_instance = copy.deepcopy(task_instance)
-            task_instance.render_templates()
-        except Exception as e:
-            logger.info(
-                f"Error rendering templates in DataHub listener. Jinja-templated variables will not be extracted correctly: {e}"
-            )
+        task_instance = _render_templates(task_instance)
 
         # The type ignore is to placate mypy on Airflow 2.1.x.
         dagrun: "DagRun" = task_instance.dag_run  # type: ignore[attr-defined]
@@ -459,8 +465,17 @@ def on_task_instance_finish(
         self, task_instance: "TaskInstance", status: InstanceRunResult
     ) -> None:
         dagrun: "DagRun" = task_instance.dag_run  # type: ignore[attr-defined]
-        task = self._task_holder.get_task(task_instance) or task_instance.task
+
+        task_instance = _render_templates(task_instance)
+
+        # We must prefer the task attribute, in case modifications to the task's inlets/outlets
+        # were made by the execute() method.
+        if getattr(task_instance, "task", None):
+            task = task_instance.task
+        else:
+            task = self._task_holder.get_task(task_instance)
         assert task is not None
+
         dag: "DAG" = task.dag  # type: ignore[assignment]
 
         datajob = AirflowGenerator.generate_datajob(

diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin.py
@@ -4,6 +4,7 @@
 
 from airflow.plugins_manager import AirflowPlugin
 
+from datahub_airflow_plugin import __package_name__
 from datahub_airflow_plugin._airflow_compat import AIRFLOW_PATCHED
 from datahub_airflow_plugin._airflow_shims import (
     HAS_AIRFLOW_LISTENER_API,
@@ -23,7 +24,10 @@
         from openlineage.airflow.utils import try_import_from_string  # noqa: F401
     except ImportError:
         # If v2 plugin dependencies are not installed, we fall back to v1.
-        logger.debug("Falling back to v1 plugin due to missing dependencies.")
+        logger.warning(
+            "Falling back to the v1 DataHub plugin due to missing dependencies. "
+            f"Please install {__package_name__}[plugin-v2] to fix this."
+        )
         _USE_AIRFLOW_LISTENER_INTERFACE = False
 
 

diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_plugin_v22.py
@@ -7,7 +7,6 @@
 from airflow.lineage import PIPELINE_OUTLETS
 from airflow.models.baseoperator import BaseOperator
 from airflow.utils.module_loading import import_string
-from cattr import structure
 from datahub.api.entities.dataprocess.dataprocess_instance import InstanceRunResult
 from datahub.telemetry import telemetry
 
@@ -52,6 +51,7 @@ def get_task_inlets_advanced(task: BaseOperator, context: Any) -> Iterable[Any]:
         )
 
         from airflow.lineage import AUTO
+        from cattr import structure
 
         # pick up unique direct upstream task_ids if AUTO is specified
         if AUTO.upper() in task_inlets or AUTO.lower() in task_inlets:

diff --git a/metadata-ingestion-modules/airflow-plugin/tox.ini b/metadata-ingestion-modules/airflow-plugin/tox.ini
@@ -23,7 +23,9 @@ deps =
 
     # Respect the Airflow constraints files.
     # We can't make ourselves work with the constraints of Airflow < 2.3.
-    py310-airflow24: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.4.3/constraints-3.10.txt
+    # The Airflow 2.4 constraints file requires a version of the sqlite provider whose
+    # hook type is missing the `conn_name_attr` property.
+    ; py310-airflow24: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.4.3/constraints-3.10.txt
     py310-airflow26: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.6.3/constraints-3.10.txt
     py310-airflow27: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.7.3/constraints-3.10.txt
     py310-airflow28: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.8.1/constraints-3.10.txt