def test_backfill_rerun_failed_tasks_without_flag(self):
dag = DAG(
dag_id='test_backfill_rerun_failed',
start_date=DEFAULT_DATE,
schedule_interval='@daily')
with dag:
DummyOperator(
task_id='test_backfill_rerun_failed_task-1',
dag=dag)
dag.clear()
executor = TestExecutor()
job = BackfillJob(dag=dag,
executor=executor,
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=2),
)
job.run()
ti = TI(task=dag.get_task('test_backfill_rerun_failed_task-1'),
execution_date=DEFAULT_DATE)
ti.refresh_from_db()
ti.set_state(State.FAILED)
job = BackfillJob(dag=dag,
executor=executor,
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=2),
rerun_failed_tasks=False
)
with self.assertRaises(AirflowException):
job.run()
def test_retry_delay(self):
"""
Test that retry delays are respected
"""
dag = models.DAG(dag_id='test_retry_handling')
task = BashOperator(
task_id='test_retry_handling_op',
bash_command='exit 1',
retries=1,
retry_delay=datetime.timedelta(seconds=3),
dag=dag,
owner='airflow',
start_date=datetime.datetime(2016, 2, 1, 0, 0, 0))
def run_with_error(ti):
try:
ti.run()
except AirflowException:
pass
ti = TI(
task=task, execution_date=datetime.datetime.now())
# first run -- up for retry
run_with_error(ti)
self.assertEqual(ti.state, State.UP_FOR_RETRY)
self.assertEqual(ti.try_number, 1)
# second run -- still up for retry because retry_delay hasn't expired
run_with_error(ti)
self.assertEqual(ti.state, State.UP_FOR_RETRY)
# third run -- failed
time.sleep(3)
run_with_error(ti)
self.assertEqual(ti.state, State.FAILED)
def test_xcom_pull_after_success(self):
"""
tests xcom set/clear relative to a task in a 'success' rerun scenario
"""
key = 'xcom_key'
value = 'xcom_value'
dag = models.DAG(dag_id='test_xcom', schedule_interval='@monthly')
task = DummyOperator(
task_id='test_xcom',
dag=dag,
pool='test_xcom',
owner='airflow',
start_date=datetime.datetime(2016, 6, 2, 0, 0, 0))
exec_date = datetime.datetime.now()
ti = TI(
task=task, execution_date=exec_date)
ti.run(mark_success=True)
ti.xcom_push(key=key, value=value)
self.assertEqual(ti.xcom_pull(task_ids='test_xcom', key=key), value)
ti.run()
# The second run and assert is to handle AIRFLOW-131 (don't clear on
# prior success)
self.assertEqual(ti.xcom_pull(task_ids='test_xcom', key=key), value)
def test_backfill_depends_on_past_backwards(self):
"""
Test that CLI respects -B argument and raises on interaction with depends_on_past
"""
dag_id = 'test_depends_on_past'
start_date = DEFAULT_DATE + datetime.timedelta(days=1)
end_date = start_date + datetime.timedelta(days=1)
kwargs = dict(
start_date=start_date,
end_date=end_date,
)
dag = self.dagbag.get_dag(dag_id)
dag.clear()
executor = MockExecutor()
job = BackfillJob(dag=dag,
executor=executor,
ignore_first_depends_on_past=True,
**kwargs)
job.run()
ti = TI(dag.get_task('test_dop_task'), end_date)
ti.refresh_from_db()
# runs fine forwards
self.assertEqual(ti.state, State.SUCCESS)
# raises backwards
expected_msg = 'You cannot backfill backwards because one or more tasks depend_on_past: {}'.format(
'test_dop_task')
with self.assertRaisesRegex(AirflowException, expected_msg):
executor = MockExecutor()
job = BackfillJob(dag=dag,
executor=executor,
run_backwards=True,
**kwargs)
job.run()
def test_cli_backfill_depends_on_past(self):
"""
Test that CLI respects -I argument
"""
dag_id = 'test_dagrun_states_deadlock'
run_date = DEFAULT_DATE + datetime.timedelta(days=1)
args = [
'backfill',
dag_id,
'-l',
'-s',
run_date.isoformat(),
]
dag = self.dagbag.get_dag(dag_id)
dag.clear()
self.assertRaisesRegexp(AirflowException, 'BackfillJob is deadlocked',
cli.backfill, self.parser.parse_args(args))
cli.backfill(self.parser.parse_args(args + ['-I']))
ti = TI(dag.get_task('test_depends_on_past'), run_date)
ti.refresh_from_db()
# task ran
self.assertEqual(ti.state, State.SUCCESS)
def test_scheduler_pooled_tasks(self):
"""
Test that the scheduler handles queued tasks correctly
See issue #1299
"""
session = settings.Session()
if not (session.query(Pool).filter(
Pool.pool == 'test_queued_pool').first()):
pool = Pool(pool='test_queued_pool', slots=5)
session.merge(pool)
session.commit()
session.close()
dag_id = 'test_scheduled_queued_tasks'
dag = self.dagbag.get_dag(dag_id)
dag.clear()
scheduler = SchedulerJob(dag_id, num_runs=1)
scheduler.run()
task_1 = dag.tasks[0]
logging.info("Trying to find task {}".format(task_1))
ti = TI(task_1, dag.start_date)
ti.refresh_from_db()
self.assertEqual(ti.state, State.QUEUED)
# now we use a DIFFERENT scheduler and executor
# to simulate the num-runs CLI arg
scheduler2 = SchedulerJob(dag_id,
num_runs=5,
executor=DEFAULT_EXECUTOR.__class__())
scheduler2.run()
ti.refresh_from_db()
self.assertEqual(ti.state, State.FAILED)
dag.clear()
def test_backfill_depends_on_past(self):
"""
Test that backfill respects ignore_depends_on_past
"""
dag = self.dagbag.get_dag('test_depends_on_past')
dag.clear()
run_date = DEFAULT_DATE + datetime.timedelta(days=5)
# backfill should deadlock
self.assertRaisesRegexp(
AirflowException,
'BackfillJob is deadlocked',
BackfillJob(dag=dag, start_date=run_date, end_date=run_date).run)
BackfillJob(
dag=dag,
start_date=run_date,
end_date=run_date,
ignore_first_depends_on_past=True).run()
# ti should have succeeded
ti = TI(dag.tasks[0], run_date)
ti.refresh_from_db()
self.assertEquals(ti.state, State.SUCCESS)
def test_localtaskjob_double_trigger(self):
dagbag = models.DagBag(
dag_folder=TEST_DAG_FOLDER,
include_examples=False,
)
dag = dagbag.dags.get('test_localtaskjob_double_trigger')
task = dag.get_task('test_localtaskjob_double_trigger_task')
session = settings.Session()
dag.clear()
dr = dag.create_dagrun(run_id="test",
state=State.SUCCESS,
execution_date=DEFAULT_DATE,
start_date=DEFAULT_DATE,
session=session)
ti = dr.get_task_instance(task_id=task.task_id, session=session)
ti.state = State.RUNNING
ti.hostname = get_hostname()
ti.pid = 1
session.commit()
ti_run = TI(task=task, execution_date=DEFAULT_DATE)
job1 = LocalTaskJob(task_instance=ti_run,
ignore_ti_state=True,
executor=SequentialExecutor())
with patch.object(BaseTaskRunner, 'start',
return_value=None) as mock_method:
job1.run()
mock_method.assert_not_called()
ti = dr.get_task_instance(task_id=task.task_id, session=session)
self.assertEqual(ti.pid, 1)
self.assertEqual(ti.state, State.RUNNING)
session.close()
def test_kill_zombies_doesn_nothing(self, mock_ti_handle_failure):
"""
Test that kill zombies does nothing when job is running and received heartbeat
"""
dagbag = models.DagBag(dag_folder=self.empty_dir, include_examples=True)
with create_session() as session:
session.query(TI).delete()
session.query(LJ).delete()
dag = dagbag.get_dag('example_branch_operator')
task = dag.get_task(task_id='run_this_first')
ti = TI(task, DEFAULT_DATE, State.RUNNING)
lj = LJ(ti)
lj.latest_heartbeat = utcnow()
lj.state = State.RUNNING
lj.id = 1
ti.job_id = lj.id
session.add(lj)
session.add(ti)
session.commit()
dagbag.kill_zombies()
mock_ti_handle_failure.assert_not_called()
def test_set_duration_empty_dates(self):
task = DummyOperator(task_id='op', email='*****@*****.**')
ti = TI(task=task, execution_date=datetime.datetime.now())
ti.set_duration()
self.assertIsNone(ti.duration)
def test_reschedule_handling(self, mock_pool_full):
"""
Test that task reschedules are handled properly
"""
# Mock the pool with a pool with slots open since the pool doesn't actually exist
mock_pool_full.return_value = False
# Return values of the python sensor callable, modified during tests
done = False
fail = False
def callable():
if fail:
raise AirflowException()
return done
dag = models.DAG(dag_id='test_reschedule_handling')
task = PythonSensor(task_id='test_reschedule_handling_sensor',
poke_interval=0,
mode='reschedule',
python_callable=callable,
retries=1,
retry_delay=datetime.timedelta(seconds=0),
dag=dag,
owner='airflow',
start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
ti = TI(task=task, execution_date=timezone.utcnow())
self.assertEqual(ti._try_number, 0)
self.assertEqual(ti.try_number, 1)
def run_ti_and_assert(run_date, expected_start_date, expected_end_date,
expected_duration, expected_state,
expected_try_number,
expected_task_reschedule_count):
with freeze_time(run_date):
try:
ti.run()
except AirflowException:
if not fail:
raise
ti.refresh_from_db()
self.assertEqual(ti.state, expected_state)
self.assertEqual(ti._try_number, expected_try_number)
self.assertEqual(ti.try_number, expected_try_number + 1)
self.assertEqual(ti.start_date, expected_start_date)
self.assertEqual(ti.end_date, expected_end_date)
self.assertEqual(ti.duration, expected_duration)
trs = TaskReschedule.find_for_task_instance(ti)
self.assertEqual(len(trs), expected_task_reschedule_count)
date1 = timezone.utcnow()
date2 = date1 + datetime.timedelta(minutes=1)
date3 = date2 + datetime.timedelta(minutes=1)
date4 = date3 + datetime.timedelta(minutes=1)
# Run with multiple reschedules.
# During reschedule the try number remains the same, but each reschedule is recorded.
# The start date is expected to remain the initial date, hence the duration increases.
# When finished the try number is incremented and there is no reschedule expected
# for this try.
done, fail = False, False
run_ti_and_assert(date1, date1, date1, 0, State.UP_FOR_RESCHEDULE, 0,
1)
done, fail = False, False
run_ti_and_assert(date2, date1, date2, 60, State.UP_FOR_RESCHEDULE, 0,
2)
done, fail = False, False
run_ti_and_assert(date3, date1, date3, 120, State.UP_FOR_RESCHEDULE, 0,
3)
done, fail = True, False
run_ti_and_assert(date4, date1, date4, 180, State.SUCCESS, 1, 0)
# Clear the task instance.
dag.clear()
ti.refresh_from_db()
self.assertEqual(ti.state, State.NONE)
self.assertEqual(ti._try_number, 1)
# Run again after clearing with reschedules and a retry.
# The retry increments the try number, and for that try no reschedule is expected.
# After the retry the start date is reset, hence the duration is also reset.
done, fail = False, False
run_ti_and_assert(date1, date1, date1, 0, State.UP_FOR_RESCHEDULE, 1,
1)
done, fail = False, True
run_ti_and_assert(date2, date1, date2, 60, State.UP_FOR_RETRY, 2, 0)
done, fail = False, False
run_ti_and_assert(date3, date3, date3, 0, State.UP_FOR_RESCHEDULE, 2,
1)
done, fail = True, False
run_ti_and_assert(date4, date3, date4, 60, State.SUCCESS, 3, 0)
def test_handle_failure_callback_with_zombies_are_correctly_passed_to_dag_file_processor(
self):
"""
Check that the same set of failure callback with zombies are passed to the dag
file processors until the next zombie detection logic is invoked.
"""
test_dag_path = os.path.join(TEST_DAG_FOLDER,
'test_example_bash_operator.py')
with conf_vars({
('scheduler', 'max_threads'): '1',
('core', 'load_examples'): 'False'
}):
dagbag = DagBag(test_dag_path)
with create_session() as session:
session.query(LJ).delete()
dag = dagbag.get_dag('test_example_bash_operator')
dag.sync_to_db()
task = dag.get_task(task_id='run_this_last')
ti = TI(task, DEFAULT_DATE, State.RUNNING)
local_job = LJ(ti)
local_job.state = State.SHUTDOWN
session.add(local_job)
session.commit()
# TODO: If there was an actual Relationshop between TI and Job
# we wouldn't need this extra commit
session.add(ti)
ti.job_id = local_job.id
session.commit()
fake_failure_callback_requests = [
FailureCallbackRequest(
full_filepath=dag.full_filepath,
simple_task_instance=SimpleTaskInstance(ti),
msg="Message")
]
test_dag_path = os.path.join(TEST_DAG_FOLDER,
'test_example_bash_operator.py')
child_pipe, parent_pipe = multiprocessing.Pipe()
async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn')
manager = DagFileProcessorManager(
dag_directory=test_dag_path,
max_runs=1,
processor_factory=FakeDagFileProcessorRunner.
_fake_dag_processor_factory,
processor_timeout=timedelta.max,
signal_conn=child_pipe,
dag_ids=[],
pickle_dags=False,
async_mode=async_mode)
parsing_result = self.run_processor_manager_one_loop(
manager, parent_pipe)
self.assertEqual(len(fake_failure_callback_requests),
len(parsing_result))
self.assertEqual(
set(zombie.simple_task_instance.key
for zombie in fake_failure_callback_requests),
set(result.simple_task_instance.key
for result in parsing_result))
child_pipe.close()
parent_pipe.close()
def test_lineage(self):
dag = DAG(dag_id='test_prepare_lineage', start_date=DEFAULT_DATE)
f1s = "/tmp/does_not_exist_1-{}"
f2s = "/tmp/does_not_exist_2-{}"
f3s = "/tmp/does_not_exist_3"
file1 = File(f1s.format("{{ execution_date }}"))
file2 = File(f2s.format("{{ execution_date }}"))
file3 = File(f3s)
with dag:
op1 = DummyOperator(
task_id='leave1',
inlets=file1,
outlets=[
file2,
],
)
op2 = DummyOperator(task_id='leave2')
op3 = DummyOperator(task_id='upstream_level_1',
inlets=AUTO,
outlets=file3)
op4 = DummyOperator(task_id='upstream_level_2')
op5 = DummyOperator(task_id='upstream_level_3',
inlets=["leave1", "upstream_level_1"])
op1.set_downstream(op3)
op2.set_downstream(op3)
op3.set_downstream(op4)
op4.set_downstream(op5)
dag.clear()
# execution_date is set in the context in order to avoid creating task instances
ctx1 = {
"ti": TI(task=op1, execution_date=DEFAULT_DATE),
"execution_date": DEFAULT_DATE
}
ctx2 = {
"ti": TI(task=op2, execution_date=DEFAULT_DATE),
"execution_date": DEFAULT_DATE
}
ctx3 = {
"ti": TI(task=op3, execution_date=DEFAULT_DATE),
"execution_date": DEFAULT_DATE
}
ctx5 = {
"ti": TI(task=op5, execution_date=DEFAULT_DATE),
"execution_date": DEFAULT_DATE
}
# prepare with manual inlets and outlets
op1.pre_execute(ctx1)
assert len(op1.inlets) == 1
assert op1.inlets[0].url == f1s.format(DEFAULT_DATE)
assert len(op1.outlets) == 1
assert op1.outlets[0].url == f2s.format(DEFAULT_DATE)
# post process with no backend
op1.post_execute(ctx1)
op2.pre_execute(ctx2)
assert len(op2.inlets) == 0
op2.post_execute(ctx2)
op3.pre_execute(ctx3)
assert len(op3.inlets) == 1
assert op3.inlets[0].url == f2s.format(DEFAULT_DATE)
assert op3.outlets[0] == file3
op3.post_execute(ctx3)
# skip 4
op5.pre_execute(ctx5)
assert len(op5.inlets) == 2
op5.post_execute(ctx5)
def test_lineage(self, _get_backend):
backend = mock.Mock()
send_mock = mock.Mock()
backend.send_lineage = send_mock
_get_backend.return_value = backend
dag = DAG(
dag_id='test_prepare_lineage',
start_date=DEFAULT_DATE
)
f1 = File("/tmp/does_not_exist_1")
f2 = File("/tmp/does_not_exist_2")
f3 = File("/tmp/does_not_exist_3")
with dag:
op1 = DummyOperator(task_id='leave1',
inlets={"datasets": [f1, ]},
outlets={"datasets": [f2, ]})
op2 = DummyOperator(task_id='leave2')
op3 = DummyOperator(task_id='upstream_level_1',
inlets={"auto": True},
outlets={"datasets": [f3, ]})
op4 = DummyOperator(task_id='upstream_level_2')
op5 = DummyOperator(task_id='upstream_level_3',
inlets={"task_ids": ["leave1", "upstream_level_1"]})
op1.set_downstream(op3)
op2.set_downstream(op3)
op3.set_downstream(op4)
op4.set_downstream(op5)
ctx1 = {"ti": TI(task=op1, execution_date=DEFAULT_DATE)}
ctx2 = {"ti": TI(task=op2, execution_date=DEFAULT_DATE)}
ctx3 = {"ti": TI(task=op3, execution_date=DEFAULT_DATE)}
ctx5 = {"ti": TI(task=op5, execution_date=DEFAULT_DATE)}
func = mock.Mock()
func.__name__ = 'foo'
# prepare with manual inlets and outlets
prep = prepare_lineage(func)
prep(op1, ctx1)
self.assertEqual(len(op1.inlets), 1)
self.assertEqual(op1.inlets[0], f1)
self.assertEqual(len(op1.outlets), 1)
self.assertEqual(op1.outlets[0], f2)
# post process with no backend
post = apply_lineage(func)
post(op1, ctx1)
self.assertEqual(send_mock.call_count, 1)
send_mock.reset_mock()
prep(op2, ctx2)
self.assertEqual(len(op2.inlets), 0)
post(op2, ctx2)
self.assertEqual(send_mock.call_count, 1)
send_mock.reset_mock()
prep(op3, ctx3)
self.assertEqual(len(op3.inlets), 1)
self.assertEqual(op3.inlets[0].qualified_name, f2.qualified_name)
post(op3, ctx3)
self.assertEqual(send_mock.call_count, 1)
send_mock.reset_mock()
# skip 4
prep(op5, ctx5)
self.assertEqual(len(op5.inlets), 2)
post(op5, ctx5)
self.assertEqual(send_mock.call_count, 1)
send_mock.reset_mock()
def test_lineage_auto_branching(self, _get_backend):
# Tests the ability for the auto feature to skip non state affecting operators
# DAG diagram:
# 1--->2---->4
# ▼ ▲
# 3-----+
backend = mock.Mock()
send_mock = mock.Mock()
backend.send_lineage = send_mock
_get_backend.return_value = backend
dag = DAG(
dag_id='test_prepare_lineage_auto_branching',
start_date=DEFAULT_DATE
)
f1 = File("/tmp/does_not_exist_1")
with dag:
op1 = DummyOperator(task_id='leave1')
op2 = DummyOperator(task_id='branch_1', outlets={"datasets": [f1, ]})
op3 = DummyOperator(task_id='branch_2')
op4 = DummyOperator(task_id='upstream_level_2', inlets={"auto": True})
op1.set_downstream(op2)
op2.set_downstream(op3)
op2.set_downstream(op4)
op3.set_downstream(op4)
ctx1 = {"ti": TI(task=op1, execution_date=DEFAULT_DATE)}
ctx2 = {"ti": TI(task=op2, execution_date=DEFAULT_DATE)}
ctx3 = {"ti": TI(task=op3, execution_date=DEFAULT_DATE)}
ctx4 = {"ti": TI(task=op4, execution_date=DEFAULT_DATE)}
func = mock.Mock()
func.__name__ = 'foo'
# prepare with manual inlets and outlets
prep = prepare_lineage(func)
prep(op1, ctx1)
self.assertEqual(len(op1.inlets), 0)
# post process with no backend
post = apply_lineage(func)
post(op1, ctx1)
send_mock.reset_mock()
prep(op2, ctx2)
self.assertEqual(len(op2.inlets), 0)
post(op2, ctx2)
self.assertEqual(send_mock.call_count, 1)
send_mock.reset_mock()
prep(op3, ctx3)
self.assertEqual(len(op3.inlets), 0)
post(op3, ctx3)
self.assertEqual(send_mock.call_count, 1)
send_mock.reset_mock()
prep(op4, ctx4)
self.assertEqual(len(op4.inlets), 1)
self.assertEqual(op4.inlets[0].name, f1.name)
post(op4, ctx4)
self.assertEqual(send_mock.call_count, 1)
send_mock.reset_mock()
def test_lineage_complicated_dag(self, _get_backend):
# Tests the ability for the auto feature to skip non state affecting operators, while still
# retrieving data from multiple outlet sources. Notice how if outlets are not specified,
# that the auto feature continues to traverse down the dag until not input sources are found.
# DAG diagram:
# 1-----------+
# |
# ▼
# 4 ----------+
# ▲ ▼
# | 5+-------->6
# 2-----------+ ▲
# |
# |
# |
# 3-----------------------+
backend = mock.Mock()
send_mock = mock.Mock()
backend.send_lineage = send_mock
_get_backend.return_value = backend
dag = DAG(
dag_id='test_prepare_lineage_auto_complicated_dag',
start_date=DEFAULT_DATE
)
f1 = File("/tmp/does_not_exist_1")
f2 = File("/tmp/does_not_exist_2")
f3 = File("/tmp/does_not_exist_3")
with dag:
op1 = DummyOperator(task_id='leave1',
outlets={"datasets": [f1, ]},
inlets={"auto": True})
op2 = DummyOperator(task_id='leave2',
outlets={"datasets": [f2, ]})
op3 = DummyOperator(task_id='leave3',
outlets={"datasets": [f3, ]})
op4 = DummyOperator(task_id='upstream_level_1')
op5 = DummyOperator(task_id='upstream_level_2', inlets={"auto": True})
op6 = DummyOperator(task_id='upstream_level_3', inlets={"auto": True})
op1.set_downstream(op4)
op2.set_downstream(op4)
op3.set_downstream(op5)
op4.set_downstream(op5)
op5.set_downstream(op6)
ctx1 = {"ti": TI(task=op1, execution_date=DEFAULT_DATE)}
ctx2 = {"ti": TI(task=op2, execution_date=DEFAULT_DATE)}
ctx3 = {"ti": TI(task=op3, execution_date=DEFAULT_DATE)}
ctx4 = {"ti": TI(task=op4, execution_date=DEFAULT_DATE)}
ctx5 = {"ti": TI(task=op5, execution_date=DEFAULT_DATE)}
ctx6 = {"ti": TI(task=op6, execution_date=DEFAULT_DATE)}
func = mock.Mock()
func.__name__ = 'foo'
# prepare with manual inlets and outlets
prep = prepare_lineage(func)
prep(op1, ctx1)
self.assertEqual(len(op1.outlets), 1)
self.assertEqual(op1.outlets[0], f1)
self.assertEqual(len(op1.inlets), 0)
# post process with no backend
post = apply_lineage(func)
post(op1, ctx1)
prep(op2, ctx2)
self.assertEqual(len(op2.outlets), 1)
post(op2, ctx2)
prep(op3, ctx3)
self.assertEqual(len(op3.outlets), 1)
post(op3, ctx3)
prep(op4, ctx4)
self.assertEqual(len(op4.inlets), 0)
post(op4, ctx4)
prep(op5, ctx5)
self.assertEqual(len(op5.inlets), 3)
self.assertEqual({file.qualified_name for file in op5.inlets}, {'file:///tmp/does_not_exist_1',
'file:///tmp/does_not_exist_2',
'file:///tmp/does_not_exist_3'})
post(op5, ctx5)
prep(op6, ctx6)
self.assertEqual(len(op6.inlets), 3)
self.assertEqual({file.qualified_name for file in op6.inlets}, {'file:///tmp/does_not_exist_1',
'file:///tmp/does_not_exist_2',
'file:///tmp/does_not_exist_3'})
post(op6, ctx6)
def test_handle_failure_callback_with_zombies_are_correctly_passed_to_dag_file_processor(self):
"""
Check that the same set of failure callback with zombies are passed to the dag
file processors until the next zombie detection logic is invoked.
"""
test_dag_path = os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py')
with conf_vars({('scheduler', 'parsing_processes'): '1', ('core', 'load_examples'): 'False'}):
dagbag = DagBag(test_dag_path, read_dags_from_db=False)
with create_session() as session:
session.query(LJ).delete()
dag = dagbag.get_dag('test_example_bash_operator')
dag.sync_to_db()
task = dag.get_task(task_id='run_this_last')
ti = TI(task, DEFAULT_DATE, State.RUNNING)
local_job = LJ(ti)
local_job.state = State.SHUTDOWN
session.add(local_job)
session.commit()
# TODO: If there was an actual Relationship between TI and Job
# we wouldn't need this extra commit
session.add(ti)
ti.job_id = local_job.id
session.commit()
expected_failure_callback_requests = [
TaskCallbackRequest(
full_filepath=dag.full_filepath,
simple_task_instance=SimpleTaskInstance(ti),
msg="Message",
)
]
test_dag_path = os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py')
child_pipe, parent_pipe = multiprocessing.Pipe()
async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn')
fake_processors = []
def fake_processor_factory(*args, **kwargs):
nonlocal fake_processors
processor = FakeDagFileProcessorRunner._fake_dag_processor_factory(*args, **kwargs)
fake_processors.append(processor)
return processor
manager = DagFileProcessorManager(
dag_directory=test_dag_path,
max_runs=1,
processor_factory=fake_processor_factory,
processor_timeout=timedelta.max,
signal_conn=child_pipe,
dag_ids=[],
pickle_dags=False,
async_mode=async_mode,
)
self.run_processor_manager_one_loop(manager, parent_pipe)
if async_mode:
# Once for initial parse, and then again for the add_callback_to_queue
assert len(fake_processors) == 2
assert fake_processors[0]._file_path == test_dag_path
assert fake_processors[0]._callback_requests == []
else:
assert len(fake_processors) == 1
assert fake_processors[-1]._file_path == test_dag_path
callback_requests = fake_processors[-1]._callback_requests
assert {zombie.simple_task_instance.key for zombie in expected_failure_callback_requests} == {
result.simple_task_instance.key for result in callback_requests
}
child_pipe.close()
parent_pipe.close()
def test_handle_failure_callback_with_zobmies_are_correctly_passed_to_dag_file_processor(
self):
"""
Check that the same set of failure callback with zombies are passed to the dag
file processors until the next zombie detection logic is invoked.
"""
test_dag_path = os.path.join(TEST_DAG_FOLDER,
'test_example_bash_operator.py')
with conf_vars({
('scheduler', 'max_threads'): '1',
('core', 'load_examples'): 'False'
}):
dagbag = DagBag(test_dag_path)
with create_session() as session:
session.query(LJ).delete()
dag = dagbag.get_dag('test_example_bash_operator')
dag.sync_to_db()
task = dag.get_task(task_id='run_this_last')
ti = TI(task, DEFAULT_DATE, State.RUNNING)
local_job = LJ(ti)
local_job.state = State.SHUTDOWN
local_job.id = 1
ti.job_id = local_job.id
session.add(local_job)
session.add(ti)
session.commit()
fake_failure_callback_requests = [
FailureCallbackRequest(
full_filepath=dag.full_filepath,
simple_task_instance=SimpleTaskInstance(ti),
msg="Message")
]
class FakeDagFileProcessorRunner(DagFileProcessorProcess):
# This fake processor will return the zombies it received in constructor
# as its processing result w/o actually parsing anything.
def __init__(self, file_path, pickle_dags, dag_id_white_list,
failure_callback_requests):
super().__init__(file_path, pickle_dags, dag_id_white_list,
failure_callback_requests)
self._result = failure_callback_requests, 0
def start(self):
pass
@property
def start_time(self):
return DEFAULT_DATE
@property
def pid(self):
return 1234
@property
def done(self):
return True
@property
def result(self):
return self._result
def processor_factory(file_path, failure_callback_requests):
return FakeDagFileProcessorRunner(file_path, False, [],
failure_callback_requests)
async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn')
processor_agent = DagFileProcessorAgent(test_dag_path, 1,
processor_factory,
timedelta.max, async_mode)
processor_agent.start()
parsing_result = []
if not async_mode:
processor_agent.run_single_parsing_loop()
while not processor_agent.done:
if not async_mode:
processor_agent.wait_until_finished()
parsing_result.extend(processor_agent.harvest_simple_dags())
self.assertEqual(len(fake_failure_callback_requests),
len(parsing_result))
self.assertEqual(
set(zombie.simple_task_instance.key
for zombie in fake_failure_callback_requests),
set(result.simple_task_instance.key
for result in parsing_result))
def test_lineage_backend(mock_emit, inlets, outlets):
DEFAULT_DATE = days_ago(2)
# Using autospec on xcom_pull and xcom_push methods fails on Python 3.6.
with mock.patch.dict(
os.environ,
{
"AIRFLOW__LINEAGE__BACKEND":
"datahub_provider.lineage.datahub.DatahubLineageBackend",
"AIRFLOW__LINEAGE__DATAHUB_CONN_ID":
datahub_rest_connection_config.conn_id,
"AIRFLOW__LINEAGE__DATAHUB_KWARGS":
json.dumps({"graceful_exceptions": False}),
},
), mock.patch("airflow.models.BaseOperator.xcom_pull"), mock.patch(
"airflow.models.BaseOperator.xcom_push"), patch_airflow_connection(
datahub_rest_connection_config):
func = mock.Mock()
func.__name__ = "foo"
dag = DAG(dag_id="test_lineage_is_sent_to_backend",
start_date=DEFAULT_DATE)
with dag:
op1 = DummyOperator(
task_id="task1_upstream",
inlets=inlets,
outlets=outlets,
)
op2 = DummyOperator(
task_id="task2",
inlets=inlets,
outlets=outlets,
)
op1 >> op2
ti = TI(task=op2, execution_date=DEFAULT_DATE)
ctx1 = {
"dag": dag,
"task": op2,
"ti": ti,
"task_instance": ti,
"execution_date": DEFAULT_DATE,
"ts": "2021-04-08T00:54:25.771575+00:00",
}
prep = prepare_lineage(func)
prep(op2, ctx1)
post = apply_lineage(func)
post(op2, ctx1)
# Verify that the inlets and outlets are registered and recognized by Airflow correctly,
# or that our lineage backend forces it to.
assert len(op2.inlets) == 1
assert len(op2.outlets) == 1
assert all(map(lambda let: isinstance(let, Dataset), op2.inlets))
assert all(map(lambda let: isinstance(let, Dataset), op2.outlets))
# Check that the right things were emitted.
mock_emit.assert_called_once()
assert len(mock_emit.call_args[0][0]) == 4
assert all(mce.validate() for mce in mock_emit.call_args[0][0])
def test_dags_clear(self):
# setup
session = settings.Session()
dags, tis = [], []
num_of_dags = 5
for i in range(num_of_dags):
dag = DAG('test_dag_clear_' + str(i),
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=10))
ti = TI(task=DummyOperator(task_id='test_task_clear_' + str(i),
owner='test',
dag=dag),
execution_date=DEFAULT_DATE)
dags.append(dag)
tis.append(ti)
# test clear all dags
for i in range(num_of_dags):
tis[i].run()
self.assertEqual(tis[i].state, State.SUCCESS)
self.assertEqual(tis[i].try_number, 2)
self.assertEqual(tis[i].max_tries, 0)
DAG.clear_dags(dags)
for i in range(num_of_dags):
tis[i].refresh_from_db()
self.assertEqual(tis[i].state, State.NONE)
self.assertEqual(tis[i].try_number, 2)
self.assertEqual(tis[i].max_tries, 1)
# test dry_run
for i in range(num_of_dags):
tis[i].run()
self.assertEqual(tis[i].state, State.SUCCESS)
self.assertEqual(tis[i].try_number, 3)
self.assertEqual(tis[i].max_tries, 1)
DAG.clear_dags(dags, dry_run=True)
for i in range(num_of_dags):
tis[i].refresh_from_db()
self.assertEqual(tis[i].state, State.SUCCESS)
self.assertEqual(tis[i].try_number, 3)
self.assertEqual(tis[i].max_tries, 1)
# test only_failed
from random import randint
failed_dag_idx = randint(0, len(tis) - 1)
tis[failed_dag_idx].state = State.FAILED
session.merge(tis[failed_dag_idx])
session.commit()
DAG.clear_dags(dags, only_failed=True)
for i in range(num_of_dags):
tis[i].refresh_from_db()
if i != failed_dag_idx:
self.assertEqual(tis[i].state, State.SUCCESS)
self.assertEqual(tis[i].try_number, 3)
self.assertEqual(tis[i].max_tries, 1)
else:
self.assertEqual(tis[i].state, State.NONE)
self.assertEqual(tis[i].try_number, 3)
self.assertEqual(tis[i].max_tries, 2)
def test_zombies_are_correctly_passed_to_dag_file_processor(self):
"""
Check that the same set of zombies are passed to the dag
file processors until the next zombie detection logic is invoked.
"""
with conf_vars({('scheduler', 'max_threads'): '1',
('core', 'load_examples'): 'False'}):
dagbag = DagBag(os.path.join(TEST_DAG_FOLDER, 'test_example_bash_operator.py'))
with create_session() as session:
session.query(LJ).delete()
dag = dagbag.get_dag('test_example_bash_operator')
task = dag.get_task(task_id='run_this_last')
ti = TI(task, DEFAULT_DATE, State.RUNNING)
lj = LJ(ti)
lj.state = State.SHUTDOWN
lj.id = 1
ti.job_id = lj.id
session.add(lj)
session.add(ti)
session.commit()
fake_zombies = [SimpleTaskInstance(ti)]
class FakeDagFIleProcessor(DagFileProcessor):
# This fake processor will return the zombies it received in constructor
# as its processing result w/o actually parsing anything.
def __init__(self, file_path, pickle_dags, dag_id_white_list, zombies):
super(FakeDagFIleProcessor, self).__init__(
file_path, pickle_dags, dag_id_white_list, zombies
)
self._result = zombies, 0
def start(self):
pass
@property
def start_time(self):
return DEFAULT_DATE
@property
def pid(self):
return 1234
@property
def done(self):
return True
@property
def result(self):
return self._result
def processor_factory(file_path, zombies):
return FakeDagFIleProcessor(file_path,
False,
[],
zombies)
test_dag_path = os.path.join(TEST_DAG_FOLDER,
'test_example_bash_operator.py')
async_mode = 'sqlite' not in conf.get('core', 'sql_alchemy_conn')
processor_agent = DagFileProcessorAgent(test_dag_path,
[],
1,
processor_factory,
timedelta.max,
async_mode)
processor_agent.start()
parsing_result = []
if not async_mode:
processor_agent.heartbeat()
while not processor_agent.done:
if not async_mode:
processor_agent.wait_until_finished()
parsing_result.extend(processor_agent.harvest_simple_dags())
self.assertEqual(len(fake_zombies), len(parsing_result))
self.assertEqual(set([zombie.key for zombie in fake_zombies]),
set([result.key for result in parsing_result]))
def test_update_counters(self):
dag = DAG(dag_id='test_manage_executor_state', start_date=DEFAULT_DATE)
task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow')
job = BackfillJob(dag=dag)
session = settings.Session()
dr = dag.create_dagrun(run_id=DagRun.ID_PREFIX,
state=State.RUNNING,
execution_date=DEFAULT_DATE,
start_date=DEFAULT_DATE,
session=session)
ti = TI(task1, dr.execution_date)
ti.refresh_from_db()
ti_status = BackfillJob._DagRunTaskStatus()
# test for success
ti.set_state(State.SUCCESS, session)
ti_status.running[ti.key] = ti
job._update_counters(ti_status=ti_status)
self.assertTrue(len(ti_status.running) == 0)
self.assertTrue(len(ti_status.succeeded) == 1)
self.assertTrue(len(ti_status.skipped) == 0)
self.assertTrue(len(ti_status.failed) == 0)
self.assertTrue(len(ti_status.to_run) == 0)
ti_status.succeeded.clear()
# test for skipped
ti.set_state(State.SKIPPED, session)
ti_status.running[ti.key] = ti
job._update_counters(ti_status=ti_status)
self.assertTrue(len(ti_status.running) == 0)
self.assertTrue(len(ti_status.succeeded) == 0)
self.assertTrue(len(ti_status.skipped) == 1)
self.assertTrue(len(ti_status.failed) == 0)
self.assertTrue(len(ti_status.to_run) == 0)
ti_status.skipped.clear()
# test for failed
ti.set_state(State.FAILED, session)
ti_status.running[ti.key] = ti
job._update_counters(ti_status=ti_status)
self.assertTrue(len(ti_status.running) == 0)
self.assertTrue(len(ti_status.succeeded) == 0)
self.assertTrue(len(ti_status.skipped) == 0)
self.assertTrue(len(ti_status.failed) == 1)
self.assertTrue(len(ti_status.to_run) == 0)
ti_status.failed.clear()
# test for retry
ti.set_state(State.UP_FOR_RETRY, session)
ti_status.running[ti.key] = ti
job._update_counters(ti_status=ti_status)
self.assertTrue(len(ti_status.running) == 0)
self.assertTrue(len(ti_status.succeeded) == 0)
self.assertTrue(len(ti_status.skipped) == 0)
self.assertTrue(len(ti_status.failed) == 0)
self.assertTrue(len(ti_status.to_run) == 1)
ti_status.to_run.clear()
# test for reschedule
ti.set_state(State.UP_FOR_RESCHEDULE, session)
ti_status.running[ti.key] = ti
job._update_counters(ti_status=ti_status)
self.assertTrue(len(ti_status.running) == 0)
self.assertTrue(len(ti_status.succeeded) == 0)
self.assertTrue(len(ti_status.skipped) == 0)
self.assertTrue(len(ti_status.failed) == 0)
self.assertTrue(len(ti_status.to_run) == 1)
ti_status.to_run.clear()
# test for none
ti.set_state(State.NONE, session)
ti_status.running[ti.key] = ti
job._update_counters(ti_status=ti_status)
self.assertTrue(len(ti_status.running) == 0)
self.assertTrue(len(ti_status.succeeded) == 0)
self.assertTrue(len(ti_status.skipped) == 0)
self.assertTrue(len(ti_status.failed) == 0)
self.assertTrue(len(ti_status.to_run) == 1)
ti_status.to_run.clear()
session.close()
def test_dags_clear(self):
# setup
session = settings.Session()
dags, tis = [], []
num_of_dags = 5
for i in range(num_of_dags):
dag = DAG(
'test_dag_clear_' + str(i),
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=10),
)
ti = TI(
task=DummyOperator(task_id='test_task_clear_' + str(i),
owner='test',
dag=dag),
execution_date=DEFAULT_DATE,
)
dag.create_dagrun(
execution_date=ti.execution_date,
state=State.RUNNING,
run_type=DagRunType.SCHEDULED,
)
dags.append(dag)
tis.append(ti)
# test clear all dags
for i in range(num_of_dags):
tis[i].run()
assert tis[i].state == State.SUCCESS
assert tis[i].try_number == 2
assert tis[i].max_tries == 0
DAG.clear_dags(dags)
for i in range(num_of_dags):
tis[i].refresh_from_db()
assert tis[i].state == State.NONE
assert tis[i].try_number == 2
assert tis[i].max_tries == 1
# test dry_run
for i in range(num_of_dags):
tis[i].run()
assert tis[i].state == State.SUCCESS
assert tis[i].try_number == 3
assert tis[i].max_tries == 1
DAG.clear_dags(dags, dry_run=True)
for i in range(num_of_dags):
tis[i].refresh_from_db()
assert tis[i].state == State.SUCCESS
assert tis[i].try_number == 3
assert tis[i].max_tries == 1
# test only_failed
from random import randint
failed_dag_idx = randint(0, len(tis) - 1)
tis[failed_dag_idx].state = State.FAILED
session.merge(tis[failed_dag_idx])
session.commit()
DAG.clear_dags(dags, only_failed=True)
for i in range(num_of_dags):
tis[i].refresh_from_db()
if i != failed_dag_idx:
assert tis[i].state == State.SUCCESS
assert tis[i].try_number == 3
assert tis[i].max_tries == 1
else:
assert tis[i].state == State.NONE
assert tis[i].try_number == 3
assert tis[i].max_tries == 2
