Convert DagFileProcessor.execute_callbacks to Internal API

airflow/api_internal/endpoints/rpc_api_endpoint.py

@@ -38,8 +38,11 @@ def _initialize_map() -> dict[str, Callable]:
     from airflow.dag_processing.manager import DagFileProcessorManager
     from airflow.dag_processing.processor import DagFileProcessor
     from airflow.models import Trigger, Variable, XCom
-    from airflow.models.dag import DagModel
+    from airflow.models.dag import DAG, DagModel
+    from airflow.models.dagrun import DagRun
     from airflow.models.dagwarning import DagWarning
+    from airflow.models.serialized_dag import SerializedDagModel
+    from airflow.models.taskinstance import TaskInstance
     from airflow.secrets.metastore import MetastoreBackend
 
     functions: list[Callable] = [
@@ -60,6 +63,16 @@ def _initialize_map() -> dict[str, Callable]:
         Variable.set,
         Variable.update,
         Variable.delete,
+        DAG.fetch_callback,
+        DAG.fetch_dagrun,
+        DagRun.fetch_task_instances,
+        DagRun.get_previous_dagrun,
+        DagRun.get_previous_scheduled_dagrun,
+        DagRun.fetch_task_instance,
+        SerializedDagModel.get_serialized_dag,
+        TaskInstance.get_task_instance,
+        TaskInstance.fetch_handle_failure_context,
+        TaskInstance.save_to_db,
         Trigger.from_object,
         Trigger.bulk_fetch,
         Trigger.clean_unused,

airflow/api_internal/internal_api_call.py

@@ -122,7 +122,10 @@ def wrapper(*args, **kwargs):
         if "cls" in arguments_dict:  # used by @classmethod
             del arguments_dict["cls"]
 
-        args_json = json.dumps(BaseSerialization.serialize(arguments_dict, use_pydantic_models=True))
+        args_json = json.dumps(
+            BaseSerialization.serialize(arguments_dict, use_pydantic_models=True),
+            default=BaseSerialization.serialize,
+        )
         method_name = f"{func.__module__}.{func.__qualname__}"
         result = make_jsonrpc_request(method_name, args_json)
         if result is None or result == b"":

airflow/cli/commands/task_command.py

@@ -73,6 +73,7 @@
     from sqlalchemy.orm.session import Session
 
     from airflow.models.operator import Operator
+    from airflow.serialization.pydantic.dag_run import DagRunPydantic
 
 log = logging.getLogger(__name__)
 
@@ -94,7 +95,7 @@ def _get_dag_run(
     create_if_necessary: CreateIfNecessary,
     exec_date_or_run_id: str | None = None,
     session: Session,
-) -> tuple[DagRun, bool]:
+) -> tuple[DagRun | DagRunPydantic, bool]:
     """Try to retrieve a DAG run from a string representing either a run ID or logical date.
 
     This checks DAG runs like this:

airflow/dag_processing/processor.py

@@ -28,7 +28,7 @@
 from typing import TYPE_CHECKING, Iterable, Iterator
 
 from setproctitle import setproctitle
-from sqlalchemy import delete, exc, func, or_, select
+from sqlalchemy import delete, func, or_, select
 
 from airflow import settings
 from airflow.api_internal.internal_api_call import internal_api_call
@@ -44,7 +44,8 @@
 from airflow.models.dagbag import DagBag
 from airflow.models.dagrun import DagRun as DR
 from airflow.models.dagwarning import DagWarning, DagWarningType
-from airflow.models.taskinstance import TaskInstance as TI
+from airflow.models.serialized_dag import SerializedDagModel
+from airflow.models.taskinstance import TaskInstance, TaskInstance as TI
 from airflow.stats import Stats
 from airflow.utils import timezone
 from airflow.utils.email import get_email_address_list, send_email
@@ -736,25 +737,28 @@ def execute_callbacks_without_dag(
     @provide_session
     def _execute_dag_callbacks(self, dagbag: DagBag, request: DagCallbackRequest, session: Session):
         dag = dagbag.dags[request.dag_id]
-        dag_run = dag.get_dagrun(run_id=request.run_id, session=session)
-        dag.handle_callback(
-            dagrun=dag_run, success=not request.is_failure_callback, reason=request.msg, session=session
-        )
+        callbacks, context = DAG.fetch_callback(
+            dag=dag,
+            dag_run_id=request.run_id,
+            success=not request.is_failure_callback,
+            reason=request.msg,
+            session=session,
+        ) or (None, None)
+
+        if callbacks and context:
+            DAG.execute_callback(callbacks, context, dag.dag_id)
 
     def _execute_task_callbacks(self, dagbag: DagBag | None, request: TaskCallbackRequest, session: Session):
         if not request.is_failure_callback:
             return
 
         simple_ti = request.simple_task_instance
-        ti: TI | None = (
-            session.query(TI)
-            .filter_by(
-                dag_id=simple_ti.dag_id,
-                run_id=simple_ti.run_id,
-                task_id=simple_ti.task_id,
-                map_index=simple_ti.map_index,
-            )
-            .one_or_none()
+        ti = TaskInstance.get_task_instance(
+            dag_id=simple_ti.dag_id,
+            run_id=simple_ti.run_id,
+            task_id=simple_ti.task_id,
+            map_index=simple_ti.map_index,
+            session=session,
         )
         if not ti:
             return
@@ -770,14 +774,10 @@ def _execute_task_callbacks(self, dagbag: DagBag | None, request: TaskCallbackRe
             # `handle_failure` so that the state of the TI gets progressed.
             #
             # Since handle_failure _really_ wants a task, we do our best effort to give it one
-            from airflow.models.serialized_dag import SerializedDagModel
+            task = SerializedDagModel.get_serialized_dag(
+                dag_id=simple_ti.dag_id, task_id=simple_ti.task_id, session=session
+            )
 
-            try:
-                model = session.get(SerializedDagModel, simple_ti.dag_id)
-                if model:
-                    task = model.dag.get_task(simple_ti.task_id)
-            except (exc.NoResultFound, TaskNotFound):
-                pass
         if task:
             ti.refresh_from_task(task)
 

airflow/models/baseoperator.py

@@ -920,7 +920,7 @@ def __init__(
             max_active_tis_per_dag = task_concurrency
         self.max_active_tis_per_dag: int | None = max_active_tis_per_dag
         self.max_active_tis_per_dagrun: int | None = max_active_tis_per_dagrun
-        self.do_xcom_push = do_xcom_push
+        self.do_xcom_push: bool = do_xcom_push
 
         self.doc_md = doc_md
         self.doc_json = doc_json

airflow/models/dag.py

@@ -88,7 +88,7 @@
     TaskNotFound,
 )
 from airflow.jobs.job import run_job
-from airflow.models.abstractoperator import AbstractOperator
+from airflow.models.abstractoperator import AbstractOperator, TaskStateChangeCallback
 from airflow.models.base import Base, StringID
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.dagcode import DagCode
@@ -139,6 +139,7 @@
     from airflow.models.operator import Operator
     from airflow.models.slamiss import SlaMiss
     from airflow.serialization.pydantic.dag import DagModelPydantic
+    from airflow.serialization.pydantic.dag_run import DagRunPydantic
     from airflow.typing_compat import Literal
     from airflow.utils.task_group import TaskGroup
 
@@ -888,7 +889,7 @@ def get_next_data_interval(self, dag_model: DagModel) -> DataInterval | None:
         # infer from the logical date.
         return self.infer_automated_data_interval(dag_model.next_dagrun)
 
-    def get_run_data_interval(self, run: DagRun) -> DataInterval:
+    def get_run_data_interval(self, run: DagRun | DagRunPydantic) -> DataInterval:
         """Get the data interval of this run.
 
         For compatibility, this method infers the data interval from the DAG's
