From 8fc45a04f9f45b0dfd8d1054db2c70d8e78f063d Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 26 Mar 2025 15:39:22 +0800 Subject: [PATCH] Move lineage mechanism to sdk The LineageBackend base class is moved to sdk.definitions. Runtime mechanism in airflow.lineage now lives in sdk.execution_time and also called in the task runner. Lineage hook and collectors are left in Core. You can make a case they can be in the sdk, but it's not necessarily so I'm keeping the work simple for now. I left a compatibility layer in Core, but will submit a PR later to remove it and rewrite related usages. --- .../administration-and-deployment/lineage.rst | 9 +- airflow-core/src/airflow/lineage/__init__.py | 109 ++------------- .../tests/unit/lineage/test_lineage.py | 7 +- .../tests/unit/models/test_baseoperator.py | 6 +- newsfragments/48379.significant.rst | 23 ++++ task-sdk/src/airflow/sdk/__init__.py | 3 + .../src/airflow/sdk/definitions/lineage.py | 13 +- .../src/airflow/sdk/execution_time/lineage.py | 124 ++++++++++++++++++ .../airflow/sdk/execution_time/task_runner.py | 3 + task-sdk/src/airflow/sdk/types.py | 2 +- 10 files changed, 185 insertions(+), 114 deletions(-) create mode 100644 newsfragments/48379.significant.rst rename airflow-core/src/airflow/lineage/backend.py => task-sdk/src/airflow/sdk/definitions/lineage.py (84%) create mode 100644 task-sdk/src/airflow/sdk/execution_time/lineage.py diff --git a/airflow-core/docs/administration-and-deployment/lineage.rst b/airflow-core/docs/administration-and-deployment/lineage.rst index 0219d7b744472..cf69dd03012ea 100644 --- a/airflow-core/docs/administration-and-deployment/lineage.rst +++ b/airflow-core/docs/administration-and-deployment/lineage.rst @@ -71,11 +71,6 @@ for the downstream task. In the example DAG task ``run_this`` (``task_id=run_me_first``) is a BashOperator that takes 3 inlets: ``CAT1``, ``CAT2``, ``CAT3``, that are generated from a list. Note that ``data_interval_start`` is a templated field and will be rendered when the task is running. -.. note:: Behind the scenes Airflow prepares the lineage metadata as part of the ``pre_execute`` method of a task. When the task - has finished execution ``post_execute`` is called and lineage metadata is pushed into XCOM. Thus if you are creating - your own operators that override this method make sure to decorate your method with ``prepare_lineage`` and ``apply_lineage`` - respectively. - Shorthand notation ------------------ @@ -143,11 +138,11 @@ It's possible to push the lineage metrics to a custom backend by providing an in [lineage] backend = my.lineage.CustomBackend -The backend should inherit from ``airflow.lineage.LineageBackend``. +The backend should inherit from ``airflow.sdk.LineageBackend``. .. code-block:: python - from airflow.lineage.backend import LineageBackend + from airflow.sdk import LineageBackend class CustomBackend(LineageBackend): diff --git a/airflow-core/src/airflow/lineage/__init__.py b/airflow-core/src/airflow/lineage/__init__.py index 8e581952e6945..1f26a6a3870da 100644 --- a/airflow-core/src/airflow/lineage/__init__.py +++ b/airflow-core/src/airflow/lineage/__init__.py @@ -19,46 +19,19 @@ from __future__ import annotations -import logging from functools import wraps -from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast +from typing import Callable, TypeVar, cast -from airflow.configuration import conf -from airflow.lineage.backend import LineageBackend -from airflow.utils.session import create_session - -if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context - -PIPELINE_OUTLETS = "pipeline_outlets" -PIPELINE_INLETS = "pipeline_inlets" -AUTO = "auto" - -log = logging.getLogger(__name__) - - -def get_backend() -> LineageBackend | None: - """Get the lineage backend if defined in the configs.""" - clazz = conf.getimport("lineage", "backend", fallback=None) - - if clazz: - if not issubclass(clazz, LineageBackend): - raise TypeError( - f"Your custom Lineage class `{clazz.__name__}` " - f"is not a subclass of `{LineageBackend.__name__}`." - ) - else: - return clazz() - - return None - - -def _render_object(obj: Any, context: Context) -> dict: - ti = context["ti"] - if TYPE_CHECKING: - assert ti.task - return ti.task.render_template(obj, context) +from airflow.sdk.execution_time.lineage import ( + AUTO, + PIPELINE_INLETS, + PIPELINE_OUTLETS, + _get_backend as get_backend, + apply_lineage as _apply_lineage, + prepare_lineage as _prepare_lineage, +) +__all__ = ["AUTO", "PIPELINE_INLETS", "PIPELINE_OUTLETS", "apply_lineage", "get_backend", "prepare_lineage"] T = TypeVar("T", bound=Callable) @@ -70,26 +43,11 @@ def apply_lineage(func: T) -> T: Saves the lineage to XCom and if configured to do so sends it to the backend. """ - _backend = get_backend() @wraps(func) def wrapper(self, context, *args, **kwargs): - self.log.debug("Lineage called with inlets: %s, outlets: %s", self.inlets, self.outlets) - ret_val = func(self, context, *args, **kwargs) - - outlets = list(self.outlets) - inlets = list(self.inlets) - - if outlets: - self.xcom_push(context, key=PIPELINE_OUTLETS, value=outlets) - - if inlets: - self.xcom_push(context, key=PIPELINE_INLETS, value=inlets) - - if _backend: - _backend.send_lineage(operator=self, inlets=self.inlets, outlets=self.outlets, context=context) - + _apply_lineage(context, self.log) return ret_val return cast(T, wrapper) @@ -105,54 +63,11 @@ def prepare_lineage(func: T) -> T: if A -> B -> C and B does not have outlets but A does, these are provided as inlets. * "list of task_ids" -> picks up outlets from the upstream task_ids * "list of datasets" -> manually defined list of dataset - """ @wraps(func) def wrapper(self, context, *args, **kwargs): - from airflow.models.abstractoperator import AbstractOperator - - self.log.debug("Preparing lineage inlets and outlets") - - if isinstance(self.inlets, (str, AbstractOperator)): - self.inlets = [self.inlets] - - if self.inlets and isinstance(self.inlets, list): - # get task_ids that are specified as parameter and make sure they are upstream - task_ids = {o for o in self.inlets if isinstance(o, str)}.union( - op.task_id for op in self.inlets if isinstance(op, AbstractOperator) - ).intersection(self.get_flat_relative_ids(upstream=True)) - - # pick up unique direct upstream task_ids if AUTO is specified - if AUTO.upper() in self.inlets or AUTO.lower() in self.inlets: - task_ids = task_ids.union(task_ids.symmetric_difference(self.upstream_task_ids)) - - # Remove auto and task_ids - self.inlets = [i for i in self.inlets if not isinstance(i, str)] - - # We manually create a session here since xcom_pull returns a - # LazySelectSequence proxy. If we do not pass a session, a new one - # will be created, but that session will not be properly closed. - # After we are done iterating, we can safely close this session. - with create_session() as session: - _inlets = self.xcom_pull( - context, task_ids=task_ids, dag_id=self.dag_id, key=PIPELINE_OUTLETS, session=session - ) - self.inlets.extend(i for it in _inlets for i in it) - - elif self.inlets: - raise AttributeError("inlets is not a list, operator, string or attr annotated object") - - if not isinstance(self.outlets, list): - self.outlets = [self.outlets] - - # render inlets and outlets - self.inlets = [_render_object(i, context) for i in self.inlets] - - self.outlets = [_render_object(i, context) for i in self.outlets] - - self.log.debug("inlets: %s, outlets: %s", self.inlets, self.outlets) - + _prepare_lineage(context, self.log) return func(self, context, *args, **kwargs) return cast(T, wrapper) diff --git a/airflow-core/tests/unit/lineage/test_lineage.py b/airflow-core/tests/unit/lineage/test_lineage.py index a45ec783b3d27..1fe47d5b3291d 100644 --- a/airflow-core/tests/unit/lineage/test_lineage.py +++ b/airflow-core/tests/unit/lineage/test_lineage.py @@ -23,11 +23,11 @@ import pytest from airflow.lineage import AUTO, apply_lineage, get_backend, prepare_lineage -from airflow.lineage.backend import LineageBackend from airflow.models import TaskInstance as TI from airflow.providers.common.compat.lineage.entities import File from airflow.providers.standard.operators.empty import EmptyOperator -from airflow.sdk.definitions.context import Context +from airflow.sdk import Context, LineageBackend +from airflow.sdk.execution_time.lineage import _get_backend from airflow.utils import timezone from airflow.utils.types import DagRunType @@ -51,6 +51,9 @@ def send_lineage(self, operator, inlets=None, outlets=None, context=None): class TestLineage: + def teardown_method(self): + _get_backend.cache_clear() + def test_lineage(self, dag_maker): f1s = "/tmp/does_not_exist_1-{}" f2s = "/tmp/does_not_exist_2-{}" diff --git a/airflow-core/tests/unit/models/test_baseoperator.py b/airflow-core/tests/unit/models/test_baseoperator.py index 0fbaadad5ae29..1573265d37dcb 100644 --- a/airflow-core/tests/unit/models/test_baseoperator.py +++ b/airflow-core/tests/unit/models/test_baseoperator.py @@ -167,7 +167,8 @@ def test_pre_execute_hook(self): op = BaseOperator(task_id="test_task", pre_execute=hook) op_copy = op.prepare_for_execution() - op_copy.pre_execute({}) + ti = mock.Mock(task=op_copy) + op_copy.pre_execute({"ti": ti}) assert hook.called def test_post_execute_hook(self): @@ -175,7 +176,8 @@ def test_post_execute_hook(self): op = BaseOperator(task_id="test_task", post_execute=hook) op_copy = op.prepare_for_execution() - op_copy.post_execute({}) + ti = mock.Mock(task=op_copy) + op_copy.post_execute({"ti": ti}) assert hook.called def test_task_naive_datetime(self): diff --git a/newsfragments/48379.significant.rst b/newsfragments/48379.significant.rst new file mode 100644 index 0000000000000..156f534f2b56c --- /dev/null +++ b/newsfragments/48379.significant.rst @@ -0,0 +1,23 @@ +LineageBackend is moved to Task SDK + +Instead of ``airflow.lineage`` or ``airflow.lineage.backend``, import it from ``airflow.sdk`` instead. + +* Types of change + + * [ ] Dag changes + * [ ] Config changes + * [ ] API changes + * [ ] CLI changes + * [ ] Behaviour changes + * [x] Plugin changes + * [ ] Dependency changes + * [x] Code interface changes + +* Migration rules needed + + * ruff + + * AIR302 + + * [ ] ``airflow.lineage.LineageBackend`` → ``airflow.sdk.LineageBackend`` + * [ ] ``airflow.lineage.backend.LineageBackend`` → ``airflow.sdk.LineageBackend`` diff --git a/task-sdk/src/airflow/sdk/__init__.py b/task-sdk/src/airflow/sdk/__init__.py index f0b011bc18e03..f78d3b29960be 100644 --- a/task-sdk/src/airflow/sdk/__init__.py +++ b/task-sdk/src/airflow/sdk/__init__.py @@ -33,6 +33,7 @@ "DAG", "EdgeModifier", "Label", + "LineageBackend", "Metadata", "Param", "TaskGroup", @@ -60,6 +61,7 @@ from airflow.sdk.definitions.context import Context, get_current_context, get_parsing_context from airflow.sdk.definitions.dag import DAG, dag from airflow.sdk.definitions.edges import EdgeModifier, Label + from airflow.sdk.definitions.lineage import LineageBackend from airflow.sdk.definitions.notifier import BaseNotifier from airflow.sdk.definitions.param import Param from airflow.sdk.definitions.taskgroup import TaskGroup @@ -81,6 +83,7 @@ "DAG": ".definitions.dag", "EdgeModifier": ".definitions.edges", "Label": ".definitions.edges", + "LineageBackend": ".definitions.lineage", "Metadata": ".definitions.asset.metadata", "Param": ".definitions.param", "TaskGroup": ".definitions.taskgroup", diff --git a/airflow-core/src/airflow/lineage/backend.py b/task-sdk/src/airflow/sdk/definitions/lineage.py similarity index 84% rename from airflow-core/src/airflow/lineage/backend.py rename to task-sdk/src/airflow/sdk/definitions/lineage.py index 25d1bd4b07eb0..d815db6fdae37 100644 --- a/airflow-core/src/airflow/lineage/backend.py +++ b/task-sdk/src/airflow/sdk/definitions/lineage.py @@ -22,7 +22,10 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from airflow.models.baseoperator import BaseOperator + from collections.abc import Sequence + + from airflow.sdk.definitions.baseoperator import BaseOperator + from airflow.sdk.definitions.context import Context class LineageBackend: @@ -31,10 +34,10 @@ class LineageBackend: def send_lineage( self, operator: BaseOperator, - inlets: list | None = None, - outlets: list | None = None, - context: dict | None = None, - ): + inlets: Sequence, + outlets: Sequence, + context: Context, + ) -> None: """ Send lineage metadata to a backend. diff --git a/task-sdk/src/airflow/sdk/execution_time/lineage.py b/task-sdk/src/airflow/sdk/execution_time/lineage.py new file mode 100644 index 0000000000000..d17828f1f64f6 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/lineage.py @@ -0,0 +1,124 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Provides lineage support functions.""" + +from __future__ import annotations + +import logging +from functools import cache +from typing import TYPE_CHECKING + +from airflow.sdk.definitions.lineage import LineageBackend + +if TYPE_CHECKING: + from structlog.typing import FilteringBoundLogger as Logger + + from airflow.sdk.definitions.context import Context + +PIPELINE_OUTLETS = "pipeline_outlets" +PIPELINE_INLETS = "pipeline_inlets" +AUTO = "auto" + + +@cache +def _get_backend() -> LineageBackend | None: + """Get the lineage backend if defined in the configs.""" + from airflow.configuration import conf + + if not (klass := conf.getimport("lineage", "backend", fallback=None)): + return None + if issubclass(klass, LineageBackend): + return klass() + raise TypeError( + f"Your custom Lineage class `{klass.__name__}` is not a subclass of `{LineageBackend.__name__}`." + ) + + +def apply_lineage(context: Context, log: Logger | logging.Logger) -> None: + """ + Conditionally send lineage to the backend. + + Saves the lineage to XCom and if configured to do so sends it to the backend. + """ + backend = _get_backend() + ti = context["ti"] + task = ti.task + + if isinstance(log, logging.Logger): + log.debug("Applying inlets: %s outlets: %s", inlets := task.inlets, outlets := task.outlets) + else: + log.debug("Applying lineage", inlets=(inlets := task.inlets), outlets=(outlets := task.outlets)) + + if outlets: + ti.xcom_push(PIPELINE_OUTLETS, outlets) + if inlets: + ti.xcom_push(PIPELINE_INLETS, inlets) + if backend: + backend.send_lineage(operator=task, inlets=inlets, outlets=outlets, context=context) + + +def prepare_lineage(context: Context, log: Logger | logging.Logger) -> None: + """ + Prepare the lineage inlets and outlets. + + Inlets can be: + + * "auto" -> picks up any outlets from direct upstream tasks that have outlets + defined, as such that if A -> B -> C and B does not have outlets but A does, + these are provided as inlets. + * "list of task_ids" -> picks up outlets from the upstream task_ids + * "list of datasets" -> manually defined list of dataset + """ + from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator + + log.debug("Preparing lineage inlets and outlets") + ti = context["ti"] + task = ti.task + + if isinstance(task.inlets, (str, AbstractOperator)): + task.inlets = [task.inlets] + if not isinstance(task.inlets, list): + raise AttributeError("inlets is not a list, operator, string or attr annotated object") + + # get task_ids that are specified as parameter and make sure they are upstream + task_ids = {o for o in task.inlets if isinstance(o, str)}.union( + op.task_id for op in task.inlets if isinstance(op, AbstractOperator) + ).intersection(task.get_flat_relative_ids(upstream=True)) + + # pick up unique direct upstream task_ids if AUTO is specified + if AUTO.upper() in task.inlets or AUTO.lower() in task.inlets: + task_ids = task_ids.union(task_ids.symmetric_difference(task.upstream_task_ids)) + + # Remove auto and task_ids + task.inlets = [i for i in task.inlets if not isinstance(i, str)] + + # Inject upstream outlets + if task_ids: + task.inlets.extend(i for it in ti.xcom_pull(task_ids=task_ids, key=PIPELINE_OUTLETS) for i in it) + + if task.outlets and not isinstance(task.outlets, list): + task.outlets = [task.outlets] + + # render inlets and outlets + task.inlets = [task.render_template(obj, context) for obj in task.inlets] + task.outlets = [task.render_template(obj, context) for obj in task.outlets] + + if isinstance(log, logging.Logger): + log.debug("Rendered inlets: %s, outlets: %s", task.inlets, task.outlets) + else: + log.debug("Rendered lineage objects", inlets=task.inlets, outlets=task.outlets) diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 31c24c3332495..73cc97d50c6f6 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -85,6 +85,7 @@ get_previous_dagrun_success, set_current_context, ) +from airflow.sdk.execution_time.lineage import apply_lineage, prepare_lineage from airflow.sdk.execution_time.xcom import XCom from airflow.utils.net import get_hostname from airflow.utils.state import TaskInstanceState @@ -883,6 +884,7 @@ def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger): if (pre_execute_hook := task._pre_execute_hook) is not None: create_executable_runner(pre_execute_hook, outlet_events, logger=log).run(context) + prepare_lineage(context, log) _run_task_state_change_callbacks(task, "on_execute_callback", context, log) @@ -906,6 +908,7 @@ def _execute_task(context: Context, ti: RuntimeTaskInstance, log: Logger): if (post_execute_hook := task._post_execute_hook) is not None: create_executable_runner(post_execute_hook, outlet_events, logger=log).run(context, result) + apply_lineage(context, log) return result diff --git a/task-sdk/src/airflow/sdk/types.py b/task-sdk/src/airflow/sdk/types.py index 6c3c40f2ab93d..63fcce0d2eeff 100644 --- a/task-sdk/src/airflow/sdk/types.py +++ b/task-sdk/src/airflow/sdk/types.py @@ -68,7 +68,7 @@ class RuntimeTaskInstanceProtocol(Protocol): def xcom_pull( self, - task_ids: str | list[str] | None = None, + task_ids: str | Iterable[str] | None = None, dag_id: str | None = None, key: str = "return_value", # TODO: `include_prior_dates` isn't yet supported in the SDK