openlineage: adjust default extractor's on_failure detection for airf…

…low 2.10 fix

Signed-off-by: Kacper Muda <mudakacper@gmail.com>

airflow/providers/openlineage/extractors/base.py

@@ -29,6 +29,7 @@
  from openlineage.client.facet import BaseFacet as BaseFacet_V1
 from openlineage.client.facet_v2 import JobFacet, RunFacet
 
+from airflow.providers.openlineage.utils.utils import IS_AIRFLOW_2_10_OR_HIGHER
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.state import TaskInstanceState
 
@@ -115,7 +116,14 @@ def _execute_extraction(self) -> OperatorLineage | None:
  return None
 
  def extract_on_complete(self, task_instance) -> OperatorLineage | None:
- if task_instance.state == TaskInstanceState.FAILED:
+ failed_states = [TaskInstanceState.FAILED, TaskInstanceState.UP_FOR_RETRY]
+ if not IS_AIRFLOW_2_10_OR_HIGHER: # todo: remove when min airflow version >= 2.10.0
+ # Before fix (#41053) implemented in Airflow 2.10 TaskInstance's state was still RUNNING when
+ # being passed to listener's on_failure method. Since `extract_on_complete()` is only called
+ # after task completion, RUNNING state means that we are dealing with FAILED task in < 2.10
+ failed_states = [TaskInstanceState.RUNNING]
+
+ if task_instance.state in failed_states:
  on_failed = getattr(self.operator, "get_openlineage_facets_on_failure", None)
  if on_failed and callable(on_failed):
  self.log.debug(

airflow/providers/openlineage/plugins/listener.py

@@ -23,15 +23,15 @@
 
 import psutil
 from openlineage.client.serde import Serde
-from packaging.version import Version
 from setproctitle import getproctitle, setproctitle
 
-from airflow import __version__ as AIRFLOW_VERSION, settings
+from airflow import settings
 from airflow.listeners import hookimpl
 from airflow.providers.openlineage import conf
 from airflow.providers.openlineage.extractors import ExtractorManager
 from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter, RunState
 from airflow.providers.openlineage.utils.utils import (
+ IS_AIRFLOW_2_10_OR_HIGHER,
  get_airflow_job_facet,
  get_airflow_mapped_task_facet,
  get_airflow_run_facet,
@@ -53,12 +53,11 @@
  from airflow.models import DagRun, TaskInstance
 
 _openlineage_listener: OpenLineageListener | None = None
-_IS_AIRFLOW_2_10_OR_HIGHER = Version(Version(AIRFLOW_VERSION).base_version) >= Version("2.10.0")
 
 
 def _get_try_number_success(val):
  # todo: remove when min airflow version >= 2.10.0
- if _IS_AIRFLOW_2_10_OR_HIGHER:
+ if IS_AIRFLOW_2_10_OR_HIGHER:
  return val.try_number
  return val.try_number - 1
 
@@ -247,7 +246,7 @@ def on_success():
 
  self._execute(on_success, "on_success", use_fork=True)
 
- if _IS_AIRFLOW_2_10_OR_HIGHER:
+ if IS_AIRFLOW_2_10_OR_HIGHER:
 
  @hookimpl
  def on_task_instance_failed(

airflow/providers/openlineage/utils/utils.py

@@ -64,7 +64,7 @@
 
 log = logging.getLogger(__name__)
 _NOMINAL_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
-_IS_AIRFLOW_2_10_OR_HIGHER = Version(Version(AIRFLOW_VERSION).base_version) >= Version("2.10.0")
+IS_AIRFLOW_2_10_OR_HIGHER = Version(Version(AIRFLOW_VERSION).base_version) >= Version("2.10.0")
 
 
 def try_import_from_string(string: str) -> Any:
@@ -663,7 +663,7 @@ def normalize_sql(sql: str | Iterable[str]):
 
 def should_use_external_connection(hook) -> bool:
  # If we're at Airflow 2.10, the execution is process-isolated, so we can safely run those again.
- if not _IS_AIRFLOW_2_10_OR_HIGHER:
+ if not IS_AIRFLOW_2_10_OR_HIGHER:
  return hook.__class__.__name__ not in ["SnowflakeHook", "SnowflakeSqlApiHook", "RedshiftSQLHook"]
  return True
 

tests/providers/amazon/aws/operators/test_redshift_sql.py

@@ -82,7 +82,7 @@ class TestRedshiftSQLOpenLineage:
  "airflow.providers.amazon.aws.hooks.redshift_sql._IS_AIRFLOW_2_10_OR_HIGHER",
  new_callable=PropertyMock,
  )
- @patch("airflow.providers.openlineage.utils.utils._IS_AIRFLOW_2_10_OR_HIGHER", new_callable=PropertyMock)
+ @patch("airflow.providers.openlineage.utils.utils.IS_AIRFLOW_2_10_OR_HIGHER", new_callable=PropertyMock)
  @patch("airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook.conn")
  def test_execute_openlineage_events(
  self,

tests/providers/openlineage/extractors/test_base.py

@@ -25,6 +25,7 @@
 from openlineage.client.facet_v2 import BaseFacet, JobFacet, parent_run, sql_job
 
 from airflow.models.baseoperator import BaseOperator
+from airflow.models.taskinstance import TaskInstanceState
 from airflow.operators.python import PythonOperator
 from airflow.providers.openlineage.extractors.base import (
  BaseExtractor,
@@ -233,6 +234,47 @@ def test_extraction_without_on_start():
  )
 
 
+@pytest.mark.parametrize(
+ "task_state, is_airflow_2_10_or_higher, should_call_on_failure",
+ (
+ # Airflow >= 2.10
+ (TaskInstanceState.FAILED, True, True),
+ (TaskInstanceState.UP_FOR_RETRY, True, True),
+ (TaskInstanceState.RUNNING, True, False),
+ (TaskInstanceState.SUCCESS, True, False),
+ # Airflow < 2.10
+ (TaskInstanceState.RUNNING, False, True),
+ (TaskInstanceState.SUCCESS, False, False),
+ (TaskInstanceState.FAILED, False, False), # should never happen, fixed in #41053
+ (TaskInstanceState.UP_FOR_RETRY, False, False), # should never happen, fixed in #41053
+ ),
+)
+def test_extract_on_failure(task_state, is_airflow_2_10_or_higher, should_call_on_failure):
+ task_instance = mock.Mock(state=task_state)
+ operator = mock.Mock()
+ operator.get_openlineage_facets_on_failure = mock.Mock(
+ return_value=OperatorLineage(run_facets={"failed": True})
+ )
+ operator.get_openlineage_facets_on_complete = mock.Mock(return_value=None)
+
+ extractor = DefaultExtractor(operator=operator)
+
+ with mock.patch(
+ "airflow.providers.openlineage.extractors.base.IS_AIRFLOW_2_10_OR_HIGHER", is_airflow_2_10_or_higher
+ ):
+ result = extractor.extract_on_complete(task_instance)
+
+ if should_call_on_failure:
+ operator.get_openlineage_facets_on_failure.assert_called_once_with(task_instance)
+ operator.get_openlineage_facets_on_complete.assert_not_called()
+ assert isinstance(result, OperatorLineage)
+ assert result.run_facets == {"failed": True}
+ else:
+ operator.get_openlineage_facets_on_failure.assert_not_called()
+ operator.get_openlineage_facets_on_complete.assert_called_once_with(task_instance)
+ assert result is None
+
+
 @mock.patch("airflow.providers.openlineage.conf.custom_extractors")
 def test_extractors_env_var(custom_extractors):
  custom_extractors.return_value = {"tests.providers.openlineage.extractors.test_base.ExampleExtractor"}