@@ -1383,8 +1384,43 @@ def normalized_schedule_interval(self) -> ScheduleInterval:
             _schedule_interval = self.schedule_interval
         return _schedule_interval
 
+    @staticmethod
+    @internal_api_call
     @provide_session
-    def handle_callback(self, dagrun, success=True, reason=None, session=NEW_SESSION):
+    def fetch_callback(
+        dag: DAG,
+        dag_run_id: str,
+        success: bool = True,
+        reason: str | None = None,
+        *,
+        session: Session = NEW_SESSION,
+    ) -> tuple[list[TaskStateChangeCallback], Context] | None:
+        """
+        Fetch the appropriate callbacks depending on the value of success.
+
+        This method gets the context of a single TaskInstance part of this DagRun and returns it along
+        the list of callbacks.
+
+        :param dag: DAG object
+        :param dag_run_id: The DAG run ID
+        :param success: Flag to specify if failure or success callback should be called
+        :param reason: Completion reason
+        :param session: Database session
+        """
+        callbacks = dag.on_success_callback if success else dag.on_failure_callback
+        if callbacks:
+            dagrun = DAG.fetch_dagrun(dag_id=dag.dag_id, run_id=dag_run_id, session=session)
+            callbacks = callbacks if isinstance(callbacks, list) else [callbacks]
+            tis = dagrun.get_task_instances(session=session)
+            ti = tis[-1]  # get first TaskInstance of DagRun
+            ti.task = dag.get_task(ti.task_id)
+            context = ti.get_template_context(session=session)
+            context["reason"] = reason
+            return callbacks, context
+        return None
+
+    @provide_session
+    def handle_callback(self, dagrun: DagRun, success=True, reason=None, session=NEW_SESSION):
         """
         Triggers on_failure_callback or on_success_callback as appropriate.
 
@@ -1400,21 +1436,29 @@ def handle_callback(self, dagrun, success=True, reason=None, session=NEW_SESSION
         :param reason: Completion reason
         :param session: Database session
         """
