You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ur...@apache.org on 2024/05/03 13:42:47 UTC
(airflow) branch main updated: Implement accessors to read dataset events defined as inlet (#39367)
This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 2294001195 Implement accessors to read dataset events defined as inlet (#39367)
2294001195 is described below
commit 2294001195276948900ffa712849cb9ff9d970ab
Author: Tzu-ping Chung <ur...@gmail.com>
AuthorDate: Fri May 3 21:42:39 2024 +0800
Implement accessors to read dataset events defined as inlet (#39367)
---
airflow/models/taskinstance.py | 2 +
airflow/utils/context.py | 88 +++++++++++++++++++++++++++++++++++
airflow/utils/context.pyi | 17 ++++++-
docs/apache-airflow/templates-ref.rst | 1 +
tests/models/test_taskinstance.py | 53 +++++++++++++++++++++
tests/operators/test_python.py | 3 +-
6 files changed, 162 insertions(+), 2 deletions(-)
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index 2fd5df910a..981f908619 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -110,6 +110,7 @@ from airflow.utils.context import (
ConnectionAccessor,
Context,
DatasetEventAccessors,
+ InletEventsAccessors,
VariableAccessor,
context_get_dataset_events,
context_merge,
@@ -804,6 +805,7 @@ def _get_template_context(
"execution_date": logical_date,
"expanded_ti_count": expanded_ti_count,
"inlets": task.inlets,
+ "inlet_events": InletEventsAccessors(task.inlets, session=session),
"logical_date": logical_date,
"macros": macros,
"map_index_template": task.map_index_template,
diff --git a/airflow/utils/context.py b/airflow/utils/context.py
index 72f3e3787e..6af778b7ff 100644
--- a/airflow/utils/context.py
+++ b/airflow/utils/context.py
@@ -32,8 +32,10 @@ from typing import (
KeysView,
Mapping,
MutableMapping,
+ Sequence,
SupportsIndex,
ValuesView,
+ overload,
)
import attrs
@@ -44,7 +46,11 @@ from airflow.exceptions import RemovedInAirflow3Warning
from airflow.utils.types import NOTSET
if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+ from sqlalchemy.sql.expression import Select
+
from airflow.models.baseoperator import BaseOperator
+ from airflow.models.dataset import DatasetEvent
# NOTE: Please keep this in sync with the following:
# * Context in airflow/utils/context.pyi.
@@ -63,6 +69,7 @@ KNOWN_CONTEXT_KEYS: set[str] = {
"expanded_ti_count",
"exception",
"inlets",
+ "inlet_events",
"logical_date",
"macros",
"map_index_template",
@@ -174,6 +181,87 @@ class DatasetEventAccessors(Mapping[str, DatasetEventAccessor]):
return self._dict[uri]
+@attrs.define()
+class InletEventsAccessor(Sequence["DatasetEvent"]):
+ """Lazy sequence to access inlet dataset events.
+
+ :meta private:
+ """
+
+ _uri: str
+ _session: Session
+
+ def _get_select_stmt(self, *, reverse: bool = False) -> Select:
+ from sqlalchemy import select
+
+ from airflow.models.dataset import DatasetEvent, DatasetModel
+
+ stmt = select(DatasetEvent).join(DatasetEvent.dataset).where(DatasetModel.uri == self._uri)
+ if reverse:
+ return stmt.order_by(DatasetEvent.timestamp.desc())
+ return stmt.order_by(DatasetEvent.timestamp.asc())
+
+ def __reversed__(self) -> Iterator[DatasetEvent]:
+ return iter(self._session.scalar(self._get_select_stmt(reverse=True)))
+
+ def __iter__(self) -> Iterator[DatasetEvent]:
+ return iter(self._session.scalar(self._get_select_stmt()))
+
+ @overload
+ def __getitem__(self, key: int) -> DatasetEvent: ...
+
+ @overload
+ def __getitem__(self, key: slice) -> Sequence[DatasetEvent]: ...
+
+ def __getitem__(self, key: int | slice) -> DatasetEvent | Sequence[DatasetEvent]:
+ if not isinstance(key, int):
+ raise ValueError("non-index access is not supported")
+ if key >= 0:
+ stmt = self._get_select_stmt().offset(key)
+ else:
+ stmt = self._get_select_stmt(reverse=True).offset(-1 - key)
+ if (event := self._session.scalar(stmt.limit(1))) is not None:
+ return event
+ raise IndexError(key)
+
+ def __len__(self) -> int:
+ from sqlalchemy import func, select
+
+ return self._session.scalar(select(func.count()).select_from(self._get_select_stmt()))
+
+
+@attrs.define(init=False)
+class InletEventsAccessors(Mapping[str, InletEventsAccessor]):
+ """Lazy mapping for inlet dataset events accessors.
+
+ :meta private:
+ """
+
+ _inlets: list[Any]
+ _datasets: dict[str, Dataset]
+ _session: Session
+
+ def __init__(self, inlets: list, *, session: Session) -> None:
+ self._inlets = inlets
+ self._datasets = {inlet.uri: inlet for inlet in inlets if isinstance(inlet, Dataset)}
+ self._session = session
+
+ def __iter__(self) -> Iterator[str]:
+ return iter(self._inlets)
+
+ def __len__(self) -> int:
+ return len(self._inlets)
+
+ def __getitem__(self, key: int | str | Dataset) -> InletEventsAccessor:
+ if isinstance(key, int): # Support index access; it's easier for trivial cases.
+ dataset = self._inlets[key]
+ if not isinstance(dataset, Dataset):
+ raise IndexError(key)
+ else:
+ dataset = self._datasets[coerce_to_uri(key)]
+ return InletEventsAccessor(dataset.uri, session=self._session)
+
+
class AirflowContextDeprecationWarning(RemovedInAirflow3Warning):
"""Warn for usage of deprecated context variables in a task."""
diff --git a/airflow/utils/context.pyi b/airflow/utils/context.pyi
index 24d31ccfc9..9cdc0a3d6f 100644
--- a/airflow/utils/context.pyi
+++ b/airflow/utils/context.pyi
@@ -26,9 +26,10 @@
# declare "these are defined, but don't error if others are accessed" someday.
from __future__ import annotations
-from typing import Any, Collection, Container, Iterable, Iterator, Mapping, overload
+from typing import Any, Collection, Container, Iterable, Iterator, Mapping, Sequence, overload
from pendulum import DateTime
+from sqlalchemy.orm import Session
from airflow.configuration import AirflowConfigParser
from airflow.datasets import Dataset
@@ -65,6 +66,19 @@ class DatasetEventAccessors(Mapping[str, DatasetEventAccessor]):
def __len__(self) -> int: ...
def __getitem__(self, key: str | Dataset) -> DatasetEventAccessor: ...
+class InletEventsAccessor(Sequence[DatasetEvent]):
+ @overload
+ def __getitem__(self, key: int) -> DatasetEvent: ...
+ @overload
+ def __getitem__(self, key: slice) -> Sequence[DatasetEvent]: ...
+ def __len__(self) -> int: ...
+
+class InletEventsAccessors(Mapping[str, InletEventsAccessor]):
+ def __init__(self, inlets: list, *, session: Session) -> None: ...
+ def __iter__(self) -> Iterator[str]: ...
+ def __len__(self) -> int: ...
+ def __getitem__(self, key: int | str | Dataset) -> InletEventsAccessor: ...
+
# NOTE: Please keep this in sync with the following:
# * KNOWN_CONTEXT_KEYS in airflow/utils/context.py
# * Table in docs/apache-airflow/templates-ref.rst
@@ -82,6 +96,7 @@ class Context(TypedDict, total=False):
execution_date: DateTime
expanded_ti_count: int | None
inlets: list
+ inlet_events: InletEventsAccessors
logical_date: DateTime
macros: Any
map_index_template: str
diff --git a/docs/apache-airflow/templates-ref.rst b/docs/apache-airflow/templates-ref.rst
index ec2333407c..0cfa14aa86 100644
--- a/docs/apache-airflow/templates-ref.rst
+++ b/docs/apache-airflow/templates-ref.rst
@@ -62,6 +62,7 @@ Variable Type Description
``{{ prev_end_date_success }}`` `pendulum.DateTime`_ End date from prior successful :class:`~airflow.models.dagrun.DagRun` (if available).
| ``None``
``{{ inlets }}`` list List of inlets declared on the task.
+``{{ inlet_events }}`` dict[str, ...] Access past events of inlet datasets. See :doc:`Datasets <authoring-and-scheduling/datasets>`. Added in version 2.10.
``{{ outlets }}`` list List of outlets declared on the task.
``{{ dag }}`` DAG The currently running :class:`~airflow.models.dag.DAG`. You can read more about DAGs in :doc:`DAGs <core-concepts/dags>`.
``{{ task }}`` BaseOperator | The currently running :class:`~airflow.models.baseoperator.BaseOperator`. You can read more about Tasks in :doc:`core-concepts/operators`
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
index 02e1836d68..9f6447b3ae 100644
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -2426,6 +2426,59 @@ class TestTaskInstance:
assert events["write2"].dataset.uri == "test_outlet_dataset_extra_2"
assert events["write2"].extra == {"x": 1}
+ def test_inlet_dataset_extra(self, dag_maker, session):
+ from airflow.datasets import Dataset
+
+ read_task_evaluated = False
+
+ with dag_maker(schedule=None, session=session):
+
+ @task(outlets=Dataset("test_inlet_dataset_extra"))
+ def write(*, ti, dataset_events):
+ dataset_events["test_inlet_dataset_extra"].extra = {"from": ti.task_id}
+
+ @task(inlets=Dataset("test_inlet_dataset_extra"))
+ def read(*, inlet_events):
+ second_event = inlet_events["test_inlet_dataset_extra"][1]
+ assert second_event.uri == "test_inlet_dataset_extra"
+ assert second_event.extra == {"from": "write2"}
+
+ last_event = inlet_events["test_inlet_dataset_extra"][-1]
+ assert last_event.uri == "test_inlet_dataset_extra"
+ assert last_event.extra == {"from": "write3"}
+
+ with pytest.raises(KeyError):
+ inlet_events["does_not_exist"]
+ with pytest.raises(IndexError):
+ inlet_events["test_inlet_dataset_extra"][5]
+
+ # TODO: Support slices.
+
+ nonlocal read_task_evaluated
+ read_task_evaluated = True
+
+ [
+ write.override(task_id="write1")(),
+ write.override(task_id="write2")(),
+ write.override(task_id="write3")(),
+ ] >> read()
+
+ dr: DagRun = dag_maker.create_dagrun()
+
+ # Run "write1", "write2", and "write3" (in this order).
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ for ti in sorted(decision.schedulable_tis, key=operator.attrgetter("task_id")):
+ ti.run(session=session)
+
+ # Run "read".
+ decision = dr.task_instance_scheduling_decisions(session=session)
+ for ti in decision.schedulable_tis:
+ ti.run(session=session)
+
+ # Should be done.
+ assert not dr.task_instance_scheduling_decisions(session=session).schedulable_tis
+ assert read_task_evaluated
+
def test_changing_of_dataset_when_ddrq_is_already_populated(self, dag_maker):
"""
Test that when a task that produces dataset has ran, that changing the consumer
diff --git a/tests/operators/test_python.py b/tests/operators/test_python.py
index 274abd6570..8a31043fd8 100644
--- a/tests/operators/test_python.py
+++ b/tests/operators/test_python.py
@@ -838,7 +838,8 @@ class BaseTestPythonVirtualenvOperator(BasePythonTest):
"ti",
"var", # Accessor for Variable; var->json and var->value.
"conn", # Accessor for Connection.
- "dataset_events", # Accessor for DatasetEvent.
+ "dataset_events", # Accessor for outlet DatasetEvent.
+ "inlet_events", # Accessor for inlet DatasetEvent.
]
ti = create_task_instance(dag_id=self.dag_id, task_id=self.task_id, schedule=None)