WIP 01 - Add deadline to DAG model

Example usage:

```
@task
def hello():
    log.info('hello world')

with DAG(
    dag_id='dag_with_deadline',
    deadline=DeadlineAlert(
        trigger=DeadlineTrigger.DagrunExecutionDate,
        interval=timedelta(hours=1),
        callback=hello,
    )
):
    hello()
```

Author: ferruzzi

airflow/dag_processing/collection.py

@@ -17,7 +17,7 @@
 # under the License.
 
 """
-Utility code that write DAGs in bulk into the database.
+Utility code that writes DAGs in bulk into the database.
 
 This should generally only be called by internal methods such as
 ``DagBag._sync_to_db``, ``DAG.bulk_write_to_db``.
@@ -453,6 +453,9 @@ def update_dags(
                     "core", "max_consecutive_failed_dag_runs_per_dag"
                 )
 
+            if dag.deadline is not None:
+                dm.deadline = dag.deadline
+
             if hasattr(dag, "has_task_concurrency_limits"):
                 dm.has_task_concurrency_limits = dag.has_task_concurrency_limits
             else:

airflow/migrations/versions/0059_3_1_0_add_deadline_to_dag.py

@@ -0,0 +1,50 @@
+#
+# 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.
+
+"""
+Remove processor_subdir.
+
+Revision ID: dfee8bd5d574
+Revises: 6a9e7a527a88
+Create Date: 2024-12-18 19:10:26.962464
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+import sqlalchemy_jsonfield
+from alembic import op
+
+from airflow.settings import json
+
+revision = "dfee8bd5d574"
+down_revision = "6a9e7a527a88"
+branch_labels = None
+depends_on = None
+airflow_version = "3.1.0"
+
+
+def upgrade():
+    op.add_column(
+        "dag",
+        sa.Column("deadline", sqlalchemy_jsonfield.JSONField(json=json), nullable=True),
+    )
+
+
+def downgrade():
+    op.drop_column("dag", "deadline")

airflow/models/__init__.py

@@ -60,7 +60,6 @@ def import_all_models():
     import airflow.models.dag_version
     import airflow.models.dagbundle
     import airflow.models.dagwarning
-    import airflow.models.deadline
     import airflow.models.errors
     import airflow.models.serialized_dag
     import airflow.models.taskinstancehistory
@@ -96,6 +95,7 @@ def __getattr__(name):
     "DagTag": "airflow.models.dag",
     "DagWarning": "airflow.models.dagwarning",
     "DbCallbackRequest": "airflow.models.db_callback_request",
+    "Deadline": "airflow.models.deadline",
     "Log": "airflow.models.log",
     "MappedOperator": "airflow.models.mappedoperator",
     "Operator": "airflow.models.operator",

airflow/models/dag.py

@@ -381,6 +381,7 @@ class DAG(TaskSDKDag, LoggingMixin):
     :param dagrun_timeout: Specify the duration a DagRun should be allowed to run before it times out or
         fails. Task instances that are running when a DagRun is timed out will be marked as skipped.
     :param sla_miss_callback: DEPRECATED - The SLA feature is removed in Airflow 3.0, to be replaced with a new implementation in 3.1
+    :param deadline: Optional Deadline Alert for the DAG.
     :param default_view: Specify DAG default view (grid, graph, duration,
                                                    gantt, landing_times), default grid
     :param orientation: Specify DAG orientation in graph view (LR, TB, RL, BT), default LR
@@ -2058,7 +2059,7 @@ class DagModel(Base):
 
     __tablename__ = "dag"
     """
-    These items are stored in the database for state related information
+    These items are stored in the database for state related information.
     """
     dag_id = Column(StringID(), primary_key=True)
     # A DAG can be paused from the UI / DB
@@ -2095,6 +2096,8 @@ class DagModel(Base):
     timetable_description = Column(String(1000), nullable=True)
     # Asset expression based on asset triggers
     asset_expression = Column(sqlalchemy_jsonfield.JSONField(json=json), nullable=True)
+    # DAG deadline information
+    deadline = Column(sqlalchemy_jsonfield.JSONField(json=json), nullable=True)
     # Tags for view filter
     tags = relationship("DagTag", cascade="all, delete, delete-orphan", backref=backref("dag"))
     # Dag owner links for DAGs view

airflow/models/deadline.py

@@ -16,8 +16,10 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import datetime
-from typing import TYPE_CHECKING
+import logging
+import sys
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING, Callable
 
 import sqlalchemy_jsonfield
 import uuid6
@@ -32,6 +34,8 @@
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
+log = logging.getLogger(__name__)
+
 
 class Deadline(Base, LoggingMixin):
     """A Deadline is a 'need-by' date which triggers a callback if the provided time has passed."""
@@ -90,3 +94,86 @@ def _determine_resource() -> tuple[str, str]:
     def add_deadline(cls, deadline: Deadline, session: Session = NEW_SESSION):
         """Add the provided deadline to the table."""
         session.add(deadline)
