Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions airflow/jobs/backfill_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand Down
9 changes: 6 additions & 3 deletions airflow/jobs/base_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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:
"""
Expand Down
3 changes: 1 addition & 2 deletions airflow/jobs/dag_processor_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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,
Expand Down
3 changes: 1 addition & 2 deletions airflow/jobs/local_task_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@
********************************************************************************************************"""


class LocalTaskJobRunner(BaseJobRunner, LoggingMixin):
class LocalTaskJobRunner(BaseJobRunner["Job | JobPydantic"], LoggingMixin):
"""LocalTaskJob runs a single task instance."""

job_type = "LocalTaskJob"
Expand All @@ -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
Expand Down
3 changes: 1 addition & 2 deletions airflow/jobs/scheduler_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions airflow/jobs/triggerer_job_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down