This is an automated email from the ASF dual-hosted git repository.
ferruzzi 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 ba20e0d0f3d AIP-86 - Remove no-longer-needed deadlines when a dagrun
succeeds (#53215)
ba20e0d0f3d is described below
commit ba20e0d0f3db4affa9aa2e841d73d4ec2e995c72
Author: D. Ferruzzi <[email protected]>
AuthorDate: Wed Jul 16 08:47:00 2025 -0700
AIP-86 - Remove no-longer-needed deadlines when a dagrun succeeds (#53215)
* AIP-86 - Remove no-longer-needed deadlines when a dagrun succeeds
When a DagRun finished successfully, if that run has any Deadlines which
have not expired, then we can safely remove them from the database. If the
Deadline has passed then it will be retained so we can still act on it in the
next Scheduler pass.
---
airflow-core/src/airflow/models/dagrun.py | 7 +++-
airflow-core/src/airflow/models/deadline.py | 54 ++++++++++++++++++++++++-
airflow-core/tests/unit/models/test_dagrun.py | 7 ++--
airflow-core/tests/unit/models/test_deadline.py | 36 +++++++++++++++++
4 files changed, 99 insertions(+), 5 deletions(-)
diff --git a/airflow-core/src/airflow/models/dagrun.py
b/airflow-core/src/airflow/models/dagrun.py
index fe6fdfb8136..cfac1936d6f 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -64,7 +64,7 @@ from airflow.callbacks.callback_requests import
DagCallbackRequest
from airflow.configuration import conf as airflow_conf
from airflow.exceptions import AirflowException, TaskNotFound
from airflow.listeners.listener import get_listener_manager
-from airflow.models import Log
+from airflow.models import Deadline, Log
from airflow.models.backfill import Backfill
from airflow.models.base import Base, StringID
from airflow.models.taskinstance import TaskInstance as TI
@@ -72,6 +72,7 @@ from airflow.models.taskinstancehistory import
TaskInstanceHistory as TIH
from airflow.models.tasklog import LogTemplate
from airflow.models.taskmap import TaskMap
from airflow.sdk.definitions._internal.abstractoperator import NotMapped
+from airflow.sdk.definitions.deadline import DeadlineReference
from airflow.stats import Stats
from airflow.ti_deps.dep_context import DepContext
from airflow.ti_deps.dependencies_states import SCHEDULEABLE_STATES
@@ -1218,6 +1219,10 @@ class DagRun(Base, LoggingMixin):
msg="success",
)
+ if (deadline := dag.deadline) and isinstance(deadline.reference,
DeadlineReference.TYPES.DAGRUN):
+ # The dagrun has succeeded. If there wre any Deadlines for it
which were not breached, they are no longer needed.
+ Deadline.prune_deadlines(session=session,
conditions={DagRun.run_id: self.run_id})
+
# if *all tasks* are deadlocked, the run failed
elif unfinished.should_schedule and not are_runnable_tasks:
self.log.error("Task deadlock (no runnable tasks); marking run %s
failed", self)
diff --git a/airflow-core/src/airflow/models/deadline.py
b/airflow-core/src/airflow/models/deadline.py
index a20f79a4cde..e2dd8b39c09 100644
--- a/airflow-core/src/airflow/models/deadline.py
+++ b/airflow-core/src/airflow/models/deadline.py
@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any
import sqlalchemy_jsonfield
import uuid6
-from sqlalchemy import Column, ForeignKey, Index, Integer, String, select
+from sqlalchemy import Column, ForeignKey, Index, Integer, String, and_, select
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import relationship
from sqlalchemy_utils import UUIDType
@@ -98,6 +98,58 @@ class Deadline(Base):
f"{self.deadline_time} or run: {self.callback}({callback_kwargs})"
)
+ @classmethod
+ @provide_session
+ def prune_deadlines(cls, *, session: Session, conditions: dict[Column,
Any]) -> int:
+ """
+ Remove deadlines from the table which match the provided conditions
and return the number removed.
+
+ NOTE: This should only be used to remove deadlines which are
associated with
+ successful DagRuns. If the deadline was missed, it will be handled
by the
+ scheduler.
+ TODO: Create the missed_deadlines table (Ramit)
+
+ :param conditions: Dictionary of conditions to evaluate against.
+ :param session: Session to use.
+ """
+ from airflow.models import DagRun # Avoids circular import
+
+ # Assemble the filter conditions.
+ filter_conditions = [column == value for column, value in
conditions.items()]
+ if not filter_conditions:
+ return 0
+
+ try:
+ # Get deadlines which match the provided conditions and their
associated DagRuns.
+ deadline_dagrun_pairs = (
+ session.query(Deadline,
DagRun).join(DagRun).filter(and_(*filter_conditions)).all()
+ )
+ except AttributeError as e:
+ logger.exception("Error resolving deadlines: %s", e)
+ raise
+
+ if not deadline_dagrun_pairs:
+ return 0
+
+ deleted_count = 0
+ dagruns_to_refresh = set()
+
+ for deadline, dagrun in deadline_dagrun_pairs:
+ if dagrun.end_date <= deadline.deadline_time:
+ # If the DagRun finished before the Deadline:
+ session.delete(deadline)
+ deleted_count += 1
+ dagruns_to_refresh.add(dagrun)
+ session.flush()
+
+ logger.debug("%d deadline records were deleted matching the conditions
%s", deleted_count, conditions)
+
+ # Refresh any affected DAG runs.
+ for dagrun in dagruns_to_refresh:
+ session.refresh(dagrun)
+
+ return deleted_count
+
class ReferenceModels:
"""
diff --git a/airflow-core/tests/unit/models/test_dagrun.py
b/airflow-core/tests/unit/models/test_dagrun.py
index ab9f0c17dda..32445dd2958 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -1256,13 +1256,14 @@ class TestDagRun:
def on_success_callable(context):
assert context["dag_run"].dag_id == "test_dagrun_success_callback"
+ future_date = datetime.datetime.now() + datetime.timedelta(days=365)
+
with dag_maker(
dag_id="test_dagrun_success_callback",
schedule=datetime.timedelta(days=1),
- start_date=datetime.datetime(2017, 1, 1),
on_success_callback=on_success_callable,
deadline=DeadlineAlert(
- reference=DeadlineReference.FIXED_DATETIME(DEFAULT_DATE),
+ reference=DeadlineReference.FIXED_DATETIME(future_date),
interval=datetime.timedelta(hours=1),
callback=test_callback_for_deadline,
),
@@ -1277,7 +1278,7 @@ class TestDagRun:
"test_state_succeeded2": TaskInstanceState.SUCCESS,
}
- # Scheduler uses Serialized DAG -- so use that instead of the Actual
DAG
+ # Scheduler uses Serialized DAG -- so use that instead of the Actual
DAG.
dag = SerializedDAG.from_dict(SerializedDAG.to_dict(dag))
dag_run = self.create_dag_run(dag=dag,
task_states=initial_task_states, session=session)
dag_run = session.merge(dag_run)
diff --git a/airflow-core/tests/unit/models/test_deadline.py
b/airflow-core/tests/unit/models/test_deadline.py
index bf0dc4cdf05..e133428a025 100644
--- a/airflow-core/tests/unit/models/test_deadline.py
+++ b/airflow-core/tests/unit/models/test_deadline.py
@@ -36,6 +36,8 @@ from unit.models import DEFAULT_DATE
DAG_ID = "dag_id_1"
RUN_ID = 1
+INVALID_DAG_ID = "invalid_dag_id"
+INVALID_RUN_ID = 2
TEST_CALLBACK_PATH = f"{__name__}.test_callback_for_deadline"
TEST_CALLBACK_KWARGS = {"arg1": "value1"}
@@ -103,6 +105,40 @@ class TestDeadline:
assert result.callback == deadline_orm.callback
assert result.callback_kwargs == deadline_orm.callback_kwargs
+ @pytest.mark.parametrize(
+ "conditions",
+ [
+ pytest.param({}, id="empty_conditions"),
+ pytest.param({Deadline.dagrun_id: INVALID_RUN_ID},
id="no_matches"),
+ pytest.param({Deadline.dagrun_id: RUN_ID}, id="single_condition"),
+ pytest.param({Deadline.dagrun_id: RUN_ID, Deadline.dag_id:
DAG_ID}, id="multiple_conditions"),
+ pytest.param(
+ {Deadline.dagrun_id: RUN_ID, Deadline.dag_id: INVALID_DAG_ID},
id="mixed_conditions"
+ ),
+ ],
+ )
+ @mock.patch("sqlalchemy.orm.Session")
+ def test_prune_deadlines(self, mock_session, conditions):
+ """Test deadline resolution with various conditions."""
+ expected_result = 1 if conditions else 0
+ # Set up the query chain to return a list of (Deadline, DagRun) pairs
+ mock_dagrun = mock.Mock(spec=DagRun, end_date=datetime.now())
+ mock_deadline = mock.Mock(spec=Deadline,
deadline_time=mock_dagrun.end_date + timedelta(days=365))
+ mock_query = mock_session.query.return_value
+ mock_query.join.return_value = mock_query
+ mock_query.filter.return_value = mock_query
+ mock_query.all.return_value = [(mock_deadline, mock_dagrun)] if
conditions else []
+
+ result = Deadline.prune_deadlines(conditions=conditions,
session=mock_session)
+
+ assert result == expected_result
+ if conditions:
+ mock_session.query.assert_called_once_with(Deadline, DagRun)
+ mock_session.query.return_value.filter.assert_called_once() #
Assert that the conditions are applied.
+ mock_session.delete.assert_called_once_with(mock_deadline)
+ else:
+ mock_session.query.assert_not_called()
+
def test_orm(self):
deadline_orm = Deadline(
deadline_time=DEFAULT_DATE,
