From e785922166b67f9f6420ef591aaabea69bd9c6d8 Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Sat, 25 Mar 2023 20:22:03 +0100 Subject: [PATCH] Rename JobRunner modules to *_job_runner and base_job* to job The #30255 introduced "JobRunner" concept and decoupled the job logic from the ORM polymorphic *Job objects. The change was implemented in the way to minimise the review effort needed, so it avoided renaming the modules for the runners (from `_job` to `_job_runner`). Also BaseJob lost its "polymorphism" properties so the package, and class name can be renamed to simply job. This PR completes the JobRunner concept introduction by applying the renames. Closes: #30296 --- .../endpoints/health_endpoint.py | 4 +- airflow/api_connexion/schemas/job_schema.py | 4 +- airflow/cli/commands/dag_command.py | 10 +- airflow/cli/commands/dag_processor_command.py | 6 +- airflow/cli/commands/jobs_command.py | 17 +- airflow/cli/commands/scheduler_command.py | 8 +- airflow/cli/commands/standalone_command.py | 8 +- airflow/cli/commands/task_command.py | 6 +- airflow/cli/commands/triggerer_command.py | 6 +- airflow/dag_processing/manager.py | 6 +- airflow/executors/celery_executor.py | 2 +- ...backfill_job.py => backfill_job_runner.py} | 6 +- .../{job_runner.py => base_job_runner.py} | 10 +- ...sor_job.py => dag_processor_job_runner.py} | 2 +- airflow/jobs/{base_job.py => job.py} | 34 +- ...l_task_job.py => local_task_job_runner.py} | 4 +- ...heduler_job.py => scheduler_job_runner.py} | 26 +- ...iggerer_job.py => triggerer_job_runner.py} | 4 +- airflow/models/__init__.py | 2 +- airflow/models/dag.py | 6 +- airflow/models/taskinstance.py | 8 +- airflow/models/trigger.py | 14 +- .../pydantic/{base_job.py => job.py} | 6 +- airflow/serialization/serialized_objects.py | 10 +- airflow/task/task_runner/__init__.py | 2 +- airflow/task/task_runner/base_task_runner.py | 6 +- .../task/task_runner/cgroup_task_runner.py | 4 +- .../task/task_runner/standard_task_runner.py | 4 +- airflow/utils/scheduler_health.py | 8 +- airflow/www/views.py | 8 +- dev/perf/scheduler_dag_execution_timing.py | 10 +- dev/perf/sql_queries.py | 6 +- .../logging-monitoring/check-health.rst | 2 +- .../endpoints/test_health_endpoint.py | 12 +- .../endpoints/test_task_instance_endpoint.py | 6 +- .../commands/test_dag_processor_command.py | 2 +- tests/cli/commands/test_jobs_command.py | 14 +- tests/core/test_impersonation_tests.py | 6 +- tests/executors/test_dask_executor.py | 6 +- tests/jobs/test_backfill_job.py | 114 +++---- tests/jobs/test_base_job.py | 44 ++- tests/jobs/test_local_task_job.py | 58 ++-- tests/jobs/test_scheduler_job.py | 312 +++++++++--------- tests/jobs/test_triggerer_job.py | 36 +- tests/jobs/test_triggerer_job_logging.py | 46 +-- tests/listeners/test_listeners.py | 4 +- tests/models/test_trigger.py | 14 +- tests/serialization/test_pydantic_models.py | 12 +- .../task/task_runner/test_base_task_runner.py | 6 +- .../task_runner/test_cgroup_task_runner.py | 10 +- .../task_runner/test_standard_task_runner.py | 64 ++-- tests/task/task_runner/test_task_runner.py | 16 +- tests/test_utils/db.py | 6 +- tests/utils/test_helpers.py | 2 +- tests/utils/test_log_handlers.py | 8 +- tests/www/views/test_views_base.py | 34 +- 56 files changed, 531 insertions(+), 560 deletions(-) rename airflow/jobs/{backfill_job.py => backfill_job_runner.py} (99%) rename airflow/jobs/{job_runner.py => base_job_runner.py} (89%) rename airflow/jobs/{dag_processor_job.py => dag_processor_job_runner.py} (96%) rename airflow/jobs/{base_job.py => job.py} (91%) rename airflow/jobs/{local_task_job.py => local_task_job_runner.py} (99%) rename airflow/jobs/{scheduler_job.py => scheduler_job_runner.py} (98%) rename airflow/jobs/{triggerer_job.py => triggerer_job_runner.py} (99%) rename airflow/serialization/pydantic/{base_job.py => job.py} (89%) diff --git a/airflow/api_connexion/endpoints/health_endpoint.py b/airflow/api_connexion/endpoints/health_endpoint.py index 83b141cdefb8a..3389343646780 100644 --- a/airflow/api_connexion/endpoints/health_endpoint.py +++ b/airflow/api_connexion/endpoints/health_endpoint.py @@ -18,8 +18,8 @@ from airflow.api_connexion.schemas.health_schema import health_schema from airflow.api_connexion.types import APIResponse -from airflow.jobs.scheduler_job import SchedulerJobRunner -from airflow.jobs.triggerer_job import TriggererJobRunner +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner +from airflow.jobs.triggerer_job_runner import TriggererJobRunner HEALTHY = "healthy" UNHEALTHY = "unhealthy" diff --git a/airflow/api_connexion/schemas/job_schema.py b/airflow/api_connexion/schemas/job_schema.py index 4d98d39c92030..d4f77eedf3b42 100644 --- a/airflow/api_connexion/schemas/job_schema.py +++ b/airflow/api_connexion/schemas/job_schema.py @@ -19,7 +19,7 @@ from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job class JobSchema(SQLAlchemySchema): @@ -28,7 +28,7 @@ class JobSchema(SQLAlchemySchema): class Meta: """Meta.""" - model = BaseJob + model = Job id = auto_field(dump_only=True) dag_id = auto_field(dump_only=True) diff --git a/airflow/cli/commands/dag_command.py b/airflow/cli/commands/dag_command.py index f1ef3abc6e7f6..8735a0e522715 100644 --- a/airflow/cli/commands/dag_command.py +++ b/airflow/cli/commands/dag_command.py @@ -35,7 +35,7 @@ from airflow.cli.simple_table import AirflowConsole from airflow.configuration import conf from airflow.exceptions import AirflowException, RemovedInAirflow3Warning -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job from airflow.models import DagBag, DagModel, DagRun, TaskInstance from airflow.models.dag import DAG from airflow.models.serialized_dag import SerializedDagModel @@ -391,15 +391,13 @@ def dag_list_jobs(args, dag: DAG | None = None, session: Session = NEW_SESSION) if not dag: raise SystemExit(f"DAG: {args.dag_id} does not exist in 'dag' table") - queries.append(BaseJob.dag_id == args.dag_id) + queries.append(Job.dag_id == args.dag_id) if args.state: - queries.append(BaseJob.state == args.state) + queries.append(Job.state == args.state) fields = ["dag_id", "state", "job_type", "start_date", "end_date"] - all_jobs = ( - session.query(BaseJob).filter(*queries).order_by(BaseJob.start_date.desc()).limit(args.limit).all() - ) + all_jobs = session.query(Job).filter(*queries).order_by(Job.start_date.desc()).limit(args.limit).all() all_jobs = [{f: str(job.__getattribute__(f)) for f in fields} for job in all_jobs] AirflowConsole().print_as( diff --git a/airflow/cli/commands/dag_processor_command.py b/airflow/cli/commands/dag_processor_command.py index 396fc746da29f..e86e6b2171fc0 100644 --- a/airflow/cli/commands/dag_processor_command.py +++ b/airflow/cli/commands/dag_processor_command.py @@ -26,14 +26,14 @@ from airflow import settings from airflow.configuration import conf -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job from airflow.utils import cli as cli_utils from airflow.utils.cli import setup_locations, setup_logging log = logging.getLogger(__name__) -def _create_dag_processor_job(args: Any) -> BaseJob: +def _create_dag_processor_job(args: Any) -> Job: """Creates DagFileProcessorProcess instance.""" from airflow.dag_processing.manager import DagFileProcessorManager @@ -47,7 +47,7 @@ def _create_dag_processor_job(args: Any) -> BaseJob: dag_ids=[], pickle_dags=args.do_pickle, ) - return BaseJob( + return Job( job_runner=processor.job_runner, ) diff --git a/airflow/cli/commands/jobs_command.py b/airflow/cli/commands/jobs_command.py index 959f7ebc4c841..79e0aad79fde5 100644 --- a/airflow/cli/commands/jobs_command.py +++ b/airflow/cli/commands/jobs_command.py @@ -18,7 +18,7 @@ from sqlalchemy.orm import Session -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job from airflow.utils.net import get_hostname from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.state import State @@ -32,22 +32,17 @@ def check(args, session: Session = NEW_SESSION) -> None: if args.hostname and args.local: raise SystemExit("You can't use --hostname and --local at the same time") - query = ( - session.query(BaseJob) - .filter(BaseJob.state == State.RUNNING) - .order_by(BaseJob.latest_heartbeat.desc()) - ) + query = session.query(Job).filter(Job.state == State.RUNNING).order_by(Job.latest_heartbeat.desc()) if args.job_type: - query = query.filter(BaseJob.job_type == args.job_type) + query = query.filter(Job.job_type == args.job_type) if args.hostname: - query = query.filter(BaseJob.hostname == args.hostname) + query = query.filter(Job.hostname == args.hostname) if args.local: - query = query.filter(BaseJob.hostname == get_hostname()) + query = query.filter(Job.hostname == get_hostname()) if args.limit > 0: query = query.limit(args.limit) - jobs: list[BaseJob] = query.all() - alive_jobs = [job for job in jobs if job.is_alive()] + alive_jobs: list[Job] = [job for job in query.all() if job.is_alive()] count_alive_jobs = len(alive_jobs) if count_alive_jobs == 0: diff --git a/airflow/cli/commands/scheduler_command.py b/airflow/cli/commands/scheduler_command.py index ba2e58d9e1e48..3a28dd2f37704 100644 --- a/airflow/cli/commands/scheduler_command.py +++ b/airflow/cli/commands/scheduler_command.py @@ -28,14 +28,14 @@ from airflow.api_internal.internal_api_call import InternalApiConfig from airflow.configuration import conf from airflow.executors.executor_loader import ExecutorLoader -from airflow.jobs.base_job import BaseJob -from airflow.jobs.scheduler_job import SchedulerJobRunner +from airflow.jobs.job import Job +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.utils import cli as cli_utils from airflow.utils.cli import process_subdir, setup_locations, setup_logging, sigint_handler, sigquit_handler from airflow.utils.scheduler_health import serve_health_check -def _run_scheduler_job(job: BaseJob, *, skip_serve_logs: bool) -> None: +def _run_scheduler_job(job: Job, *, skip_serve_logs: bool) -> None: InternalApiConfig.force_database_direct_access() enable_health_check = conf.getboolean("scheduler", "ENABLE_HEALTH_CHECK") with _serve_logs(skip_serve_logs), _serve_health_check(enable_health_check): @@ -47,7 +47,7 @@ def scheduler(args): """Starts Airflow Scheduler.""" print(settings.HEADER) - job = BaseJob( + job = Job( job_runner=SchedulerJobRunner( subdir=process_subdir(args.subdir), num_runs=args.num_runs, diff --git a/airflow/cli/commands/standalone_command.py b/airflow/cli/commands/standalone_command.py index 1f5dcc6216cf5..35bd49ab68ab2 100644 --- a/airflow/cli/commands/standalone_command.py +++ b/airflow/cli/commands/standalone_command.py @@ -30,10 +30,10 @@ from airflow.configuration import AIRFLOW_HOME, conf, make_group_other_inaccessible from airflow.executors import executor_constants from airflow.executors.executor_loader import ExecutorLoader -from airflow.jobs.base_job import most_recent_job -from airflow.jobs.job_runner import BaseJobRunner -from airflow.jobs.scheduler_job import SchedulerJobRunner -from airflow.jobs.triggerer_job import TriggererJobRunner +from airflow.jobs.base_job_runner import BaseJobRunner +from airflow.jobs.job import most_recent_job +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner +from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.utils import db diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py index 65945ffdc7cb2..c318c48034f6c 100644 --- a/airflow/cli/commands/task_command.py +++ b/airflow/cli/commands/task_command.py @@ -37,8 +37,8 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException, DagRunNotFound, TaskInstanceNotFound from airflow.executors.executor_loader import ExecutorLoader -from airflow.jobs.base_job import BaseJob -from airflow.jobs.local_task_job import LocalTaskJobRunner +from airflow.jobs.job import Job +from airflow.jobs.local_task_job_runner import LocalTaskJobRunner from airflow.listeners.listener import get_listener_manager from airflow.models import DagPickle, TaskInstance from airflow.models.dag import DAG @@ -260,7 +260,7 @@ def _run_task_by_local_task_job(args, ti: TaskInstance) -> TaskReturnCode | None pool=args.pool, external_executor_id=_extract_external_executor_id(args), ) - run_job = BaseJob( + run_job = Job( job_runner=local_task_job_runner, dag_id=ti.dag_id, ) diff --git a/airflow/cli/commands/triggerer_command.py b/airflow/cli/commands/triggerer_command.py index e4e1f5fd029f5..604695b92eed9 100644 --- a/airflow/cli/commands/triggerer_command.py +++ b/airflow/cli/commands/triggerer_command.py @@ -28,8 +28,8 @@ from airflow import settings from airflow.configuration import conf -from airflow.jobs.base_job import BaseJob -from airflow.jobs.triggerer_job import TriggererJobRunner +from airflow.jobs.job import Job +from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.utils import cli as cli_utils from airflow.utils.cli import setup_locations, setup_logging, sigint_handler, sigquit_handler from airflow.utils.serve_logs import serve_logs @@ -56,7 +56,7 @@ def triggerer(args): settings.MASK_SECRETS_IN_LOGS = True print(settings.HEADER) triggerer_job_runner = TriggererJobRunner(capacity=args.capacity) - job = BaseJob(job_runner=triggerer_job_runner) + job = Job(job_runner=triggerer_job_runner) if args.daemon: pid, stdout, stderr, log_file = setup_locations( diff --git a/airflow/dag_processing/manager.py b/airflow/dag_processing/manager.py index 18fcb41012e95..b45d833e715be 100644 --- a/airflow/dag_processing/manager.py +++ b/airflow/dag_processing/manager.py @@ -46,7 +46,7 @@ from airflow.callbacks.callback_requests import CallbackRequest, SlaCallbackRequest from airflow.configuration import conf from airflow.dag_processing.processor import DagFileProcessorProcess -from airflow.jobs.base_job import perform_heartbeat +from airflow.jobs.job import perform_heartbeat from airflow.models import errors from airflow.models.dag import DagModel from airflow.models.dagwarning import DagWarning @@ -68,7 +68,7 @@ from airflow.utils.sqlalchemy import prohibit_commit, skip_locked, with_row_locks if TYPE_CHECKING: - from airflow.jobs.dag_processor_job import DagProcessorJobRunner + from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner class DagParsingStat(NamedTuple): @@ -385,7 +385,7 @@ def __init__( signal_conn: MultiprocessingConnection | None = None, async_mode: bool = True, ): - from airflow.jobs.dag_processor_job import DagProcessorJobRunner + from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner super().__init__() # known files; this will be updated every `dag_dir_list_interval` and stuff added/removed accordingly diff --git a/airflow/executors/celery_executor.py b/airflow/executors/celery_executor.py index 5be3f7fc60c2d..63502f61ed872 100644 --- a/airflow/executors/celery_executor.py +++ b/airflow/executors/celery_executor.py @@ -197,7 +197,7 @@ def on_celery_import_modules(*args, **kwargs): """ import jinja2.ext # noqa: F401 - import airflow.jobs.local_task_job + import airflow.jobs.local_task_job_runner import airflow.macros import airflow.operators.bash import airflow.operators.python diff --git a/airflow/jobs/backfill_job.py b/airflow/jobs/backfill_job_runner.py similarity index 99% rename from airflow/jobs/backfill_job.py rename to airflow/jobs/backfill_job_runner.py index c170ced21449a..435dfe40df8b8 100644 --- a/airflow/jobs/backfill_job.py +++ b/airflow/jobs/backfill_job_runner.py @@ -37,8 +37,8 @@ TaskConcurrencyLimitReached, ) from airflow.executors.executor_loader import ExecutorLoader -from airflow.jobs.base_job import BaseJob -from airflow.jobs.job_runner import BaseJobRunner +from airflow.jobs.base_job_runner import BaseJobRunner +from airflow.jobs.job import Job from airflow.models import DAG, DagPickle from airflow.models.dagrun import DagRun from airflow.models.taskinstance import TaskInstance, TaskInstanceKey @@ -69,7 +69,7 @@ class BackfillJobRunner(BaseJobRunner, LoggingMixin): STATES_COUNT_AS_RUNNING = (State.RUNNING, State.QUEUED) - job: BaseJob # backfill_job can only run with BaseJob class not the Pydantic serialized version + job: Job # backfill_job can only run with Job class not the Pydantic serialized version @attr.define class _DagRunTaskStatus: diff --git a/airflow/jobs/job_runner.py b/airflow/jobs/base_job_runner.py similarity index 89% rename from airflow/jobs/job_runner.py rename to airflow/jobs/base_job_runner.py index 27fff814a9b9e..096a929d1e6ef 100644 --- a/airflow/jobs/job_runner.py +++ b/airflow/jobs/base_job_runner.py @@ -24,8 +24,8 @@ if TYPE_CHECKING: from sqlalchemy.orm import Session - from airflow.jobs.base_job import BaseJob - from airflow.serialization.pydantic.base_job import BaseJobPydantic + from airflow.jobs.job import Job + from airflow.serialization.pydantic.job import JobPydantic class BaseJobRunner: @@ -33,7 +33,7 @@ class BaseJobRunner: job_type = "undefined" - job: BaseJob | BaseJobPydantic + job: Job | JobPydantic def _execute(self) -> int | None: """ @@ -51,8 +51,8 @@ def heartbeat_callback(self, session: Session = NEW_SESSION) -> None: @classmethod @provide_session - def most_recent_job(cls, session: Session = NEW_SESSION) -> BaseJob | None: + def most_recent_job(cls, session: Session = NEW_SESSION) -> Job | None: """Returns the most recent job of this type, if any, based on last heartbeat received.""" - from airflow.jobs.base_job import most_recent_job + from airflow.jobs.job import most_recent_job return most_recent_job(cls.job_type, session=session) diff --git a/airflow/jobs/dag_processor_job.py b/airflow/jobs/dag_processor_job_runner.py similarity index 96% rename from airflow/jobs/dag_processor_job.py rename to airflow/jobs/dag_processor_job_runner.py index 056fdc2d7d820..a6a4548e40ec1 100644 --- a/airflow/jobs/dag_processor_job.py +++ b/airflow/jobs/dag_processor_job_runner.py @@ -18,7 +18,7 @@ from __future__ import annotations from airflow.dag_processing.manager import DagFileProcessorManager -from airflow.jobs.job_runner import BaseJobRunner +from airflow.jobs.base_job_runner import BaseJobRunner from airflow.utils.log.logging_mixin import LoggingMixin diff --git a/airflow/jobs/base_job.py b/airflow/jobs/job.py similarity index 91% rename from airflow/jobs/base_job.py rename to airflow/jobs/job.py index c01eeb4800274..5390ce04fe75e 100644 --- a/airflow/jobs/base_job.py +++ b/airflow/jobs/job.py @@ -29,10 +29,10 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.executors.executor_loader import ExecutorLoader -from airflow.jobs.job_runner import BaseJobRunner +from airflow.jobs.base_job_runner import BaseJobRunner from airflow.listeners.listener import get_listener_manager from airflow.models.base import ID_LEN, Base -from airflow.serialization.pydantic.base_job import BaseJobPydantic +from airflow.serialization.pydantic.job import JobPydantic from airflow.stats import Stats from airflow.utils import timezone from airflow.utils.helpers import convert_camel_to_snake @@ -50,9 +50,9 @@ def _resolve_dagrun_model(): return DagRun -class BaseJob(Base, LoggingMixin): +class Job(Base, LoggingMixin): """ - Abstract class to be derived for jobs. + The ORM class representing Job stored in the database. Jobs are processing items with state and duration that aren't task instances. For instance a BackfillJob is a collection of task instance runs, @@ -82,13 +82,13 @@ class BaseJob(Base, LoggingMixin): task_instances_enqueued = relationship( "TaskInstance", - primaryjoin="BaseJob.id == foreign(TaskInstance.queued_by_job_id)", + primaryjoin="Job.id == foreign(TaskInstance.queued_by_job_id)", backref=backref("queued_by_job", uselist=False), ) dag_runs = relationship( "DagRun", - primaryjoin=lambda: BaseJob.id == foreign(_resolve_dagrun_model().creating_job_id), + primaryjoin=lambda: Job.id == foreign(_resolve_dagrun_model().creating_job_id), backref="creating_job", ) @@ -143,7 +143,7 @@ def is_alive(self, grace_multiplier=2.1): @provide_session def kill(self, session: Session = NEW_SESSION) -> NoReturn: """Handles on_kill callback and updates state in database.""" - job = session.query(BaseJob).filter(BaseJob.id == self.id).first() + job = session.query(Job).filter(Job.id == self.id).first() job.end_date = timezone.utcnow() try: self.on_kill() @@ -257,24 +257,22 @@ def job_runner(self) -> BaseJobRunner: return self._job_runner @provide_session - def most_recent_job(self, session: Session = NEW_SESSION) -> BaseJob | None: + def most_recent_job(self, session: Session = NEW_SESSION) -> Job | None: """Returns the most recent job of this type, if any, based on last heartbeat received.""" return most_recent_job(self.job_type, session=session) @provide_session -def perform_heartbeat( - job: BaseJob | BaseJobPydantic, only_if_necessary: bool, session: Session = NEW_SESSION -): - if isinstance(job, BaseJob): +def perform_heartbeat(job: Job | JobPydantic, only_if_necessary: bool, session: Session = NEW_SESSION): + if isinstance(job, Job): job.heartbeat() else: # TODO (potiuk): Make it works over internal API as a follow up - BaseJob.get(job.id, session=session).heartbeat(only_if_necessary=only_if_necessary) + Job.get(job.id, session=session).heartbeat(only_if_necessary=only_if_necessary) @provide_session -def most_recent_job(job_type: str, session: Session = NEW_SESSION) -> BaseJob | None: +def most_recent_job(job_type: str, session: Session = NEW_SESSION) -> Job | None: """ Return the most recent job of this type, if any, based on last heartbeat received. @@ -285,12 +283,12 @@ def most_recent_job(job_type: str, session: Session = NEW_SESSION) -> BaseJob | :param session: Database session """ return ( - session.query(BaseJob) - .filter(BaseJob.job_type == job_type) + session.query(Job) + .filter(Job.job_type == job_type) .order_by( # Put "running" jobs at the front. - case({State.RUNNING: 0}, value=BaseJob.state, else_=1), - BaseJob.latest_heartbeat.desc(), + case({State.RUNNING: 0}, value=Job.state, else_=1), + Job.latest_heartbeat.desc(), ) .first() ) diff --git a/airflow/jobs/local_task_job.py b/airflow/jobs/local_task_job_runner.py similarity index 99% rename from airflow/jobs/local_task_job.py rename to airflow/jobs/local_task_job_runner.py index 1ba7ef9fb97f9..d5becaa243903 100644 --- a/airflow/jobs/local_task_job.py +++ b/airflow/jobs/local_task_job_runner.py @@ -24,8 +24,8 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException -from airflow.jobs.base_job import perform_heartbeat -from airflow.jobs.job_runner import BaseJobRunner +from airflow.jobs.base_job_runner import BaseJobRunner +from airflow.jobs.job import perform_heartbeat from airflow.models.taskinstance import TaskInstance, TaskReturnCode from airflow.stats import Stats from airflow.utils import timezone diff --git a/airflow/jobs/scheduler_job.py b/airflow/jobs/scheduler_job_runner.py similarity index 98% rename from airflow/jobs/scheduler_job.py rename to airflow/jobs/scheduler_job_runner.py index af790f266e66a..98fdfdb82d6a7 100644 --- a/airflow/jobs/scheduler_job.py +++ b/airflow/jobs/scheduler_job_runner.py @@ -42,8 +42,8 @@ from airflow.configuration import conf from airflow.exceptions import RemovedInAirflow3Warning from airflow.executors.executor_loader import ExecutorLoader -from airflow.jobs.base_job import BaseJob -from airflow.jobs.job_runner import BaseJobRunner +from airflow.jobs.base_job_runner import BaseJobRunner +from airflow.jobs.job import Job from airflow.models.dag import DAG, DagModel from airflow.models.dagbag import DagBag from airflow.models.dagrun import DagRun @@ -121,7 +121,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin): job_type = "SchedulerJob" heartrate: int = conf.getint("scheduler", "SCHEDULER_HEARTBEAT_SEC") - job: BaseJob # scheduler can only run with BaseJob class not the Pydantic serialized version + job: Job # scheduler can only run with Job class not the Pydantic serialized version def __init__( self, @@ -1457,11 +1457,11 @@ def adopt_or_reset_orphaned_tasks(self, session: Session = NEW_SESSION) -> int: self.log.debug("Calling SchedulerJob.adopt_or_reset_orphaned_tasks method") try: num_failed = ( - session.query(BaseJob) + session.query(Job) .filter( - BaseJob.job_type == "SchedulerJob", - BaseJob.state == State.RUNNING, - BaseJob.latest_heartbeat < (timezone.utcnow() - timedelta(seconds=timeout)), + Job.job_type == "SchedulerJob", + Job.state == State.RUNNING, + Job.latest_heartbeat < (timezone.utcnow() - timedelta(seconds=timeout)), ) .update({"state": State.FAILED}) ) @@ -1479,7 +1479,7 @@ def adopt_or_reset_orphaned_tasks(self, session: Session = NEW_SESSION) -> int: # "or queued_by_job_id IS NONE") can go as soon as scheduler HA is # released. .outerjoin(TI.queued_by_job) - .filter(or_(TI.queued_by_job_id.is_(None), BaseJob.state != State.RUNNING)) + .filter(or_(TI.queued_by_job_id.is_(None), Job.state != State.RUNNING)) .join(TI.dag_run) .filter( DagRun.run_type != DagRunType.BACKFILL_JOB, @@ -1555,7 +1555,7 @@ def _find_zombies(self) -> None: or have a no-longer-running LocalTaskJob, and create a TaskCallbackRequest to be handled by the DAG processor. """ - from airflow.jobs.base_job import BaseJob + from airflow.jobs.job import Job self.log.debug("Finding 'running' jobs without a recent heartbeat") limit_dttm = timezone.utcnow() - timedelta(seconds=self._zombie_threshold_secs) @@ -1564,16 +1564,16 @@ def _find_zombies(self) -> None: zombies: list[tuple[TI, str, str]] = ( session.query(TI, DM.fileloc, DM.processor_subdir) .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql") - .join(BaseJob, TI.job_id == BaseJob.id) + .join(Job, TI.job_id == Job.id) .join(DM, TI.dag_id == DM.dag_id) .filter(TI.state == TaskInstanceState.RUNNING) .filter( or_( - BaseJob.state != State.RUNNING, - BaseJob.latest_heartbeat < limit_dttm, + Job.state != State.RUNNING, + Job.latest_heartbeat < limit_dttm, ) ) - .filter(BaseJob.job_type == "LocalTaskJob") + .filter(Job.job_type == "LocalTaskJob") .filter(TI.queued_by_job_id == self.job.id) .all() ) diff --git a/airflow/jobs/triggerer_job.py b/airflow/jobs/triggerer_job_runner.py similarity index 99% rename from airflow/jobs/triggerer_job.py rename to airflow/jobs/triggerer_job_runner.py index cd2d9854f436f..da5c1b969ebf1 100644 --- a/airflow/jobs/triggerer_job.py +++ b/airflow/jobs/triggerer_job_runner.py @@ -33,8 +33,8 @@ from sqlalchemy import func from airflow.configuration import conf -from airflow.jobs.base_job import perform_heartbeat -from airflow.jobs.job_runner import BaseJobRunner +from airflow.jobs.base_job_runner import BaseJobRunner +from airflow.jobs.job import perform_heartbeat from airflow.models.trigger import Trigger from airflow.stats import Stats from airflow.triggers.base import BaseTrigger, TriggerEvent diff --git a/airflow/models/__init__.py b/airflow/models/__init__.py index d98a60ce6c8db..aa1e2c98ef4c6 100644 --- a/airflow/models/__init__.py +++ b/airflow/models/__init__.py @@ -59,7 +59,7 @@ def import_all_models(): for name in __lazy_imports: __getattr__(name) - import airflow.jobs.base_job + import airflow.jobs.job import airflow.models.dagwarning import airflow.models.dataset import airflow.models.serialized_dag diff --git a/airflow/models/dag.py b/airflow/models/dag.py index a182be157f9f2..d77644b76a7c4 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -2455,7 +2455,7 @@ def run( :param run_at_least_once: If true, always run the DAG at least once even if no logical run exists within the time range. """ - from airflow.jobs.backfill_job import BackfillJobRunner + from airflow.jobs.backfill_job_runner import BackfillJobRunner if not executor and local: from airflow.executors.local_executor import LocalExecutor @@ -2465,9 +2465,9 @@ def run( from airflow.executors.executor_loader import ExecutorLoader executor = ExecutorLoader.get_default_executor() - from airflow.jobs.base_job import BaseJob + from airflow.jobs.job import Job - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( self, start_date=start_date, diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 24b81de03103d..c2ed7b786864c 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -265,9 +265,9 @@ def clear_task_instances( session.execute(delete_qry) if job_ids: - from airflow.jobs.base_job import BaseJob + from airflow.jobs.job import Job - for job in session.query(BaseJob).filter(BaseJob.id.in_(job_ids)).all(): + for job in session.query(Job).filter(Job.id.in_(job_ids)).all(): job.state = TaskInstanceState.RESTARTING if activate_dag_runs is not None: @@ -2955,6 +2955,6 @@ def __repr__(self): STATICA_HACK = True globals()["kcah_acitats"[::-1].upper()] = False if STATICA_HACK: # pragma: no cover - from airflow.jobs.base_job import BaseJob + from airflow.jobs.job import Job - TaskInstance.queued_by_job = relationship(BaseJob) + TaskInstance.queued_by_job = relationship(Job) diff --git a/airflow/models/trigger.py b/airflow/models/trigger.py index f6195fa3d6cb8..adad7ba1c34aa 100644 --- a/airflow/models/trigger.py +++ b/airflow/models/trigger.py @@ -60,8 +60,8 @@ class Trigger(Base): triggerer_id = Column(Integer, nullable=True) triggerer_job = relationship( - "BaseJob", - primaryjoin="BaseJob.id == Trigger.triggerer_id", + "Job", + primaryjoin="Job.id == Trigger.triggerer_id", foreign_keys=triggerer_id, uselist=False, ) @@ -201,7 +201,7 @@ def assign_unassigned(cls, triggerer_id, capacity, session: Session = NEW_SESSIO Takes a triggerer_id and the capacity for that triggerer and assigns unassigned triggers until that capacity is reached, or there are no more unassigned triggers. """ - from airflow.jobs.base_job import BaseJob # To avoid circular import + from airflow.jobs.job import Job # To avoid circular import count = session.query(func.count(cls.id)).filter(cls.triggerer_id == triggerer_id).scalar() capacity -= count @@ -211,10 +211,10 @@ def assign_unassigned(cls, triggerer_id, capacity, session: Session = NEW_SESSIO alive_triggerer_ids = [ row[0] - for row in session.query(BaseJob.id).filter( - BaseJob.end_date.is_(None), - BaseJob.latest_heartbeat > timezone.utcnow() - datetime.timedelta(seconds=30), - BaseJob.job_type == "TriggererJob", + for row in session.query(Job.id).filter( + Job.end_date.is_(None), + Job.latest_heartbeat > timezone.utcnow() - datetime.timedelta(seconds=30), + Job.job_type == "TriggererJob", ) ] diff --git a/airflow/serialization/pydantic/base_job.py b/airflow/serialization/pydantic/job.py similarity index 89% rename from airflow/serialization/pydantic/base_job.py rename to airflow/serialization/pydantic/job.py index e55a8ac5ed465..511dc5720a674 100644 --- a/airflow/serialization/pydantic/base_job.py +++ b/airflow/serialization/pydantic/job.py @@ -19,7 +19,7 @@ from pydantic import BaseModel as BaseModelPydantic -from airflow.jobs.job_runner import BaseJobRunner +from airflow.jobs.base_job_runner import BaseJobRunner def check_runner_initialized(job_runner: Optional[BaseJobRunner], job_type: str) -> BaseJobRunner: @@ -28,8 +28,8 @@ def check_runner_initialized(job_runner: Optional[BaseJobRunner], job_type: str) return job_runner -class BaseJobPydantic(BaseModelPydantic): - """Serializable representation of the BaseJob ORM SqlAlchemyModel used by internal API""" +class JobPydantic(BaseModelPydantic): + """Serializable representation of the Job ORM SqlAlchemyModel used by internal API""" id: Optional[int] dag_id: Optional[str] diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index 3fc85d93b3006..0f319a1afb363 100644 --- a/airflow/serialization/serialized_objects.py +++ b/airflow/serialization/serialized_objects.py @@ -38,7 +38,7 @@ from airflow.configuration import conf from airflow.datasets import Dataset from airflow.exceptions import AirflowException, RemovedInAirflow3Warning, SerializationError -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job from airflow.models.baseoperator import BaseOperator, BaseOperatorLink from airflow.models.connection import Connection from airflow.models.dag import DAG, create_timetable @@ -54,9 +54,9 @@ from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding from airflow.serialization.helpers import serialize_template_field from airflow.serialization.json_schema import Validator, load_dag_schema -from airflow.serialization.pydantic.base_job import BaseJobPydantic from airflow.serialization.pydantic.dag_run import DagRunPydantic from airflow.serialization.pydantic.dataset import DatasetPydantic +from airflow.serialization.pydantic.job import JobPydantic from airflow.serialization.pydantic.taskinstance import TaskInstancePydantic from airflow.settings import _ENABLE_AIP_44, DAGS_FOLDER, json from airflow.timetables.base import Timetable @@ -479,8 +479,8 @@ def serialize( type_=DAT.SIMPLE_TASK_INSTANCE, ) elif use_pydantic_models and _ENABLE_AIP_44: - if isinstance(var, BaseJob): - return cls._encode(BaseJobPydantic.from_orm(var).dict(), type_=DAT.BASE_JOB) + if isinstance(var, Job): + return cls._encode(JobPydantic.from_orm(var).dict(), type_=DAT.BASE_JOB) elif isinstance(var, TaskInstance): return cls._encode(TaskInstancePydantic.from_orm(var).dict(), type_=DAT.TASK_INSTANCE) elif isinstance(var, DagRun): @@ -554,7 +554,7 @@ def deserialize(cls, encoded_var: Any, use_pydantic_models=False) -> Any: return SimpleTaskInstance(**cls.deserialize(var)) elif use_pydantic_models and _ENABLE_AIP_44: if type_ == DAT.BASE_JOB: - return BaseJobPydantic.parse_obj(var) + return JobPydantic.parse_obj(var) elif type_ == DAT.TASK_INSTANCE: return TaskInstancePydantic.parse_obj(var) elif type_ == DAT.DAG_RUN: diff --git a/airflow/task/task_runner/__init__.py b/airflow/task/task_runner/__init__.py index baaf628bb386b..fa57d52b5cc89 100644 --- a/airflow/task/task_runner/__init__.py +++ b/airflow/task/task_runner/__init__.py @@ -21,7 +21,7 @@ from airflow.configuration import conf from airflow.exceptions import AirflowConfigException -from airflow.jobs.local_task_job import LocalTaskJobRunner +from airflow.jobs.local_task_job_runner import LocalTaskJobRunner from airflow.task.task_runner.base_task_runner import BaseTaskRunner from airflow.utils.module_loading import import_string diff --git a/airflow/task/task_runner/base_task_runner.py b/airflow/task/task_runner/base_task_runner.py index 28773bff3cc98..570ce1b1b8181 100644 --- a/airflow/task/task_runner/base_task_runner.py +++ b/airflow/task/task_runner/base_task_runner.py @@ -23,8 +23,8 @@ import threading from typing import cast -from airflow.jobs.base_job import BaseJob -from airflow.jobs.local_task_job import LocalTaskJobRunner +from airflow.jobs.job import Job +from airflow.jobs.local_task_job_runner import LocalTaskJobRunner from airflow.utils.dag_parsing_context import _airflow_parsing_context_manager from airflow.utils.platform import IS_WINDOWS @@ -53,7 +53,7 @@ class BaseTaskRunner(LoggingMixin): should be LocalTaskJobRunner """ - def __init__(self, base_job: BaseJob): + def __init__(self, base_job: Job): self.job_runner: LocalTaskJobRunner = cast(LocalTaskJobRunner, base_job.job_runner) if not hasattr(self.job_runner, "task_instance"): raise ValueError( diff --git a/airflow/task/task_runner/cgroup_task_runner.py b/airflow/task/task_runner/cgroup_task_runner.py index 63fcd78591961..8b6b03e5a559d 100644 --- a/airflow/task/task_runner/cgroup_task_runner.py +++ b/airflow/task/task_runner/cgroup_task_runner.py @@ -25,7 +25,7 @@ import psutil from cgroupspy import trees -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job from airflow.task.task_runner.base_task_runner import BaseTaskRunner from airflow.utils.operator_resources import Resources from airflow.utils.platform import getuser @@ -62,7 +62,7 @@ class CgroupTaskRunner(BaseTaskRunner): airflow ALL= (root) NOEXEC: !/bin/chmod /CGROUPS_FOLDER/cpu/airflow/* * """ - def __init__(self, base_job: BaseJob): + def __init__(self, base_job: Job): super().__init__(base_job=base_job) self.process = None self._finished_running = False diff --git a/airflow/task/task_runner/standard_task_runner.py b/airflow/task/task_runner/standard_task_runner.py index eb90d4819ee66..b4233d20ea7dd 100644 --- a/airflow/task/task_runner/standard_task_runner.py +++ b/airflow/task/task_runner/standard_task_runner.py @@ -24,7 +24,7 @@ import psutil from setproctitle import setproctitle -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job from airflow.models.taskinstance import TaskReturnCode from airflow.settings import CAN_FORK from airflow.task.task_runner.base_task_runner import BaseTaskRunner @@ -35,7 +35,7 @@ class StandardTaskRunner(BaseTaskRunner): """Standard runner for all tasks.""" - def __init__(self, base_job: BaseJob): + def __init__(self, base_job: Job): super().__init__(base_job=base_job) self._rc = None self.dag = self.job_runner.task_instance.task.dag diff --git a/airflow/utils/scheduler_health.py b/airflow/utils/scheduler_health.py index d79cc0eca0cd3..4facd0a75c40d 100644 --- a/airflow/utils/scheduler_health.py +++ b/airflow/utils/scheduler_health.py @@ -20,8 +20,8 @@ from http.server import BaseHTTPRequestHandler, HTTPServer from airflow.configuration import conf -from airflow.jobs.base_job import BaseJob -from airflow.jobs.scheduler_job import SchedulerJobRunner +from airflow.jobs.job import Job +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.utils.net import get_hostname from airflow.utils.session import create_session @@ -36,10 +36,10 @@ def do_GET(self): try: with create_session() as session: scheduler_job = ( - session.query(BaseJob) + session.query(Job) .filter_by(job_type=SchedulerJobRunner.job_type) .filter_by(hostname=get_hostname()) - .order_by(BaseJob.latest_heartbeat.desc()) + .order_by(Job.latest_heartbeat.desc()) .first() ) if scheduler_job and scheduler_job.is_alive(): diff --git a/airflow/www/views.py b/airflow/www/views.py index b39769bd5c1dd..bc38f358950df 100644 --- a/airflow/www/views.py +++ b/airflow/www/views.py @@ -88,9 +88,9 @@ from airflow.datasets import Dataset from airflow.exceptions import AirflowException, ParamValidationError, RemovedInAirflow3Warning from airflow.executors.executor_loader import ExecutorLoader -from airflow.jobs.base_job import BaseJob -from airflow.jobs.scheduler_job import SchedulerJobRunner -from airflow.jobs.triggerer_job import TriggererJobRunner +from airflow.jobs.job import Job +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner +from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.models import Connection, DagModel, DagTag, Log, SlaMiss, TaskFail, XCom, errors from airflow.models.abstractoperator import AbstractOperator from airflow.models.dag import DAG, get_dataset_triggered_next_run_info @@ -4945,7 +4945,7 @@ class JobModelView(AirflowModelView): route_base = "/job" - datamodel = AirflowModelView.CustomSQLAInterface(BaseJob) # type: ignore + datamodel = AirflowModelView.CustomSQLAInterface(Job) # type: ignore class_permission_name = permissions.RESOURCE_JOB method_permission_name = { diff --git a/dev/perf/scheduler_dag_execution_timing.py b/dev/perf/scheduler_dag_execution_timing.py index cf90612cb8188..9f5e6335e9849 100755 --- a/dev/perf/scheduler_dag_execution_timing.py +++ b/dev/perf/scheduler_dag_execution_timing.py @@ -240,8 +240,8 @@ def main(num_runs, repeat, pre_create_dag_runs, executor_class, dag_ids): if pre_create_dag_runs: os.environ["AIRFLOW__SCHEDULER__USE_JOB_SCHEDULE"] = "False" - from airflow.jobs.base_job import BaseJob - from airflow.jobs.scheduler_job import SchedulerJobRunner + from airflow.jobs.job import Job + from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.models.dagbag import DagBag from airflow.utils import db @@ -277,9 +277,7 @@ def main(num_runs, repeat, pre_create_dag_runs, executor_class, dag_ids): ShortCircuitExecutor = get_executor_under_test(executor_class) executor = ShortCircuitExecutor(dag_ids_to_watch=dag_ids, num_runs=num_runs) - scheduler_job = BaseJob( - job_runner=SchedulerJobRunner(dag_ids=dag_ids, do_pickle=False), executor=executor - ) + scheduler_job = Job(job_runner=SchedulerJobRunner(dag_ids=dag_ids, do_pickle=False), executor=executor) executor.scheduler_job = scheduler_job total_tasks = sum(len(dag.tasks) for dag in dags) @@ -310,7 +308,7 @@ def main(num_runs, repeat, pre_create_dag_runs, executor_class, dag_ids): reset_dag(dag, session) executor.reset(dag_ids) - scheduler_job = BaseJob( + scheduler_job = Job( job_runner=SchedulerJobRunner(dag_ids=dag_ids, do_pickle=False), executor=executor ) executor.scheduler_job = scheduler_job diff --git a/dev/perf/sql_queries.py b/dev/perf/sql_queries.py index 2d0c0f6e7ee07..22c1da8e6e035 100644 --- a/dev/perf/sql_queries.py +++ b/dev/perf/sql_queries.py @@ -23,7 +23,7 @@ import pandas as pd -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job # Setup environment before any Airflow import DAG_FOLDER = os.path.join(os.path.dirname(__file__), "dags") @@ -120,11 +120,11 @@ def run_scheduler_job(with_db_reset=False) -> None: """ Run the scheduler job, selectively resetting the db before creating a ScheduleJob instance """ - from airflow.jobs.scheduler_job import SchedulerJobRunner + from airflow.jobs.scheduler_job_runner import SchedulerJobRunner if with_db_reset: reset_db() - BaseJob(job_runner=SchedulerJobRunner(subdir=DAG_FOLDER, do_pickle=False, num_runs=3)).run() + Job(job_runner=SchedulerJobRunner(subdir=DAG_FOLDER, do_pickle=False, num_runs=3)).run() def is_query(line: str) -> bool: diff --git a/docs/apache-airflow/administration-and-deployment/logging-monitoring/check-health.rst b/docs/apache-airflow/administration-and-deployment/logging-monitoring/check-health.rst index 2aa4386df5bd7..798343f9d1e0e 100644 --- a/docs/apache-airflow/administration-and-deployment/logging-monitoring/check-health.rst +++ b/docs/apache-airflow/administration-and-deployment/logging-monitoring/check-health.rst @@ -99,7 +99,7 @@ using `http.server.BaseHTTPRequestHandler None: self.app = minimal_app_for_api self.client = self.app.test_client() # type:ignore with create_session() as session: - session.query(BaseJob).delete() + session.query(Job).delete() def teardown_method(self): with create_session() as session: - session.query(BaseJob).delete() + session.query(Job).delete() class TestGetHealth(TestHealthTestBase): @@ -49,7 +49,7 @@ class TestGetHealth(TestHealthTestBase): def test_healthy_scheduler_status(self, session): last_scheduler_heartbeat_for_testing_1 = timezone.utcnow() session.add( - BaseJob( + Job( job_type="SchedulerJob", state=State.RUNNING, latest_heartbeat=last_scheduler_heartbeat_for_testing_1, @@ -69,7 +69,7 @@ def test_healthy_scheduler_status(self, session): def test_unhealthy_scheduler_is_slow(self, session): last_scheduler_heartbeat_for_testing_2 = timezone.utcnow() - timedelta(minutes=1) session.add( - BaseJob( + Job( job_type="SchedulerJob", state=State.RUNNING, latest_heartbeat=last_scheduler_heartbeat_for_testing_2, diff --git a/tests/api_connexion/endpoints/test_task_instance_endpoint.py b/tests/api_connexion/endpoints/test_task_instance_endpoint.py index db8108972b435..ed3327bca4de2 100644 --- a/tests/api_connexion/endpoints/test_task_instance_endpoint.py +++ b/tests/api_connexion/endpoints/test_task_instance_endpoint.py @@ -24,8 +24,8 @@ import pytest from sqlalchemy.orm import contains_eager -from airflow.jobs.base_job import BaseJob -from airflow.jobs.triggerer_job import TriggererJobRunner +from airflow.jobs.job import Job +from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.models import DagRun, SlaMiss, TaskInstance, Trigger from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF from airflow.security import permissions @@ -232,7 +232,7 @@ def test_should_respond_200_with_task_state_in_deferred(self, session): )[0] ti.trigger = Trigger("none", {}) ti.trigger.created_date = now - ti.triggerer_job = BaseJob(job_runner=TriggererJobRunner()) + ti.triggerer_job = Job(job_runner=TriggererJobRunner()) ti.triggerer_job.state = "running" session.commit() response = self.client.get( diff --git a/tests/cli/commands/test_dag_processor_command.py b/tests/cli/commands/test_dag_processor_command.py index e151b5aff1c6e..0fbcfe37ea580 100644 --- a/tests/cli/commands/test_dag_processor_command.py +++ b/tests/cli/commands/test_dag_processor_command.py @@ -42,7 +42,7 @@ def setup_class(cls): ("core", "load_examples"): "False", } ) - @mock.patch("airflow.jobs.dag_processor_job.DagProcessorJobRunner") + @mock.patch("airflow.jobs.dag_processor_job_runner.DagProcessorJobRunner") @pytest.mark.skipif( conf.get_mandatory_value("database", "sql_alchemy_conn").lower().startswith("sqlite"), reason="Standalone Dag Processor doesn't support sqlite.", diff --git a/tests/cli/commands/test_jobs_command.py b/tests/cli/commands/test_jobs_command.py index ab576e02db4eb..cf3a276414440 100644 --- a/tests/cli/commands/test_jobs_command.py +++ b/tests/cli/commands/test_jobs_command.py @@ -23,8 +23,8 @@ from airflow.cli import cli_parser from airflow.cli.commands import jobs_command -from airflow.jobs.base_job import BaseJob -from airflow.jobs.scheduler_job import SchedulerJobRunner +from airflow.jobs.job import Job +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.utils.session import create_session from airflow.utils.state import State from tests.test_utils.db import clear_db_jobs @@ -46,7 +46,7 @@ def teardown_method(self) -> None: def test_should_report_success_for_one_working_scheduler(self): with create_session() as session: - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner()) + self.scheduler_job = Job(job_runner=SchedulerJobRunner()) self.scheduler_job.state = State.RUNNING session.add(self.scheduler_job) session.commit() @@ -58,7 +58,7 @@ def test_should_report_success_for_one_working_scheduler(self): def test_should_report_success_for_one_working_scheduler_with_hostname(self): with create_session() as session: - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner()) + self.scheduler_job = Job(job_runner=SchedulerJobRunner()) self.scheduler_job.state = State.RUNNING self.scheduler_job.hostname = "HOSTNAME" session.add(self.scheduler_job) @@ -77,7 +77,7 @@ def test_should_report_success_for_ha_schedulers(self): scheduler_jobs = [] with create_session() as session: for _ in range(3): - scheduler_job = BaseJob(job_runner=SchedulerJobRunner()) + scheduler_job = Job(job_runner=SchedulerJobRunner()) scheduler_job.state = State.RUNNING session.add(scheduler_job) scheduler_jobs.append(scheduler_job) @@ -99,7 +99,7 @@ def test_should_ignore_not_running_jobs(self): scheduler_jobs = [] with create_session() as session: for _ in range(3): - scheduler_job = BaseJob(job_runner=SchedulerJobRunner()) + scheduler_job = Job(job_runner=SchedulerJobRunner()) scheduler_job.state = State.SHUTDOWN session.add(scheduler_job) scheduler_jobs.append(scheduler_job) @@ -115,7 +115,7 @@ def test_should_raise_exception_for_multiple_scheduler_on_one_host(self): scheduler_jobs = [] with create_session() as session: for _ in range(3): - scheduler_job = BaseJob(job_runner=SchedulerJobRunner()) + scheduler_job = Job(job_runner=SchedulerJobRunner()) scheduler_job.state = State.RUNNING scheduler_job.hostname = "HOSTNAME" session.add(scheduler_job) diff --git a/tests/core/test_impersonation_tests.py b/tests/core/test_impersonation_tests.py index d82845113dd4d..c70d761870d6b 100644 --- a/tests/core/test_impersonation_tests.py +++ b/tests/core/test_impersonation_tests.py @@ -24,8 +24,8 @@ import pytest -from airflow.jobs.backfill_job import BackfillJobRunner -from airflow.jobs.base_job import BaseJob +from airflow.jobs.backfill_job_runner import BackfillJobRunner +from airflow.jobs.job import Job from airflow.models import DagBag, DagRun, TaskInstance from airflow.utils.db import add_default_pool_if_not_exists from airflow.utils.state import State @@ -113,7 +113,7 @@ def run_backfill(self, dag_id, task_id): dag = self.dagbag.get_dag(dag_id) dag.clear() - BaseJob(job_runner=BackfillJobRunner(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)).run() + Job(job_runner=BackfillJobRunner(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)).run() run_id = DagRun.generate_run_id(DagRunType.BACKFILL_JOB, execution_date=DEFAULT_DATE) ti = TaskInstance(task=dag.get_task(task_id), run_id=run_id) ti.refresh_from_db() diff --git a/tests/executors/test_dask_executor.py b/tests/executors/test_dask_executor.py index af903a44dec62..fdd0c3e3ad3c1 100644 --- a/tests/executors/test_dask_executor.py +++ b/tests/executors/test_dask_executor.py @@ -25,8 +25,8 @@ from airflow.exceptions import AirflowException from airflow.executors.dask_executor import DaskExecutor -from airflow.jobs.backfill_job import BackfillJobRunner -from airflow.jobs.base_job import BaseJob +from airflow.jobs.backfill_job_runner import BackfillJobRunner +from airflow.jobs.job import Job from airflow.models import DagBag from airflow.utils import timezone from tests.test_utils.config import conf_vars @@ -108,7 +108,7 @@ def test_backfill_integration(self): """ dag = self.dagbag.get_dag("example_bash_operator") - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, diff --git a/tests/jobs/test_backfill_job.py b/tests/jobs/test_backfill_job.py index 9621b7872b7a7..b6298bd5f0f59 100644 --- a/tests/jobs/test_backfill_job.py +++ b/tests/jobs/test_backfill_job.py @@ -38,8 +38,8 @@ ) from airflow.executors.executor_constants import MOCK_EXECUTOR from airflow.executors.executor_loader import ExecutorLoader -from airflow.jobs.backfill_job import BackfillJobRunner -from airflow.jobs.base_job import BaseJob +from airflow.jobs.backfill_job_runner import BackfillJobRunner +from airflow.jobs.job import Job from airflow.listeners.listener import get_listener_manager from airflow.models import DagBag, Pool, TaskInstance as TI from airflow.models.dagrun import DagRun @@ -121,7 +121,7 @@ def test_unfinished_dag_runs_set_to_failed(self, dag_maker): dag = self._get_dummy_dag(dag_maker) dag_run = dag_maker.create_dagrun(state=None) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -143,7 +143,7 @@ def test_dag_run_with_finished_tasks_set_to_success(self, dag_maker): for ti in dag_run.get_task_instances(): ti.set_state(State.SUCCESS) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -167,7 +167,7 @@ def test_trigger_controller_dag(self, session): target_dag_run = session.query(DagRun).filter(DagRun.dag_id == target_dag.dag_id).one_or_none() assert target_dag_run is None - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True ) @@ -188,7 +188,7 @@ def test_backfill_multi_dates(self): end_date = DEFAULT_DATE + datetime.timedelta(days=1) executor = MockExecutor(parallelism=16) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -282,7 +282,7 @@ def test_backfill_examples(self, dag_id, expected_execution_order): logger.info("*** Running example DAG: %s", dag.dag_id) executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -305,7 +305,7 @@ def test_backfill_conf(self, dag_maker): executor = MockExecutor() conf_ = json.loads("""{"key": "value"}""") - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -323,7 +323,7 @@ def test_backfill_conf(self, dag_maker): assert conf_ == dr[0].conf - @patch("airflow.jobs.backfill_job.BackfillJobRunner.log") + @patch("airflow.jobs.backfill_job_runner.BackfillJobRunner.log") def test_backfill_respect_max_active_tis_per_dag_limit(self, mock_log, dag_maker): max_active_tis_per_dag = 2 dag = self._get_dummy_dag( @@ -335,7 +335,7 @@ def test_backfill_respect_max_active_tis_per_dag_limit(self, mock_log, dag_maker executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -379,7 +379,7 @@ def test_backfill_respect_max_active_tis_per_dag_limit(self, mock_log, dag_maker assert 0 == times_dag_concurrency_limit_reached_in_debug assert times_task_concurrency_limit_reached_in_debug > 0 - @patch("airflow.jobs.backfill_job.BackfillJobRunner.log") + @patch("airflow.jobs.backfill_job_runner.BackfillJobRunner.log") def test_backfill_respect_dag_concurrency_limit(self, mock_log, dag_maker): dag = self._get_dummy_dag(dag_maker, dag_id="test_backfill_respect_concurrency_limit") dag_maker.create_dagrun(state=None) @@ -387,7 +387,7 @@ def test_backfill_respect_dag_concurrency_limit(self, mock_log, dag_maker): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -432,7 +432,7 @@ def test_backfill_respect_dag_concurrency_limit(self, mock_log, dag_maker): assert 0 == times_task_concurrency_limit_reached_in_debug assert times_dag_concurrency_limit_reached_in_debug > 0 - @patch("airflow.jobs.backfill_job.BackfillJobRunner.log") + @patch("airflow.jobs.backfill_job_runner.BackfillJobRunner.log") def test_backfill_respect_default_pool_limit(self, mock_log, dag_maker): default_pool_slots = 2 set_default_pool_slots(default_pool_slots) @@ -442,7 +442,7 @@ def test_backfill_respect_default_pool_limit(self, mock_log, dag_maker): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -500,7 +500,7 @@ def test_backfill_pool_not_found(self, dag_maker): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -514,7 +514,7 @@ def test_backfill_pool_not_found(self, dag_maker): except AirflowException: return - @patch("airflow.jobs.backfill_job.BackfillJobRunner.log") + @patch("airflow.jobs.backfill_job_runner.BackfillJobRunner.log") def test_backfill_respect_pool_limit(self, mock_log, dag_maker): session = settings.Session() @@ -535,7 +535,7 @@ def test_backfill_respect_pool_limit(self, mock_log, dag_maker): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -587,7 +587,7 @@ def test_backfill_run_rescheduled(self, dag_maker): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -601,7 +601,7 @@ def test_backfill_run_rescheduled(self, dag_maker): ti.refresh_from_db() ti.set_state(State.UP_FOR_RESCHEDULE) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -626,7 +626,7 @@ def test_backfill_override_conf(self, dag_maker): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -658,7 +658,7 @@ def test_backfill_skip_active_scheduled_dagrun(self, dag_maker, caplog): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -666,7 +666,7 @@ def test_backfill_skip_active_scheduled_dagrun(self, dag_maker, caplog): ), executor=executor, ) - with caplog.at_level(logging.ERROR, logger="airflow.jobs.backfill_job.BackfillJob"): + with caplog.at_level(logging.ERROR, logger="airflow.jobs.backfill_job_runner.BackfillJob"): caplog.clear() job.run() assert "Backfill cannot be created for DagRun" in caplog.messages[0] @@ -686,7 +686,7 @@ def test_backfill_rerun_failed_tasks(self, dag_maker): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -700,7 +700,7 @@ def test_backfill_rerun_failed_tasks(self, dag_maker): ti.refresh_from_db() ti.set_state(State.FAILED) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -724,7 +724,7 @@ def test_backfill_rerun_upstream_failed_tasks(self, dag_maker): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -738,7 +738,7 @@ def test_backfill_rerun_upstream_failed_tasks(self, dag_maker): ti.refresh_from_db() ti.set_state(State.UPSTREAM_FAILED) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -760,7 +760,7 @@ def test_backfill_rerun_failed_tasks_without_flag(self, dag_maker): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -774,7 +774,7 @@ def test_backfill_rerun_failed_tasks_without_flag(self, dag_maker): ti.refresh_from_db() ti.set_state(State.FAILED) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -806,7 +806,7 @@ def test_backfill_retry_intermittent_failed_task(self, dag_maker): executor.mock_task_results[ TaskInstanceKey(dag.dag_id, task1.task_id, DEFAULT_DATE, try_number=2) ] = State.UP_FOR_RETRY - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -833,7 +833,7 @@ def test_backfill_retry_always_failed_task(self, dag_maker): TaskInstanceKey(dag.dag_id, task1.task_id, dr.run_id, try_number=1) ] = State.UP_FOR_RETRY executor.mock_task_fail(dag.dag_id, task1.task_id, dr.run_id, try_number=2) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -864,7 +864,7 @@ def test_backfill_ordered_concurrent_execute(self, dag_maker): dag_maker.create_dagrun(run_id=runid0) executor = MockExecutor(parallelism=16) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -908,7 +908,7 @@ def test_backfill_pooled_tasks(self): dag.clear() executor = MockExecutor(do_update=True) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE), executor=executor, ) @@ -932,7 +932,7 @@ def test_backfill_depends_on_past_works_independently_on_ignore_depends_on_past( dag.clear() run_date = DEFAULT_DATE + datetime.timedelta(days=5) - BaseJob( + Job( job_runner=BackfillJobRunner( dag=dag, start_date=run_date, @@ -962,7 +962,7 @@ def test_backfill_depends_on_past_backwards(self): dag.clear() executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner(dag=dag, ignore_first_depends_on_past=True, **kwargs), executor=executor, ) @@ -977,7 +977,7 @@ def test_backfill_depends_on_past_backwards(self): expected_msg = "You cannot backfill backwards because one or more tasks depend_on_past: test_dop_task" with pytest.raises(AirflowException, match=expected_msg): executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner(dag=dag, run_backwards=True, **kwargs), executor=executor, ) @@ -1028,7 +1028,7 @@ def test_backfill_max_limit_check_within_limit(self, dag_maker): end_date = DEFAULT_DATE executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=start_date, @@ -1053,7 +1053,7 @@ def test_backfill_notifies_dagrun_listener(self, dag_maker): end_date = DEFAULT_DATE executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=start_date, @@ -1107,7 +1107,7 @@ def run_backfill(cond): thread_session.close() executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=start_date, @@ -1160,7 +1160,7 @@ def test_backfill_max_limit_check_no_count_existing(self, dag_maker): dag_maker.create_dagrun(state=None) executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=start_date, end_date=end_date, donot_pickle=True ), @@ -1188,7 +1188,7 @@ def test_backfill_max_limit_check_complete_loop(self, dag_maker): # backfill job 3 times success_expected = 2 executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=start_date, @@ -1226,7 +1226,7 @@ def test_sub_set_subdag(self, dag_maker): sub_dag = dag.partial_subset( task_ids_or_regex="leave*", include_downstream=False, include_upstream=False ) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner(dag=sub_dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE), executor=executor, ) @@ -1273,7 +1273,7 @@ def test_backfill_fill_blanks(self, dag_maker): session.commit() session.close() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE), executor=executor, ) @@ -1304,7 +1304,7 @@ def test_backfill_execute_subdag(self): start_date = timezone.utcnow() executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=subdag, start_date=start_date, @@ -1349,7 +1349,7 @@ def test_subdag_clear_parentdag_downstream_clear(self): subdag = subdag_op_task.subdag executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -1403,7 +1403,7 @@ def test_backfill_execute_subdag_with_removed_task(self): session = settings.Session() executor = MockExecutor() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=subdag, start_date=DEFAULT_DATE, @@ -1450,7 +1450,7 @@ def test_update_counters(self, dag_maker, session): with dag_maker(dag_id="test_manage_executor_state", start_date=DEFAULT_DATE, session=session) as dag: task1 = EmptyOperator(task_id="dummy", owner="airflow") dr = dag_maker.create_dagrun(state=None) - job = BaseJob(job_runner=BackfillJobRunner(dag=dag)) + job = Job(job_runner=BackfillJobRunner(dag=dag)) ti = TI(task1, dr.execution_date) ti.refresh_from_db() @@ -1616,7 +1616,7 @@ def test_backfill_run_backwards(self): executor = MockExecutor(parallelism=16) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -1657,7 +1657,7 @@ def test_reset_orphaned_tasks_with_orphans(self, dag_maker): tasks.append(task) session = settings.Session() - job = BaseJob(job_runner=BackfillJobRunner(dag=dag)) + job = Job(job_runner=BackfillJobRunner(dag=dag)) # create dagruns dr1 = dag_maker.create_dagrun(state=State.RUNNING) @@ -1717,7 +1717,7 @@ def test_reset_orphaned_tasks_specified_dagrun(self, session, dag_maker): ) as dag: EmptyOperator(task_id=task_id, dag=dag) - job = BaseJob(job_runner=BackfillJobRunner(dag=dag)) + job = Job(job_runner=BackfillJobRunner(dag=dag)) # make two dagruns, only reset for one dr1 = dag_maker.create_dagrun(state=State.SUCCESS) dr2 = dag.create_dagrun(run_id="test2", state=State.RUNNING, session=session) @@ -1744,7 +1744,7 @@ def test_job_id_is_assigned_to_dag_run(self, dag_maker): with dag_maker(dag_id=dag_id, start_date=DEFAULT_DATE, schedule="@daily") as dag: EmptyOperator(task_id="dummy_task", dag=dag) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner(dag=dag, start_date=timezone.utcnow() - datetime.timedelta(days=1)), executor=MockExecutor(), ) @@ -1759,7 +1759,7 @@ def test_backfill_has_job_id_int(self): executor = MockExecutor(parallelism=16) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -1790,7 +1790,7 @@ def test_backfilling_dags(self, dag_id, executor_name, session): when = timezone.datetime(2022, 1, 1) - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=when, @@ -1850,7 +1850,7 @@ def consumer(value): ti_status.active_runs.append(dr) ti_status.to_run = {ti.key: ti for ti in dr.task_instances} - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=dr.execution_date, @@ -1923,7 +1923,7 @@ def consumer(a, b): executor = MockExecutor() when = timezone.datetime(2022, 1, 1) - BaseJob( + Job( job_runner=BackfillJobRunner(dag=dag, start_date=when, end_date=when, donot_pickle=True), executor=executor, ).run() @@ -1948,7 +1948,7 @@ def test_start_date_set_for_resetted_dagruns(self, dag_maker, session, caplog): session.merge(dr) session.flush() dag.clear() - BaseJob( + Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, @@ -1978,7 +1978,7 @@ def test_task_instances_are_not_set_to_scheduled_when_dagrun_reset(self, dag_mak dag.clear() - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE + datetime.timedelta(days=1), @@ -2020,7 +2020,7 @@ def test_backfill_disable_retry(self, dag_maker, disable_retry, try_number, exce TaskInstanceKey(dag.dag_id, task1.task_id, dag_run.run_id, try_number=2) ] = TaskInstanceState.FAILED - job = BaseJob( + job = Job( job_runner=BackfillJobRunner( dag=dag, start_date=DEFAULT_DATE, diff --git a/tests/jobs/test_base_job.py b/tests/jobs/test_base_job.py index d6e86a5922561..b66ec76cc585b 100644 --- a/tests/jobs/test_base_job.py +++ b/tests/jobs/test_base_job.py @@ -25,7 +25,7 @@ from sqlalchemy.exc import OperationalError from airflow.executors.sequential_executor import SequentialExecutor -from airflow.jobs.base_job import BaseJob, most_recent_job +from airflow.jobs.job import Job, most_recent_job from airflow.listeners.listener import get_listener_manager from airflow.utils import timezone from airflow.utils.session import create_session @@ -35,9 +35,9 @@ from tests.utils.test_helpers import MockJobRunner -class TestBaseJob: +class TestJob: def test_state_success(self): - job = BaseJob(job_runner=MockJobRunner()) + job = Job(job_runner=MockJobRunner()) job.run() assert job.state == State.SUCCESS @@ -46,7 +46,7 @@ def test_state_success(self): def test_state_sysexit(self): import sys - job = BaseJob(job_runner=MockJobRunner(lambda: sys.exit(0))) + job = Job(job_runner=MockJobRunner(lambda: sys.exit(0))) job.run() assert job.state == State.SUCCESS @@ -56,7 +56,7 @@ def test_base_job_respects_plugin_hooks(self): import sys - job = BaseJob(job_runner=MockJobRunner(lambda: sys.exit(0))) + job = Job(job_runner=MockJobRunner(lambda: sys.exit(0))) job.run() assert job.state == State.SUCCESS @@ -68,7 +68,7 @@ def test_base_job_respects_plugin_lifecycle(self, dag_maker): """ get_listener_manager().add_listener(lifecycle_listener) - job = BaseJob(job_runner=MockJobRunner(lambda: sys.exit(0))) + job = Job(job_runner=MockJobRunner(lambda: sys.exit(0))) job.run() assert lifecycle_listener.started_component is job @@ -78,7 +78,7 @@ def test_state_failed(self): def abort(): raise RuntimeError("fail") - job = BaseJob(job_runner=MockJobRunner(abort)) + job = Job(job_runner=MockJobRunner(abort)) with raises(RuntimeError): job.run() @@ -87,9 +87,9 @@ def abort(): def test_most_recent_job(self): with create_session() as session: - old_job = BaseJob(job_runner=MockJobRunner(), heartrate=10) + old_job = Job(job_runner=MockJobRunner(), heartrate=10) old_job.latest_heartbeat = old_job.latest_heartbeat - datetime.timedelta(seconds=20) - job = BaseJob(job_runner=MockJobRunner(), heartrate=10) + job = Job(job_runner=MockJobRunner(), heartrate=10) session.add(job) session.add(old_job) session.flush() @@ -101,13 +101,13 @@ def test_most_recent_job(self): def test_most_recent_job_running_precedence(self): with create_session() as session: - old_running_state_job = BaseJob(job_runner=MockJobRunner(), heartrate=10) + old_running_state_job = Job(job_runner=MockJobRunner(), heartrate=10) old_running_state_job.latest_heartbeat = timezone.utcnow() old_running_state_job.state = State.RUNNING - new_failed_state_job = BaseJob(job_runner=MockJobRunner(), heartrate=10) + new_failed_state_job = Job(job_runner=MockJobRunner(), heartrate=10) new_failed_state_job.latest_heartbeat = timezone.utcnow() new_failed_state_job.state = State.FAILED - new_null_state_job = BaseJob(job_runner=MockJobRunner(), heartrate=10) + new_null_state_job = Job(job_runner=MockJobRunner(), heartrate=10) new_null_state_job.latest_heartbeat = timezone.utcnow() new_null_state_job.state = None session.add(old_running_state_job) @@ -120,7 +120,7 @@ def test_most_recent_job_running_precedence(self): session.rollback() def test_is_alive(self): - job = BaseJob(job_runner=MockJobRunner(), heartrate=10, state=State.RUNNING) + job = Job(job_runner=MockJobRunner(), heartrate=10, state=State.RUNNING) assert job.is_alive() is True job.latest_heartbeat = timezone.utcnow() - datetime.timedelta(seconds=20) @@ -138,14 +138,14 @@ def test_is_alive(self): job.latest_heartbeat = timezone.utcnow() - datetime.timedelta(seconds=10) assert job.is_alive() is False, "Completed jobs even with recent heartbeat should not be alive" - @patch("airflow.jobs.base_job.create_session") + @patch("airflow.jobs.job.create_session") def test_heartbeat_failed(self, mock_create_session): when = timezone.utcnow() - datetime.timedelta(seconds=60) with create_session() as session: mock_session = Mock(spec_set=session, name="MockSession") mock_create_session.return_value.__enter__.return_value = mock_session - job = BaseJob(job_runner=MockJobRunner(), heartrate=10, state=State.RUNNING) + job = Job(job_runner=MockJobRunner(), heartrate=10, state=State.RUNNING) job.latest_heartbeat = when mock_session.commit.side_effect = OperationalError("Force fail", {}, None) @@ -160,18 +160,16 @@ def test_heartbeat_failed(self, mock_create_session): ("core", "executor"): "SequentialExecutor", } ) - @patch("airflow.jobs.base_job.ExecutorLoader.get_default_executor") - @patch("airflow.jobs.base_job.get_hostname") - @patch("airflow.jobs.base_job.getuser") + @patch("airflow.jobs.job.ExecutorLoader.get_default_executor") + @patch("airflow.jobs.job.get_hostname") + @patch("airflow.jobs.job.getuser") def test_essential_attr(self, mock_getuser, mock_hostname, mock_default_executor): mock_sequential_executor = SequentialExecutor() mock_hostname.return_value = "test_hostname" mock_getuser.return_value = "testuser" mock_default_executor.return_value = mock_sequential_executor - test_job = BaseJob( - job_runner=MockJobRunner(), heartrate=10, dag_id="example_dag", state=State.RUNNING - ) + test_job = Job(job_runner=MockJobRunner(), heartrate=10, dag_id="example_dag", state=State.RUNNING) assert test_job.executor_class == "SequentialExecutor" assert test_job.heartrate == 10 assert test_job.dag_id == "example_dag" @@ -182,9 +180,9 @@ def test_essential_attr(self, mock_getuser, mock_hostname, mock_default_executor assert test_job.executor == mock_sequential_executor def test_heartbeat(self, frozen_sleep, monkeypatch): - monkeypatch.setattr("airflow.jobs.base_job.sleep", frozen_sleep) + monkeypatch.setattr("airflow.jobs.job.sleep", frozen_sleep) with create_session() as session: - job = BaseJob(job_runner=MockJobRunner(), heartrate=10) + job = Job(job_runner=MockJobRunner(), heartrate=10) job.latest_heartbeat = timezone.utcnow() session.add(job) session.commit() diff --git a/tests/jobs/test_local_task_job.py b/tests/jobs/test_local_task_job.py index 6c218587f46e2..23c8a7a23758a 100644 --- a/tests/jobs/test_local_task_job.py +++ b/tests/jobs/test_local_task_job.py @@ -36,9 +36,9 @@ from airflow import DAG, settings from airflow.exceptions import AirflowException from airflow.executors.sequential_executor import SequentialExecutor -from airflow.jobs.base_job import BaseJob -from airflow.jobs.local_task_job import SIGSEGV_MESSAGE, LocalTaskJobRunner -from airflow.jobs.scheduler_job import SchedulerJobRunner +from airflow.jobs.job import Job +from airflow.jobs.local_task_job_runner import SIGSEGV_MESSAGE, LocalTaskJobRunner +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.models.dagbag import DagBag from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance @@ -91,7 +91,7 @@ class TestLocalTaskJob: @pytest.fixture(autouse=True) def set_instance_attrs(self, dagbag): self.dagbag = dagbag - with patch("airflow.jobs.base_job.sleep") as self.mock_base_job_sleep: + with patch("airflow.jobs.job.sleep") as self.mock_base_job_sleep: yield def validate_ti_states(self, dag_run, ti_state_mapping, error_message): @@ -113,7 +113,7 @@ def test_localtaskjob_essential_attr(self, dag_maker): ti = dr.get_task_instance(task_id=op1.task_id) - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id, executor=SequentialExecutor(), @@ -138,7 +138,7 @@ def test_localtaskjob_heartbeat(self, dag_maker): ti.hostname = "blablabla" session.commit() - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id, executor=SequentialExecutor(), @@ -177,7 +177,7 @@ def test_localtaskjob_heartbeat(self, dag_maker): job1.job_runner.heartbeat_callback() @mock.patch("subprocess.check_call") - @mock.patch("airflow.jobs.local_task_job.psutil") + @mock.patch("airflow.jobs.local_task_job_runner.psutil") def test_localtaskjob_heartbeat_with_run_as_user(self, psutil_mock, _, dag_maker): session = settings.Session() with dag_maker("test_localtaskjob_heartbeat"): @@ -189,7 +189,7 @@ def test_localtaskjob_heartbeat_with_run_as_user(self, psutil_mock, _, dag_maker ti.hostname = get_hostname() session.commit() - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id, executor=SequentialExecutor(), @@ -235,7 +235,7 @@ def test_localtaskjob_heartbeat_with_run_as_user(self, psutil_mock, _, dag_maker @conf_vars({("core", "default_impersonation"): "testuser"}) @mock.patch("subprocess.check_call") - @mock.patch("airflow.jobs.local_task_job.psutil") + @mock.patch("airflow.jobs.local_task_job_runner.psutil") def test_localtaskjob_heartbeat_with_default_impersonation(self, psutil_mock, _, dag_maker): session = settings.Session() with dag_maker("test_localtaskjob_heartbeat"): @@ -247,7 +247,7 @@ def test_localtaskjob_heartbeat_with_default_impersonation(self, psutil_mock, _, ti.hostname = get_hostname() session.commit() - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id, executor=SequentialExecutor(), @@ -320,7 +320,7 @@ def test_heartbeat_failed_fast(self): ti.pid = 1 session.commit() - job = BaseJob( + job = Job( job_runner=LocalTaskJobRunner(task_instance=ti), dag_id=ti.dag_id, executor=MockExecutor(do_update=False), @@ -354,9 +354,7 @@ def test_mark_success_no_kill(self, caplog, get_test_dag, session): ti = dr.get_task_instance(task.task_id) ti.refresh_from_task(task) - job1 = BaseJob( - job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id - ) + job1 = Job(job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id) with timeout(30): job1.run() @@ -391,7 +389,7 @@ def test_localtaskjob_double_trigger(self): ti_run = TaskInstance(task=task, run_id=dr.run_id) ti_run.refresh_from_db() - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner(task_instance=ti_run), dag_id=ti_run.dag_id, executor=SequentialExecutor(), @@ -407,14 +405,14 @@ def test_localtaskjob_double_trigger(self): session.close() @patch.object(StandardTaskRunner, "return_code") - @mock.patch("airflow.jobs.scheduler_job.Stats.incr", autospec=True) + @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr", autospec=True) def test_local_task_return_code_metric(self, mock_stats_incr, mock_return_code, create_dummy_dag): _, task = create_dummy_dag("test_localtaskjob_code") ti_run = TaskInstance(task=task, execution_date=DEFAULT_DATE) ti_run.refresh_from_db() - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner(task_instance=ti_run), dag_id=ti_run.dag_id, executor=SequentialExecutor(), @@ -439,7 +437,7 @@ def test_localtaskjob_maintain_heart_rate(self, mock_return_code, caplog, create ti_run = TaskInstance(task=task, execution_date=DEFAULT_DATE) ti_run.refresh_from_db() - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner(task_instance=ti_run), dag_id=ti_run.dag_id, executor=SequentialExecutor(), @@ -485,7 +483,7 @@ def test_mark_failure_on_failure_callback(self, caplog, get_test_dag): ti = dr.get_task_instance(task.task_id) ti.refresh_from_task(task) - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner( task_instance=ti, ignore_ti_state=True, @@ -523,7 +521,7 @@ def test_dagrun_timeout_logged_in_task_logs(self, caplog, get_test_dag): ti = dr.get_task_instance(task.task_id) ti.refresh_from_task(task) - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner( task_instance=ti, ignore_ti_state=True, @@ -559,7 +557,7 @@ def test_failure_callback_called_by_airflow_run_raw_process(self, monkeypatch, t ti = TaskInstance(task=task, execution_date=DEFAULT_DATE) ti.refresh_from_db() - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), executor=SequentialExecutor(), dag_id=ti.dag_id, @@ -594,7 +592,7 @@ def test_mark_success_on_success_callback(self, caplog, get_test_dag): ti = dr.get_task_instance(task.task_id) ti.refresh_from_task(task) - job = BaseJob( + job = Job( job_runner=LocalTaskJobRunner( task_instance=ti, ignore_ti_state=True, @@ -672,7 +670,7 @@ def send_signal(ti, signal_sent, sig): thread.daemon = True thread.start() - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner( task_instance=ti, ignore_ti_state=True, @@ -760,7 +758,7 @@ def test_fast_follow( "test_dagrun_fast_follow", ) - scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) scheduler_job.job_runner.dagbag.bag_dag(dag, root_dag=dag) dag_run = dag.create_dagrun(run_id="test_dagrun_fast_follow", state=State.RUNNING) @@ -774,7 +772,7 @@ def test_fast_follow( ti = TaskInstance(task=dag.get_task(task_ids_to_run[0]), execution_date=dag_run.execution_date) ti.refresh_from_db() - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner( task_instance=ti, ignore_ti_state=True, @@ -791,7 +789,7 @@ def test_fast_follow( task=dag.get_task(task_ids_to_run[1]), execution_date=dag_run.execution_date ) ti.refresh_from_db() - job2 = BaseJob( + job2 = Job( job_runner=LocalTaskJobRunner( task_instance=ti, ignore_ti_state=True, @@ -830,7 +828,7 @@ def test_mini_scheduler_works_with_wait_for_upstream(self, caplog, get_test_dag) session.merge(ti2_k) session.merge(ti2_l) - job1 = BaseJob( + job1 = Job( job_runner=LocalTaskJobRunner( task_instance=ti2_k, ignore_ti_state=True, @@ -878,7 +876,7 @@ def task_function(ti): dag_run = dag_maker.create_dagrun() ti = TaskInstance(task=task, run_id=dag_run.run_id) ti.refresh_from_db() - job = BaseJob( + job = Job( job_runner=LocalTaskJobRunner( task_instance=ti, ignore_ti_state=True, @@ -915,7 +913,7 @@ def test_number_of_queries_single_loop(mock_get_task_runner, dag_maker): ti = dr.task_instances[0] ti.refresh_from_task(task) - job = BaseJob(job_runner=LocalTaskJobRunner(task_instance=ti), dag_id=ti.dag_id, executor=MockExecutor()) + job = Job(job_runner=LocalTaskJobRunner(task_instance=ti), dag_id=ti.dag_id, executor=MockExecutor()) with assert_queries_count(18): job.run() @@ -1030,7 +1028,7 @@ def task_function(): dag.create_dagrun(state=State.RUNNING, run_id=run_id, execution_date=execution_date) ti = TaskInstance(task=task, execution_date=execution_date) ti.refresh_from_db() - job = BaseJob( + job = Job( job_runner=LocalTaskJobRunner( task_instance=ti, ignore_ti_state=True, diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py index cd3df11b3a908..cc478d8ea9a07 100644 --- a/tests/jobs/test_scheduler_job.py +++ b/tests/jobs/test_scheduler_job.py @@ -44,10 +44,10 @@ from airflow.executors.base_executor import BaseExecutor from airflow.executors.executor_constants import MOCK_EXECUTOR from airflow.executors.executor_loader import ExecutorLoader -from airflow.jobs.backfill_job import BackfillJobRunner -from airflow.jobs.base_job import BaseJob -from airflow.jobs.local_task_job import LocalTaskJobRunner -from airflow.jobs.scheduler_job import SchedulerJobRunner +from airflow.jobs.backfill_job_runner import BackfillJobRunner +from airflow.jobs.job import Job +from airflow.jobs.local_task_job_runner import LocalTaskJobRunner +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.models import DAG, DagBag, DagModel, DbCallbackRequest, Pool, TaskInstance from airflow.models.dagrun import DagRun from airflow.models.dataset import DatasetDagRunQueue, DatasetEvent, DatasetModel @@ -171,9 +171,7 @@ def set_instance_attrs(self, dagbag) -> Generator: ) def test_is_alive(self, configs): with conf_vars(configs): - self.scheduler_job = BaseJob( - job_runner=SchedulerJobRunner(None), heartrate=10, state=State.RUNNING - ) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(None), heartrate=10, state=State.RUNNING) assert self.scheduler_job.is_alive() self.scheduler_job.latest_heartbeat = timezone.utcnow() - datetime.timedelta(seconds=20) @@ -202,7 +200,7 @@ def run_single_scheduler_loop_with_no_dags(self, dags_folder): :param dags_folder: the directory to traverse """ - self.scheduler_job = BaseJob( + self.scheduler_job = Job( job_runner=SchedulerJobRunner( executor=self.null_exec, num_times_parse_dags=1, subdir=os.path.join(dags_folder) ) @@ -214,7 +212,7 @@ def test_no_orphan_process_will_be_left(self): empty_dir = mkdtemp() current_process = psutil.Process() old_children = current_process.children(recursive=True) - self.scheduler_job = BaseJob( + self.scheduler_job = Job( job_runner=SchedulerJobRunner(subdir=empty_dir, num_runs=1), executor=MockExecutor(do_update=False), ) @@ -225,8 +223,8 @@ def test_no_orphan_process_will_be_left(self): current_children = set(current_process.children(recursive=True)) - set(old_children) assert not current_children - @mock.patch("airflow.jobs.scheduler_job.TaskCallbackRequest") - @mock.patch("airflow.jobs.scheduler_job.Stats.incr") + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") + @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr") def test_process_executor_events(self, mock_stats_incr, mock_task_callback, dag_maker): dag_id = "test_process_executor_events" task_id_1 = "dummy_task" @@ -241,7 +239,7 @@ def test_process_executor_events(self, mock_stats_incr, mock_task_callback, dag_ executor = MockExecutor(do_update=False) task_callback = mock.MagicMock() mock_task_callback.return_value = task_callback - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() ti1.state = State.QUEUED session.merge(ti1) @@ -279,8 +277,8 @@ def test_process_executor_events(self, mock_stats_incr, mock_task_callback, dag_ any_order=True, ) - @mock.patch("airflow.jobs.scheduler_job.TaskCallbackRequest") - @mock.patch("airflow.jobs.scheduler_job.Stats.incr") + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") + @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr") def test_process_executor_events_with_no_callback(self, mock_stats_incr, mock_task_callback, dag_maker): dag_id = "test_process_executor_events_with_no_callback" task_id_1 = "dummy_task" @@ -289,7 +287,7 @@ def test_process_executor_events_with_no_callback(self, mock_stats_incr, mock_ta executor = MockExecutor(do_update=False) task_callback = mock.MagicMock() mock_task_callback.return_value = task_callback - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() session = settings.Session() @@ -304,7 +302,7 @@ def test_process_executor_events_with_no_callback(self, mock_stats_incr, mock_ta executor = MockExecutor(do_update=False) task_callback = mock.MagicMock() mock_task_callback.return_value = task_callback - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() ti1.state = State.QUEUED session.merge(ti1) @@ -339,8 +337,8 @@ def test_process_executor_events_with_no_callback(self, mock_stats_incr, mock_ta any_order=True, ) - @mock.patch("airflow.jobs.scheduler_job.TaskCallbackRequest") - @mock.patch("airflow.jobs.scheduler_job.Stats.incr") + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") + @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr") def test_process_executor_events_with_callback(self, mock_stats_incr, mock_task_callback, dag_maker): dag_id = "test_process_executor_events_with_callback" task_id_1 = "dummy_task" @@ -354,7 +352,7 @@ def test_process_executor_events_with_callback(self, mock_stats_incr, mock_task_ executor = MockExecutor(do_update=False) task_callback = mock.MagicMock() mock_task_callback.return_value = task_callback - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() session = settings.Session() @@ -389,8 +387,8 @@ def test_process_executor_events_with_callback(self, mock_stats_incr, mock_task_ }, ) - @mock.patch("airflow.jobs.scheduler_job.TaskCallbackRequest") - @mock.patch("airflow.jobs.scheduler_job.Stats.incr") + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") + @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr") def test_process_executor_event_missing_dag(self, mock_stats_incr, mock_task_callback, dag_maker, caplog): dag_id = "test_process_executor_events_with_callback" task_id_1 = "dummy_task" @@ -404,7 +402,7 @@ def test_process_executor_event_missing_dag(self, mock_stats_incr, mock_task_cal executor = MockExecutor(do_update=False) task_callback = mock.MagicMock() mock_task_callback.return_value = task_callback - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) self.scheduler_job.job_runner.dagbag = mock.MagicMock() self.scheduler_job.job_runner.dagbag.get_dag.side_effect = Exception("failed") self.scheduler_job.job_runner.processor_agent = mock.MagicMock() @@ -419,8 +417,8 @@ def test_process_executor_event_missing_dag(self, mock_stats_incr, mock_task_cal ti1.refresh_from_db() assert ti1.state == State.FAILED - @mock.patch("airflow.jobs.scheduler_job.TaskCallbackRequest") - @mock.patch("airflow.jobs.scheduler_job.Stats.incr") + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") + @mock.patch("airflow.jobs.scheduler_job_runner.Stats.incr") def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_callback, dag_maker): dag_id = "test_process_executor_events_ti_requeued" task_id_1 = "dummy_task" @@ -435,7 +433,7 @@ def test_process_executor_events_ti_requeued(self, mock_stats_incr, mock_task_ca executor = MockExecutor(do_update=False) task_callback = mock.MagicMock() mock_task_callback.return_value = task_callback - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) self.scheduler_job.id = 1 self.scheduler_job.job_runner.processor_agent = mock.MagicMock() @@ -489,7 +487,7 @@ def test_execute_task_instances_is_paused_wont_execute(self, session, dag_maker) EmptyOperator(task_id=task_id_1) assert isinstance(dag, SerializedDAG) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) dr1 = dag_maker.create_dagrun(run_type=DagRunType.BACKFILL_JOB) (ti1,) = dr1.task_instances @@ -511,7 +509,7 @@ def test_execute_task_instances_backfill_tasks_wont_execute(self, dag_maker): with dag_maker(dag_id=dag_id): task1 = EmptyOperator(task_id=task_id_1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dr1 = dag_maker.create_dagrun(run_type=DagRunType.BACKFILL_JOB) @@ -532,7 +530,7 @@ def test_execute_task_instances_backfill_tasks_wont_execute(self, dag_maker): @conf_vars({("scheduler", "standalone_dag_processor"): "False"}) def test_setup_callback_sink_not_standalone_dag_processor(self): - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull, num_runs=1)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull, num_runs=1)) self.scheduler_job.job_runner._execute() @@ -540,7 +538,7 @@ def test_setup_callback_sink_not_standalone_dag_processor(self): @conf_vars({("scheduler", "standalone_dag_processor"): "True"}) def test_setup_callback_sink_standalone_dag_processor(self): - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull, num_runs=1)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull, num_runs=1)) self.scheduler_job.job_runner._execute() @@ -552,7 +550,7 @@ def test_find_executable_task_instances_backfill(self, dag_maker): with dag_maker(dag_id=dag_id, max_active_tasks=16): task1 = EmptyOperator(task_id=task_id_1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) @@ -584,7 +582,7 @@ def test_find_executable_task_instances_pool(self, dag_maker): EmptyOperator(task_id=task_id_1, pool="a", priority_weight=2) EmptyOperator(task_id=task_id_2, pool="b", priority_weight=1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) dr2 = dag_maker.create_dagrun_after(dr1, run_type=DagRunType.SCHEDULED) @@ -637,7 +635,7 @@ def test_find_executable_task_instances_only_running_dagruns( EmptyOperator(task_id=task_id_1) EmptyOperator(task_id=task_id_2) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) dr = dag_maker.create_dagrun(state=state) @@ -670,7 +668,7 @@ def test_find_executable_task_instances_order_execution_date(self, dag_maker): dr1 = session.merge(dr1, load=False) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) tis = dr1.task_instances + dr2.task_instances for ti in tis: @@ -698,7 +696,7 @@ def test_find_executable_task_instances_order_priority(self, dag_maker): dr1 = session.merge(dr1, load=False) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) tis = dr1.task_instances + dr2.task_instances for ti in tis: @@ -717,7 +715,7 @@ def test_find_executable_task_instances_order_priority_with_pools(self, dag_make even if different pools are involved. """ - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dag_id = "SchedulerJobTest.test_find_executable_task_instances_order_priority_with_pools" @@ -764,7 +762,7 @@ def test_find_executable_task_instances_order_execution_date_and_priority(self, dr2 = dag_maker.create_dagrun(execution_date=DEFAULT_DATE + timedelta(hours=1)) dr1 = session.merge(dr1, load=False) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) tis = dr1.task_instances + dr2.task_instances for ti in tis: @@ -786,7 +784,7 @@ def test_find_executable_task_instances_in_default_pool(self, dag_maker): op2 = EmptyOperator(task_id="dummy2") executor = MockExecutor(do_update=True) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) session = settings.Session() dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) @@ -823,7 +821,7 @@ def test_queued_task_instances_fails_with_missing_dag(self, dag_maker, session): EmptyOperator(task_id=task_id_1) EmptyOperator(task_id=task_id_2) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.dagbag = mock.MagicMock() self.scheduler_job.job_runner.dagbag.get_dag.return_value = None @@ -846,7 +844,7 @@ def test_nonexistent_pool(self, dag_maker): with dag_maker(dag_id=dag_id, max_active_tasks=16): EmptyOperator(task_id="dummy_wrong_pool", pool="this_pool_doesnt_exist") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dr = dag_maker.create_dagrun() @@ -866,7 +864,7 @@ def test_infinite_pool(self, dag_maker): with dag_maker(dag_id=dag_id, concurrency=16): EmptyOperator(task_id="dummy", pool="infinite_pool") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dr = dag_maker.create_dagrun() @@ -888,7 +886,7 @@ def test_not_enough_pool_slots(self, caplog, dag_maker): EmptyOperator(task_id="cannot_run", pool="some_pool", pool_slots=4) EmptyOperator(task_id="can_run", pool="some_pool", pool_slots=1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dr = dag_maker.create_dagrun() ti = dr.task_instances[0] @@ -931,7 +929,7 @@ def test_find_executable_task_instances_none(self, dag_maker): with dag_maker(dag_id=dag_id, max_active_tasks=16): EmptyOperator(task_id=task_id_1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() assert 0 == len( @@ -950,7 +948,7 @@ def test_tis_for_queued_dagruns_are_not_run(self, dag_maker): task1 = EmptyOperator(task_id=task_id_1) dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED, state=State.QUEUED) dr2 = dag_maker.create_dagrun_after(dr1, run_type=DagRunType.SCHEDULED) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() ti1 = TaskInstance(task1, run_id=dr1.run_id) ti2 = TaskInstance(task1, run_id=dr2.run_id) @@ -974,7 +972,7 @@ def test_find_executable_task_instances_concurrency(self, dag_maker): with dag_maker(dag_id=dag_id, max_active_tasks=2, session=session): EmptyOperator(task_id="dummy") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) dr2 = dag_maker.create_dagrun_after(dr1, run_type=DagRunType.SCHEDULED) @@ -1014,7 +1012,7 @@ def test_find_executable_task_instances_concurrency_queued(self, dag_maker): task2 = EmptyOperator(task_id="dummy2") task3 = EmptyOperator(task_id="dummy3") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dag_run = dag_maker.create_dagrun() @@ -1048,7 +1046,7 @@ def test_find_executable_task_instances_max_active_tis_per_dag(self, dag_maker): task2 = EmptyOperator(task_id=task_id_2) executor = MockExecutor(do_update=True) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) session = settings.Session() dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) @@ -1123,7 +1121,7 @@ def test_change_state_for_executable_task_instances_no_tis_with_state(self, dag_ with dag_maker(dag_id=dag_id, max_active_tasks=2): task1 = EmptyOperator(task_id=task_id_1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dr1 = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) @@ -1150,7 +1148,7 @@ def test_change_state_for_executable_task_instances_no_tis_with_state(self, dag_ def test_find_executable_task_instances_not_enough_pool_slots_for_first(self, dag_maker): set_default_pool_slots(1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dag_id = "SchedulerJobTest.test_find_executable_task_instances_not_enough_pool_slots_for_first" @@ -1175,7 +1173,7 @@ def test_find_executable_task_instances_not_enough_pool_slots_for_first(self, da session.rollback() def test_find_executable_task_instances_not_enough_dag_concurrency_for_first(self, dag_maker): - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dag_id_1 = ( @@ -1211,7 +1209,7 @@ def test_find_executable_task_instances_not_enough_dag_concurrency_for_first(sel session.rollback() def test_find_executable_task_instances_not_enough_task_concurrency_for_first(self, dag_maker): - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dag_id = "SchedulerJobTest.test_find_executable_task_instances_not_enough_task_concurrency_for_first" @@ -1246,7 +1244,7 @@ def test_find_executable_task_instances_negative_open_pool_slots(self, dag_maker """ set_default_pool_slots(0) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() pool1 = Pool(pool="pool1", slots=1) @@ -1274,9 +1272,9 @@ def test_find_executable_task_instances_negative_open_pool_slots(self, dag_maker session.rollback() - @mock.patch("airflow.jobs.scheduler_job.Stats.gauge") + @mock.patch("airflow.jobs.scheduler_job_runner.Stats.gauge") def test_emit_pool_starving_tasks_metrics(self, mock_stats_gauge, dag_maker): - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dag_id = "SchedulerJobTest.test_emit_pool_starving_tasks_metrics" @@ -1327,7 +1325,7 @@ def test_enqueue_task_instances_with_queued_state(self, dag_maker, session): with dag_maker(dag_id=dag_id, start_date=DEFAULT_DATE, session=session): task1 = EmptyOperator(task_id=task_id_1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) dr1 = dag_maker.create_dagrun() ti1 = dr1.get_task_instance(task1.task_id, session) @@ -1347,7 +1345,7 @@ def test_enqueue_task_instances_sets_ti_state_to_None_if_dagrun_in_finish_state( with dag_maker(dag_id=dag_id, start_date=DEFAULT_DATE, session=session): task1 = EmptyOperator(task_id=task_id_1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) dr1 = dag_maker.create_dagrun(state=state) ti = dr1.get_task_instance(task1.task_id, session) @@ -1375,7 +1373,7 @@ def test_critical_section_enqueue_task_instances(self, dag_maker): task1 = EmptyOperator(task_id=task_id_1) task2 = EmptyOperator(task_id=task_id_2) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) # create first dag run with 1 running and 1 queued @@ -1429,7 +1427,7 @@ def test_execute_task_instances_limit(self, dag_maker): task1 = EmptyOperator(task_id=task_id_1) task2 = EmptyOperator(task_id=task_id_2) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) def _create_dagruns(): dagrun = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED, state=State.RUNNING) @@ -1480,7 +1478,7 @@ def test_execute_task_instances_unlimited(self, dag_maker): task1 = EmptyOperator(task_id=task_id_1) task2 = EmptyOperator(task_id=task_id_2) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) def _create_dagruns(): dagrun = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED, state=State.RUNNING) @@ -1528,7 +1526,7 @@ def test_adopt_or_reset_orphaned_tasks(self, dag_maker): processor = mock.MagicMock() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(num_runs=0)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(num_runs=0)) self.scheduler_job.job_runner.processor_agent = processor self.scheduler_job.job_runner.adopt_or_reset_orphaned_tasks() @@ -1544,7 +1542,7 @@ def test_executor_end_called(self, mock_processor_agent): """ Test to make sure executor.end gets called with a successful scheduler loop run """ - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull, num_runs=1)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull, num_runs=1)) self.scheduler_job.executor = mock.MagicMock(slots_available=8) self.scheduler_job.run() @@ -1557,7 +1555,7 @@ def test_cleanup_methods_all_called(self, mock_processor_agent): """ Test to make sure all cleanup methods are called when the scheduler loop has an exception """ - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull, num_runs=1)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull, num_runs=1)) self.scheduler_job.executor = mock.MagicMock(slots_available=8) self.scheduler_job.job_runner._run_scheduler_loop = mock.MagicMock(side_effect=Exception("oops")) mock_processor_agent.return_value.end.side_effect = Exception("double oops") @@ -1576,7 +1574,7 @@ def test_queued_dagruns_stops_creating_when_max_active_is_reached(self, dag_make EmptyOperator(task_id="mytask") session = settings.Session() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner.processor_agent = mock.MagicMock() @@ -1606,7 +1604,7 @@ def test_runs_are_created_after_max_active_runs_was_reached(self, dag_maker, ses """ Test that when creating runs once max_active_runs is reached the runs does not stick """ - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor(do_update=True) self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -1654,7 +1652,7 @@ def test_dagrun_timeout_verify_max_active_runs(self, dag_maker): ) as dag: EmptyOperator(task_id="dummy") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.dagbag = dag_maker.dagbag session = settings.Session() @@ -1722,7 +1720,7 @@ def test_dagrun_timeout_fails_run(self, dag_maker): dr = dag_maker.create_dagrun(start_date=timezone.utcnow() - datetime.timedelta(days=1)) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.dagbag = dag_maker.dagbag # Mock that processor_agent is started @@ -1764,7 +1762,7 @@ def test_dagrun_timeout_fails_run_and_update_next_dagrun(self, dag_maker): dr = dag_maker.create_dagrun(start_date=timezone.utcnow() - datetime.timedelta(days=1)) # check that next_dagrun is dr.execution_date dag_maker.dag_model.next_dagrun == dr.execution_date - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.dagbag = dag_maker.dagbag self.scheduler_job.executor = MockExecutor() @@ -1798,7 +1796,7 @@ def test_dagrun_callbacks_are_called(self, state, expected_callback_msg, dag_mak ) as dag: EmptyOperator(task_id="dummy") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner.dagbag = dag_maker.dagbag self.scheduler_job.job_runner.processor_agent = mock.Mock() @@ -1844,7 +1842,7 @@ def test_dagrun_plugins_are_notified(self, state, expected_callback_msg, dag_mak dag_listener.clear() get_listener_manager().add_listener(dag_listener) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner.dagbag = dag_maker.dagbag self.scheduler_job.job_runner.processor_agent = mock.Mock() @@ -1874,7 +1872,7 @@ def test_dagrun_timeout_callbacks_are_stored_in_database(self, dag_maker, sessio ) as dag: EmptyOperator(task_id="empty") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.executor.callback_sink = DatabaseCallbackSink() self.scheduler_job.job_runner.dagbag = dag_maker.dagbag @@ -1911,7 +1909,7 @@ def test_dagrun_callbacks_commited_before_sent(self, dag_maker): with dag_maker(dag_id="test_dagrun_callbacks_commited_before_sent"): EmptyOperator(task_id="dummy") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.Mock() self.scheduler_job.job_runner._send_dag_callbacks_to_processor = mock.Mock() self.scheduler_job.job_runner._schedule_dag_run = mock.Mock() @@ -1923,7 +1921,7 @@ def test_dagrun_callbacks_commited_before_sent(self, dag_maker): ti.set_state(State.SUCCESS, session) with mock.patch.object(settings, "USE_JOB_SCHEDULE", False), mock.patch( - "airflow.jobs.scheduler_job.prohibit_commit" + "airflow.jobs.scheduler_job_runner.prohibit_commit" ) as mock_guard: mock_guard.return_value.__enter__.return_value.commit.side_effect = session.commit @@ -1958,7 +1956,7 @@ def test_dagrun_callbacks_are_not_added_when_callbacks_are_not_defined(self, sta ): BashOperator(task_id="test_task", bash_command="echo hi") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.Mock() self.scheduler_job.job_runner._send_dag_callbacks_to_processor = mock.Mock() @@ -1991,7 +1989,7 @@ def test_dagrun_callbacks_are_added_when_callbacks_are_defined(self, state, msg, ): BashOperator(task_id="test_task", bash_command="echo hi") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.Mock() self.scheduler_job.job_runner._send_dag_callbacks_to_processor = mock.Mock() @@ -2026,7 +2024,7 @@ def test_dagrun_notify_called_success(self, dag_maker): executor = MockExecutor(do_update=False) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) self.scheduler_job.job_runner.dagbag = dag_maker.dagbag self.scheduler_job.job_runner.processor_agent = mock.MagicMock() @@ -2066,7 +2064,7 @@ def test_do_not_schedule_removed_task(self, dag_maker): ): pass - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) res = self.scheduler_job.job_runner._executable_task_instances_to_queued(max_tis=32, session=session) assert [] == res @@ -2214,7 +2212,7 @@ def test_dagrun_root_after_dagrun_unfinished(self): dag_id = "test_dagrun_states_root_future" dag = self.dagbag.get_dag(dag_id) dag.sync_to_db() - self.scheduler_job = BaseJob( + self.scheduler_job = Job( job_runner=SchedulerJobRunner(num_runs=1, subdir=dag.fileloc), executor=self.null_exec, ) @@ -2285,7 +2283,7 @@ def test_scheduler_start_date(self, configs): other_dag.is_paused_upon_creation = True other_dag.sync_to_db() - self.scheduler_job = BaseJob( + self.scheduler_job = Job( job_runner=SchedulerJobRunner(subdir=dag.fileloc, num_runs=1), executor=self.null_exec, ) @@ -2301,7 +2299,7 @@ def test_scheduler_start_date(self, configs): # That behavior still exists, but now it will only do so if after the # start date bf_exec = MockExecutor() - backfill = BaseJob( + backfill = Job( BackfillJobRunner(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE), executor=bf_exec, ) @@ -2317,7 +2315,7 @@ def test_scheduler_start_date(self, configs): ] == bf_exec.sorted_tasks session.commit() - self.scheduler_job = BaseJob( + self.scheduler_job = Job( job_runner=SchedulerJobRunner(dag.fileloc, num_runs=1), executor=self.null_exec, ) @@ -2356,7 +2354,7 @@ def test_scheduler_task_start_date(self, configs): dagbag.sync_to_db() - self.scheduler_job = BaseJob( + self.scheduler_job = Job( job_runner=SchedulerJobRunner(subdir=dag.fileloc, num_runs=3), executor=self.null_exec, ) @@ -2388,7 +2386,7 @@ def test_scheduler_multiprocessing(self, configs): dag = self.dagbag.get_dag(dag_id) dag.clear() - self.scheduler_job = BaseJob( + self.scheduler_job = Job( job_runner=SchedulerJobRunner( subdir=os.path.join(TEST_DAG_FOLDER, "test_scheduler_dags.py"), num_runs=1, @@ -2426,7 +2424,7 @@ def test_scheduler_verify_pool_full(self, dag_maker, configs): session.add(pool) session.flush() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=self.null_exec) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=self.null_exec) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() # Create 2 dagruns, which will create 2 task instances. @@ -2466,7 +2464,7 @@ def test_scheduler_verify_pool_full_2_slots_per_task(self, dag_maker, session): session.add(pool) session.flush() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=self.null_exec) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=self.null_exec) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() # Create 5 dagruns, which will create 5 task instances. @@ -2520,7 +2518,7 @@ def test_scheduler_keeps_scheduling_pool_full(self, dag_maker): session.add(pool_p2) session.flush() - scheduler = BaseJob(job_runner=SchedulerJobRunner(), executor=self.null_exec) + scheduler = Job(job_runner=SchedulerJobRunner(), executor=self.null_exec) scheduler.job_runner.processor_agent = mock.MagicMock() def _create_dagruns(dag: DAG): @@ -2595,7 +2593,7 @@ def test_scheduler_verify_priority_and_slots(self, dag_maker): session.add(pool) session.flush() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=self.null_exec) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=self.null_exec) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() dr = dag_maker.create_dagrun() @@ -2642,13 +2640,13 @@ def test_verify_integrity_if_dag_not_changed(self, dag_maker): with dag_maker(dag_id="test_verify_integrity_if_dag_not_changed") as dag: BashOperator(task_id="dummy", bash_command="echo hi") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() orm_dag = dag_maker.dag_model assert orm_dag is not None - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() dag = self.scheduler_job.job_runner.dagbag.get_dag( "test_verify_integrity_if_dag_not_changed", session=session @@ -2660,7 +2658,7 @@ def test_verify_integrity_if_dag_not_changed(self, dag_maker): dr = drs[0] # Verify that DagRun.verify_integrity is not called - with mock.patch("airflow.jobs.scheduler_job.DagRun.verify_integrity") as mock_verify_integrity: + with mock.patch("airflow.jobs.scheduler_job_runner.DagRun.verify_integrity") as mock_verify_integrity: self.scheduler_job.job_runner._schedule_dag_run(dr, session) mock_verify_integrity.assert_not_called() session.flush() @@ -2693,13 +2691,13 @@ def test_verify_integrity_if_dag_changed(self, dag_maker): with dag_maker(dag_id="test_verify_integrity_if_dag_changed") as dag: BashOperator(task_id="dummy", bash_command="echo hi") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() orm_dag = dag_maker.dag_model assert orm_dag is not None - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() dag = self.scheduler_job.job_runner.dagbag.get_dag( "test_verify_integrity_if_dag_changed", session=session @@ -2765,13 +2763,13 @@ def test_verify_integrity_if_dag_disappeared(self, dag_maker, caplog): with dag_maker(dag_id="test_verify_integrity_if_dag_disappeared") as dag: BashOperator(task_id="dummy", bash_command="echo hi") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() orm_dag = dag_maker.dag_model assert orm_dag is not None - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() dag = self.scheduler_job.job_runner.dagbag.get_dag( "test_verify_integrity_if_dag_disappeared", session=session @@ -2834,7 +2832,7 @@ def do_schedule(session): # Use a empty file since the above mock will return the # expected DAGs. Also specify only a single file so that it doesn't # try to schedule the above DAG repeatedly. - self.scheduler_job = BaseJob( + self.scheduler_job = Job( job_runner=SchedulerJobRunner(num_runs=1, subdir=os.devnull), executor=executor, ) @@ -2895,7 +2893,7 @@ def test_retry_handling_job(self): dag_task1 = dag.get_task("test_retry_handling_op") dag.clear() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(dag_id=dag.dag_id, num_runs=1)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(dag_id=dag.dag_id, num_runs=1)) self.scheduler_job.heartrate = 0 self.scheduler_job.run() @@ -2989,7 +2987,7 @@ def test_list_py_file_paths(self): def test_adopt_or_reset_orphaned_tasks_nothing(self): """Try with nothing.""" - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner()) + self.scheduler_job = Job(job_runner=SchedulerJobRunner()) session = settings.Session() assert 0 == self.scheduler_job.job_runner.adopt_or_reset_orphaned_tasks(session=session) @@ -2999,7 +2997,7 @@ def test_adopt_or_reset_orphaned_tasks_external_triggered_dag(self, dag_maker): task_id = dag_id + "_task" EmptyOperator(task_id=task_id) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() dr1 = dag_maker.create_dagrun(external_trigger=True) @@ -3018,7 +3016,7 @@ def test_adopt_or_reset_orphaned_tasks_backfill_dag(self, dag_maker): task_id = dag_id + "_task" EmptyOperator(task_id=task_id) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() session.add(self.scheduler_job) session.flush() @@ -3041,7 +3039,7 @@ def test_reset_orphaned_tasks_no_orphans(self, dag_maker): task_id = dag_id + "_task" EmptyOperator(task_id=task_id) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() session.add(self.scheduler_job) session.flush() @@ -3065,7 +3063,7 @@ def test_reset_orphaned_tasks_non_running_dagruns(self, dag_maker): task_id = dag_id + "_task" EmptyOperator(task_id=task_id) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() session.add(self.scheduler_job) session.flush() @@ -3088,13 +3086,13 @@ def test_adopt_or_reset_orphaned_tasks_stale_scheduler_jobs(self, dag_maker): EmptyOperator(task_id="task1") EmptyOperator(task_id="task2") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) session = settings.Session() self.scheduler_job.state = State.RUNNING self.scheduler_job.latest_heartbeat = timezone.utcnow() session.add(self.scheduler_job) - old_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + old_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) old_job.state = State.RUNNING old_job.latest_heartbeat = timezone.utcnow() - timedelta(minutes=15) session.add(old_job) @@ -3134,24 +3132,24 @@ def test_adopt_or_reset_orphaned_tasks_only_fails_scheduler_jobs(self, caplog): """Make sure we only set SchedulerJobs to failed, not all jobs""" session = settings.Session() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.state = State.RUNNING self.scheduler_job.latest_heartbeat = timezone.utcnow() session.add(self.scheduler_job) session.flush() - old_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + old_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) old_job.state = State.RUNNING old_job.latest_heartbeat = timezone.utcnow() - timedelta(minutes=15) session.add(old_job) session.flush() - old_task_job = BaseJob(job_runner=MockJobRunner(), state=State.RUNNING) + old_task_job = Job(job_runner=MockJobRunner(), state=State.RUNNING) old_task_job.latest_heartbeat = timezone.utcnow() - timedelta(minutes=15) session.add(old_task_job) session.flush() - with caplog.at_level("INFO", logger="airflow.jobs.scheduler_job"): + with caplog.at_level("INFO", logger="airflow.jobs.scheduler_job_runner"): self.scheduler_job.job_runner.adopt_or_reset_orphaned_tasks(session=session) session.expire_all() @@ -3166,7 +3164,7 @@ def test_send_sla_callbacks_to_processor_sla_disabled(self, dag_maker): EmptyOperator(task_id="task1") with patch.object(settings, "CHECK_SLAS", False): - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner._send_sla_callbacks_to_processor(dag) @@ -3179,7 +3177,7 @@ def test_send_sla_callbacks_to_processor_sla_no_task_slas(self, dag_maker): EmptyOperator(task_id="task1") with patch.object(settings, "CHECK_SLAS", True): - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner._send_sla_callbacks_to_processor(dag) @@ -3204,7 +3202,7 @@ def test_send_sla_callbacks_to_processor_sla_with_task_slas(self, schedule, dag_ EmptyOperator(task_id="task1", sla=timedelta(seconds=60)) with patch.object(settings, "CHECK_SLAS", True): - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner._send_sla_callbacks_to_processor(dag) @@ -3230,7 +3228,7 @@ def test_send_sla_callbacks_to_processor_sla_dag_not_scheduled(self, schedule, d EmptyOperator(task_id="task1", sla=timedelta(seconds=5)) with patch.object(settings, "CHECK_SLAS", True): - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner._send_sla_callbacks_to_processor(dag) @@ -3249,7 +3247,7 @@ def test_create_dag_runs(self, dag_maker): dag_model = dag_maker.dag_model - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=self.null_exec) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=self.null_exec) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() with create_session() as session: @@ -3326,7 +3324,7 @@ def test_create_dag_runs_datasets(self, session, dag_maker): ) session.flush() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=self.null_exec) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=self.null_exec) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() with create_session() as session: @@ -3358,7 +3356,7 @@ def dict_from_obj(obj): assert dag3.get_last_dagrun().creating_job_id == self.scheduler_job.id @time_machine.travel(DEFAULT_DATE + datetime.timedelta(days=1, seconds=9), tick=False) - @mock.patch("airflow.jobs.scheduler_job.Stats.timing") + @mock.patch("airflow.jobs.scheduler_job_runner.Stats.timing") def test_start_dagruns(self, stats_timing, dag_maker): """ Test that _start_dagrun: @@ -3373,7 +3371,7 @@ def test_start_dagruns(self, stats_timing, dag_maker): dag_model = dag_maker.dag_model - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=self.null_exec) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=self.null_exec) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() with create_session() as session: @@ -3416,7 +3414,7 @@ def test_extra_operator_links_not_loaded_in_scheduler_loop(self, dag_maker): assert custom_task.operator_extra_links session = settings.Session() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=self.null_exec) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=self.null_exec) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() self.scheduler_job.job_runner._start_queued_dagruns(session) @@ -3437,16 +3435,14 @@ def test_scheduler_create_dag_runs_does_not_raise_error(self, caplog, dag_maker) task_id="dummy", ) - self.scheduler_job = BaseJob( - job_runner=SchedulerJobRunner(subdir=os.devnull), executor=self.null_exec - ) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull), executor=self.null_exec) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() caplog.set_level("FATAL") caplog.clear() with create_session() as session, caplog.at_level( "ERROR", - logger="airflow.jobs.scheduler_job", + logger="airflow.jobs.scheduler_job_runner", ): self.scheduler_job.job_runner._create_dag_runs([dag_maker.dag_model], session) assert caplog.messages == [ @@ -3474,7 +3470,7 @@ def test_bulk_write_to_db_external_trigger_dont_skip_scheduled_run(self, dag_mak assert dag_model.next_dagrun_data_interval_start == DEFAULT_DATE assert dag_model.next_dagrun_data_interval_end == DEFAULT_DATE + timedelta(minutes=1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -3548,9 +3544,7 @@ def test_scheduler_create_dag_runs_check_existing_run(self, dag_maker): assert dag.get_last_dagrun(session) == dagrun - self.scheduler_job = BaseJob( - job_runner=SchedulerJobRunner(subdir=os.devnull), executor=self.null_exec - ) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull), executor=self.null_exec) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() # Test that this does not raise any error @@ -3598,7 +3592,7 @@ def test_do_schedule_max_active_runs_dag_timed_out(self, dag_maker): session=session, ) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -3639,7 +3633,7 @@ def test_do_schedule_max_active_runs_task_removed(self, session, dag_maker): state=State.RUNNING, ) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -3658,7 +3652,7 @@ def test_more_runs_are_not_created_when_max_active_runs_is_reached(self, dag_mak """ with dag_maker(max_active_runs=1): EmptyOperator(task_id="task") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) session = settings.Session() @@ -3723,7 +3717,7 @@ def complete_one_dagrun(): # Need to use something that doesn't immediately get marked as success by the scheduler BashOperator(task_id="task", bash_command="true") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor(do_update=True) self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -3781,7 +3775,7 @@ def test_do_schedule_max_active_runs_and_manual_trigger(self, dag_maker): dag.sync_to_db(session=session) # Update the date fields - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -3836,7 +3830,7 @@ def test_max_active_runs_in_a_dag_doesnt_stop_running_dagruns_in_otherdags(self, for _ in range(9): dr = dag_maker.create_dagrun_after(dr, run_type=DagRunType.SCHEDULED, state=State.QUEUED) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -3864,7 +3858,7 @@ def test_start_queued_dagruns_do_follow_execution_date_order(self, dag_maker): run_id=f"dagrun_{i}", run_type=DagRunType.SCHEDULED, state=State.QUEUED, execution_date=date ) date = dr.execution_date + timedelta(hours=1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -3933,7 +3927,7 @@ def test_no_dagruns_would_stuck_in_running(self, dag_maker): dr = dag_maker.create_dagrun(run_id=f"dr2_run_{i+1}", state=State.QUEUED, execution_date=date) date = dr.execution_date + timedelta(hours=1) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -3974,7 +3968,7 @@ def test_dag_file_processor_process_task_instances(self, state, start_date, end_ with dag_maker(dag_id="test_scheduler_process_execute_task"): BashOperator(task_id="dummy", bash_command="echo hi") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() dr = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) @@ -4018,7 +4012,7 @@ def test_dag_file_processor_process_task_instances_with_max_active_tis_per_dag( with dag_maker(dag_id="test_scheduler_process_execute_task_with_max_active_tis_per_dag"): BashOperator(task_id="dummy", max_active_tis_per_dag=2, bash_command="echo Hi") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() dr = dag_maker.create_dagrun( @@ -4070,7 +4064,7 @@ def test_dag_file_processor_process_task_instances_depends_on_past( BashOperator(task_id="dummy1", bash_command="echo hi") BashOperator(task_id="dummy2", bash_command="echo hi") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() dr = dag_maker.create_dagrun( run_type=DagRunType.SCHEDULED, @@ -4099,7 +4093,7 @@ def test_scheduler_job_add_new_task(self, dag_maker): with dag_maker(dag_id="test_scheduler_add_new_task") as dag: BashOperator(task_id="dummy", bash_command="echo test") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.dagbag = dag_maker.dagbag session = settings.Session() @@ -4108,7 +4102,7 @@ def test_scheduler_job_add_new_task(self, dag_maker): if self.scheduler_job.job_runner.processor_agent: self.scheduler_job.job_runner.processor_agent.end() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() dag = self.scheduler_job.job_runner.dagbag.get_dag("test_scheduler_add_new_task", session=session) self.scheduler_job.job_runner._create_dag_runs([orm_dag], session) @@ -4145,7 +4139,7 @@ def test_runs_respected_after_clear(self, dag_maker): ) as dag: BashOperator(task_id="dummy", bash_command="echo Hi") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() session = settings.Session() @@ -4198,7 +4192,7 @@ def test_timeout_triggers(self, dag_maker): session.flush() # Boot up the scheduler and make it check timeouts - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.check_trigger_timeouts(session=session) # Make sure that TI1 is now scheduled to fail, and 2 wasn't touched @@ -4210,7 +4204,7 @@ def test_timeout_triggers(self, dag_maker): def test_find_zombies_nothing(self): executor = MockExecutor(do_update=False) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(), executor=executor) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(), executor=executor) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() self.scheduler_job.job_runner._find_zombies() @@ -4220,7 +4214,7 @@ def test_find_zombies_nothing(self): def test_find_zombies(self, load_examples): dagbag = DagBag(TEST_DAG_FOLDER, read_dags_from_db=False) with create_session() as session: - session.query(BaseJob).delete() + session.query(Job).delete() dag = dagbag.get_dag("example_branch_operator") dag.sync_to_db() @@ -4231,7 +4225,7 @@ def test_find_zombies(self, load_examples): session=session, ) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner.processor_agent = mock.MagicMock() @@ -4243,7 +4237,7 @@ def test_find_zombies(self, load_examples): ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING) ti.queued_by_job_id = 999 - local_job = BaseJob(job_runner=LocalTaskJobRunner(ti), dag_id=ti.dag_id) + local_job = Job(job_runner=LocalTaskJobRunner(ti), dag_id=ti.dag_id) local_job.state = State.SHUTDOWN session.add(local_job) @@ -4274,7 +4268,7 @@ def test_find_zombies(self, load_examples): with create_session() as session: session.query(TaskInstance).delete() - session.query(BaseJob).delete() + session.query(Job).delete() def test_zombie_message(self, load_examples): """ @@ -4283,7 +4277,7 @@ def test_zombie_message(self, load_examples): dagbag = DagBag(TEST_DAG_FOLDER, read_dags_from_db=False) with create_session() as session: - session.query(BaseJob).delete() + session.query(Job).delete() dag = dagbag.get_dag("example_branch_operator") dag.sync_to_db() @@ -4294,7 +4288,7 @@ def test_zombie_message(self, load_examples): session=session, ) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner.processor_agent = mock.MagicMock() @@ -4306,7 +4300,7 @@ def test_zombie_message(self, load_examples): ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING) ti.queued_by_job_id = 999 - local_job = BaseJob(job_runner=LocalTaskJobRunner(ti), dag_id=ti.dag_id) + local_job = Job(job_runner=LocalTaskJobRunner(ti), dag_id=ti.dag_id) local_job.state = State.SHUTDOWN session.add(local_job) @@ -4352,7 +4346,7 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce dag_folder=os.path.join(settings.DAGS_FOLDER, "test_example_bash_operator.py"), read_dags_from_db=False, ) - session.query(BaseJob).delete() + session.query(Job).delete() dag = dagbag.get_dag("test_example_bash_operator") dag.sync_to_db(processor_subdir=TEST_DAG_FOLDER) @@ -4365,7 +4359,7 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce task = dag.get_task(task_id="run_this_last") ti = TaskInstance(task, run_id=dag_run.run_id, state=State.RUNNING) - local_job = BaseJob(job_runner=LocalTaskJobRunner(ti), dag_id=ti.dag_id) + local_job = Job(job_runner=LocalTaskJobRunner(ti), dag_id=ti.dag_id) local_job.state = State.SHUTDOWN session.add(local_job) session.flush() @@ -4376,7 +4370,7 @@ def test_find_zombies_handle_failure_callbacks_are_correctly_passed_to_dag_proce ti.job_id = local_job.id session.flush() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner.processor_agent = mock.MagicMock() @@ -4417,7 +4411,7 @@ def test_cleanup_stale_dags(self): session.flush() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner.processor_agent = mock.MagicMock() @@ -4511,7 +4505,7 @@ def test_mapped_dag(self, dag_id, session): executor = SequentialExecutor() - job = BaseJob(job_runner=SchedulerJobRunner(subdir=dag.fileloc, executor=executor)) + job = Job(job_runner=SchedulerJobRunner(subdir=dag.fileloc, executor=executor)) self.run_scheduler_until_dagrun_terminal(job) @@ -4527,7 +4521,7 @@ def test_should_mark_empty_task_as_success(self): dagbag = DagBag(dag_folder=dag_file, include_examples=False, read_dags_from_db=False) dagbag.sync_to_db() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner.processor_agent = mock.MagicMock() dag = self.scheduler_job.job_runner.dagbag.get_dag("test_only_empty_tasks") @@ -4607,7 +4601,7 @@ def test_catchup_works_correctly(self, dag_maker): ) as dag: EmptyOperator(task_id="dummy") - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner.processor_agent = mock.MagicMock(spec=DagFileProcessorAgent) @@ -4663,7 +4657,7 @@ def test_update_dagrun_state_for_paused_dag(self, dag_maker, session): assert scheduled_run.state == State.RUNNING - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner._update_dag_run_state_for_paused_dags(session=session) session.flush() @@ -4705,7 +4699,7 @@ def test_update_dagrun_state_for_paused_dag_not_for_backfill(self, dag_maker, se assert backfill_run.state == State.RUNNING - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.executor = MockExecutor() self.scheduler_job.job_runner._update_dag_run_state_for_paused_dags() session.flush() @@ -4731,7 +4725,7 @@ def test_dataset_orphaning(self, dag_maker, session): with dag_maker(dag_id="datasets-1", schedule=[dataset1], session=session): BashOperator(task_id="task", bash_command="echo 1", outlets=[dataset3]) - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) self.scheduler_job.job_runner._orphan_unreferenced_datasets(session=session) session.flush() @@ -4780,7 +4774,7 @@ def test_schedule_dag_run_with_upstream_skip(dag_maker, session): # dag_runs = DagRun.find(dag_id='test_task_with_upstream_skip_dag') # dag_file_processor._process_task_instances(dag, dag_runs=dag_runs) - scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=os.devnull)) + scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=os.devnull)) scheduler_job.job_runner._schedule_dag_run(dr, session) session.flush() tis = {ti.task_id: ti for ti in dr.get_task_instances(session=session)} @@ -4797,7 +4791,7 @@ class TestSchedulerJobQueriesCount: made that affects the performance of the SchedulerJob. """ - scheduler_job: BaseJob | None + scheduler_job: Job | None @staticmethod def clean_db(): @@ -4869,7 +4863,7 @@ def test_execute_queries_count_with_harvested_dags(self, expected_query_count, d mock_agent = mock.MagicMock() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=PERF_DAGS_FOLDER, num_runs=1)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=PERF_DAGS_FOLDER, num_runs=1)) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.heartbeat = mock.MagicMock() self.scheduler_job.job_runner.processor_agent = mock_agent @@ -4944,7 +4938,7 @@ def test_process_dags_queries_count( mock_agent = mock.MagicMock() - self.scheduler_job = BaseJob(job_runner=SchedulerJobRunner(subdir=PERF_DAGS_FOLDER, num_runs=1)) + self.scheduler_job = Job(job_runner=SchedulerJobRunner(subdir=PERF_DAGS_FOLDER, num_runs=1)) self.scheduler_job.executor = MockExecutor(do_update=False) self.scheduler_job.heartbeat = mock.MagicMock() self.scheduler_job.job_runner.processor_agent = mock_agent diff --git a/tests/jobs/test_triggerer_job.py b/tests/jobs/test_triggerer_job.py index c3816ddd9df51..d22a371a9f52b 100644 --- a/tests/jobs/test_triggerer_job.py +++ b/tests/jobs/test_triggerer_job.py @@ -29,8 +29,8 @@ from airflow import DAG from airflow.config_templates import airflow_local_settings -from airflow.jobs.base_job import BaseJob -from airflow.jobs.triggerer_job import TriggererJobRunner, TriggerRunner, setup_queue_listener +from airflow.jobs.job import Job +from airflow.jobs.triggerer_job_runner import TriggererJobRunner, TriggerRunner, setup_queue_listener from airflow.logging_config import configure_logging from airflow.models import DagModel, DagRun, TaskInstance, Trigger from airflow.models.baseoperator import BaseOperator @@ -125,7 +125,7 @@ def __init__(self, password, **kwargs): trigger = SuccessTrigger() op = SensitiveArgOperator(task_id="sensitive_arg_task", password="some_password") create_trigger_in_db(session, trigger, operator=op) - job = BaseJob(job_runner=TriggererJobRunner()) + job = Job(job_runner=TriggererJobRunner()) job.job_runner.load_triggers() # Now, start TriggerRunner up (and set it as a daemon thread during tests) job.job_runner.daemon = True @@ -150,7 +150,7 @@ def __init__(self, password, **kwargs): def test_is_alive(): """Checks the heartbeat logic""" # Current time - triggerer_job = BaseJob(job_runner=TriggererJobRunner(None), heartrate=10, state=State.RUNNING) + triggerer_job = Job(job_runner=TriggererJobRunner(None), heartrate=10, state=State.RUNNING) assert triggerer_job.is_alive() # Slightly old, but still fresh @@ -170,7 +170,7 @@ def test_is_alive(): def test_is_needed(session): """Checks the triggerer-is-needed logic""" # No triggers, no need - triggerer_job = BaseJob(job_runner=TriggererJobRunner(None), heartrate=10, state=State.RUNNING) + triggerer_job = Job(job_runner=TriggererJobRunner(None), heartrate=10, state=State.RUNNING) assert triggerer_job.job_runner.is_needed() is False # Add a trigger, it's needed trigger = TimeDeltaTrigger(datetime.timedelta(days=7)) @@ -192,7 +192,7 @@ def test_capacity_decode(): None, ] for input_str in variants: - job = BaseJob(job_runner=TriggererJobRunner(capacity=input_str)) + job = Job(job_runner=TriggererJobRunner(capacity=input_str)) assert job.job_runner.capacity == input_str or job.job_runner.capacity == 1000 # Negative cases @@ -217,7 +217,7 @@ def test_trigger_lifecycle(session): trigger = TimeDeltaTrigger(datetime.timedelta(days=7)) dag_model, run, trigger_orm, task_instance = create_trigger_in_db(session, trigger) # Make a TriggererJobRunner and have it retrieve DB tasks - job = BaseJob(job_runner=TriggererJobRunner()) + job = Job(job_runner=TriggererJobRunner()) job.job_runner.load_triggers() # Make sure it turned up in TriggerRunner's queue assert [x for x, y in job.job_runner.trigger_runner.to_create] == [1] @@ -339,7 +339,7 @@ def handle_events(self): session.commit() - job = BaseJob(job_runner=TriggererJob_()) + job = Job(job_runner=TriggererJob_()) job.job_runner.trigger_runner = TriggerRunner_() thread = Thread(target=job.job_runner._execute) thread.start() @@ -372,7 +372,7 @@ def test_trigger_from_dead_triggerer(session): session.add(trigger_orm) session.commit() # Make a TriggererJobRunner and have it retrieve DB tasks - job = BaseJob(job_runner=TriggererJobRunner()) + job = Job(job_runner=TriggererJobRunner()) job.job_runner.load_triggers() # Make sure it turned up in TriggerRunner's queue assert [x for x, y in job.job_runner.trigger_runner.to_create] == [1] @@ -390,7 +390,7 @@ def test_trigger_from_expired_triggerer(session): trigger_orm.triggerer_id = 42 session.add(trigger_orm) # Use a TriggererJobRunner with an expired heartbeat - triggerer_job_orm = BaseJob(job_runner=TriggererJobRunner()) + triggerer_job_orm = Job(job_runner=TriggererJobRunner()) triggerer_job_orm.id = 42 triggerer_job_orm.start_date = timezone.utcnow() - datetime.timedelta(hours=1) triggerer_job_orm.end_date = None @@ -398,7 +398,7 @@ def test_trigger_from_expired_triggerer(session): session.add(triggerer_job_orm) session.commit() # Make a TriggererJobRunner and have it retrieve DB tasks - job = BaseJob(job_runner=TriggererJobRunner()) + job = Job(job_runner=TriggererJobRunner()) job.job_runner.load_triggers() # Make sure it turned up in TriggerRunner's queue assert [x for x, y in job.job_runner.trigger_runner.to_create] == [1] @@ -413,7 +413,7 @@ def test_trigger_firing(session): trigger = SuccessTrigger() create_trigger_in_db(session, trigger) # Make a TriggererJobRunner and have it retrieve DB tasks - job = BaseJob(job_runner=TriggererJobRunner()) + job = Job(job_runner=TriggererJobRunner()) job.job_runner.load_triggers() # Now, start TriggerRunner up (and set it as a daemon thread during tests) job.job_runner.daemon = True @@ -441,7 +441,7 @@ def test_trigger_failing(session): trigger = FailureTrigger() create_trigger_in_db(session, trigger) # Make a TriggererJobRunner and have it retrieve DB tasks - job = BaseJob(job_runner=TriggererJobRunner()) + job = Job(job_runner=TriggererJobRunner()) job.job_runner.load_triggers() # Now, start TriggerRunner up (and set it as a daemon thread during tests) job.job_runner.daemon = True @@ -506,7 +506,7 @@ def test_invalid_trigger(session, dag_maker): session.commit() # Make a TriggererJobRunner and have it retrieve DB tasks - job = BaseJob(job_runner=TriggererJobRunner()) + job = Job(job_runner=TriggererJobRunner()) job.job_runner.load_triggers() # Make sure it turned up in the failed queue @@ -525,11 +525,11 @@ def test_invalid_trigger(session, dag_maker): @pytest.mark.parametrize("should_wrap", (True, False)) -@patch("airflow.jobs.triggerer_job.configure_trigger_log_handler") +@patch("airflow.jobs.triggerer_job_runner.configure_trigger_log_handler") def test_handler_config_respects_donot_wrap(mock_configure, should_wrap): - from airflow.jobs import triggerer_job + from airflow.jobs import triggerer_job_runner - triggerer_job.DISABLE_WRAPPER = not should_wrap + triggerer_job_runner.DISABLE_WRAPPER = not should_wrap TriggererJobRunner() if should_wrap: mock_configure.assert_called() @@ -537,7 +537,7 @@ def test_handler_config_respects_donot_wrap(mock_configure, should_wrap): mock_configure.assert_not_called() -@patch("airflow.jobs.triggerer_job.setup_queue_listener") +@patch("airflow.jobs.triggerer_job_runner.setup_queue_listener") def test_triggerer_job_always_creates_listener(mock_setup): mock_setup.assert_not_called() TriggererJobRunner() diff --git a/tests/jobs/test_triggerer_job_logging.py b/tests/jobs/test_triggerer_job_logging.py index 27a641d8e22df..d11f4bb3686f7 100644 --- a/tests/jobs/test_triggerer_job_logging.py +++ b/tests/jobs/test_triggerer_job_logging.py @@ -24,7 +24,7 @@ import pytest from airflow.config_templates import airflow_local_settings -from airflow.jobs import triggerer_job +from airflow.jobs import triggerer_job_runner from airflow.logging_config import configure_logging from airflow.providers.amazon.aws.log.s3_task_handler import S3TaskHandler from airflow.utils.log.file_task_handler import FileTaskHandler @@ -51,7 +51,7 @@ def clear_logger_handlers(log): @pytest.fixture(autouse=True) def reload_triggerer_job(): - importlib.reload(triggerer_job) + importlib.reload(triggerer_job_runner) def test_configure_trigger_log_handler_file(): @@ -74,11 +74,11 @@ def test_configure_trigger_log_handler_file(): task_handlers = assert_handlers(task_logger, FileTaskHandler) # not yet configured to use wrapper - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is False + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is False - triggerer_job.configure_trigger_log_handler() + triggerer_job_runner.configure_trigger_log_handler() # after config - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is True + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is True root_handlers = assert_handlers(root_logger, RedirectStdHandler, TriggererHandlerWrapper) assert root_handlers[1].base_handler == task_handlers[0] # other handlers have DropTriggerLogsFilter @@ -112,11 +112,11 @@ def test_configure_trigger_log_handler_s3(): task_logger = logging.getLogger("airflow.task") task_handlers = assert_handlers(task_logger, S3TaskHandler) # not yet configured to use wrapper - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is False + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is False - triggerer_job.configure_trigger_log_handler() + triggerer_job_runner.configure_trigger_log_handler() # after config - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is True + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is True handlers = assert_handlers(root_logger, RedirectStdHandler, TriggererHandlerWrapper) assert handlers[1].base_handler == task_handlers[0] # other handlers have DropTriggerLogsFilter @@ -196,16 +196,16 @@ def test_configure_trigger_log_handler_not_file_task_handler(cfg, cls, msg): assert_handlers(task_logger, cls) # not yet configured to use wrapper - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is False + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is False with warnings.catch_warnings(record=True) as captured: - triggerer_job.configure_trigger_log_handler() + triggerer_job_runner.configure_trigger_log_handler() assert [x.message.args[0] for x in captured] == msg # after config # doesn't use TriggererHandlerWrapper, no change in handler - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is False + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is False # still no root handlers assert_handlers(root_logger) @@ -246,12 +246,12 @@ def test_configure_trigger_log_handler_fallback_task(): # before config root_logger = logging.getLogger() assert_handlers(root_logger) - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is False + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is False - triggerer_job.configure_trigger_log_handler() + triggerer_job_runner.configure_trigger_log_handler() # after config - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is True + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is True handlers = assert_handlers(root_logger, TriggererHandlerWrapper) assert handlers[0].base_handler == task_logger.handlers[0] @@ -297,13 +297,13 @@ def test_configure_trigger_log_handler_root_has_task_handler(): # before config root_logger = logging.getLogger() assert_handlers(root_logger, FileTaskHandler) - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is False + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is False # configure - triggerer_job.configure_trigger_log_handler() + triggerer_job_runner.configure_trigger_log_handler() # after config - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is True + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is True handlers = assert_handlers(root_logger, TriggererHandlerWrapper) # no filters on wrapper handler assert handlers[0].filters == [] @@ -354,15 +354,15 @@ def test_configure_trigger_log_handler_root_not_file_task(): # before config root_logger = logging.getLogger() assert_handlers(root_logger, logging.Handler) - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is False + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is False # configure with warnings.catch_warnings(record=True) as captured: - triggerer_job.configure_trigger_log_handler() + triggerer_job_runner.configure_trigger_log_handler() assert captured == [] # after config - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is True + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is True handlers = assert_handlers(root_logger, logging.Handler, TriggererHandlerWrapper) # other handlers have DropTriggerLogsFilter assert handlers[0].filters[0].__class__ == DropTriggerLogsFilter @@ -420,10 +420,10 @@ def test_configure_trigger_log_handler_root_old_file_task(): root_logger = logging.getLogger() assert_handlers(root_logger, OldFileTaskHandler) - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is False + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is False with warnings.catch_warnings(record=True) as captured: - triggerer_job.configure_trigger_log_handler() + triggerer_job_runner.configure_trigger_log_handler() # since a root logger is explicitly configured with an old FileTaskHandler which doesn't # work properly with individual trigger logging, warn @@ -435,7 +435,7 @@ def test_configure_trigger_log_handler_root_old_file_task(): ] # after config - assert triggerer_job.HANDLER_SUPPORTS_TRIGGERER is True + assert triggerer_job_runner.HANDLER_SUPPORTS_TRIGGERER is True handlers = assert_handlers(root_logger, OldFileTaskHandler, TriggererHandlerWrapper) # other handlers have DropTriggerLogsFilter assert handlers[0].filters[0].__class__ == DropTriggerLogsFilter diff --git a/tests/listeners/test_listeners.py b/tests/listeners/test_listeners.py index 19594ea7a6157..c900f761c7134 100644 --- a/tests/listeners/test_listeners.py +++ b/tests/listeners/test_listeners.py @@ -20,7 +20,7 @@ import pytest as pytest from airflow import AirflowException -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job from airflow.listeners.listener import get_listener_manager from airflow.operators.bash import BashOperator from airflow.utils import timezone @@ -80,7 +80,7 @@ def test_multiple_listeners(create_task_instance, session=None): lm.add_listener(full_listener) lm.add_listener(lifecycle_listener) - job = BaseJob(job_runner=MockJobRunner()) + job = Job(job_runner=MockJobRunner()) try: job.run() except NotImplementedError: diff --git a/tests/models/test_trigger.py b/tests/models/test_trigger.py index 3756100fe6ff3..827312d9f23f6 100644 --- a/tests/models/test_trigger.py +++ b/tests/models/test_trigger.py @@ -20,8 +20,8 @@ import pytest -from airflow.jobs.base_job import BaseJob -from airflow.jobs.triggerer_job import TriggererJobRunner +from airflow.jobs.job import Job +from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.models import TaskInstance, Trigger from airflow.operators.empty import EmptyOperator from airflow.triggers.base import TriggerEvent @@ -41,11 +41,11 @@ def session(): def clear_db(session): session.query(TaskInstance).delete() session.query(Trigger).delete() - session.query(BaseJob).delete() + session.query(Job).delete() yield session session.query(TaskInstance).delete() session.query(Trigger).delete() - session.query(BaseJob).delete() + session.query(Job).delete() session.commit() @@ -140,14 +140,14 @@ def test_assign_unassigned(session, create_task_instance): """ Tests that unassigned triggers of all appropriate states are assigned. """ - finished_triggerer = BaseJob(job_runner=TriggererJobRunner(None), heartrate=10, state=State.SUCCESS) + finished_triggerer = Job(job_runner=TriggererJobRunner(None), heartrate=10, state=State.SUCCESS) finished_triggerer.end_date = timezone.utcnow() - datetime.timedelta(hours=1) session.add(finished_triggerer) assert not finished_triggerer.is_alive() - healthy_triggerer = BaseJob(job_runner=TriggererJobRunner(None), heartrate=10, state=State.RUNNING) + healthy_triggerer = Job(job_runner=TriggererJobRunner(None), heartrate=10, state=State.RUNNING) session.add(healthy_triggerer) assert healthy_triggerer.is_alive() - new_triggerer = BaseJob(job_runner=TriggererJobRunner(None), heartrate=10, state=State.RUNNING) + new_triggerer = Job(job_runner=TriggererJobRunner(None), heartrate=10, state=State.RUNNING) session.add(new_triggerer) assert new_triggerer.is_alive() session.commit() diff --git a/tests/serialization/test_pydantic_models.py b/tests/serialization/test_pydantic_models.py index 2a050383ec624..8947afea35634 100644 --- a/tests/serialization/test_pydantic_models.py +++ b/tests/serialization/test_pydantic_models.py @@ -19,17 +19,17 @@ from pydantic import parse_raw_as -from airflow.jobs.base_job import BaseJob -from airflow.jobs.local_task_job import LocalTaskJobRunner +from airflow.jobs.job import Job +from airflow.jobs.local_task_job_runner import LocalTaskJobRunner from airflow.models.dataset import ( DagScheduleDatasetReference, DatasetEvent, DatasetModel, TaskOutletDatasetReference, ) -from airflow.serialization.pydantic.base_job import BaseJobPydantic from airflow.serialization.pydantic.dag_run import DagRunPydantic from airflow.serialization.pydantic.dataset import DatasetEventPydantic +from airflow.serialization.pydantic.job import JobPydantic from airflow.serialization.pydantic.taskinstance import TaskInstancePydantic from airflow.utils import timezone from airflow.utils.state import State @@ -76,15 +76,15 @@ def test_serializing_pydantic_dagrun(session, create_task_instance): def test_serializing_pydantic_local_task_job(session, create_task_instance): dag_id = "test-dag" ti = create_task_instance(dag_id=dag_id, session=session) - ltj = BaseJob(job_runner=LocalTaskJobRunner(task_instance=ti), dag_id=ti.dag_id) + ltj = Job(job_runner=LocalTaskJobRunner(task_instance=ti), dag_id=ti.dag_id) ltj.state = State.RUNNING session.commit() - pydantic_job = BaseJobPydantic.from_orm(ltj) + pydantic_job = JobPydantic.from_orm(ltj) json_string = pydantic_job.json() print(json_string) - deserialized_model = parse_raw_as(BaseJobPydantic, json_string) + deserialized_model = parse_raw_as(JobPydantic, json_string) assert deserialized_model.dag_id == dag_id assert deserialized_model.job_type == "LocalTaskJob" assert deserialized_model.state == State.RUNNING diff --git a/tests/task/task_runner/test_base_task_runner.py b/tests/task/task_runner/test_base_task_runner.py index 10188ffa677e8..fe5dcc099ba5c 100644 --- a/tests/task/task_runner/test_base_task_runner.py +++ b/tests/task/task_runner/test_base_task_runner.py @@ -21,8 +21,8 @@ import pytest -from airflow.jobs.base_job import BaseJob -from airflow.jobs.local_task_job import LocalTaskJobRunner +from airflow.jobs.job import Job +from airflow.jobs.local_task_job_runner import LocalTaskJobRunner from airflow.models.baseoperator import BaseOperator from airflow.task.task_runner.base_task_runner import BaseTaskRunner @@ -40,7 +40,7 @@ def test_config_copy_mode(tmp_configuration_copy, subprocess_call, dag_maker, im ti = dr.task_instances[0] task_runner = LocalTaskJobRunner(ti) - job = BaseJob(job_runner=task_runner, dag_id=ti.dag_id) + job = Job(job_runner=task_runner, dag_id=ti.dag_id) runner = BaseTaskRunner(job) # So we don't try to delete it -- cos the file won't exist del runner._cfg_path diff --git a/tests/task/task_runner/test_cgroup_task_runner.py b/tests/task/task_runner/test_cgroup_task_runner.py index 786ab222d8296..0d40f0c00282d 100644 --- a/tests/task/task_runner/test_cgroup_task_runner.py +++ b/tests/task/task_runner/test_cgroup_task_runner.py @@ -32,12 +32,12 @@ def test_cgroup_task_runner_super_calls(self, mock_super_on_finish, mock_super_i and when task finishes, CgroupTaskRunner.on_finish() calls super().on_finish() to delete the temp cfg file. """ - local_task_job = mock.Mock() - local_task_job.task_instance = mock.MagicMock() - local_task_job.task_instance.run_as_user = None - local_task_job.task_instance.command_as_list.return_value = ["sleep", "1000"] + base_job = mock.Mock() + base_job.task_instance = mock.MagicMock() + base_job.task_instance.run_as_user = None + base_job.task_instance.command_as_list.return_value = ["sleep", "1000"] - runner = CgroupTaskRunner(local_task_job) + runner = CgroupTaskRunner(base_job) assert mock_super_init.called runner.on_finish() diff --git a/tests/task/task_runner/test_standard_task_runner.py b/tests/task/task_runner/test_standard_task_runner.py index 321912935a586..d642c5418cecf 100644 --- a/tests/task/task_runner/test_standard_task_runner.py +++ b/tests/task/task_runner/test_standard_task_runner.py @@ -28,8 +28,8 @@ import psutil import pytest -from airflow.jobs.base_job import BaseJob -from airflow.jobs.local_task_job import LocalTaskJobRunner +from airflow.jobs.job import Job +from airflow.jobs.local_task_job_runner import LocalTaskJobRunner from airflow.listeners.listener import get_listener_manager from airflow.models.dagbag import DagBag from airflow.models.taskinstance import TaskInstance @@ -94,10 +94,10 @@ def setup_class(self): @patch("airflow.utils.log.file_task_handler.FileTaskHandler._init_file") def test_start_and_terminate(self, mock_init): mock_init.return_value = "/tmp/any" - local_task_job = mock.Mock() - local_task_job.task_instance = mock.MagicMock() - local_task_job.task_instance.run_as_user = None - local_task_job.task_instance.command_as_list.return_value = [ + base_job = mock.Mock() + base_job.task_instance = mock.MagicMock() + base_job.task_instance.run_as_user = None + base_job.task_instance.command_as_list.return_value = [ "airflow", "tasks", "run", @@ -105,9 +105,9 @@ def test_start_and_terminate(self, mock_init): "task1", "2016-01-01", ] - local_task_job.job_runner = LocalTaskJobRunner(local_task_job.task_instance) + base_job.job_runner = LocalTaskJobRunner(base_job.task_instance) - runner = StandardTaskRunner(local_task_job) + runner = StandardTaskRunner(base_job) runner.start() # Wait until process sets its pgid to be equal to pid with timeout(seconds=1): @@ -150,9 +150,7 @@ def test_notifies_about_start_and_stop(self): start_date=DEFAULT_DATE, ) ti = TaskInstance(task=task, run_id="test") - job1 = BaseJob( - job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id - ) + job1 = Job(job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id) runner = StandardTaskRunner(job1) runner.start() @@ -195,9 +193,7 @@ def test_notifies_about_fail(self): start_date=DEFAULT_DATE, ) ti = TaskInstance(task=task, run_id="test") - job1 = BaseJob( - job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id - ) + job1 = Job(job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id) runner = StandardTaskRunner(job1) runner.start() @@ -220,12 +216,12 @@ def test_notifies_about_fail(self): @patch("airflow.utils.log.file_task_handler.FileTaskHandler._init_file") def test_start_and_terminate_run_as_user(self, mock_init): mock_init.return_value = "/tmp/any" - local_task_job = mock.Mock() - local_task_job.task_instance = mock.MagicMock() - local_task_job.task_instance.task_id = "task_id" - local_task_job.task_instance.dag_id = "dag_id" - local_task_job.task_instance.run_as_user = getuser() - local_task_job.task_instance.command_as_list.return_value = [ + base_job = mock.Mock() + base_job.task_instance = mock.MagicMock() + base_job.task_instance.task_id = "task_id" + base_job.task_instance.dag_id = "dag_id" + base_job.task_instance.run_as_user = getuser() + base_job.task_instance.command_as_list.return_value = [ "airflow", "tasks", "test", @@ -233,8 +229,8 @@ def test_start_and_terminate_run_as_user(self, mock_init): "task1", "2016-01-01", ] - local_task_job.job_runner = LocalTaskJobRunner(local_task_job.task_instance) - runner = StandardTaskRunner(local_task_job) + base_job.job_runner = LocalTaskJobRunner(base_job.task_instance) + runner = StandardTaskRunner(base_job) runner.start() time.sleep(0.5) @@ -261,12 +257,12 @@ def test_early_reap_exit(self, mock_init, caplog): -9 and a log message. """ mock_init.return_value = "/tmp/any" - local_task_job = mock.Mock() - local_task_job.task_instance = mock.MagicMock() - local_task_job.task_instance.task_id = "task_id" - local_task_job.task_instance.dag_id = "dag_id" - local_task_job.task_instance.run_as_user = getuser() - local_task_job.task_instance.command_as_list.return_value = [ + base_job = mock.Mock() + base_job.task_instance = mock.MagicMock() + base_job.task_instance.task_id = "task_id" + base_job.task_instance.dag_id = "dag_id" + base_job.task_instance.run_as_user = getuser() + base_job.task_instance.command_as_list.return_value = [ "airflow", "tasks", "test", @@ -274,10 +270,10 @@ def test_early_reap_exit(self, mock_init, caplog): "task1", "2016-01-01", ] - local_task_job.job_runner = LocalTaskJobRunner(local_task_job.task_instance) + base_job.job_runner = LocalTaskJobRunner(base_job.task_instance) # Kick off the runner - runner = StandardTaskRunner(local_task_job) + runner = StandardTaskRunner(base_job) runner.start() time.sleep(0.2) @@ -323,9 +319,7 @@ def test_on_kill(self): start_date=DEFAULT_DATE, ) ti = TaskInstance(task=task, run_id="test") - job1 = BaseJob( - job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id - ) + job1 = Job(job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id) runner = StandardTaskRunner(job1) runner.start() @@ -383,9 +377,7 @@ def test_parsing_context(self): start_date=DEFAULT_DATE, ) ti = TaskInstance(task=task, run_id="test") - job1 = BaseJob( - job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id - ) + job1 = Job(job_runner=LocalTaskJobRunner(task_instance=ti, ignore_ti_state=True), dag_id=ti.dag_id) runner = StandardTaskRunner(job1) runner.start() diff --git a/tests/task/task_runner/test_task_runner.py b/tests/task/task_runner/test_task_runner.py index 9241d43f21b2b..e95d47c473022 100644 --- a/tests/task/task_runner/test_task_runner.py +++ b/tests/task/task_runner/test_task_runner.py @@ -20,7 +20,7 @@ import pytest -from airflow.jobs.local_task_job import LocalTaskJobRunner +from airflow.jobs.local_task_job_runner import LocalTaskJobRunner from airflow.task.task_runner import CORE_TASK_RUNNERS, get_task_runner from airflow.utils.module_loading import import_string @@ -38,10 +38,10 @@ def test_should_support_core_task_runner(self, mock_subprocess): ti = mock.MagicMock(map_index=-1, run_as_user=None) ti.get_template_context.return_value = {"ti": ti} ti.get_dagrun.return_value.get_log_template.return_value.filename = "blah" - local_task_job = mock.MagicMock(task_instance=ti) - local_task_job.job_runner = LocalTaskJobRunner(ti) - local_task_job.job_runner.job = local_task_job - task_runner = get_task_runner(local_task_job.job_runner) + base_job = mock.MagicMock(task_instance=ti) + base_job.job_runner = LocalTaskJobRunner(ti) + base_job.job_runner.job = base_job + task_runner = get_task_runner(base_job.job_runner) assert "StandardTaskRunner" == task_runner.__class__.__name__ @@ -50,13 +50,13 @@ def test_should_support_core_task_runner(self, mock_subprocess): "tests.task.task_runner.test_task_runner.custom_task_runner", ) def test_should_support_custom_legacy_task_runner(self): - local_task_job = mock.MagicMock( + base_job = mock.MagicMock( **{"task_instance.get_template_context.return_value": {"ti": mock.MagicMock()}} ) custom_task_runner.reset_mock() - task_runner = get_task_runner(local_task_job) + task_runner = get_task_runner(base_job) - custom_task_runner.assert_called_once_with(local_task_job.job) + custom_task_runner.assert_called_once_with(base_job.job) assert custom_task_runner.return_value == task_runner diff --git a/tests/test_utils/db.py b/tests/test_utils/db.py index f9128e035e9a4..0ac1d7bb8ef0c 100644 --- a/tests/test_utils/db.py +++ b/tests/test_utils/db.py @@ -17,7 +17,7 @@ # under the License. from __future__ import annotations -from airflow.jobs.base_job import BaseJob +from airflow.jobs.job import Job from airflow.models import ( Connection, DagModel, @@ -55,7 +55,7 @@ def clear_db_runs(): with create_session() as session: - session.query(BaseJob).delete() + session.query(Job).delete() session.query(Trigger).delete() session.query(DagRun).delete() session.query(TaskInstance).delete() @@ -156,7 +156,7 @@ def clear_db_logs(): def clear_db_jobs(): with create_session() as session: - session.query(BaseJob).delete() + session.query(Job).delete() def clear_db_task_fail(): diff --git a/tests/utils/test_helpers.py b/tests/utils/test_helpers.py index a976d51292cbc..0adbd43b210e4 100644 --- a/tests/utils/test_helpers.py +++ b/tests/utils/test_helpers.py @@ -23,7 +23,7 @@ import pytest from airflow import AirflowException -from airflow.jobs.job_runner import BaseJobRunner +from airflow.jobs.base_job_runner import BaseJobRunner from airflow.utils import helpers, timezone from airflow.utils.helpers import ( at_most_one, diff --git a/tests/utils/test_log_handlers.py b/tests/utils/test_log_handlers.py index e6f9492ebc0fa..dc66c4a4e204f 100644 --- a/tests/utils/test_log_handlers.py +++ b/tests/utils/test_log_handlers.py @@ -31,8 +31,8 @@ from kubernetes.client import models as k8s from airflow.config_templates.airflow_local_settings import DEFAULT_LOGGING_CONFIG -from airflow.jobs.base_job import BaseJob -from airflow.jobs.triggerer_job import TriggererJobRunner +from airflow.jobs.job import Job +from airflow.jobs.triggerer_job_runner import TriggererJobRunner from airflow.models import DAG, DagRun, TaskInstance, Trigger from airflow.operators.python import PythonOperator from airflow.utils.log.file_task_handler import ( @@ -403,7 +403,7 @@ def test_set_context_trigger(self, create_dummy_dag, dag_maker, is_a_trigger, se assert isinstance(ti, TaskInstance) if is_a_trigger: ti.is_trigger_log_context = True - job = BaseJob(job_runner=TriggererJobRunner()) + job = Job(job_runner=TriggererJobRunner()) t = Trigger("", {}) t.triggerer_job = job ti.triggerer = t @@ -475,7 +475,7 @@ def test_log_retrieval_valid_trigger(self, create_task_instance): ) ti.hostname = "hostname" trigger = Trigger("", {}) - job = BaseJob(job_runner=TriggererJobRunner()) + job = Job(job_runner=TriggererJobRunner()) job.id = 123 trigger.triggerer_job = job ti.trigger = trigger diff --git a/tests/www/views/test_views_base.py b/tests/www/views/test_views_base.py index 69b0cc757a003..0b4ee78c7d310 100644 --- a/tests/www/views/test_views_base.py +++ b/tests/www/views/test_views_base.py @@ -22,8 +22,8 @@ import pytest -from airflow.jobs.base_job import BaseJob -from airflow.jobs.scheduler_job import SchedulerJobRunner +from airflow.jobs.job import Job +from airflow.jobs.scheduler_job_runner import SchedulerJobRunner from airflow.utils import timezone from airflow.utils.session import create_session from airflow.www import app as application @@ -62,7 +62,7 @@ def test_doc_urls(admin_client, monkeypatch): def heartbeat_healthy(): # case-1: healthy scheduler status last_heartbeat = timezone.utcnow() - job = BaseJob( + job = Job( job_type="SchedulerJob", state="running", latest_heartbeat=last_heartbeat, @@ -72,10 +72,10 @@ def heartbeat_healthy(): session.add(job) yield "healthy", last_heartbeat.isoformat() with create_session() as session: - session.query(BaseJob).filter( - BaseJob.job_type == "SchedulerJob", - BaseJob.state == "running", - BaseJob.latest_heartbeat == last_heartbeat, + session.query(Job).filter( + Job.job_type == "SchedulerJob", + Job.state == "running", + Job.latest_heartbeat == last_heartbeat, ).delete() @@ -83,23 +83,23 @@ def heartbeat_healthy(): def heartbeat_too_slow(): # case-2: unhealthy scheduler status - scenario 1 (SchedulerJob is running too slowly) last_heartbeat = timezone.utcnow() - datetime.timedelta(minutes=1) - job = BaseJob( + job = Job( job_type="SchedulerJob", state="running", latest_heartbeat=last_heartbeat, job_runner=SchedulerJobRunner(), ) with create_session() as session: - session.query(BaseJob).filter( - BaseJob.job_type == "SchedulerJob", + session.query(Job).filter( + Job.job_type == "SchedulerJob", ).update({"latest_heartbeat": last_heartbeat - datetime.timedelta(seconds=1)}) session.add(job) yield "unhealthy", last_heartbeat.isoformat() with create_session() as session: - session.query(BaseJob).filter( - BaseJob.job_type == "SchedulerJob", - BaseJob.state == "running", - BaseJob.latest_heartbeat == last_heartbeat, + session.query(Job).filter( + Job.job_type == "SchedulerJob", + Job.state == "running", + Job.latest_heartbeat == last_heartbeat, ).delete() @@ -107,9 +107,9 @@ def heartbeat_too_slow(): def heartbeat_not_running(): # case-3: unhealthy scheduler status - scenario 2 (no running SchedulerJob) with create_session() as session: - session.query(BaseJob).filter( - BaseJob.job_type == "SchedulerJob", - BaseJob.state == "running", + session.query(Job).filter( + Job.job_type == "SchedulerJob", + Job.state == "running", ).delete() yield "unhealthy", None