-        callbacks = self.on_success_callback if success else self.on_failure_callback
-        if callbacks:
-            callbacks = callbacks if isinstance(callbacks, list) else [callbacks]
-            tis = dagrun.get_task_instances(session=session)
-            ti = tis[-1]  # get first TaskInstance of DagRun
-            ti.task = self.get_task(ti.task_id)
-            context = ti.get_template_context(session=session)
-            context.update({"reason": reason})
+        callbacks, context = DAG.fetch_callback(
+            dag=self, dag_run_id=dagrun.run_id, success=success, reason=reason, session=session
+        ) or (None, None)
+
+        DAG.execute_callback(callbacks, context, self.dag_id)
+
+    @classmethod
+    def execute_callback(cls, callbacks: list[Callable] | None, context: Context | None, dag_id: str):
+        """
+        Triggers the callbacks with the given context.
+
+        :param callbacks: List of callbacks to call
+        :param context: Context to pass to all callbacks
+        :param dag_id: The dag_id of the DAG to find.
+        """
+        if callbacks and context:
             for callback in callbacks:
-                self.log.info("Executing dag callback function: %s", callback)
+                cls.logger().info("Executing dag callback function: %s", callback)
                 try:
                     callback(context)
                 except Exception:
-                    self.log.exception("failed to invoke dag state update callback")
-                    Stats.incr("dag.callback_exceptions", tags={"dag_id": dagrun.dag_id})
+                    cls.logger().exception("failed to invoke dag state update callback")
+                    Stats.incr("dag.callback_exceptions", tags={"dag_id": dag_id})
 
     def get_active_runs(self):
         """
@@ -1452,16 +1496,19 @@ def get_num_active_runs(self, external_trigger=None, only_running=True, session=
 
         return session.scalar(query)
 
+    @staticmethod
+    @internal_api_call
     @provide_session
-    def get_dagrun(
-        self,
+    def fetch_dagrun(
+        dag_id: str,
         execution_date: datetime | None = None,
         run_id: str | None = None,
         session: Session = NEW_SESSION,
-    ):
+    ) -> DagRun | DagRunPydantic:
         """
         Return the dag run for a given execution date or run_id if it exists, otherwise none.
 
+        :param dag_id: The dag_id of the DAG to find.
         :param execution_date: The execution date of the DagRun to find.
         :param run_id: The run_id of the DagRun to find.
         :param session:
@@ -1471,11 +1518,22 @@ def get_dagrun(
             raise TypeError("You must provide either the execution_date or the run_id")
         query = select(DagRun)
         if execution_date:
-            query = query.where(DagRun.dag_id == self.dag_id, DagRun.execution_date == execution_date)
+            query = query.where(DagRun.dag_id == dag_id, DagRun.execution_date == execution_date)
         if run_id:
-            query = query.where(DagRun.dag_id == self.dag_id, DagRun.run_id == run_id)
+            query = query.where(DagRun.dag_id == dag_id, DagRun.run_id == run_id)
         return session.scalar(query)
 
+    @provide_session
+    def get_dagrun(
+        self,
+        execution_date: datetime | None = None,
+        run_id: str | None = None,
+        session: Session = NEW_SESSION,
+    ) -> DagRun | DagRunPydantic:
+        return DAG.fetch_dagrun(
+            dag_id=self.dag_id, execution_date=execution_date, run_id=run_id, session=session
+        )
+
     @provide_session
     def get_dagruns_between(self, start_date, end_date, session=NEW_SESSION):
         """