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