+
+
+class DeadlineTrigger:
+    """
+    Store the calculation methods for the various SDeadline Alert triggers.
+
+    TODO:  Embetter this docstring muchly.
+
+    usage:
+
+    In the DAG define a deadline as
+
+    deadline=DeadlineAlert(
+        trigger=DeadlineTrigger.DAGRUN_EXECUTION_DATE,
+        interval=timedelta(hours=1),
+        callback=hello,
+    )
+
+    to parse the deadline trigger use DeadlineTrigger.evaluate(dag.deadline.trigger)
+    """
+
+    DAGRUN_EXECUTION_DATE = "dagrun_execution_date"
+
+    @staticmethod
+    def evaluate(trigger: str):
+        return eval(f"DeadlineTrigger().{trigger}()")
+
+    @staticmethod
+    def get_from_db(table_name, column_name):
+        # TODO:
+        #   fetch appropriate timestamp from db
+        #   cast to datetime
+        #   return
+        log.info("MOCKED Getting %s :: %s", table_name, column_name)
+        return datetime(2024, 1, 1)
+
+    def dagrun_execution_date(self) -> datetime:
+        return self.get_from_db("dagrun", "execution_date")
+
+
+class DeadlineAlert(LoggingMixin):
+    """Store Deadline values needed to calculate the need-by timestamp and the callback information."""
+
+    def __init__(
+        self,
+        trigger: type[DeadlineTrigger] | datetime,
+        interval: timedelta,
+        callback: Callable | str,
+        callback_kwargs: dict | None = None,
+    ):
+        super().__init__()
+        self.trigger = trigger
+        self.interval = interval
+        self.callback_kwargs = callback_kwargs
+        self.callback = self.get_callback_path(callback)
+
+    @staticmethod
+    def get_callback_path(_callback: str | Callable) -> str:
+        if callable(_callback):
+            # Get the reference path to the callable in the form `airflow.models.deadline.get_from_db`
+            return f"{_callback.__module__}.{_callback.__qualname__}"
+
+        # Check if the dotpath can resolve to a callable; store it or raise a ValueError
+        try:
+            _callback_module, _callback_name = _callback.rsplit(".", 1)
+            getattr(sys.modules[_callback_module], _callback_name)
+            return _callback
+        except (KeyError, AttributeError):
+            # KeyError if the path is not valid
+            # AttributeError if the provided value can't be rsplit
+            raise ValueError("callback is not a path to a callable")
+
+    def serialize_deadline_alert(self):
+        from airflow.serialization.serialized_objects import BaseSerialization
+
+        return BaseSerialization.serialize(
+            {
+                "trigger": self.trigger,
+                "interval": self.interval,
+                "callback": self.callback,
+                "callback_kwargs": self.callback_kwargs,
+            }
+        )

airflow/serialization/enums.py

@@ -75,3 +75,4 @@ class DagAttributeTypes(str, Enum):
     DAG_CALLBACK_REQUEST = "dag_callback_request"
     TASK_INSTANCE_KEY = "task_instance_key"
     TRIGGER = "trigger"
+    DEADLINE_ALERT = "deadline_alert"

airflow/serialization/schema.json

@@ -186,6 +186,12 @@
         },
         "dag_display_name": { "type" : "string"},
         "description": { "type" : "string"},
+        "deadline": {
+            "anyOf": [
+                { "$ref": "#/definitions/dict" },
+                { "type": "null" }
+            ]
+        },
         "_concurrency": { "type" : "number"},
         "max_active_tasks": { "type" : "number"},
         "max_active_runs": { "type" : "number"},

airflow/serialization/serialized_objects.py

@@ -43,6 +43,7 @@
 from airflow.models.baseoperatorlink import BaseOperatorLink, XComOperatorLink
 from airflow.models.connection import Connection
 from airflow.models.dag import DAG, _get_model_data_interval
+from airflow.models.deadline import DeadlineAlert
 from airflow.models.expandinput import (
     EXPAND_INPUT_EMPTY,
     create_expand_input,
@@ -720,6 +721,8 @@ def serialize(
             )
         elif isinstance(var, DAG):
             return cls._encode(SerializedDAG.serialize_dag(var), type_=DAT.DAG)
+        elif isinstance(var, DeadlineAlert):
+            return cls._encode(DeadlineAlert.serialize_deadline_alert(var), type_=DAT.DEADLINE_ALERT)
         elif isinstance(var, Resources):
             return var.to_dict()
         elif isinstance(var, MappedOperator):
@@ -1666,6 +1669,8 @@ def serialize_dag(cls, dag: DAG) -> dict:
             serialized_dag["dag_dependencies"] = [x.__dict__ for x in sorted(dag_deps)]
             serialized_dag["task_group"] = TaskGroupSerialization.serialize_task_group(dag.task_group)
 
+            serialized_dag["deadline"] = dag.deadline.serialize_deadline_alert() if dag.deadline else None
+
             # Edge info in the JSON exactly matches our internal structure
             serialized_dag["edge_info"] = dag.edge_info
             serialized_dag["params"] = cls._serialize_params_dict(dag.params)

airflow/utils/db.py

@@ -95,6 +95,7 @@ class MappedClassProtocol(Protocol):
     "2.10.0": "22ed7efa9da2",
     "2.10.3": "5f2621c13b39",
     "3.0.0": "6a9e7a527a88",
+    "3.1.0": "dfee8bd5d574",
 }
 
 

docs/apache-airflow/img/airflow_erd.sha256

@@ -1 +1 @@
-617116c74735faa69a297fe665664b691f341487bc8dcbef3e3ec7e76cdea799
+b4c271a3e9913d4d0a00ae7ec0744a5071d593dff92447872a5cf2bed89000c9