Move lineage mechanism to sdk

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):

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)

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-{}"

airflow-core/tests/unit/models/test_baseoperator.py

@@ -167,15 +167,17 @@ 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):
         hook = mock.MagicMock()
 
         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):

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``

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",

airflow-core/src/airflow/lineage/backend.py → 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.