From 522421b7954853ffd3c569fdbaeac10cf712df36 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung 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: