Move get_serialized_dag and get_task_instance methods to model classes

airflow/dag_processing/processor.py

@@ -32,8 +32,6 @@
 from sqlalchemy.orm.session import Session
 
 from airflow import settings
-from airflow.api_internal.actions.dag import InternalApiDagActions
-from airflow.api_internal.actions.taskinstance import InternalApiTaskInstanceActions
 from airflow.api_internal.internal_api_call import internal_api_call
 from airflow.callbacks.callback_requests import (
     CallbackRequest,
@@ -48,7 +46,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
@@ -691,7 +690,7 @@ def _execute_task_callbacks(self, dagbag: DagBag | None, request: TaskCallbackRe
             return
 
         simple_ti = request.simple_task_instance
-        ti = InternalApiTaskInstanceActions.get_task_instance(
+        ti = TaskInstance.get_task_instance(
             dag_id=simple_ti.dag_id,
             run_id=simple_ti.run_id,
             task_id=simple_ti.task_id,
@@ -714,7 +713,7 @@ 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
-            task = InternalApiDagActions.get_serialized_dag(
+            task = SerializedDagModel.get_serialized_dag(
                 dag_id=simple_ti.dag_id, task_id=simple_ti.task_id, session=session
             )
 

airflow/models/serialized_dag.py

@@ -24,18 +24,21 @@
 from datetime import datetime, timedelta
 
 import sqlalchemy_jsonfield
-from sqlalchemy import BigInteger, Column, Index, LargeBinary, String, and_, or_
+from sqlalchemy import BigInteger, Column, Index, LargeBinary, String, and_, exc, or_
 from sqlalchemy.orm import Session, backref, foreign, relationship
 from sqlalchemy.sql.expression import func, literal
 
+from airflow.api_internal.internal_api_call import internal_api_call
+from airflow.exceptions import TaskNotFound
+from airflow.models import Operator
 from airflow.models.base import ID_LEN, Base
 from airflow.models.dag import DAG, DagModel
 from airflow.models.dagcode import DagCode
 from airflow.models.dagrun import DagRun
 from airflow.serialization.serialized_objects import DagDependency, SerializedDAG
 from airflow.settings import COMPRESS_SERIALIZED_DAGS, MIN_SERIALIZED_DAG_UPDATE_INTERVAL, json
 from airflow.utils import timezone
-from airflow.utils.session import provide_session
+from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.sqlalchemy import UtcDateTime
 
 log = logging.getLogger(__name__)
@@ -368,3 +371,18 @@ def get_dag_dependencies(cls, session: Session = None) -> dict[str, list[DagDepe
         else:
             iterator = session.query(cls.dag_id, func.json_extract_path(cls._data, "dag", "dag_dependencies"))
         return {dag_id: [DagDependency(**d) for d in (deps_data or [])] for dag_id, deps_data in iterator}
+
+    @staticmethod
+    @internal_api_call
+    @provide_session
+    def get_serialized_dag(dag_id: str, task_id: str, session: Session = NEW_SESSION) -> Operator | None:
+        from airflow.models.serialized_dag import SerializedDagModel
+
+        try:
+            model = session.query(SerializedDagModel).get(dag_id)
+            if model:
+                return model.dag.get_task(task_id)
+        except (exc.NoResultFound, TaskNotFound):
+            pass
+
+        return None

airflow/models/taskinstance.py

@@ -64,7 +64,7 @@
 from sqlalchemy.sql.expression import ColumnOperators, case
 
 from airflow import settings
-from airflow.api_internal.actions.taskinstance import InternalApiTaskInstanceActions
+from airflow.api_internal.internal_api_call import internal_api_call
 from airflow.compat.functools import cache
 from airflow.configuration import conf
 from airflow.datasets import Dataset
@@ -84,6 +84,7 @@
     UnmappableXComTypePushed,
     XComForMappingNotPushed,
 )
+from airflow.jobs.scheduler_job import TI
 from airflow.models.base import Base, StringID
 from airflow.models.log import Log
 from airflow.models.mappedoperator import MappedOperator
@@ -109,6 +110,7 @@
 from airflow.utils.net import get_hostname
 from airflow.utils.operator_helpers import context_to_airflow_vars
 from airflow.utils.platform import getuser
+from airflow.utils.retries import run_with_db_retries
 from airflow.utils.session import NEW_SESSION, create_session, provide_session
 from airflow.utils.sqlalchemy import (
     ExecutorConfigType,
@@ -766,6 +768,33 @@ def error(self, session: Session = NEW_SESSION) -> None:
         session.merge(self)
         session.commit()
 
+    @staticmethod
+    @internal_api_call
+    @provide_session
+    def get_task_instance(
+        dag_id: str,
+        run_id: str,
+        task_id: str,
+        map_index: int,
+        lock_for_update: bool = False,
+        session: Session = NEW_SESSION,
+    ) -> TI | None:
+        # TODO: need to convert SQLAlchemy objects to internal API objects
+        query = session.query(TI).filter_by(
+            dag_id=dag_id,
+            run_id=run_id,
+            task_id=task_id,
+            map_index=map_index,
+        )
+
+        if lock_for_update:
+            # TODO: pass the logger once resolved here https://github.com/apache/airflow/pull/28502
+            for attempt in run_with_db_retries(logger=None):
+                with attempt:
+                    return query.with_for_update().one_or_none()
+        else:
+            return query.one_or_none()
+
     @provide_session
     def refresh_from_db(self, session: Session = NEW_SESSION, lock_for_update: bool = False) -> None:
         """
@@ -776,7 +805,7 @@ def refresh_from_db(self, session: Session = NEW_SESSION, lock_for_update: bool
             lock the TaskInstance (issuing a FOR UPDATE clause) until the
             session is committed.
         """
-        ti = InternalApiTaskInstanceActions.get_task_instance(
+        ti = TaskInstance.get_task_instance(
             dag_id=self.dag_id,
             task_id=self.task_id,
             run_id=self.run_id,