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