AIP-44 Migrate Job to Internal API (#34026)

* AIP-44 Migrate Job to Internal API

* Update airflow/jobs/base_job_runner.py

---------

Co-authored-by: Tzu-ping Chung <[email protected]>

airflow/api_internal/endpoints/rpc_api_endpoint.py

@@ -24,6 +24,7 @@
 
 from flask import Response
 
+from airflow.jobs.job import Job, most_recent_job
 from airflow.serialization.serialized_objects import BaseSerialization
 from airflow.utils.session import create_session
 
@@ -51,6 +52,12 @@ def _initialize_map() -> dict[str, Callable]:
         DagModel.get_current,
         DagFileProcessorManager.clear_nonexistent_import_errors,
         DagWarning.purge_inactive_dag_warnings,
+        Job._add_to_db,
+        Job._fetch_from_db,
+        Job._kill,
+        Job._update_heartbeat,
+        Job._update_in_db,
+        most_recent_job,
         MetastoreBackend._fetch_connection,
         MetastoreBackend._fetch_variable,
         XCom.get_value,

airflow/jobs/backfill_job_runner.py

@@ -62,7 +62,7 @@
     from airflow.models.taskinstance import TaskInstanceKey
 
 
-class BackfillJobRunner(BaseJobRunner[Job], LoggingMixin):
+class BackfillJobRunner(BaseJobRunner, LoggingMixin):
     """
     A backfill job runner consists of a dag or subdag for a specific time range.
 

airflow/jobs/base_job_runner.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING
 
 from airflow.utils.session import NEW_SESSION, provide_session
 
@@ -27,22 +27,20 @@
     from airflow.jobs.job import Job
     from airflow.serialization.pydantic.job import JobPydantic
 
-J = TypeVar("J", "Job", "JobPydantic", "Job | JobPydantic")
 
-
-class BaseJobRunner(Generic[J]):
+class BaseJobRunner:
     """Abstract class for job runners to derive from."""
 
     job_type = "undefined"
 
-    def __init__(self, job: J) -> None:
+    def __init__(self, job: Job) -> None:
         if job.job_type and job.job_type != self.job_type:
             raise Exception(
                 f"The job is already assigned a different job_type: {job.job_type}."
                 f"This is a bug and should be reported."
             )
         job.job_type = self.job_type
-        self.job: J = job
+        self.job: Job = job
 
     def _execute(self) -> int | None:
         """
@@ -65,7 +63,7 @@ def heartbeat_callback(self, session: Session = NEW_SESSION) -> None:
 
     @classmethod
     @provide_session
-    def most_recent_job(cls, session: Session = NEW_SESSION) -> Job | None:
+    def most_recent_job(cls, session: Session = NEW_SESSION) -> Job | JobPydantic | None:
         """Return the most recent job of this type, if any, based on last heartbeat received."""
         from airflow.jobs.job import most_recent_job
 

airflow/jobs/dag_processor_job_runner.py

@@ -31,7 +31,7 @@ def empty_callback(_: Any) -> None:
     pass
 
 
-class DagProcessorJobRunner(BaseJobRunner[Job], LoggingMixin):
+class DagProcessorJobRunner(BaseJobRunner, LoggingMixin):
     """
     DagProcessorJobRunner is a job runner that runs a DagFileProcessorManager processor.