From 522421b7954853ffd3c569fdbaeac10cf712df36 Mon Sep 17 00:00:00 2001
From: Tzu-ping Chung <uranusjr@gmail.com>
Date: Mon, 15 May 2023 17:27:09 +0800
Subject: [PATCH] Make BaseJobRunner a generic on the job class

---
 airflow/jobs/backfill_job_runner.py      | 3 +--
 airflow/jobs/base_job_runner.py          | 9 ++++++---
 airflow/jobs/dag_processor_job_runner.py | 3 +--
 airflow/jobs/local_task_job_runner.py    | 3 +--
 airflow/jobs/scheduler_job_runner.py     | 3 +--
 airflow/jobs/triggerer_job_runner.py     | 3 +--
 6 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/airflow/jobs/backfill_job_runner.py b/airflow/jobs/backfill_job_runner.py
index ebc92187f9551..28986b2bede84 100644
--- a/airflow/jobs/backfill_job_runner.py
+++ b/airflow/jobs/backfill_job_runner.py
@@ -57,7 +57,7 @@
     from airflow.models.abstractoperator import AbstractOperator
 
 
-class BackfillJobRunner(BaseJobRunner, LoggingMixin):
+class BackfillJobRunner(BaseJobRunner[Job], LoggingMixin):
     """
     A backfill job runner consists of a dag or subdag for a specific time range.
 
@@ -149,7 +149,6 @@ def __init__(
         :param kwargs:
         """
         super().__init__(job)
-        self.job = job
         self.dag = dag
         self.dag_id = dag.dag_id
         self.bf_start_date = start_date
diff --git a/airflow/jobs/base_job_runner.py b/airflow/jobs/base_job_runner.py
index 5b0af0dbecd01..0a3222bc14b2a 100644
--- a/airflow/jobs/base_job_runner.py
+++ b/airflow/jobs/base_job_runner.py
@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Generic, TypeVar
 
 from airflow.utils.session import NEW_SESSION, provide_session
 
@@ -27,19 +27,22 @@
     from airflow.jobs.job import Job
     from airflow.serialization.pydantic.job import JobPydantic
 
+J = TypeVar("J", "Job", "JobPydantic", "Job | JobPydantic")
 
-class BaseJobRunner:
+
+class BaseJobRunner(Generic[J]):
     """Abstract class for job runners to derive from."""
 
     job_type = "undefined"
 
-    def __init__(self, job: Job | JobPydantic) -> None:
+    def __init__(self, job: J) -> 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
 
     def _execute(self) -> int | None:
         """
diff --git a/airflow/jobs/dag_processor_job_runner.py b/airflow/jobs/dag_processor_job_runner.py
index b2420726c6a99..466b5caa7d19a 100644
--- a/airflow/jobs/dag_processor_job_runner.py
+++ b/airflow/jobs/dag_processor_job_runner.py
@@ -29,7 +29,7 @@ def empty_callback(_: Any) -> None:
     pass
 
 
-class DagProcessorJobRunner(BaseJobRunner, LoggingMixin):
+class DagProcessorJobRunner(BaseJobRunner[Job], LoggingMixin):
     """
     DagProcessorJobRunner is a job runner that runs a DagFileProcessorManager processor.
 
@@ -47,7 +47,6 @@ def __init__(
         **kwargs,
     ):
         super().__init__(job)
-        self.job = job
         self.processor = processor
         self.processor.heartbeat = lambda: perform_heartbeat(
             job=self.job,
diff --git a/airflow/jobs/local_task_job_runner.py b/airflow/jobs/local_task_job_runner.py
index 2fb4eaaba3ef6..93046696e69d3 100644
--- a/airflow/jobs/local_task_job_runner.py
+++ b/airflow/jobs/local_task_job_runner.py
@@ -68,7 +68,7 @@
 ********************************************************************************************************"""
 
 
-class LocalTaskJobRunner(BaseJobRunner, LoggingMixin):
+class LocalTaskJobRunner(BaseJobRunner["Job | JobPydantic"], LoggingMixin):
     """LocalTaskJob runs a single task instance."""
 
     job_type = "LocalTaskJob"
@@ -88,7 +88,6 @@ def __init__(
         external_executor_id: str | None = None,
     ):
         super().__init__(job)
-        self.job = job
         LoggingMixin.__init__(self, context=task_instance)
         self.task_instance = task_instance
         self.ignore_all_deps = ignore_all_deps
diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py
index 080b8bf372fcf..e7ed5ce00ee22 100644
--- a/airflow/jobs/scheduler_job_runner.py
+++ b/airflow/jobs/scheduler_job_runner.py
@@ -119,7 +119,7 @@ def _is_parent_process() -> bool:
     return multiprocessing.current_process().name == "MainProcess"
 
 
-class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
+class SchedulerJobRunner(BaseJobRunner[Job], LoggingMixin):
     """
     SchedulerJobRunner runs for a specific time interval and schedules jobs that are ready to run.
 
@@ -157,7 +157,6 @@ def __init__(
         processor_poll_interval: float | None = None,
     ):
         super().__init__(job)
-        self.job = job
         self.subdir = subdir
         self.num_runs = num_runs
         # In specific tests, we want to stop the parse loop after the _files_ have been parsed a certain
diff --git a/airflow/jobs/triggerer_job_runner.py b/airflow/jobs/triggerer_job_runner.py
index 4651b6cf9210b..c417be1539340 100644
--- a/airflow/jobs/triggerer_job_runner.py
+++ b/airflow/jobs/triggerer_job_runner.py
@@ -234,7 +234,7 @@ def setup_queue_listener():
         return None
 
 
-class TriggererJobRunner(BaseJobRunner, LoggingMixin):
+class TriggererJobRunner(BaseJobRunner["Job | JobPydantic"], LoggingMixin):
     """
     TriggererJobRunner continuously runs active triggers in asyncio, watching
     for them to fire off their events and then dispatching that information
@@ -253,7 +253,6 @@ def __init__(
         capacity=None,
     ):
         super().__init__(job)
-        self.job = job
         if capacity is None:
             self.capacity = conf.getint("triggerer", "default_capacity", fallback=1000)
         elif isinstance(capacity, int) and capacity > 0:
