def execute(self, context): if self.execution_date is not None: run_id = 'trig__{}'.format(self.execution_date) self.execution_date = timezone.parse(self.execution_date) else: run_id = 'trig__' + timezone.utcnow().isoformat() dro = DagRunOrder(run_id=run_id) if self.python_callable is not None: dro = self.python_callable(context, dro) if dro: trigger_dag(dag_id=self.trigger_dag_id, run_id=dro.run_id, conf=json.dumps(dro.payload), execution_date=self.execution_date, replace_microseconds=False) DagRun.refresh_from_db this_dag_run = DagRun.find(dag_id=self.trigger_dag_id, run_id=dro.run_id) DagRun.refresh_from_db while (this_dag_run[0].state == 'running'): DagRun.refresh_from_db this_dag_run = DagRun.find(dag_id=self.trigger_dag_id, run_id=dro.run_id) if (this_dag_run[0].state == 'failed'): raise ValueError( self.trigger_dag_id + ' is throwing an exception, it is in failed status') else: self.log.info("Criteria not met, moving on")
def test_trigger_controller_dag(self): dag = self.dagbag.get_dag('example_trigger_controller_dag') target_dag = self.dagbag.get_dag('example_trigger_target_dag') target_dag.sync_to_db() dag_file_processor = DagFileProcessor(dag_ids=[], log=Mock()) task_instances_list = dag_file_processor._process_task_instances( target_dag, dag_runs=DagRun.find(dag_id='example_trigger_target_dag') ) self.assertFalse(task_instances_list) job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True ) job.run() task_instances_list = dag_file_processor._process_task_instances( target_dag, dag_runs=DagRun.find(dag_id='example_trigger_target_dag') ) self.assertTrue(task_instances_list)
def _create_dagruns(dag, execution_dates, state, run_type): """ Infers from the dates which dag runs need to be created and does so. :param dag: the dag to create dag runs for :param execution_dates: list of execution dates to evaluate :param state: the state to set the dag run to :param run_type: The prefix will be used to construct dag run id: {run_id_prefix}__{execution_date} :return: newly created and existing dag runs for the execution dates supplied """ # find out if we need to create any dag runs dag_runs = DagRun.find(dag_id=dag.dag_id, execution_date=execution_dates) dates_to_create = list( set(execution_dates) - {dag_run.execution_date for dag_run in dag_runs}) for date in dates_to_create: dag_run = dag.create_dagrun( execution_date=date, start_date=timezone.utcnow(), external_trigger=False, state=state, run_type=run_type, ) dag_runs.append(dag_run) return dag_runs
def _get_dag_run(self, run_date: datetime, dag: DAG, session: Session = None): """ Returns a dag run for the given run date, which will be matched to an existing dag run if available or create a new dag run otherwise. If the max_active_runs limit is reached, this function will return None. :param run_date: the execution date for the dag run :param dag: DAG :param session: the database session object :return: a DagRun in state RUNNING or None """ run_id = f"{DagRunType.BACKFILL_JOB.value}__{run_date.isoformat()}" # consider max_active_runs but ignore when running subdags respect_dag_max_active_limit = bool(dag.schedule_interval and not dag.is_subdag) current_active_dag_count = dag.get_num_active_runs( external_trigger=False) # check if we are scheduling on top of a already existing dag_run # we could find a "scheduled" run instead of a "backfill" run = DagRun.find(dag_id=dag.dag_id, execution_date=run_date, session=session) if run is not None and len(run) > 0: run = run[0] if run.state == State.RUNNING: respect_dag_max_active_limit = False else: run = None # enforce max_active_runs limit for dag, special cases already # handled by respect_dag_max_active_limit if (respect_dag_max_active_limit and current_active_dag_count >= dag.max_active_runs): return None run = run or dag.create_dagrun( run_id=run_id, execution_date=run_date, start_date=timezone.utcnow(), state=State.RUNNING, external_trigger=False, session=session, conf=self.conf, ) # set required transient field run.dag = dag # explicitly mark as backfill and running run.state = State.RUNNING run.run_id = run_id run.verify_integrity(session=session) return run
def _get_dep_statuses(self, ti, session, dep_context): dag = ti.task.dag dagrun = ti.get_dagrun(session) if not dagrun: # The import is needed here to avoid a circular dependency from airflow.models.dagrun import DagRun running_dagruns = DagRun.find( dag_id=dag.dag_id, state=State.RUNNING, external_trigger=False, session=session ) if len(running_dagruns) >= dag.max_active_runs: reason = ( "The maximum number of active dag runs ({}) for this task " "instance's DAG '{}' has been reached.".format(dag.max_active_runs, ti.dag_id) ) else: reason = "Unknown reason" yield self._failing_status(reason=f"Task instance's dagrun did not exist: {reason}.") else: if dagrun.state != State.RUNNING: yield self._failing_status( reason="Task instance's dagrun was not in the 'running' state but in " "the state '{}'.".format(dagrun.state) )
def test_backfill_max_limit_check_no_count_existing(self): dag = self._get_dag_test_max_active_limits( 'test_backfill_max_limit_check_no_count_existing') start_date = DEFAULT_DATE end_date = DEFAULT_DATE # Existing dagrun that is within the backfill range dag.create_dagrun(run_id="test_existing_backfill", state=State.RUNNING, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE) executor = MockExecutor() job = BackfillJob(dag=dag, start_date=start_date, end_date=end_date, executor=executor, donot_pickle=True) job.run() # BackfillJob will run since the existing DagRun does not count for the max # active limit since it's within the backfill date range. dagruns = DagRun.find(dag_id=dag.dag_id) # will only be able to run 1 (the existing one) since there's just # one dag run slot left given the max_active_runs limit self.assertEqual(1, len(dagruns)) self.assertEqual(State.SUCCESS, dagruns[0].state)
def get_run_ids(dag: DAG, run_id: str, future: bool, past: bool, session: SASession = NEW_SESSION): """Returns run_ids of DAG execution""" last_dagrun = dag.get_last_dagrun(include_externally_triggered=True) current_dagrun = dag.get_dagrun(run_id=run_id) first_dagrun = ( session.query(DagRun) .filter(DagRun.dag_id == dag.dag_id) .order_by(DagRun.execution_date.asc()) .first() ) if last_dagrun is None: raise ValueError(f'DagRun for {dag.dag_id} not found') # determine run_id range of dag runs and tasks to consider end_date = last_dagrun.logical_date if future else current_dagrun.logical_date start_date = current_dagrun.logical_date if not past else first_dagrun.logical_date if not dag.timetable.can_run: # If the DAG never schedules, need to look at existing DagRun if the user wants future or # past runs. dag_runs = dag.get_dagruns_between(start_date=start_date, end_date=end_date) run_ids = sorted({d.run_id for d in dag_runs}) elif not dag.timetable.periodic: run_ids = [run_id] else: dates = [ info.logical_date for info in dag.iter_dagrun_infos_between(start_date, end_date, align=False) ] run_ids = [dr.run_id for dr in DagRun.find(dag_id=dag.dag_id, execution_date=dates)] return run_ids
def _iter_existing_dag_run_infos(dag: DAG, run_ids: List[str]) -> Iterator[_DagRunInfo]: for dag_run in DagRun.find(dag_id=dag.dag_id, run_id=run_ids): dag_run.dag = dag dag_run.verify_integrity() yield _DagRunInfo(dag_run.logical_date, dag.get_run_data_interval(dag_run))
def _create_dagruns( dag: DAG, infos: Iterable[_DagRunInfo], state: DagRunState, run_type: DagRunType, ) -> Iterable[DagRun]: """Infers from data intervals which DAG runs need to be created and does so. :param dag: The DAG to create runs for. :param infos: List of logical dates and data intervals to evaluate. :param state: The state to set the dag run to :param run_type: The prefix will be used to construct dag run id: ``{run_id_prefix}__{execution_date}``. :return: Newly created and existing dag runs for the execution dates supplied. """ # Find out existing DAG runs that we don't need to create. dag_runs = { run.logical_date: run for run in DagRun.find( dag_id=dag.dag_id, execution_date=[info.logical_date for info in infos]) } for info in infos: if info.logical_date in dag_runs: continue dag_runs[info.logical_date] = dag.create_dagrun( execution_date=info.logical_date, data_interval=info.data_interval, start_date=timezone.utcnow(), external_trigger=False, state=state, run_type=run_type, ) return dag_runs.values()
def schedule(self): self.log.info("Starting the scheduler.") self._restore_unfinished_dag_run() while True: identified_message = self.mailbox.get_identified_message() origin_event = identified_message.deserialize() self.log.debug("Event: {}".format(origin_event)) if SchedulerInnerEventUtil.is_inner_event(origin_event): event = SchedulerInnerEventUtil.to_inner_event(origin_event) else: event = origin_event with create_session() as session: if isinstance(event, BaseEvent): dagruns = self._find_dagruns_by_event(event, session) for dagrun in dagruns: dag_run_id = DagRunId(dagrun.dag_id, dagrun.run_id) self.task_event_manager.handle_event(dag_run_id, event) elif isinstance(event, RequestEvent): self._process_request_event(event) elif isinstance(event, ResponseEvent): continue elif isinstance(event, TaskSchedulingEvent): self._schedule_task(event) elif isinstance(event, TaskStatusChangedEvent): dagrun = self._find_dagrun(event.dag_id, event.execution_date, session) tasks = self._find_schedulable_tasks(dagrun, session) self._send_scheduling_task_events(tasks, SchedulingAction.START) elif isinstance(event, DagExecutableEvent): dagrun = self._create_dag_run(event.dag_id, session=session) tasks = self._find_schedulable_tasks(dagrun, session) self._send_scheduling_task_events(tasks, SchedulingAction.START) elif isinstance(event, EventHandleEvent): dag_runs = DagRun.find(dag_id=event.dag_id, run_id=event.dag_run_id) assert len(dag_runs) == 1 ti = dag_runs[0].get_task_instance(event.task_id) self._send_scheduling_task_event(ti, event.action) elif isinstance(event, StopDagEvent): self._stop_dag(event.dag_id, session) elif isinstance(event, ParseDagRequestEvent) or isinstance( event, ParseDagResponseEvent): pass elif isinstance(event, StopSchedulerEvent): self.log.info("{} {}".format(self.id, event.job_id)) if self.id == event.job_id or 0 == event.job_id: self.log.info("break the scheduler event loop.") identified_message.remove_handled_message() session.expunge_all() break else: self.log.error( "can not handler the event {}".format(event)) identified_message.remove_handled_message() session.expunge_all() self._stop_timer()
def _remove_periodic_events(self, run_id, session=None): dagruns = DagRun.find(run_id=run_id) dag = self.dagbag.get_dag(dag_id=dagruns[0].dag_id, session=session) for task in dag.tasks: if task.executor_config is not None and 'periodic_config' in task.executor_config: self.log.debug('remove periodic task {} {}'.format( run_id, task.task_id)) self.periodic_manager.remove_task(run_id, task.task_id)
def _stop_dag(self, dag_id, session: Session): """ Stop the dag. Pause the dag and cancel all running dag_runs and task_instances. """ DagModel.get_dagmodel(dag_id, session)\ .set_is_paused(is_paused=True, including_subdags=True, session=session) active_runs = DagRun.find(dag_id=dag_id, state=State.RUNNING) for dag_run in active_runs: self._stop_dag_run(dag_run)
def _get_dag_run(self, dagrun_info: DagRunInfo, dag: DAG, session: Session = None): """ Returns a dag run for the given run date, which will be matched to an existing dag run if available or create a new dag run otherwise. If the max_active_runs limit is reached, this function will return None. :param dagrun_info: Schedule information for the dag run :param dag: DAG :param session: the database session object :return: a DagRun in state RUNNING or None """ run_date = dagrun_info.logical_date # consider max_active_runs but ignore when running subdags respect_dag_max_active_limit = bool(dag.timetable.can_run and not dag.is_subdag) current_active_dag_count = dag.get_num_active_runs(external_trigger=False) # check if we are scheduling on top of a already existing dag_run # we could find a "scheduled" run instead of a "backfill" runs = DagRun.find(dag_id=dag.dag_id, execution_date=run_date, session=session) run: Optional[DagRun] if runs: run = runs[0] if run.state == DagRunState.RUNNING: respect_dag_max_active_limit = False # Fixes --conf overwrite for backfills with already existing DagRuns run.conf = self.conf or {} else: run = None # enforce max_active_runs limit for dag, special cases already # handled by respect_dag_max_active_limit if respect_dag_max_active_limit and current_active_dag_count >= dag.max_active_runs: return None run = run or dag.create_dagrun( execution_date=run_date, data_interval=dagrun_info.data_interval, start_date=timezone.utcnow(), state=DagRunState.RUNNING, external_trigger=False, session=session, conf=self.conf, run_type=DagRunType.BACKFILL_JOB, creating_job_id=self.id, ) # set required transient field run.dag = dag # explicitly mark as backfill and running run.state = DagRunState.RUNNING run.run_type = DagRunType.BACKFILL_JOB run.verify_integrity(session=session) return run
def test_backfill_max_limit_check_complete_loop(self): dag = self._get_dag_test_max_active_limits( 'test_backfill_max_limit_check_complete_loop') start_date = DEFAULT_DATE - datetime.timedelta(hours=1) end_date = DEFAULT_DATE # Given the max limit to be 1 in active dag runs, we need to run the # backfill job 3 times success_expected = 2 executor = MockExecutor() job = BackfillJob(dag=dag, start_date=start_date, end_date=end_date, executor=executor, donot_pickle=True) job.run() success_dagruns = len(DagRun.find(dag_id=dag.dag_id, state=State.SUCCESS)) running_dagruns = len(DagRun.find(dag_id=dag.dag_id, state=State.RUNNING)) self.assertEqual(success_expected, success_dagruns) self.assertEqual(0, running_dagruns) # no dag_runs in running state are left
def verify_dag_run_integrity(dag, dates): """ Verify the integrity of the dag runs in case a task was added or removed set the confirmed execution dates as they might be different from what was provided """ confirmed_dates = [] dag_runs = DagRun.find(dag_id=dag.dag_id, execution_date=dates) for dag_run in dag_runs: dag_run.dag = dag dag_run.verify_integrity() confirmed_dates.append(dag_run.execution_date) return confirmed_dates
def test_backfill_conf(self): dag = self._get_dummy_dag('test_backfill_conf') executor = MockExecutor() conf_ = json.loads("""{"key": "value"}""") job = BackfillJob(dag=dag, executor=executor, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE + datetime.timedelta(days=2), conf=conf_) job.run() dr = DagRun.find(dag_id='test_backfill_conf') self.assertEqual(conf_, dr[0].conf)
def test_sub_set_subdag(self): dag = DAG('test_sub_set_subdag', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}) with dag: op1 = DummyOperator(task_id='leave1') op2 = DummyOperator(task_id='leave2') op3 = DummyOperator(task_id='upstream_level_1') op4 = DummyOperator(task_id='upstream_level_2') op5 = DummyOperator(task_id='upstream_level_3') # order randomly op2.set_downstream(op3) op1.set_downstream(op3) op4.set_downstream(op5) op3.set_downstream(op4) dag.clear() dr = dag.create_dagrun(run_id="test", state=State.RUNNING, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE) executor = MockExecutor() sub_dag = dag.sub_dag(task_regex="leave*", include_downstream=False, include_upstream=False) job = BackfillJob(dag=sub_dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, executor=executor) job.run() self.assertRaises(sqlalchemy.orm.exc.NoResultFound, dr.refresh_from_db) # the run_id should have changed, so a refresh won't work drs = DagRun.find(dag_id=dag.dag_id, execution_date=DEFAULT_DATE) dr = drs[0] self.assertEqual( DagRun.generate_run_id(DagRunType.BACKFILL_JOB, DEFAULT_DATE), dr.run_id) for ti in dr.get_task_instances(): if ti.task_id == 'leave1' or ti.task_id == 'leave2': self.assertEqual(State.SUCCESS, ti.state) else: self.assertEqual(State.NONE, ti.state)
def test_backfill_max_limit_check_within_limit(self): dag = self._get_dag_test_max_active_limits( 'test_backfill_max_limit_check_within_limit', max_active_runs=16) start_date = DEFAULT_DATE - datetime.timedelta(hours=1) end_date = DEFAULT_DATE executor = MockExecutor() job = BackfillJob(dag=dag, start_date=start_date, end_date=end_date, executor=executor, donot_pickle=True) job.run() dagruns = DagRun.find(dag_id=dag.dag_id) self.assertEqual(2, len(dagruns)) self.assertTrue(all(run.state == State.SUCCESS for run in dagruns))
def getDagRunTasks(): dag_run = DagRun.find(dag_id='dentaway_ETL', execution_date=dag.latest_execution_date) instances = dag_run[0].get_task_instances() inst = [] for x in instances: if x.task_id != 'get_dag_run_tasks' and x.task_id.startswith( 'save_dagrun') == False: duration = x.duration if x.duration == None: duration = 0.0 inst.append({ 'task_id': x.task_id, 'duration': duration, 'state': x.state }) Variable.set("instances", inst) return inst
def poke(self, context): """ Implementation of the Airflow interface to check if the DAG should proceed. """ dag_runs = DagRun.find(dag_id=self.dag_name) length = len(dag_runs) # Query the latest start date of the DAG last_start_date = dag_runs[length - 1].start_date.replace(tzinfo=None) update = False dvc = DVCHook(self.dvc_repo) # Check modification dates of the given files for file in self.files: print( f"Current date = {last_start_date} vs. file modified date {dvc.modified_date(file)}" ) if dvc.modified_date(file) >= last_start_date: update = True break return update
def _create_task_instances(self, dag_run, session=None): """ This method schedules the tasks for a single DAG by looking at the active DAG runs and adding task instances that should run to the queue. """ # update the state of the previously active dag runs dag_runs = DagRun.find(dag_id=dag_run.dag_id, state=State.RUNNING, session=session) active_dag_runs = [] for run in dag_runs: self.log.info("Examining DAG run %s", run) # don't consider runs that are executed in the future unless # specified by config and schedule_interval is None if run.execution_date > timezone.utcnow( ) and not dag_run.dag.allow_future_exec_dates: self.log.error("Execution date is in future: %s", run.execution_date) continue if len(active_dag_runs) >= dag_run.dag.max_active_runs: self.log.info( "Number of active dag runs reached max_active_run.") break # skip backfill dagruns for now as long as they are not really scheduled if run.is_backfill: continue run.dag = dag_run.dag # todo: preferably the integrity check happens at dag collection time run.verify_integrity(session=session) run.update_state(session=session) if run.state == State.RUNNING: make_transient(run) active_dag_runs.append(run)
def _execute(self, session=None): """ Initializes all components required to run a dag for a specified date range and calls helper method to execute the tasks. """ ti_status = BackfillJob._DagRunTaskStatus() start_date = self.bf_start_date # Get DagRun schedule between the start/end dates, which will turn into dag runs. dagrun_start_date = timezone.coerce_datetime(start_date) if self.bf_end_date is None: dagrun_end_date = pendulum.now(timezone.utc) else: dagrun_end_date = pendulum.instance(self.bf_end_date) dagrun_infos = list( self.dag.iter_dagrun_infos_between(dagrun_start_date, dagrun_end_date)) if self.run_backwards: tasks_that_depend_on_past = [ t.task_id for t in self.dag.task_dict.values() if t.depends_on_past ] if tasks_that_depend_on_past: raise AirflowException( f'You cannot backfill backwards because one or more ' f'tasks depend_on_past: {",".join(tasks_that_depend_on_past)}' ) dagrun_infos = dagrun_infos[::-1] if not dagrun_infos: if not self.run_at_least_once: self.log.info( "No run dates were found for the given dates and dag interval." ) return dagrun_infos = [ DagRunInfo.interval(dagrun_start_date, dagrun_end_date) ] dag_with_subdags_ids = [d.dag_id for d in self._get_dag_with_subdags()] running_dagruns = DagRun.find( dag_id=dag_with_subdags_ids, execution_start_date=self.bf_start_date, execution_end_date=self.bf_end_date, no_backfills=True, state=DagRunState.RUNNING, ) if running_dagruns: for run in running_dagruns: self.log.error( "Backfill cannot be created for DagRun %s in %s, as there's already %s in a RUNNING " "state.", run.run_id, run.execution_date.strftime("%Y-%m-%dT%H:%M:%S"), run.run_type, ) self.log.error( "Changing DagRun into BACKFILL would cause scheduler to lose track of executing " "tasks. Not changing DagRun type into BACKFILL, and trying insert another DagRun into " "database would cause database constraint violation for dag_id + execution_date " "combination. Please adjust backfill dates or wait for this DagRun to finish.", ) return # picklin' pickle_id = None if not self.donot_pickle and self.executor_class not in ( executor_constants.LOCAL_EXECUTOR, executor_constants.SEQUENTIAL_EXECUTOR, executor_constants.DASK_EXECUTOR, ): pickle = DagPickle(self.dag) session.add(pickle) session.commit() pickle_id = pickle.id executor = self.executor executor.job_id = "backfill" executor.start() ti_status.total_runs = len(dagrun_infos) # total dag runs in backfill try: remaining_dates = ti_status.total_runs while remaining_dates > 0: dagrun_infos_to_process = [ dagrun_info for dagrun_info in dagrun_infos if dagrun_info.logical_date not in ti_status.executed_dag_run_dates ] self._execute_dagruns( dagrun_infos=dagrun_infos_to_process, ti_status=ti_status, executor=executor, pickle_id=pickle_id, start_date=start_date, session=session, ) remaining_dates = ti_status.total_runs - len( ti_status.executed_dag_run_dates) err = self._collect_errors(ti_status=ti_status, session=session) if err: if not self.continue_on_failures or ti_status.deadlocked: raise BackfillUnfinished(err, ti_status) if remaining_dates > 0: self.log.info( "max_active_runs limit for dag %s has been reached " " - waiting for other dag runs to finish", self.dag_id, ) time.sleep(self.delay_on_limit_secs) except (KeyboardInterrupt, SystemExit): self.log.warning("Backfill terminated by user.") # TODO: we will need to terminate running task instances and set the # state to failed. self._set_unfinished_dag_runs_to_failed(ti_status.active_runs) finally: session.commit() executor.end() self.log.info("Backfill done for DAG %s. Exiting.", self.dag)
def get_dag_run_state(dag_id: str, run_id: str): return DagRun.find(dag_id=dag_id, run_id=run_id)[0].state
def schedule(self) -> bool: identified_message = self.mailbox.get_identified_message() if not identified_message: return True origin_event = identified_message.deserialize() self.log.debug("Event: {}".format(origin_event)) if SchedulerInnerEventUtil.is_inner_event(origin_event): event = SchedulerInnerEventUtil.to_inner_event(origin_event) else: event = origin_event with create_session() as session: if isinstance(event, BaseEvent): dagruns = self._find_dagruns_by_event(event, session) for dagrun in dagruns: dag_run_id = DagRunId(dagrun.dag_id, dagrun.run_id) self.task_event_manager.handle_event(dag_run_id, event) elif isinstance(event, RequestEvent): self._process_request_event(event) elif isinstance(event, TaskSchedulingEvent): self._schedule_task(event) elif isinstance(event, TaskStateChangedEvent): dagrun = self._find_dagrun(event.dag_id, event.execution_date, session) if dagrun is not None: dag_run_id = DagRunId(dagrun.dag_id, dagrun.run_id) self.task_event_manager.handle_event( dag_run_id, origin_event) tasks = self._find_scheduled_tasks(dagrun, session) self._send_scheduling_task_events(tasks, SchedulingAction.START) if dagrun.state in State.finished: self.mailbox.send_message( DagRunFinishedEvent(dagrun.run_id).to_event()) else: self.log.warning( "dagrun is None for dag_id:{} execution_date: {}". format(event.dag_id, event.execution_date)) elif isinstance(event, DagExecutableEvent): dagrun = self._create_dag_run(event.dag_id, session=session) tasks = self._find_scheduled_tasks(dagrun, session) self._send_scheduling_task_events(tasks, SchedulingAction.START) elif isinstance(event, EventHandleEvent): dag_runs = DagRun.find(dag_id=event.dag_id, run_id=event.dag_run_id) assert len(dag_runs) == 1 ti = dag_runs[0].get_task_instance(event.task_id) self._send_scheduling_task_event(ti, event.action) elif isinstance(event, StopDagEvent): self._stop_dag(event.dag_id, session) elif isinstance(event, DagRunFinishedEvent): self._remove_periodic_events(event.run_id) elif isinstance(event, PeriodicEvent): dag_runs = DagRun.find(run_id=event.run_id) assert len(dag_runs) == 1 ti = dag_runs[0].get_task_instance(event.task_id) self._send_scheduling_task_event(ti, SchedulingAction.RESTART) elif isinstance(event, StopSchedulerEvent): self.log.info("{} {}".format(self.id, event.job_id)) if self.id == event.job_id or 0 == event.job_id: self.log.info("break the scheduler event loop.") identified_message.remove_handled_message() session.expunge_all() return False elif isinstance(event, ParseDagRequestEvent) or isinstance( event, ParseDagResponseEvent): pass elif isinstance(event, ResponseEvent): pass else: self.log.error("can not handler the event {}".format(event)) identified_message.remove_handled_message() session.expunge_all() return True
def test_backfill_max_limit_check(self): dag_id = 'test_backfill_max_limit_check' run_id = 'test_dagrun' start_date = DEFAULT_DATE - datetime.timedelta(hours=1) end_date = DEFAULT_DATE dag_run_created_cond = threading.Condition() def run_backfill(cond): cond.acquire() # this session object is different than the one in the main thread with create_session() as thread_session: try: dag = self._get_dag_test_max_active_limits(dag_id) # Existing dagrun that is not within the backfill range dag.create_dagrun( run_id=run_id, state=State.RUNNING, execution_date=DEFAULT_DATE + datetime.timedelta(hours=1), start_date=DEFAULT_DATE, ) thread_session.commit() cond.notify() finally: cond.release() thread_session.close() executor = MockExecutor() job = BackfillJob(dag=dag, start_date=start_date, end_date=end_date, executor=executor, donot_pickle=True) job.run() backfill_job_thread = threading.Thread(target=run_backfill, name="run_backfill", args=(dag_run_created_cond, )) dag_run_created_cond.acquire() with create_session() as session: backfill_job_thread.start() try: # at this point backfill can't run since the max_active_runs has been # reached, so it is waiting dag_run_created_cond.wait(timeout=1.5) dagruns = DagRun.find(dag_id=dag_id) dr = dagruns[0] self.assertEqual(1, len(dagruns)) self.assertEqual(dr.run_id, run_id) # allow the backfill to execute # by setting the existing dag run to SUCCESS, # backfill will execute dag runs 1 by 1 dr.set_state(State.SUCCESS) session.merge(dr) session.commit() backfill_job_thread.join() dagruns = DagRun.find(dag_id=dag_id) self.assertEqual(3, len(dagruns)) # 2 from backfill + 1 existing self.assertEqual(dagruns[-1].run_id, dr.run_id) finally: dag_run_created_cond.release()
def test_backfill_fill_blanks(self): dag = DAG( 'test_backfill_fill_blanks', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'}, ) with dag: op1 = DummyOperator(task_id='op1') op2 = DummyOperator(task_id='op2') op3 = DummyOperator(task_id='op3') op4 = DummyOperator(task_id='op4') op5 = DummyOperator(task_id='op5') op6 = DummyOperator(task_id='op6') dag.clear() dr = dag.create_dagrun(run_id='test', state=State.RUNNING, execution_date=DEFAULT_DATE, start_date=DEFAULT_DATE) executor = MockExecutor() session = settings.Session() tis = dr.get_task_instances() for ti in tis: if ti.task_id == op1.task_id: ti.state = State.UP_FOR_RETRY ti.end_date = DEFAULT_DATE elif ti.task_id == op2.task_id: ti.state = State.FAILED elif ti.task_id == op3.task_id: ti.state = State.SKIPPED elif ti.task_id == op4.task_id: ti.state = State.SCHEDULED elif ti.task_id == op5.task_id: ti.state = State.UPSTREAM_FAILED # op6 = None session.merge(ti) session.commit() session.close() job = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, executor=executor) self.assertRaisesRegex(AirflowException, 'Some task instances failed', job.run) self.assertRaises(sqlalchemy.orm.exc.NoResultFound, dr.refresh_from_db) # the run_id should have changed, so a refresh won't work drs = DagRun.find(dag_id=dag.dag_id, execution_date=DEFAULT_DATE) dr = drs[0] self.assertEqual(dr.state, State.FAILED) tis = dr.get_task_instances() for ti in tis: if ti.task_id in (op1.task_id, op4.task_id, op6.task_id): self.assertEqual(ti.state, State.SUCCESS) elif ti.task_id == op2.task_id: self.assertEqual(ti.state, State.FAILED) elif ti.task_id == op3.task_id: self.assertEqual(ti.state, State.SKIPPED) elif ti.task_id == op5.task_id: self.assertEqual(ti.state, State.UPSTREAM_FAILED)
def execute(self, context: Context): if isinstance(self.execution_date, datetime.datetime): parsed_execution_date = self.execution_date elif isinstance(self.execution_date, str): parsed_execution_date = timezone.parse(self.execution_date) else: parsed_execution_date = timezone.utcnow() if self.trigger_run_id: run_id = self.trigger_run_id else: run_id = DagRun.generate_run_id(DagRunType.MANUAL, parsed_execution_date) try: dag_run = trigger_dag( dag_id=self.trigger_dag_id, run_id=run_id, conf=self.conf, execution_date=parsed_execution_date, replace_microseconds=False, ) except DagRunAlreadyExists as e: if self.reset_dag_run: self.log.info("Clearing %s on %s", self.trigger_dag_id, parsed_execution_date) # Get target dag object and call clear() dag_model = DagModel.get_current(self.trigger_dag_id) if dag_model is None: raise DagNotFound( f"Dag id {self.trigger_dag_id} not found in DagModel") dag_bag = DagBag(dag_folder=dag_model.fileloc, read_dags_from_db=True) dag = dag_bag.get_dag(self.trigger_dag_id) dag.clear(start_date=parsed_execution_date, end_date=parsed_execution_date) dag_run = DagRun.find(dag_id=dag.dag_id, run_id=run_id)[0] else: raise e if dag_run is None: raise RuntimeError("The dag_run should be set here!") # Store the execution date from the dag run (either created or found above) to # be used when creating the extra link on the webserver. ti = context['task_instance'] ti.xcom_push(key=XCOM_EXECUTION_DATE_ISO, value=dag_run.execution_date.isoformat()) ti.xcom_push(key=XCOM_RUN_ID, value=dag_run.run_id) if self.wait_for_completion: # wait for dag to complete while True: self.log.info( 'Waiting for %s on %s to become allowed state %s ...', self.trigger_dag_id, dag_run.execution_date, self.allowed_states, ) time.sleep(self.poke_interval) dag_run.refresh_from_db() state = dag_run.state if state in self.failed_states: raise AirflowException( f"{self.trigger_dag_id} failed with failed states {state}" ) if state in self.allowed_states: self.log.info("%s finished with allowed state %s", self.trigger_dag_id, state) return
def _get_dagrun(self, execution_date): dag_runs = DagRun.find( dag_id=self.subdag.dag_id, execution_date=execution_date, ) return dag_runs[0] if dag_runs else None