def get_task_instance(dag_id, task_id, execution_date): """Return the task object identified by the given dag_id and task_id.""" dagbag = DagBag() # Check DAG exists. if dag_id not in dagbag.dags: error_message = "Dag id {} not found".format(dag_id) raise AirflowException(error_message) # Get DAG object and check Task Exists dag = dagbag.get_dag(dag_id) if not dag.has_task(task_id): error_message = 'Task {} not found in dag {}'.format(task_id, dag_id) raise AirflowException(error_message) # Get DagRun object and check that it exists dagrun = dag.get_dagrun(execution_date=execution_date) if not dagrun: error_message = ('Dag Run for date {} not found in dag {}' .format(execution_date, dag_id)) raise AirflowException(error_message) # Get task instance object and check that it exists task_instance = dagrun.get_task_instance(task_id) if not task_instance: error_message = ('Task {} instance for date {} not found' .format(task_id, execution_date)) raise AirflowException(error_message) return task_instance
def poke(self, context, session=None): if self.execution_delta: dttm = context['execution_date'] - self.execution_delta elif self.execution_date_fn: dttm = self.execution_date_fn(context['execution_date']) else: dttm = context['execution_date'] dttm_filter = dttm if isinstance(dttm, list) else [dttm] serialized_dttm_filter = ','.join( [datetime.isoformat() for datetime in dttm_filter]) self.log.info( 'Poking for ' '{self.external_dag_id}.' '{self.external_task_id} on ' '{} ... '.format(serialized_dttm_filter, **locals())) DM = DagModel TI = TaskInstance DR = DagRun # we only do the check for 1st time, no need for subsequent poke if self.check_existence and not self.has_checked_existence: dag_to_wait = session.query(DM).filter( DM.dag_id == self.external_dag_id ).first() if not dag_to_wait: raise AirflowException('The external DAG ' '{} does not exist.'.format(self.external_dag_id)) else: if not os.path.exists(dag_to_wait.fileloc): raise AirflowException('The external DAG ' '{} was deleted.'.format(self.external_dag_id)) if self.external_task_id: refreshed_dag_info = DagBag(dag_to_wait.fileloc).get_dag(self.external_dag_id) if not refreshed_dag_info.has_task(self.external_task_id): raise AirflowException('The external task' '{} in DAG {} does not exist.'.format(self.external_task_id, self.external_dag_id)) self.has_checked_existence = True if self.external_task_id: count = session.query(TI).filter( TI.dag_id == self.external_dag_id, TI.task_id == self.external_task_id, TI.state.in_(self.allowed_states), TI.execution_date.in_(dttm_filter), ).count() else: count = session.query(DR).filter( DR.dag_id == self.external_dag_id, DR.state.in_(self.allowed_states), DR.execution_date.in_(dttm_filter), ).count() session.commit() return count == len(dttm_filter)
def trigger_dag(dag_id, run_id=None, conf=None, execution_date=None): dagbag = DagBag() if dag_id not in dagbag.dags: raise AirflowException("Dag id {} not found".format(dag_id)) dag = dagbag.get_dag(dag_id) if not execution_date: execution_date = datetime.now() if not run_id: run_id = "manual__{0}".format(execution_date.isoformat()) dr = DagRun.find(dag_id=dag_id, run_id=run_id) if dr: raise AirflowException("Run id {} already exists for dag id {}".format( run_id, dag_id )) run_conf = None if conf: run_conf = json.loads(conf) trigger = dag.create_dagrun( run_id=run_id, execution_date=execution_date, state=State.RUNNING, conf=run_conf, external_trigger=True ) return trigger
def downgrade(): engine = settings.engine if engine.dialect.has_table(engine, 'task_instance'): connection = op.get_bind() sessionmaker = sa.orm.sessionmaker() session = sessionmaker(bind=connection) dagbag = DagBag(settings.DAGS_FOLDER) query = session.query(sa.func.count(TaskInstance.max_tries)).filter( TaskInstance.max_tries != -1 ) while query.scalar(): tis = session.query(TaskInstance).filter( TaskInstance.max_tries != -1 ).limit(BATCH_SIZE).all() for ti in tis: dag = dagbag.get_dag(ti.dag_id) if not dag or not dag.has_task(ti.task_id): ti.try_number = 0 else: task = dag.get_task(ti.task_id) # max_tries - try_number is number of times a task instance # left to retry by itself. So the current try_number should be # max number of self retry (task.retries) minus number of # times left for task instance to try the task. ti.try_number = max(0, task.retries - (ti.max_tries - ti.try_number)) ti.max_tries = -1 session.merge(ti) session.commit() session.commit() op.drop_column('task_instance', 'max_tries')
def test_dag_with_system_exit(self): """ Test to check that a DAG with a system.exit() doesn't break the scheduler. """ dag_id = 'exit_test_dag' dag_ids = [dag_id] dag_directory = os.path.join(models.DAGS_FOLDER, "..", "dags_with_system_exit") dag_file = os.path.join(dag_directory, 'b_test_scheduler_dags.py') dagbag = DagBag(dag_folder=dag_file) for dag_id in dag_ids: dag = dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, subdir= dag_directory, num_runs=1, **self.default_scheduler_args) scheduler.run() session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1)
def test_subdag_deadlock(self): dagbag = DagBag() dag = dagbag.get_dag('test_subdag_deadlock') dag.clear() subdag = dagbag.get_dag('test_subdag_deadlock.subdag') subdag.clear() # first make sure subdag has failed self.assertRaises(AirflowException, subdag.run, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # now make sure dag picks up the subdag error self.assertRaises(AirflowException, dag.run, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
def list_dags(args): dagbag = DagBag(process_subdir(args.subdir)) s = textwrap.dedent("""\n ------------------------------------------------------------------- DAGS ------------------------------------------------------------------- {dag_list} """) dag_list = "\n".join(sorted(dagbag.dags)) print(s.format(dag_list=dag_list)) if args.report: print(dagbag.dagbag_report())
def test_scheduler_reschedule(self): """ Checks if tasks that are not taken up by the executor get rescheduled """ executor = TestExecutor() dagbag = DagBag(executor=executor) dagbag.dags.clear() dagbag.executor = executor dag = DAG( dag_id='test_scheduler_reschedule', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') dag.clear() dag.is_subdag = False session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.is_paused = False session.merge(orm_dag) session.commit() dagbag.bag_dag(dag=dag, root_dag=dag, parent_dag=dag) @mock.patch('airflow.models.DagBag', return_value=dagbag) @mock.patch('airflow.models.DagBag.collect_dags') def do_schedule(function, function2): # Use a empty file since the above mock will return the # expected DAGs. Also specify only a single file so that it doesn't # try to schedule the above DAG repeatedly. scheduler = SchedulerJob(num_runs=1, executor=executor, subdir=os.path.join(models.DAGS_FOLDER, "no_dags.py")) scheduler.heartrate = 0 scheduler.run() do_schedule() self.assertEquals(1, len(executor.queued_tasks)) executor.queued_tasks.clear() do_schedule() self.assertEquals(2, len(executor.queued_tasks))
def execute(self, context): dro = DagRunOrder(run_id='trig__' + timezone.utcnow().isoformat()) dro = self.python_callable(context, dro) if dro: with create_session() as session: dbag = DagBag(settings.DAGS_FOLDER) trigger_dag = dbag.get_dag(self.trigger_dag_id) dr = trigger_dag.create_dagrun( run_id=dro.run_id, state=State.RUNNING, conf=dro.payload, external_trigger=True) self.log.info("Creating DagRun %s", dr) session.add(dr) session.commit() else: self.log.info("Criteria not met, moving on")
def get_task(dag_id, task_id): """Return the task object identified by the given dag_id and task_id.""" dagbag = DagBag() # Check DAG exists. if dag_id not in dagbag.dags: error_message = "Dag id {} not found".format(dag_id) raise AirflowException(error_message) # Get DAG object and check Task Exists dag = dagbag.get_dag(dag_id) if not dag.has_task(task_id): error_message = 'Task {} not found in dag {}'.format(task_id, dag_id) raise AirflowException(error_message) # Return the task. return dag.get_task(task_id)
def upgrade(): op.add_column('task_instance', sa.Column('max_tries', sa.Integer, server_default="-1")) # Check if table task_instance exist before data migration. This check is # needed for database that does not create table until migration finishes. # Checking task_instance table exists prevent the error of querying # non-existing task_instance table. connection = op.get_bind() inspector = Inspector.from_engine(connection) tables = inspector.get_table_names() if 'task_instance' in tables: # Get current session sessionmaker = sa.orm.sessionmaker() session = sessionmaker(bind=connection) dagbag = DagBag(settings.DAGS_FOLDER) query = session.query(sa.func.count(TaskInstance.max_tries)).filter( TaskInstance.max_tries == -1 ) # Separate db query in batch to prevent loading entire table # into memory and cause out of memory error. while query.scalar(): tis = session.query(TaskInstance).filter( TaskInstance.max_tries == -1 ).limit(BATCH_SIZE).all() for ti in tis: dag = dagbag.get_dag(ti.dag_id) if not dag or not dag.has_task(ti.task_id): # task_instance table might not have the up-to-date # information, i.e dag or task might be modified or # deleted in dagbag but is reflected in task instance # table. In this case we do not retry the task that can't # be parsed. ti.max_tries = ti.try_number else: task = dag.get_task(ti.task_id) if task.retries: ti.max_tries = task.retries else: ti.max_tries = ti.try_number session.merge(ti) session.commit() # Commit the current session. session.commit()
def execute(self, context): dro = DagRunOrder(run_id='trig__' + datetime.now().isoformat()) dro = self.python_callable(context, dro) if dro: session = settings.Session() dbag = DagBag(os.path.expanduser(conf.get('core', 'DAGS_FOLDER'))) trigger_dag = dbag.get_dag(self.trigger_dag_id) dr = trigger_dag.create_dagrun( run_id=dro.run_id, state=State.RUNNING, conf=dro.payload, external_trigger=True) logging.info("Creating DagRun {}".format(dr)) session.add(dr) session.commit() session.close() else: logging.info("Criteria not met, moving on")
def test_scheduler_reschedule(self): """ Checks if tasks that are not taken up by the executor get rescheduled """ executor = TestExecutor() dagbag = DagBag(executor=executor) dagbag.dags.clear() dagbag.executor = executor dag = DAG( dag_id='test_scheduler_reschedule', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') dag.clear() dag.is_subdag = False session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.is_paused = False session.merge(orm_dag) session.commit() dagbag.bag_dag(dag=dag, root_dag=dag, parent_dag=dag) @mock.patch('airflow.models.DagBag', return_value=dagbag) @mock.patch('airflow.models.DagBag.collect_dags') def do_schedule(function, function2): scheduler = SchedulerJob(num_runs=1, executor=executor,) scheduler.heartrate = 0 scheduler.run() do_schedule() self.assertEquals(1, len(executor.queued_tasks)) executor.queued_tasks.clear() do_schedule() self.assertEquals(2, len(executor.queued_tasks))
def test_trigger_dag_for_date(self): url_template = '/api/experimental/dags/{}/dag_runs' dag_id = 'example_bash_operator' hour_from_now = utcnow() + timedelta(hours=1) execution_date = datetime(hour_from_now.year, hour_from_now.month, hour_from_now.day, hour_from_now.hour) datetime_string = execution_date.isoformat() # Test Correct execution response = self.client.post( url_template.format(dag_id), data=json.dumps({'execution_date': datetime_string}), content_type="application/json" ) self.assertEqual(200, response.status_code) dagbag = DagBag() dag = dagbag.get_dag(dag_id) dag_run = dag.get_dagrun(execution_date) self.assertTrue(dag_run, 'Dag Run not found for execution date {}' .format(execution_date)) # Test error for nonexistent dag response = self.client.post( url_template.format('does_not_exist_dag'), data=json.dumps({'execution_date': execution_date.isoformat()}), content_type="application/json" ) self.assertEqual(404, response.status_code) # Test error for bad datetime format response = self.client.post( url_template.format(dag_id), data=json.dumps({'execution_date': 'not_a_datetime'}), content_type="application/json" ) self.assertEqual(400, response.status_code)
def test_trigger_dag_for_date(self): url_template = '/api/experimental/dags/{}/dag_runs/{}' dag_id = 'example_bash_operator' now = datetime.now() execution_date = datetime(now.year, now.month, now.day, now.hour + 1) datetime_string = execution_date.isoformat() # Test Correct execution response = self.app.post( quote(url_template.format(dag_id, datetime_string)), data=json.dumps(dict(run_id='my_run'.format(datetime_string))), content_type="application/json" ) self.assertEqual(200, response.status_code) dagbag = DagBag() dag = dagbag.get_dag(dag_id) dag_run = dag.get_dagrun(execution_date) self.assertTrue(dag_run, 'Dag Run not found for execution date {}' .format(execution_date)) # Test error for nonexistent dag response = self.app.post( quote(url_template.format('does_not_exist_dag', datetime_string)), data=json.dumps(dict()), content_type="application/json" ) self.assertEqual(404, response.status_code) # Test error for bad datetime format response = self.app.post( quote( url_template.format(dag_id, 'not_a_datetime')), data=json.dumps(dict(run_id='my_run'.format(datetime_string))), content_type="application/json" ) self.assertEqual(400, response.status_code)
def test_find_zombies(self): manager = DagFileProcessorManager( dag_directory='directory', file_paths=['abc.txt'], max_runs=1, processor_factory=MagicMock().return_value, signal_conn=MagicMock(), stat_queue=MagicMock(), result_queue=MagicMock, async_mode=True) dagbag = DagBag(TEST_DAG_FOLDER) with create_session() as session: 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.state = State.SHUTDOWN lj.id = 1 ti.job_id = lj.id session.add(lj) session.add(ti) session.commit() manager._last_zombie_query_time = timezone.utcnow() - timedelta( seconds=manager._zombie_threshold_secs + 1) zombies = manager._find_zombies() self.assertEqual(1, len(zombies)) self.assertIsInstance(zombies[0], SimpleTaskInstance) self.assertEqual(ti.dag_id, zombies[0].dag_id) self.assertEqual(ti.task_id, zombies[0].task_id) self.assertEqual(ti.execution_date, zombies[0].execution_date) session.query(TI).delete() session.query(LJ).delete()
class SchedulerJobTest(unittest.TestCase): # These defaults make the test faster to run default_scheduler_args = {"file_process_interval": 0, "processor_poll_interval": 0.5} def setUp(self): self.dagbag = DagBag() @provide_session def evaluate_dagrun( self, dag_id, expected_task_states, # dict of task_id: state dagrun_state, run_kwargs=None, advance_execution_date=False, session=None): """ Helper for testing DagRun states with simple two-task DAGS. This is hackish: a dag run is created but its tasks are run by a backfill. """ if run_kwargs is None: run_kwargs = {} scheduler = SchedulerJob(**self.default_scheduler_args) dag = self.dagbag.get_dag(dag_id) dag.clear() dr = scheduler.create_dag_run(dag) if advance_execution_date: # run a second time to schedule a dagrun after the start_date dr = scheduler.create_dag_run(dag) ex_date = dr.execution_date try: dag.run(start_date=ex_date, end_date=ex_date, **run_kwargs) except AirflowException: pass # test tasks for task_id, expected_state in expected_task_states.items(): task = dag.get_task(task_id) ti = TI(task, ex_date) ti.refresh_from_db() self.assertEqual(ti.state, expected_state) # load dagrun dr = DagRun.find(dag_id=dag_id, execution_date=ex_date) dr = dr[0] dr.dag = dag # dagrun is running self.assertEqual(dr.state, State.RUNNING) dr.update_state() # dagrun failed self.assertEqual(dr.state, dagrun_state) def test_dagrun_fail(self): """ DagRuns with one failed and one incomplete root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_fail', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.UPSTREAM_FAILED, }, dagrun_state=State.FAILED) def test_dagrun_success(self): """ DagRuns with one failed and one successful root task -> SUCCESS """ self.evaluate_dagrun( dag_id='test_dagrun_states_success', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.SUCCESS, }, dagrun_state=State.SUCCESS) def test_dagrun_root_fail(self): """ DagRuns with one successful and one failed root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_root_fail', expected_task_states={ 'test_dagrun_succeed': State.SUCCESS, 'test_dagrun_fail': State.FAILED, }, dagrun_state=State.FAILED) 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, executor=TestExecutor(), **self.default_scheduler_args) 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() logging.error("TI is: {}".format(ti)) 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__(), **self.default_scheduler_args) scheduler2.run() ti.refresh_from_db() self.assertEqual(ti.state, State.FAILED) dag.clear() def test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date(self): """ DagRun is marked a success if ignore_first_depends_on_past=True Test that an otherwise-deadlocked dagrun is marked as a success if ignore_first_depends_on_past=True and the dagrun execution_date is after the start_date. """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, advance_execution_date=True, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_dagrun_deadlock_ignore_depends_on_past(self): """ Test that ignore_first_depends_on_past doesn't affect results (this is the same test as test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date except that start_date == execution_date so depends_on_past is irrelevant). """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_scheduler_start_date(self): """ Test that the scheduler respects start_dates, even when DAGS have run """ dag_id = 'test_start_date_scheduling' dag = self.dagbag.get_dag(dag_id) dag.clear() self.assertTrue(dag.start_date > DEFAULT_DATE) scheduler = SchedulerJob(dag_id, num_runs=2, **self.default_scheduler_args) scheduler.run() # zero tasks ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) # previously, running this backfill would kick off the Scheduler # because it would take the most recent run and start from there # That behavior still exists, but now it will only do so if after the # start date backfill = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) backfill.run() # one task ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) scheduler = SchedulerJob(dag_id, num_runs=2, **self.default_scheduler_args) scheduler.run() # still one task session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) def test_scheduler_multiprocessing(self): """ Test that the scheduler can successfully queue multiple dags in parallel """ dag_ids = ['test_start_date_scheduling', 'test_dagrun_states_success'] for dag_id in dag_ids: dag = self.dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, file_process_interval=0, processor_poll_interval=0.5, num_runs=2) scheduler.run() # zero tasks ran dag_id = 'test_start_date_scheduling' session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) def test_scheduler_dagrun_once(self): """ Test if the scheduler does not create multiple dagruns if a dag is scheduled with @once and a start_date """ dag = DAG( 'test_scheduler_dagrun_once', start_date=datetime.datetime(2015, 1, 1), schedule_interval="@once") scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dr = scheduler.create_dag_run(dag) self.assertIsNone(dr) def test_scheduler_process_task_instances(self): """ Test if _process_task_instances puts the right task instances into the queue. """ dag = DAG( dag_id='test_scheduler_process_execute_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.append.assert_called_with( (dag.dag_id, dag_task1.task_id, DEFAULT_DATE) ) def test_scheduler_do_not_schedule_removed_task(self): dag = DAG( dag_id='test_scheduler_do_not_schedule_removed_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dag = DAG( dag_id='test_scheduler_do_not_schedule_removed_task', start_date=DEFAULT_DATE) queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_do_not_schedule_too_early(self): dag = DAG( dag_id='test_scheduler_do_not_schedule_too_early', start_date=datetime.datetime(2200, 1, 1)) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNone(dr) queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_do_not_run_finished(self): dag = DAG( dag_id='test_scheduler_do_not_run_finished', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) tis = dr.get_task_instances(session=session) for ti in tis: ti.state = State.SUCCESS session.commit() session.close() queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_add_new_task(self): """ Test if a task instance will be added if the dag is updated """ dag = DAG( dag_id='test_scheduler_add_new_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) tis = dr.get_task_instances() self.assertEquals(len(tis), 1) dag_task2 = DummyOperator( task_id='dummy2', dag=dag, owner='airflow') queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) tis = dr.get_task_instances() self.assertEquals(len(tis), 2) def test_scheduler_verify_max_active_runs(self): """ Test if a a dagrun will not be scheduled if max_dag_runs has been reached """ dag = DAG( dag_id='test_scheduler_verify_max_active_runs', start_date=DEFAULT_DATE) dag.max_active_runs = 1 dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dr = scheduler.create_dag_run(dag) self.assertIsNone(dr) def test_scheduler_fail_dagrun_timeout(self): """ Test if a a dagrun wil be set failed if timeout """ dag = DAG( dag_id='test_scheduler_fail_dagrun_timeout', start_date=DEFAULT_DATE) dag.dagrun_timeout = datetime.timedelta(seconds=60) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dr.start_date = datetime.datetime.now() - datetime.timedelta(days=1) session.merge(dr) session.commit() dr2 = scheduler.create_dag_run(dag) self.assertIsNotNone(dr2) dr.refresh_from_db(session=session) self.assertEquals(dr.state, State.FAILED) def test_scheduler_verify_max_active_runs_and_dagrun_timeout(self): """ Test if a a dagrun will not be scheduled if max_dag_runs has been reached and dagrun_timeout is not reached Test if a a dagrun will be scheduled if max_dag_runs has been reached but dagrun_timeout is also reached """ dag = DAG( dag_id='test_scheduler_verify_max_active_runs_and_dagrun_timeout', start_date=DEFAULT_DATE) dag.max_active_runs = 1 dag.dagrun_timeout = datetime.timedelta(seconds=60) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) # Should not be scheduled as DagRun has not timedout and max_active_runs is reached new_dr = scheduler.create_dag_run(dag) self.assertIsNone(new_dr) # Should be scheduled as dagrun_timeout has passed dr.start_date = datetime.datetime.now() - datetime.timedelta(days=1) session.merge(dr) session.commit() new_dr = scheduler.create_dag_run(dag) self.assertIsNotNone(new_dr) def test_scheduler_auto_align(self): """ Test if the schedule_interval will be auto aligned with the start_date such that if the start_date coincides with the schedule the first execution_date will be start_date, otherwise it will be start_date + interval. """ dag = DAG( dag_id='test_scheduler_auto_align_1', start_date=datetime.datetime(2016, 1, 1, 10, 10, 0), schedule_interval="4 5 * * *" ) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) self.assertEquals(dr.execution_date, datetime.datetime(2016, 1, 2, 5, 4)) dag = DAG( dag_id='test_scheduler_auto_align_2', start_date=datetime.datetime(2016, 1, 1, 10, 10, 0), schedule_interval="10 10 * * *" ) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) self.assertEquals(dr.execution_date, datetime.datetime(2016, 1, 1, 10, 10)) def test_scheduler_reschedule(self): """ Checks if tasks that are not taken up by the executor get rescheduled """ executor = TestExecutor() dagbag = DagBag(executor=executor) dagbag.dags.clear() dagbag.executor = executor dag = DAG( dag_id='test_scheduler_reschedule', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') dag.clear() dag.is_subdag = False session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.is_paused = False session.merge(orm_dag) session.commit() dagbag.bag_dag(dag=dag, root_dag=dag, parent_dag=dag) @mock.patch('airflow.models.DagBag', return_value=dagbag) @mock.patch('airflow.models.DagBag.collect_dags') def do_schedule(function, function2): # Use a empty file since the above mock will return the # expected DAGs. Also specify only a single file so that it doesn't # try to schedule the above DAG repeatedly. scheduler = SchedulerJob(num_runs=1, executor=executor, subdir=os.path.join(models.DAGS_FOLDER, "no_dags.py")) scheduler.heartrate = 0 scheduler.run() do_schedule() self.assertEquals(1, len(executor.queued_tasks)) executor.queued_tasks.clear() do_schedule() self.assertEquals(2, len(executor.queued_tasks)) def test_scheduler_run_duration(self): """ Verifies that the scheduler run duration limit is followed. """ dag_id = 'test_start_date_scheduling' dag = self.dagbag.get_dag(dag_id) dag.clear() self.assertTrue(dag.start_date > DEFAULT_DATE) expected_run_duration = 5 start_time = datetime.datetime.now() scheduler = SchedulerJob(dag_id, run_duration=expected_run_duration, **self.default_scheduler_args) scheduler.run() end_time = datetime.datetime.now() run_duration = (end_time - start_time).total_seconds() logging.info("Test ran in %.2fs, expected %.2fs", run_duration, expected_run_duration) assert run_duration - expected_run_duration < 5.0 def test_dag_with_system_exit(self): """ Test to check that a DAG with a system.exit() doesn't break the scheduler. """ dag_id = 'exit_test_dag' dag_ids = [dag_id] dag_directory = os.path.join(models.DAGS_FOLDER, "..", "dags_with_system_exit") dag_file = os.path.join(dag_directory, 'b_test_scheduler_dags.py') dagbag = DagBag(dag_folder=dag_file) for dag_id in dag_ids: dag = dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, subdir= dag_directory, num_runs=1, **self.default_scheduler_args) scheduler.run() session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1)
def setUp(self): self.dagbag = DagBag(dag_folder="/Users/oscar.barlow/Projects/work/infinityworks/airflow-hacknight-coke/dags")
class SchedulerJobTest(unittest.TestCase): def setUp(self): self.dagbag = DagBag() @provide_session def evaluate_dagrun( self, dag_id, expected_task_states, # dict of task_id: state dagrun_state, run_kwargs=None, advance_execution_date=False, session=None): """ Helper for testing DagRun states with simple two-task DAGS """ if run_kwargs is None: run_kwargs = {} scheduler = SchedulerJob() dag = self.dagbag.get_dag(dag_id) dag.clear() dr = scheduler.schedule_dag(dag) if advance_execution_date: # run a second time to schedule a dagrun after the start_date dr = scheduler.schedule_dag(dag) ex_date = dr.execution_date try: dag.run(start_date=ex_date, end_date=ex_date, **run_kwargs) except AirflowException: pass # test tasks for task_id, expected_state in expected_task_states.items(): task = dag.get_task(task_id) ti = TI(task, ex_date) ti.refresh_from_db() self.assertEqual(ti.state, expected_state) # load dagrun dr = session.query(DagRun).filter( DagRun.dag_id == dag.dag_id, DagRun.execution_date == ex_date ).first() # dagrun is running self.assertEqual(dr.state, State.RUNNING) dag.get_active_runs() # dagrun failed self.assertEqual(dr.state, dagrun_state) def test_dagrun_fail(self): """ DagRuns with one failed and one incomplete root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_fail', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.UPSTREAM_FAILED, }, dagrun_state=State.FAILED) def test_dagrun_success(self): """ DagRuns with one failed and one successful root task -> SUCCESS """ self.evaluate_dagrun( dag_id='test_dagrun_states_success', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.SUCCESS, }, dagrun_state=State.SUCCESS) def test_dagrun_root_fail(self): """ DagRuns with one successful and one failed root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_root_fail', expected_task_states={ 'test_dagrun_succeed': State.SUCCESS, 'test_dagrun_fail': State.FAILED, }, dagrun_state=State.FAILED) def test_dagrun_deadlock(self): """ Deadlocked DagRun is marked a failure Test that a deadlocked dagrun is marked as a failure by having depends_on_past and an execution_date after the start_date """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': None, 'test_depends_on_past_2': None, }, dagrun_state=State.FAILED, advance_execution_date=True) 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=10) 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.FAILED) dag.clear() def test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date(self): """ DagRun is marked a success if ignore_first_depends_on_past=True Test that an otherwise-deadlocked dagrun is marked as a success if ignore_first_depends_on_past=True and the dagrun execution_date is after the start_date. """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, advance_execution_date=True, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_dagrun_deadlock_ignore_depends_on_past(self): """ Test that ignore_first_depends_on_past doesn't affect results (this is the same test as test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date except that start_date == execution_date so depends_on_past is irrelevant). """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_scheduler_start_date(self): """ Test that the scheduler respects start_dates, even when DAGS have run """ dag_id = 'test_start_date_scheduling' dag = self.dagbag.get_dag(dag_id) dag.clear() self.assertTrue(dag.start_date > DEFAULT_DATE) scheduler = SchedulerJob(dag_id, num_runs=2) scheduler.run() # zero tasks ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) # previously, running this backfill would kick off the Scheduler # because it would take the most recent run and start from there # That behavior still exists, but now it will only do so if after the # start date backfill = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) backfill.run() # one task ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) scheduler = SchedulerJob(dag_id, num_runs=2) scheduler.run() # still one task session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) def test_scheduler_multiprocessing(self): """ Test that the scheduler can successfully queue multiple dags in parallel """ dag_ids = ['test_start_date_scheduling', 'test_dagrun_states_success'] for dag_id in dag_ids: dag = self.dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, num_runs=2) scheduler.run() # zero tasks ran dag_id = 'test_start_date_scheduling' session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0)
def run(args): utils.pessimistic_connection_handling() # Setting up logging log = os.path.expanduser(conf.get('core', 'BASE_LOG_FOLDER')) directory = log + "/{args.dag_id}/{args.task_id}".format(args=args) if not os.path.exists(directory): os.makedirs(directory) args.execution_date = dateutil.parser.parse(args.execution_date) iso = args.execution_date.isoformat() filename = "{directory}/{iso}".format(**locals()) subdir = None if args.subdir: subdir = args.subdir.replace( "DAGS_FOLDER", conf.get("core", "DAGS_FOLDER")) subdir = os.path.expanduser(subdir) logging.basicConfig( filename=filename, level=settings.LOGGING_LEVEL, format=settings.LOG_FORMAT) if not args.pickle: dagbag = DagBag(subdir) if args.dag_id not in dagbag.dags: msg = 'DAG [{0}] could not be found'.format(args.dag_id) logging.error(msg) raise AirflowException(msg) dag = dagbag.dags[args.dag_id] task = dag.get_task(task_id=args.task_id) else: session = settings.Session() logging.info('Loading pickle id {args.pickle}'.format(**locals())) dag_pickle = session.query( DagPickle).filter(DagPickle.id == args.pickle).first() if not dag_pickle: raise AirflowException("Who hid the pickle!? [missing pickle]") dag = dag_pickle.pickle task = dag.get_task(task_id=args.task_id) task_start_date = None if args.task_start_date: task_start_date = dateutil.parser.parse(args.task_start_date) task.start_date = task_start_date ti = TaskInstance(task, args.execution_date) if args.local: print("Logging into: " + filename) run_job = jobs.LocalTaskJob( task_instance=ti, mark_success=args.mark_success, force=args.force, pickle_id=args.pickle, task_start_date=task_start_date, ignore_dependencies=args.ignore_dependencies) run_job.run() elif args.raw: ti.run( mark_success=args.mark_success, force=args.force, ignore_dependencies=args.ignore_dependencies, job_id=args.job_id, ) else: pickle_id = None if args.ship_dag: try: # Running remotely, so pickling the DAG session = settings.Session() pickle = DagPickle(dag) session.add(pickle) session.commit() pickle_id = pickle.id print( 'Pickled dag {dag} ' 'as pickle_id:{pickle_id}').format(**locals()) except Exception as e: print('Could not pickle the DAG') print(e) raise e executor = DEFAULT_EXECUTOR executor.start() print("Sending to executor.") executor.queue_task_instance( ti, mark_success=args.mark_success, pickle_id=pickle_id, ignore_dependencies=args.ignore_dependencies, force=args.force) executor.heartbeat() executor.end()
def get_dagbag(): return DagBag()
def test_retry_still_in_executor(self): """ Checks if the scheduler does not put a task in limbo, when a task is retried but is still present in the executor. """ executor = TestExecutor() dagbag = DagBag(executor=executor) dagbag.dags.clear() dagbag.executor = executor dag = DAG(dag_id='test_retry_still_in_executor', start_date=DEFAULT_DATE) dag_task1 = BashOperator(task_id='test_retry_handling_op', bash_command='exit 1', retries=1, dag=dag, owner='airflow') dag.clear() dag.is_subdag = False session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.is_paused = False session.merge(orm_dag) session.commit() dagbag.bag_dag(dag=dag, root_dag=dag, parent_dag=dag) @mock.patch('airflow.models.DagBag', return_value=dagbag) @mock.patch('airflow.models.DagBag.collect_dags') def do_schedule(function, function2): # Use a empty file since the above mock will return the # expected DAGs. Also specify only a single file so that it doesn't # try to schedule the above DAG repeatedly. scheduler = SchedulerJob(num_runs=1, executor=executor, subdir=os.path.join( models.DAGS_FOLDER, "no_dags.py")) scheduler.heartrate = 0 scheduler.run() do_schedule() self.assertEquals(1, len(executor.queued_tasks)) def run_with_error(task): try: task.run() except AirflowException: pass ti_tuple = six.next(six.itervalues(executor.queued_tasks)) (command, priority, queue, ti) = ti_tuple ti.task = dag_task1 # fail execution run_with_error(ti) self.assertEqual(ti.state, State.UP_FOR_RETRY) self.assertEqual(ti.try_number, 1) # do not schedule do_schedule() self.assertTrue(executor.has_task(ti)) ti.refresh_from_db() self.assertEqual(ti.state, State.UP_FOR_RETRY) # now the executor has cleared and it should be allowed the re-queue executor.queued_tasks.clear() do_schedule() ti.refresh_from_db() self.assertEqual(ti.state, State.QUEUED)
def setUp(self): self.dagbag = DagBag()
class SchedulerJobTest(unittest.TestCase): # These defaults make the test faster to run default_scheduler_args = { "file_process_interval": 0, "processor_poll_interval": 0.5 } def setUp(self): self.dagbag = DagBag() @provide_session def evaluate_dagrun( self, dag_id, expected_task_states, # dict of task_id: state dagrun_state, run_kwargs=None, advance_execution_date=False, session=None): """ Helper for testing DagRun states with simple two-task DAGS. This is hackish: a dag run is created but its tasks are run by a backfill. """ if run_kwargs is None: run_kwargs = {} scheduler = SchedulerJob(**self.default_scheduler_args) dag = self.dagbag.get_dag(dag_id) dag.clear() dr = scheduler.create_dag_run(dag) if advance_execution_date: # run a second time to schedule a dagrun after the start_date dr = scheduler.create_dag_run(dag) ex_date = dr.execution_date try: dag.run(start_date=ex_date, end_date=ex_date, **run_kwargs) except AirflowException: pass # test tasks for task_id, expected_state in expected_task_states.items(): task = dag.get_task(task_id) ti = TI(task, ex_date) ti.refresh_from_db() self.assertEqual(ti.state, expected_state) # load dagrun dr = DagRun.find(dag_id=dag_id, execution_date=ex_date) dr = dr[0] dr.dag = dag self.assertEqual(dr.state, dagrun_state) def test_dagrun_fail(self): """ DagRuns with one failed and one incomplete root task -> FAILED """ self.evaluate_dagrun(dag_id='test_dagrun_states_fail', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.UPSTREAM_FAILED, }, dagrun_state=State.FAILED) def test_dagrun_success(self): """ DagRuns with one failed and one successful root task -> SUCCESS """ self.evaluate_dagrun(dag_id='test_dagrun_states_success', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.SUCCESS, }, dagrun_state=State.SUCCESS) def test_dagrun_root_fail(self): """ DagRuns with one successful and one failed root task -> FAILED """ self.evaluate_dagrun(dag_id='test_dagrun_states_root_fail', expected_task_states={ 'test_dagrun_succeed': State.SUCCESS, 'test_dagrun_fail': State.FAILED, }, dagrun_state=State.FAILED) def test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date(self): """ DagRun is marked a success if ignore_first_depends_on_past=True Test that an otherwise-deadlocked dagrun is marked as a success if ignore_first_depends_on_past=True and the dagrun execution_date is after the start_date. """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, advance_execution_date=True, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_dagrun_deadlock_ignore_depends_on_past(self): """ Test that ignore_first_depends_on_past doesn't affect results (this is the same test as test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date except that start_date == execution_date so depends_on_past is irrelevant). """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_scheduler_start_date(self): """ Test that the scheduler respects start_dates, even when DAGS have run """ dag_id = 'test_start_date_scheduling' dag = self.dagbag.get_dag(dag_id) dag.clear() self.assertTrue(dag.start_date > DEFAULT_DATE) scheduler = SchedulerJob(dag_id, num_runs=2, **self.default_scheduler_args) scheduler.run() # zero tasks ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) # previously, running this backfill would kick off the Scheduler # because it would take the most recent run and start from there # That behavior still exists, but now it will only do so if after the # start date backfill = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) backfill.run() # one task ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) scheduler = SchedulerJob(dag_id, num_runs=2, **self.default_scheduler_args) scheduler.run() # still one task session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) def test_scheduler_multiprocessing(self): """ Test that the scheduler can successfully queue multiple dags in parallel """ dag_ids = ['test_start_date_scheduling', 'test_dagrun_states_success'] for dag_id in dag_ids: dag = self.dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, file_process_interval=0, processor_poll_interval=0.5, num_runs=2) scheduler.run() # zero tasks ran dag_id = 'test_start_date_scheduling' session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) def test_scheduler_dagrun_once(self): """ Test if the scheduler does not create multiple dagruns if a dag is scheduled with @once and a start_date """ dag = DAG('test_scheduler_dagrun_once', start_date=datetime.datetime(2015, 1, 1), schedule_interval="@once") scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dr = scheduler.create_dag_run(dag) self.assertIsNone(dr) def test_scheduler_process_task_instances(self): """ Test if _process_task_instances puts the right task instances into the queue. """ dag = DAG(dag_id='test_scheduler_process_execute_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.append.assert_called_with( (dag.dag_id, dag_task1.task_id, DEFAULT_DATE)) def test_scheduler_do_not_schedule_removed_task(self): dag = DAG(dag_id='test_scheduler_do_not_schedule_removed_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dag = DAG(dag_id='test_scheduler_do_not_schedule_removed_task', start_date=DEFAULT_DATE) queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_do_not_schedule_too_early(self): dag = DAG(dag_id='test_scheduler_do_not_schedule_too_early', start_date=datetime.datetime(2200, 1, 1)) dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNone(dr) queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_do_not_run_finished(self): dag = DAG(dag_id='test_scheduler_do_not_run_finished', start_date=DEFAULT_DATE) dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) tis = dr.get_task_instances(session=session) for ti in tis: ti.state = State.SUCCESS session.commit() session.close() queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_add_new_task(self): """ Test if a task instance will be added if the dag is updated """ dag = DAG(dag_id='test_scheduler_add_new_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) tis = dr.get_task_instances() self.assertEquals(len(tis), 1) dag_task2 = DummyOperator(task_id='dummy2', dag=dag, owner='airflow') queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) tis = dr.get_task_instances() self.assertEquals(len(tis), 2) def test_scheduler_verify_max_active_runs(self): """ Test if a a dagrun will not be scheduled if max_dag_runs has been reached """ dag = DAG(dag_id='test_scheduler_verify_max_active_runs', start_date=DEFAULT_DATE) dag.max_active_runs = 1 dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dr = scheduler.create_dag_run(dag) self.assertIsNone(dr) def test_scheduler_fail_dagrun_timeout(self): """ Test if a a dagrun wil be set failed if timeout """ dag = DAG(dag_id='test_scheduler_fail_dagrun_timeout', start_date=DEFAULT_DATE) dag.dagrun_timeout = datetime.timedelta(seconds=60) dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dr.start_date = datetime.datetime.now() - datetime.timedelta(days=1) session.merge(dr) session.commit() dr2 = scheduler.create_dag_run(dag) self.assertIsNotNone(dr2) dr.refresh_from_db(session=session) self.assertEquals(dr.state, State.FAILED) def test_scheduler_verify_max_active_runs_and_dagrun_timeout(self): """ Test if a a dagrun will not be scheduled if max_dag_runs has been reached and dagrun_timeout is not reached Test if a a dagrun will be scheduled if max_dag_runs has been reached but dagrun_timeout is also reached """ dag = DAG( dag_id='test_scheduler_verify_max_active_runs_and_dagrun_timeout', start_date=DEFAULT_DATE) dag.max_active_runs = 1 dag.dagrun_timeout = datetime.timedelta(seconds=60) dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) # Should not be scheduled as DagRun has not timedout and max_active_runs is reached new_dr = scheduler.create_dag_run(dag) self.assertIsNone(new_dr) # Should be scheduled as dagrun_timeout has passed dr.start_date = datetime.datetime.now() - datetime.timedelta(days=1) session.merge(dr) session.commit() new_dr = scheduler.create_dag_run(dag) self.assertIsNotNone(new_dr) def test_scheduler_max_active_runs_respected_after_clear(self): """ Test if _process_task_instances only schedules ti's up to max_active_runs (related to issue AIRFLOW-137) """ dag = DAG( dag_id='test_scheduler_max_active_runs_respected_after_clear', start_date=DEFAULT_DATE) dag.max_active_runs = 3 dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() # First create up to 3 dagruns in RUNNING state. scheduler.create_dag_run(dag) # Reduce max_active_runs to 1 dag.max_active_runs = 1 queue = mock.Mock() # and schedule them in, so we can check how many # tasks are put on the queue (should be one, not 3) scheduler._process_task_instances(dag, queue=queue) queue.append.assert_called_with( (dag.dag_id, dag_task1.task_id, DEFAULT_DATE)) @patch.object(TI, 'pool_full') def test_scheduler_verify_pool_full(self, mock_pool_full): """ Test task instances not queued when pool is full """ mock_pool_full.return_value = False dag = DAG(dag_id='test_scheduler_verify_pool_full', start_date=DEFAULT_DATE) DummyOperator(task_id='dummy', dag=dag, owner='airflow', pool='test_scheduler_verify_pool_full') session = settings.Session() pool = Pool(pool='test_scheduler_verify_pool_full', slots=1) session.add(pool) orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.is_paused = False session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() # Create 2 dagruns, which will create 2 task instances. dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) self.assertEquals(dr.execution_date, DEFAULT_DATE) dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) queue = [] scheduler._process_task_instances(dag, queue=queue) self.assertEquals(len(queue), 2) dagbag = SimpleDagBag([dag]) # Recreated part of the scheduler here, to kick off tasks -> executor for ti_key in queue: task = dag.get_task(ti_key[1]) ti = models.TaskInstance(task, ti_key[2]) # Task starts out in the scheduled state. All tasks in the # scheduled state will be sent to the executor ti.state = State.SCHEDULED # Also save this task instance to the DB. session.merge(ti) session.commit() scheduler._execute_task_instances( dagbag, (State.SCHEDULED, State.UP_FOR_RETRY)) self.assertEquals(len(scheduler.executor.queued_tasks), 1) def test_scheduler_auto_align(self): """ Test if the schedule_interval will be auto aligned with the start_date such that if the start_date coincides with the schedule the first execution_date will be start_date, otherwise it will be start_date + interval. """ dag = DAG(dag_id='test_scheduler_auto_align_1', start_date=datetime.datetime(2016, 1, 1, 10, 10, 0), schedule_interval="4 5 * * *") dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) self.assertEquals(dr.execution_date, datetime.datetime(2016, 1, 2, 5, 4)) dag = DAG(dag_id='test_scheduler_auto_align_2', start_date=datetime.datetime(2016, 1, 1, 10, 10, 0), schedule_interval="10 10 * * *") dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) self.assertEquals(dr.execution_date, datetime.datetime(2016, 1, 1, 10, 10)) def test_scheduler_reschedule(self): """ Checks if tasks that are not taken up by the executor get rescheduled """ executor = TestExecutor() dagbag = DagBag(executor=executor) dagbag.dags.clear() dagbag.executor = executor dag = DAG(dag_id='test_scheduler_reschedule', start_date=DEFAULT_DATE) dag_task1 = DummyOperator(task_id='dummy', dag=dag, owner='airflow') dag.clear() dag.is_subdag = False session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.is_paused = False session.merge(orm_dag) session.commit() dagbag.bag_dag(dag=dag, root_dag=dag, parent_dag=dag) @mock.patch('airflow.models.DagBag', return_value=dagbag) @mock.patch('airflow.models.DagBag.collect_dags') def do_schedule(function, function2): # Use a empty file since the above mock will return the # expected DAGs. Also specify only a single file so that it doesn't # try to schedule the above DAG repeatedly. scheduler = SchedulerJob(num_runs=1, executor=executor, subdir=os.path.join( models.DAGS_FOLDER, "no_dags.py")) scheduler.heartrate = 0 scheduler.run() do_schedule() self.assertEquals(1, len(executor.queued_tasks)) executor.queued_tasks.clear() do_schedule() self.assertEquals(2, len(executor.queued_tasks)) def test_retry_still_in_executor(self): """ Checks if the scheduler does not put a task in limbo, when a task is retried but is still present in the executor. """ executor = TestExecutor() dagbag = DagBag(executor=executor) dagbag.dags.clear() dagbag.executor = executor dag = DAG(dag_id='test_retry_still_in_executor', start_date=DEFAULT_DATE) dag_task1 = BashOperator(task_id='test_retry_handling_op', bash_command='exit 1', retries=1, dag=dag, owner='airflow') dag.clear() dag.is_subdag = False session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.is_paused = False session.merge(orm_dag) session.commit() dagbag.bag_dag(dag=dag, root_dag=dag, parent_dag=dag) @mock.patch('airflow.models.DagBag', return_value=dagbag) @mock.patch('airflow.models.DagBag.collect_dags') def do_schedule(function, function2): # Use a empty file since the above mock will return the # expected DAGs. Also specify only a single file so that it doesn't # try to schedule the above DAG repeatedly. scheduler = SchedulerJob(num_runs=1, executor=executor, subdir=os.path.join( models.DAGS_FOLDER, "no_dags.py")) scheduler.heartrate = 0 scheduler.run() do_schedule() self.assertEquals(1, len(executor.queued_tasks)) def run_with_error(task): try: task.run() except AirflowException: pass ti_tuple = six.next(six.itervalues(executor.queued_tasks)) (command, priority, queue, ti) = ti_tuple ti.task = dag_task1 # fail execution run_with_error(ti) self.assertEqual(ti.state, State.UP_FOR_RETRY) self.assertEqual(ti.try_number, 1) # do not schedule do_schedule() self.assertTrue(executor.has_task(ti)) ti.refresh_from_db() self.assertEqual(ti.state, State.UP_FOR_RETRY) # now the executor has cleared and it should be allowed the re-queue executor.queued_tasks.clear() do_schedule() ti.refresh_from_db() self.assertEqual(ti.state, State.QUEUED) @unittest.skipUnless("INTEGRATION" in os.environ, "Can only run end to end") def test_retry_handling_job(self): """ Integration test of the scheduler not accidentally resetting the try_numbers for a task """ dag = self.dagbag.get_dag('test_retry_handling_job') dag_task1 = dag.get_task("test_retry_handling_op") dag.clear() scheduler = SchedulerJob(dag_id=dag.dag_id, num_runs=1) scheduler.heartrate = 0 scheduler.run() session = settings.Session() ti = session.query(TI).filter(TI.dag_id == dag.dag_id, TI.task_id == dag_task1.task_id).first() # make sure the counter has increased self.assertEqual(ti.try_number, 2) self.assertEqual(ti.state, State.UP_FOR_RETRY) def test_scheduler_run_duration(self): """ Verifies that the scheduler run duration limit is followed. """ dag_id = 'test_start_date_scheduling' dag = self.dagbag.get_dag(dag_id) dag.clear() self.assertTrue(dag.start_date > DEFAULT_DATE) expected_run_duration = 5 start_time = datetime.datetime.now() scheduler = SchedulerJob(dag_id, run_duration=expected_run_duration, **self.default_scheduler_args) scheduler.run() end_time = datetime.datetime.now() run_duration = (end_time - start_time).total_seconds() logging.info("Test ran in %.2fs, expected %.2fs", run_duration, expected_run_duration) assert run_duration - expected_run_duration < 5.0 def test_dag_with_system_exit(self): """ Test to check that a DAG with a system.exit() doesn't break the scheduler. """ dag_id = 'exit_test_dag' dag_ids = [dag_id] dag_directory = os.path.join(models.DAGS_FOLDER, "..", "dags_with_system_exit") dag_file = os.path.join(dag_directory, 'b_test_scheduler_dags.py') dagbag = DagBag(dag_folder=dag_file) for dag_id in dag_ids: dag = dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, subdir=dag_directory, num_runs=1, **self.default_scheduler_args) scheduler.run() session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) def test_dag_get_active_runs(self): """ Test to check that a DAG returns it's active runs """ now = datetime.datetime.now() six_hours_ago_to_the_hour = (now - datetime.timedelta(hours=6)).replace( minute=0, second=0, microsecond=0) START_DATE = six_hours_ago_to_the_hour DAG_NAME1 = 'get_active_runs_test' default_args = { 'owner': 'airflow', 'depends_on_past': False, 'start_date': START_DATE } dag1 = DAG(DAG_NAME1, schedule_interval='* * * * *', max_active_runs=1, default_args=default_args) run_this_1 = DummyOperator(task_id='run_this_1', dag=dag1) run_this_2 = DummyOperator(task_id='run_this_2', dag=dag1) run_this_2.set_upstream(run_this_1) run_this_3 = DummyOperator(task_id='run_this_3', dag=dag1) run_this_3.set_upstream(run_this_2) session = settings.Session() orm_dag = DagModel(dag_id=dag1.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag1.clear() dr = scheduler.create_dag_run(dag1) # We had better get a dag run self.assertIsNotNone(dr) execution_date = dr.execution_date running_dates = dag1.get_active_runs() try: running_date = running_dates[0] except: running_date = 'Except' self.assertEqual(execution_date, running_date, 'Running Date must match Execution Date')
class BackfillJobTest(unittest.TestCase): def setUp(self): self.parser = cli.CLIFactory.get_parser() self.dagbag = DagBag(include_examples=True) @unittest.skipIf('sqlite' in configuration.get('core', 'sql_alchemy_conn'), "concurrent access not supported in sqlite") 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') dag.clear() target_dag.clear() scheduler = SchedulerJob() queue = mock.Mock() scheduler._process_task_instances(target_dag, queue=queue) self.assertFalse(queue.append.called) job = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True) job.run() scheduler = SchedulerJob() queue = mock.Mock() scheduler._process_task_instances(target_dag, queue=queue) self.assertTrue(queue.append.called) target_dag.clear() dag.clear() @unittest.skipIf('sqlite' in configuration.get('core', 'sql_alchemy_conn'), "concurrent access not supported in sqlite") def test_backfill_multi_dates(self): dag = self.dagbag.get_dag('example_bash_operator') dag.clear() job = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE + datetime.timedelta(days=1), ignore_first_depends_on_past=True) job.run() session = settings.Session() drs = session.query(DagRun).filter( DagRun.dag_id == 'example_bash_operator').order_by( DagRun.execution_date).all() self.assertTrue(drs[0].execution_date == DEFAULT_DATE) self.assertTrue(drs[0].state == State.SUCCESS) self.assertTrue(drs[1].execution_date == DEFAULT_DATE + datetime.timedelta(days=1)) self.assertTrue(drs[1].state == State.SUCCESS) dag.clear() session.close() @unittest.skipIf('sqlite' in configuration.get('core', 'sql_alchemy_conn'), "concurrent access not supported in sqlite") def test_backfill_examples(self): """ Test backfilling example dags """ # some DAGs really are just examples... but try to make them work! skip_dags = [ 'example_http_operator', 'example_twitter_dag', 'example_trigger_target_dag', 'example_trigger_controller_dag', # tested above 'test_utils', # sleeps forever ] logger = logging.getLogger('BackfillJobTest.test_backfill_examples') dags = [ dag for dag in self.dagbag.dags.values() if 'example_dags' in dag.full_filepath and dag.dag_id not in skip_dags ] for dag in dags: dag.clear(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) for i, dag in enumerate(sorted(dags, key=lambda d: d.dag_id)): logger.info('*** Running example DAG #{}: {}'.format( i, dag.dag_id)) job = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True) job.run() def test_backfill_pooled_tasks(self): """ Test that queued tasks are executed by BackfillJob Test for https://github.com/airbnb/airflow/pull/1225 """ session = settings.Session() pool = Pool(pool='test_backfill_pooled_task_pool', slots=1) session.add(pool) session.commit() dag = self.dagbag.get_dag('test_backfill_pooled_task_dag') dag.clear() job = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # run with timeout because this creates an infinite loop if not # caught with timeout(seconds=30): job.run() ti = TI(task=dag.get_task('test_backfill_pooled_task'), execution_date=DEFAULT_DATE) ti.refresh_from_db() self.assertEqual(ti.state, State.SUCCESS) 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_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) dag.clear()
def dag_bag(): return DagBag(dag_folder="dags/", include_examples=False)
def execute(self, context: Context): if isinstance(self.execution_date, datetime.datetime): execution_date = self.execution_date elif isinstance(self.execution_date, str): execution_date = timezone.parse(self.execution_date) self.execution_date = execution_date else: execution_date = timezone.utcnow() if self.trigger_run_id: run_id = self.trigger_run_id else: run_id = DagRun.generate_run_id(DagRunType.MANUAL, execution_date) try: dag_run = trigger_dag( dag_id=self.trigger_dag_id, run_id=run_id, conf=self.conf, execution_date=self.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, self.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=self.execution_date, end_date=self.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 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, 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 setUp(self): configuration.load_test_config() self.dagbag = DagBag(dag_folder=DEV_NULL, include_examples=True) self.args = {'owner': 'airflow', 'start_date': DEFAULT_DATE} self.dag = DAG(TEST_DAG_ID, default_args=self.args)
class SchedulerJobTest(unittest.TestCase): def setUp(self): self.dagbag = DagBag() @provide_session def evaluate_dagrun( self, dag_id, expected_task_states, # dict of task_id: state dagrun_state, run_kwargs=None, advance_execution_date=False, session=None): """ Helper for testing DagRun states with simple two-task DAGS """ if run_kwargs is None: run_kwargs = {} scheduler = SchedulerJob() dag = self.dagbag.get_dag(dag_id) dag.clear() dr = scheduler.schedule_dag(dag) if advance_execution_date: # run a second time to schedule a dagrun after the start_date dr = scheduler.schedule_dag(dag) ex_date = dr.execution_date try: dag.run(start_date=ex_date, end_date=ex_date, **run_kwargs) except AirflowException: pass # test tasks for task_id, expected_state in expected_task_states.items(): task = dag.get_task(task_id) ti = TI(task, ex_date) ti.refresh_from_db() self.assertEqual(ti.state, expected_state) # load dagrun dr = session.query(DagRun).filter( DagRun.dag_id == dag.dag_id, DagRun.execution_date == ex_date).first() # dagrun is running self.assertEqual(dr.state, State.RUNNING) dag.get_active_runs() # dagrun failed self.assertEqual(dr.state, dagrun_state) def test_dagrun_fail(self): """ DagRuns with one failed and one incomplete root task -> FAILED """ self.evaluate_dagrun(dag_id='test_dagrun_states_fail', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.UPSTREAM_FAILED, }, dagrun_state=State.FAILED) def test_dagrun_success(self): """ DagRuns with one failed and one successful root task -> SUCCESS """ self.evaluate_dagrun(dag_id='test_dagrun_states_success', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.SUCCESS, }, dagrun_state=State.SUCCESS) def test_dagrun_root_fail(self): """ DagRuns with one successful and one failed root task -> FAILED """ self.evaluate_dagrun(dag_id='test_dagrun_states_root_fail', expected_task_states={ 'test_dagrun_succeed': State.SUCCESS, 'test_dagrun_fail': State.FAILED, }, dagrun_state=State.FAILED) def test_dagrun_deadlock(self): """ Deadlocked DagRun is marked a failure Test that a deadlocked dagrun is marked as a failure by having depends_on_past and an execution_date after the start_date """ self.evaluate_dagrun(dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': None, 'test_depends_on_past_2': None, }, dagrun_state=State.FAILED, advance_execution_date=True) 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=10) 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.FAILED) dag.clear() def test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date(self): """ DagRun is marked a success if ignore_first_depends_on_past=True Test that an otherwise-deadlocked dagrun is marked as a success if ignore_first_depends_on_past=True and the dagrun execution_date is after the start_date. """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, advance_execution_date=True, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_dagrun_deadlock_ignore_depends_on_past(self): """ Test that ignore_first_depends_on_past doesn't affect results (this is the same test as test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date except that start_date == execution_date so depends_on_past is irrelevant). """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_scheduler_start_date(self): """ Test that the scheduler respects start_dates, even when DAGS have run """ dag_id = 'test_start_date_scheduling' dag = self.dagbag.get_dag(dag_id) dag.clear() self.assertTrue(dag.start_date > DEFAULT_DATE) scheduler = SchedulerJob(dag_id, num_runs=2) scheduler.run() # zero tasks ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) # previously, running this backfill would kick off the Scheduler # because it would take the most recent run and start from there # That behavior still exists, but now it will only do so if after the # start date backfill = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) backfill.run() # one task ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) scheduler = SchedulerJob(dag_id, num_runs=2) scheduler.run() # still one task session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) def test_scheduler_multiprocessing(self): """ Test that the scheduler can successfully queue multiple dags in parallel """ dag_ids = ['test_start_date_scheduling', 'test_dagrun_states_success'] for dag_id in dag_ids: dag = self.dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, num_runs=2) scheduler.run() # zero tasks ran dag_id = 'test_start_date_scheduling' session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0)
def setUp(self): os.environ['AIRFLOW_VAR_ENTERPRISE_PROJECT'] = 'test' self.dagbag = DagBag() self.dag_id = 'Copy_Table_BigQuery'
def setUp(self): self.parser = cli.CLIFactory.get_parser() self.dagbag = DagBag(include_examples=True)
class BackfillJobTest(unittest.TestCase): def setUp(self): self.parser = cli.CLIFactory.get_parser() self.dagbag = DagBag(include_examples=True) def test_backfill_examples(self): """ Test backfilling example dags """ # some DAGs really are just examples... but try to make them work! skip_dags = [ 'example_http_operator', 'example_twitter_dag', ] logger = logging.getLogger('BackfillJobTest.test_backfill_examples') dags = [ dag for dag in self.dagbag.dags.values() if 'example_dags' in dag.full_filepath and dag.dag_id not in skip_dags ] for dag in dags: dag.clear(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) for i, dag in enumerate(sorted(dags, key=lambda d: d.dag_id)): logger.info('*** Running example DAG #{}: {}'.format( i, dag.dag_id)) job = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True) job.run() def test_backfill_pooled_tasks(self): """ Test that queued tasks are executed by BackfillJob Test for https://github.com/airbnb/airflow/pull/1225 """ session = settings.Session() pool = Pool(pool='test_backfill_pooled_task_pool', slots=1) session.add(pool) session.commit() dag = self.dagbag.get_dag('test_backfill_pooled_task_dag') dag.clear() job = BackfillJob(dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # run with timeout because this creates an infinite loop if not # caught with timeout(seconds=30): job.run() ti = TI(task=dag.get_task('test_backfill_pooled_task'), execution_date=DEFAULT_DATE) ti.refresh_from_db() self.assertEqual(ti.state, State.SUCCESS) def test_backfill_depends_on_past(self): """ Test that backfill resects 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_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 _trigger_dag( dag_id: str, dag_bag: DagBag, run_id: Optional[str] = None, conf: Optional[Union[dict, str]] = None, execution_date: Optional[datetime] = None, replace_microseconds: bool = True, ) -> List[DagRun]: # pylint: disable=too-many-arguments """Triggers DAG run. :param dag_id: DAG ID :param dag_bag: DAG Bag model :param run_id: ID of the dag_run :param conf: configuration :param execution_date: date of execution :param replace_microseconds: whether microseconds should be zeroed :return: list of triggered dags """ dag = dag_bag.get_dag(dag_id) # prefetch dag if it is stored serialized if dag_id not in dag_bag.dags: raise DagNotFound(f"Dag id {dag_id} not found") execution_date = execution_date if execution_date else timezone.utcnow() if not timezone.is_localized(execution_date): raise ValueError("The execution_date should be localized") if replace_microseconds: execution_date = execution_date.replace(microsecond=0) if dag.default_args and 'start_date' in dag.default_args: min_dag_start_date = dag.default_args["start_date"] if min_dag_start_date and execution_date < min_dag_start_date: raise ValueError( "The execution_date [{}] should be >= start_date [{}] from DAG's default_args" .format(execution_date.isoformat(), min_dag_start_date.isoformat())) run_id = run_id or DagRun.generate_run_id(DagRunType.MANUAL, execution_date) dag_run = DagRun.find(dag_id=dag_id, run_id=run_id) if dag_run: raise DagRunAlreadyExists( f"Run id {run_id} already exists for dag id {dag_id}") run_conf = None if conf: run_conf = conf if isinstance(conf, dict) else json.loads(conf) triggers = [] dags_to_trigger = [dag] + dag.subdags for _dag in dags_to_trigger: trigger = _dag.create_dagrun( run_id=run_id, execution_date=execution_date, state=State.RUNNING, conf=run_conf, external_trigger=True, dag_hash=dag_bag.dags_hash.get(dag_id), ) triggers.append(trigger) return triggers
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 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, zombies): super().__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 FakeDagFileProcessorRunner(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]))
class SchedulerJobTest(unittest.TestCase): def setUp(self): self.dagbag = DagBag() @provide_session def evaluate_dagrun( self, dag_id, first_task_state, second_task_state, dagrun_state, run_kwargs=None, advance_execution_date=False, session=None): """ Helper for testing DagRun states with simple two-task DAGS """ if run_kwargs is None: run_kwargs = {} scheduler = SchedulerJob() dag = self.dagbag.get_dag(dag_id) dag.clear() dr = scheduler.schedule_dag(dag) if advance_execution_date: # run a second time to schedule a dagrun after the start_date dr = scheduler.schedule_dag(dag) ex_date = dr.execution_date try: dag.run(start_date=ex_date, end_date=ex_date, **run_kwargs) except AirflowException: pass # test tasks task_1, task_2 = dag.tasks ti = TI(task_1, ex_date) ti.refresh_from_db() self.assertEqual(ti.state, first_task_state) ti = TI(task_2, ex_date) ti.refresh_from_db() self.assertEqual(ti.state, second_task_state) # load dagrun dr = session.query(DagRun).filter( DagRun.dag_id == dag.dag_id, DagRun.execution_date == ex_date ).first() # dagrun is running self.assertEqual(dr.state, State.RUNNING) dag.get_active_runs() # dagrun failed self.assertEqual(dr.state, dagrun_state) def test_dagrun_fail(self): """ DagRuns with one failed and one incomplete root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_fail', first_task_state=State.FAILED, second_task_state=State.UPSTREAM_FAILED, dagrun_state=State.FAILED) def test_dagrun_success(self): """ DagRuns with one failed and one successful root task -> SUCCESS """ self.evaluate_dagrun( dag_id='test_dagrun_states_success', first_task_state=State.FAILED, second_task_state=State.SUCCESS, dagrun_state=State.SUCCESS) def test_dagrun_root_fail(self): """ DagRuns with one successful and one failed root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_root_fail', first_task_state=State.SUCCESS, second_task_state=State.FAILED, dagrun_state=State.FAILED) def test_dagrun_deadlock(self): """ Deadlocked DagRun is marked a failure Test that a deadlocked dagrun is marked as a failure by having depends_on_past and an execution_date after the start_date """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', first_task_state=None, second_task_state=None, dagrun_state=State.FAILED, advance_execution_date=True) 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=10) scheduler.run() task_1 = dag.tasks[0] ti = TI(task_1, dag.start_date) ti.refresh_from_db() self.assertEqual(ti.state, State.FAILED) dag.clear() def test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date(self): """ DagRun is marked a success if ignore_first_depends_on_past=True Test that an otherwise-deadlocked dagrun is marked as a success if ignore_first_depends_on_past=True and the dagrun execution_date is after the start_date. """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', first_task_state=State.SUCCESS, second_task_state=State.SUCCESS, dagrun_state=State.SUCCESS, advance_execution_date=True, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_dagrun_deadlock_ignore_depends_on_past(self): """ Test that ignore_first_depends_on_past doesn't affect results (this is the same test as test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date except that start_date == execution_date so depends_on_past is irrelevant). """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', first_task_state=State.SUCCESS, second_task_state=State.SUCCESS, dagrun_state=State.SUCCESS, run_kwargs=dict(ignore_first_depends_on_past=True))
class SchedulerJobTest(unittest.TestCase): def setUp(self): self.dagbag = DagBag() @provide_session def evaluate_dagrun( self, dag_id, expected_task_states, # dict of task_id: state dagrun_state, run_kwargs=None, advance_execution_date=False, session=None): """ Helper for testing DagRun states with simple two-task DAGS. This is hackish: a dag run is created but its tasks are run by a backfill. """ if run_kwargs is None: run_kwargs = {} scheduler = SchedulerJob() dag = self.dagbag.get_dag(dag_id) dag.clear() dr = scheduler.schedule_dag(dag) if advance_execution_date: # run a second time to schedule a dagrun after the start_date dr = scheduler.schedule_dag(dag) ex_date = dr.execution_date try: dag.run(start_date=ex_date, end_date=ex_date, **run_kwargs) except AirflowException: pass # test tasks for task_id, expected_state in expected_task_states.items(): task = dag.get_task(task_id) ti = TI(task, ex_date) ti.refresh_from_db() self.assertEqual(ti.state, expected_state) # load dagrun dr = DagRun.find(dag_id=dag_id, execution_date=ex_date) dr = dr[0] dr.dag = dag # dagrun is running self.assertEqual(dr.state, State.RUNNING) dr.update_state() # dagrun failed self.assertEqual(dr.state, dagrun_state) def test_dagrun_fail(self): """ DagRuns with one failed and one incomplete root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_fail', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.UPSTREAM_FAILED, }, dagrun_state=State.FAILED) def test_dagrun_success(self): """ DagRuns with one failed and one successful root task -> SUCCESS """ self.evaluate_dagrun( dag_id='test_dagrun_states_success', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.SUCCESS, }, dagrun_state=State.SUCCESS) def test_dagrun_root_fail(self): """ DagRuns with one successful and one failed root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_root_fail', expected_task_states={ 'test_dagrun_succeed': State.SUCCESS, 'test_dagrun_fail': State.FAILED, }, dagrun_state=State.FAILED) def test_dagrun_deadlock(self): """ Deadlocked DagRun is marked a failure Test that a deadlocked dagrun is marked as a failure by having depends_on_past and an execution_date after the start_date """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': None, 'test_depends_on_past_2': None, }, dagrun_state=State.FAILED, advance_execution_date=True) 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_dagrun_deadlock_ignore_depends_on_past_advance_ex_date(self): """ DagRun is marked a success if ignore_first_depends_on_past=True Test that an otherwise-deadlocked dagrun is marked as a success if ignore_first_depends_on_past=True and the dagrun execution_date is after the start_date. """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, advance_execution_date=True, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_dagrun_deadlock_ignore_depends_on_past(self): """ Test that ignore_first_depends_on_past doesn't affect results (this is the same test as test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date except that start_date == execution_date so depends_on_past is irrelevant). """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_scheduler_start_date(self): """ Test that the scheduler respects start_dates, even when DAGS have run """ dag_id = 'test_start_date_scheduling' dag = self.dagbag.get_dag(dag_id) dag.clear() self.assertTrue(dag.start_date > DEFAULT_DATE) scheduler = SchedulerJob(dag_id, num_runs=2) scheduler.run() # zero tasks ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) # previously, running this backfill would kick off the Scheduler # because it would take the most recent run and start from there # That behavior still exists, but now it will only do so if after the # start date backfill = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) backfill.run() # one task ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) scheduler = SchedulerJob(dag_id, num_runs=2) scheduler.run() # still one task session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) def test_scheduler_multiprocessing(self): """ Test that the scheduler can successfully queue multiple dags in parallel """ dag_ids = ['test_start_date_scheduling', 'test_dagrun_states_success'] for dag_id in dag_ids: dag = self.dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, num_runs=2) scheduler.run() # zero tasks ran dag_id = 'test_start_date_scheduling' session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) def test_scheduler_dagrun_once(self): """ Test if the scheduler does not create multiple dagruns if a dag is scheduled with @once and a start_date """ dag = DAG( 'test_scheduler_dagrun_once', start_date=datetime.datetime(2015, 1, 1), schedule_interval="@once") scheduler = SchedulerJob() dag.clear() dr = scheduler.schedule_dag(dag) self.assertIsNotNone(dr) dr = scheduler.schedule_dag(dag) self.assertIsNone(dr) def test_scheduler_process_execute_task(self): """ Test if process dag sends a task to the executor """ dag = DAG( dag_id='test_scheduler_process_execute_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.schedule_dag(dag) self.assertIsNotNone(dr) queue = mock.Mock() scheduler.process_dag(dag, queue=queue) queue.put.assert_called_with( ((dag.dag_id, dag_task1.task_id, DEFAULT_DATE), None) ) tis = dr.get_task_instances(state=State.SCHEDULED) self.assertIsNotNone(tis) def test_scheduler_process_check_heartrate(self): """ Test if process dag honors the heartrate """ dag = DAG( dag_id='test_scheduler_process_check_heartrate', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.last_scheduler_run = datetime.datetime.now() session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() scheduler.heartrate = 1000 dag.clear() dr = scheduler.schedule_dag(dag) self.assertIsNotNone(dr) queue = mock.Mock() scheduler.process_dag(dag, queue=queue) queue.put.assert_not_called()
def setUp(self): self.parser = cli.CLIFactory.get_parser() self.dagbag = DagBag()
def make_example_dags(module_path): """Loads DAGs from a module for test.""" dagbag = DagBag(module_path) return dagbag.dags
def setUpClass(cls): cls.dagbag = DagBag()
class BackfillJobTest(unittest.TestCase): def setUp(self): self.parser = cli.CLIFactory.get_parser() self.dagbag = DagBag() def test_backfill_examples(self): """ Test backfilling example dags """ dags = [ dag for dag in self.dagbag.dags.values() if dag.dag_id in ('example_bash_operator',)] for dag in dags: dag.clear( start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) for dag in dags: job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) job.run() def test_trap_executor_error(self): """ Test that errors setting up tasks (before tasks run) are caught Test for https://github.com/airbnb/airflow/pull/1220 """ dag = self.dagbag.get_dag('test_raise_executor_error') dag.clear() job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # run with timeout because this creates an infinite loop if not # caught def run_with_timeout(): with timeout(seconds=30): job.run() self.assertRaises(AirflowException, run_with_timeout) def test_backfill_pooled_tasks(self): """ Test that queued tasks are executed by BackfillJob Test for https://github.com/airbnb/airflow/pull/1225 """ session = settings.Session() pool = Pool(pool='test_backfill_pooled_task_pool', slots=1) session.add(pool) session.commit() dag = self.dagbag.get_dag('test_backfill_pooled_task_dag') dag.clear() job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # run with timeout because this creates an infinite loop if not # caught with timeout(seconds=30): job.run() ti = TI( task=dag.get_task('test_backfill_pooled_task'), execution_date=DEFAULT_DATE) ti.refresh_from_db() self.assertEqual(ti.state, State.SUCCESS) def test_backfill_depends_on_past(self): """ Test that backfill resects 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_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 setUp(self): self.dagbag = DagBag(include_examples=True)
def list_dags(args): dagbag = DagBag(args.subdir) print("\n".join(sorted(dagbag.dags)))
def _trigger_dag( dag_id: str, dag_bag: DagBag, dag_run: DagModel, run_id: Optional[str], conf: Optional[Union[dict, str]], execution_date: Optional[datetime], replace_microseconds: bool, ) -> List[DagRun]: # pylint: disable=too-many-arguments """Triggers DAG run. :param dag_id: DAG ID :param dag_bag: DAG Bag model :param dag_run: DAG Run model :param run_id: ID of the dag_run :param conf: configuration :param execution_date: date of execution :param replace_microseconds: whether microseconds should be zeroed :return: list of triggered dags """ if dag_id not in dag_bag.dags: raise DagNotFound("Dag id {} not found".format(dag_id)) dag = dag_bag.get_dag(dag_id) execution_date = execution_date if execution_date else timezone.utcnow() assert timezone.is_localized(execution_date) if replace_microseconds: execution_date = execution_date.replace(microsecond=0) if not run_id: run_id = "manual__{0}".format(execution_date.isoformat()) dag_run_id = dag_run.find(dag_id=dag_id, run_id=run_id) if dag_run_id: raise DagRunAlreadyExists("Run id {} already exists for dag id {}".format( run_id, dag_id )) run_conf = None if conf: if isinstance(conf, dict): run_conf = conf else: run_conf = json.loads(conf) triggers = list() dags_to_trigger = list() dags_to_trigger.append(dag) while dags_to_trigger: dag = dags_to_trigger.pop() trigger = dag.create_dagrun( run_id=run_id, execution_date=execution_date, state=State.RUNNING, conf=run_conf, external_trigger=True, ) triggers.append(trigger) if dag.subdags: dags_to_trigger.extend(dag.subdags) return triggers
class BackfillJobTest(unittest.TestCase): def setUp(self): self.parser = cli.CLIFactory.get_parser() self.dagbag = DagBag(include_examples=True) def test_backfill_examples(self): """ Test backfilling example dags """ # some DAGs really are just examples... but try to make them work! skip_dags = [ 'example_http_operator', 'example_twitter_dag', ] logger = logging.getLogger('BackfillJobTest.test_backfill_examples') dags = [ dag for dag in self.dagbag.dags.values() if 'example_dags' in dag.full_filepath and dag.dag_id not in skip_dags ] for dag in dags: dag.clear( start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) for i, dag in enumerate(sorted(dags, key=lambda d: d.dag_id)): logger.info('*** Running example DAG #{}: {}'.format(i, dag.dag_id)) job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True) job.run() def test_backfill_pooled_tasks(self): """ Test that queued tasks are executed by BackfillJob Test for https://github.com/airbnb/airflow/pull/1225 """ session = settings.Session() pool = Pool(pool='test_backfill_pooled_task_pool', slots=1) session.add(pool) session.commit() dag = self.dagbag.get_dag('test_backfill_pooled_task_dag') dag.clear() job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # run with timeout because this creates an infinite loop if not # caught with timeout(seconds=30): job.run() ti = TI( task=dag.get_task('test_backfill_pooled_task'), execution_date=DEFAULT_DATE) ti.refresh_from_db() self.assertEqual(ti.state, State.SUCCESS) def test_backfill_depends_on_past(self): """ Test that backfill resects 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_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) dag.clear()
def make_example_dags(module, dag_ids): """Loads DAGs from a module for test.""" dagbag = DagBag(module.__path__[0]) return {k: dagbag.dags[k] for k in dag_ids}
def test_dag_cluster_policy_obeyed(self): dag_file = os.path.join(TEST_DAGS_FOLDER, "test_dag_with_no_tags.py") dagbag = DagBag(dag_folder=dag_file, include_examples=False, include_smart_sensor=False) assert len(dagbag.dag_ids) == 0 assert "has no tags" in dagbag.import_errors[dag_file]
def setUpClass(cls): super().setUpClass() DagBag(example_bash_operator.__file__).get_dag( "example_bash_operator").sync_to_db()
class SchedulerJobTest(unittest.TestCase): # These defaults make the test faster to run default_scheduler_args = {"file_process_interval": 0, "processor_poll_interval": 0.5} def setUp(self): self.dagbag = DagBag() @provide_session def evaluate_dagrun( self, dag_id, expected_task_states, # dict of task_id: state dagrun_state, run_kwargs=None, advance_execution_date=False, session=None): """ Helper for testing DagRun states with simple two-task DAGS. This is hackish: a dag run is created but its tasks are run by a backfill. """ if run_kwargs is None: run_kwargs = {} scheduler = SchedulerJob(**self.default_scheduler_args) dag = self.dagbag.get_dag(dag_id) dag.clear() dr = scheduler.create_dag_run(dag) if advance_execution_date: # run a second time to schedule a dagrun after the start_date dr = scheduler.create_dag_run(dag) ex_date = dr.execution_date try: dag.run(start_date=ex_date, end_date=ex_date, **run_kwargs) except AirflowException: pass # test tasks for task_id, expected_state in expected_task_states.items(): task = dag.get_task(task_id) ti = TI(task, ex_date) ti.refresh_from_db() self.assertEqual(ti.state, expected_state) # load dagrun dr = DagRun.find(dag_id=dag_id, execution_date=ex_date) dr = dr[0] dr.dag = dag self.assertEqual(dr.state, dagrun_state) def test_dagrun_fail(self): """ DagRuns with one failed and one incomplete root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_fail', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.UPSTREAM_FAILED, }, dagrun_state=State.FAILED) def test_dagrun_success(self): """ DagRuns with one failed and one successful root task -> SUCCESS """ self.evaluate_dagrun( dag_id='test_dagrun_states_success', expected_task_states={ 'test_dagrun_fail': State.FAILED, 'test_dagrun_succeed': State.SUCCESS, }, dagrun_state=State.SUCCESS) def test_dagrun_root_fail(self): """ DagRuns with one successful and one failed root task -> FAILED """ self.evaluate_dagrun( dag_id='test_dagrun_states_root_fail', expected_task_states={ 'test_dagrun_succeed': State.SUCCESS, 'test_dagrun_fail': State.FAILED, }, dagrun_state=State.FAILED) def test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date(self): """ DagRun is marked a success if ignore_first_depends_on_past=True Test that an otherwise-deadlocked dagrun is marked as a success if ignore_first_depends_on_past=True and the dagrun execution_date is after the start_date. """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, advance_execution_date=True, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_dagrun_deadlock_ignore_depends_on_past(self): """ Test that ignore_first_depends_on_past doesn't affect results (this is the same test as test_dagrun_deadlock_ignore_depends_on_past_advance_ex_date except that start_date == execution_date so depends_on_past is irrelevant). """ self.evaluate_dagrun( dag_id='test_dagrun_states_deadlock', expected_task_states={ 'test_depends_on_past': State.SUCCESS, 'test_depends_on_past_2': State.SUCCESS, }, dagrun_state=State.SUCCESS, run_kwargs=dict(ignore_first_depends_on_past=True)) def test_scheduler_start_date(self): """ Test that the scheduler respects start_dates, even when DAGS have run """ dag_id = 'test_start_date_scheduling' dag = self.dagbag.get_dag(dag_id) dag.clear() self.assertTrue(dag.start_date > DEFAULT_DATE) scheduler = SchedulerJob(dag_id, num_runs=2, **self.default_scheduler_args) scheduler.run() # zero tasks ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) # previously, running this backfill would kick off the Scheduler # because it would take the most recent run and start from there # That behavior still exists, but now it will only do so if after the # start date backfill = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) backfill.run() # one task ran session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) scheduler = SchedulerJob(dag_id, num_runs=2, **self.default_scheduler_args) scheduler.run() # still one task session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) def test_scheduler_multiprocessing(self): """ Test that the scheduler can successfully queue multiple dags in parallel """ dag_ids = ['test_start_date_scheduling', 'test_dagrun_states_success'] for dag_id in dag_ids: dag = self.dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, file_process_interval=0, processor_poll_interval=0.5, num_runs=2) scheduler.run() # zero tasks ran dag_id = 'test_start_date_scheduling' session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 0) def test_scheduler_dagrun_once(self): """ Test if the scheduler does not create multiple dagruns if a dag is scheduled with @once and a start_date """ dag = DAG( 'test_scheduler_dagrun_once', start_date=datetime.datetime(2015, 1, 1), schedule_interval="@once") scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dr = scheduler.create_dag_run(dag) self.assertIsNone(dr) def test_scheduler_process_task_instances(self): """ Test if _process_task_instances puts the right task instances into the queue. """ dag = DAG( dag_id='test_scheduler_process_execute_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.append.assert_called_with( (dag.dag_id, dag_task1.task_id, DEFAULT_DATE) ) def test_scheduler_do_not_schedule_removed_task(self): dag = DAG( dag_id='test_scheduler_do_not_schedule_removed_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dag = DAG( dag_id='test_scheduler_do_not_schedule_removed_task', start_date=DEFAULT_DATE) queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_do_not_schedule_too_early(self): dag = DAG( dag_id='test_scheduler_do_not_schedule_too_early', start_date=datetime.datetime(2200, 1, 1)) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNone(dr) queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_do_not_run_finished(self): dag = DAG( dag_id='test_scheduler_do_not_run_finished', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) tis = dr.get_task_instances(session=session) for ti in tis: ti.state = State.SUCCESS session.commit() session.close() queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_add_new_task(self): """ Test if a task instance will be added if the dag is updated """ dag = DAG( dag_id='test_scheduler_add_new_task', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) tis = dr.get_task_instances() self.assertEquals(len(tis), 1) dag_task2 = DummyOperator( task_id='dummy2', dag=dag, owner='airflow') queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) tis = dr.get_task_instances() self.assertEquals(len(tis), 2) def test_scheduler_does_not_run_excluded(self): dag = DAG( dag_id='test_scheduler_does_not_run_excluded', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) tis = dr.get_task_instances(session=session) for ti in tis: ti.state = State.EXCLUDED session.commit() session.close() queue = mock.Mock() scheduler._process_task_instances(dag, queue=queue) queue.put.assert_not_called() def test_scheduler_verify_max_active_runs(self): """ Test if a a dagrun will not be scheduled if max_dag_runs has been reached """ dag = DAG( dag_id='test_scheduler_verify_max_active_runs', start_date=DEFAULT_DATE) dag.max_active_runs = 1 dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dr = scheduler.create_dag_run(dag) self.assertIsNone(dr) def test_scheduler_fail_dagrun_timeout(self): """ Test if a a dagrun wil be set failed if timeout """ dag = DAG( dag_id='test_scheduler_fail_dagrun_timeout', start_date=DEFAULT_DATE) dag.dagrun_timeout = datetime.timedelta(seconds=60) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) dr.start_date = datetime.datetime.now() - datetime.timedelta(days=1) session.merge(dr) session.commit() dr2 = scheduler.create_dag_run(dag) self.assertIsNotNone(dr2) dr.refresh_from_db(session=session) self.assertEquals(dr.state, State.FAILED) def test_scheduler_verify_max_active_runs_and_dagrun_timeout(self): """ Test if a a dagrun will not be scheduled if max_dag_runs has been reached and dagrun_timeout is not reached Test if a a dagrun will be scheduled if max_dag_runs has been reached but dagrun_timeout is also reached """ dag = DAG( dag_id='test_scheduler_verify_max_active_runs_and_dagrun_timeout', start_date=DEFAULT_DATE) dag.max_active_runs = 1 dag.dagrun_timeout = datetime.timedelta(seconds=60) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) # Should not be scheduled as DagRun has not timedout and max_active_runs is reached new_dr = scheduler.create_dag_run(dag) self.assertIsNone(new_dr) # Should be scheduled as dagrun_timeout has passed dr.start_date = datetime.datetime.now() - datetime.timedelta(days=1) session.merge(dr) session.commit() new_dr = scheduler.create_dag_run(dag) self.assertIsNotNone(new_dr) def test_scheduler_max_active_runs_respected_after_clear(self): """ Test if _process_task_instances only schedules ti's up to max_active_runs (related to issue AIRFLOW-137) """ dag = DAG( dag_id='test_scheduler_max_active_runs_respected_after_clear', start_date=DEFAULT_DATE) dag.max_active_runs = 3 dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag.clear() # First create up to 3 dagruns in RUNNING state. scheduler.create_dag_run(dag) # Reduce max_active_runs to 1 dag.max_active_runs = 1 queue = mock.Mock() # and schedule them in, so we can check how many # tasks are put on the queue (should be one, not 3) scheduler._process_task_instances(dag, queue=queue) queue.append.assert_called_with( (dag.dag_id, dag_task1.task_id, DEFAULT_DATE) ) @patch.object(TI, 'pool_full') def test_scheduler_verify_pool_full(self, mock_pool_full): """ Test task instances not queued when pool is full """ mock_pool_full.return_value = False dag = DAG( dag_id='test_scheduler_verify_pool_full', start_date=DEFAULT_DATE) DummyOperator( task_id='dummy', dag=dag, owner='airflow', pool='test_scheduler_verify_pool_full') session = settings.Session() pool = Pool(pool='test_scheduler_verify_pool_full', slots=1) session.add(pool) orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.is_paused = False session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() # Create 2 dagruns, which will create 2 task instances. dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) self.assertEquals(dr.execution_date, DEFAULT_DATE) dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) queue = [] scheduler._process_task_instances(dag, queue=queue) self.assertEquals(len(queue), 2) dagbag = SimpleDagBag([dag]) # Recreated part of the scheduler here, to kick off tasks -> executor for ti_key in queue: task = dag.get_task(ti_key[1]) ti = models.TaskInstance(task, ti_key[2]) # Task starts out in the scheduled state. All tasks in the # scheduled state will be sent to the executor ti.state = State.SCHEDULED # Also save this task instance to the DB. session.merge(ti) session.commit() scheduler._execute_task_instances(dagbag, (State.SCHEDULED, State.UP_FOR_RETRY)) self.assertEquals(len(scheduler.executor.queued_tasks), 1) def test_scheduler_auto_align(self): """ Test if the schedule_interval will be auto aligned with the start_date such that if the start_date coincides with the schedule the first execution_date will be start_date, otherwise it will be start_date + interval. """ dag = DAG( dag_id='test_scheduler_auto_align_1', start_date=datetime.datetime(2016, 1, 1, 10, 10, 0), schedule_interval="4 5 * * *" ) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) self.assertEquals(dr.execution_date, datetime.datetime(2016, 1, 2, 5, 4)) dag = DAG( dag_id='test_scheduler_auto_align_2', start_date=datetime.datetime(2016, 1, 1, 10, 10, 0), schedule_interval="10 10 * * *" ) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) session.merge(orm_dag) session.commit() scheduler = SchedulerJob() dag.clear() dr = scheduler.create_dag_run(dag) self.assertIsNotNone(dr) self.assertEquals(dr.execution_date, datetime.datetime(2016, 1, 1, 10, 10)) def test_scheduler_reschedule(self): """ Checks if tasks that are not taken up by the executor get rescheduled """ executor = TestExecutor() dagbag = DagBag(executor=executor) dagbag.dags.clear() dagbag.executor = executor dag = DAG( dag_id='test_scheduler_reschedule', start_date=DEFAULT_DATE) dag_task1 = DummyOperator( task_id='dummy', dag=dag, owner='airflow') dag.clear() dag.is_subdag = False session = settings.Session() orm_dag = DagModel(dag_id=dag.dag_id) orm_dag.is_paused = False session.merge(orm_dag) session.commit() dagbag.bag_dag(dag=dag, root_dag=dag, parent_dag=dag) @mock.patch('airflow.models.DagBag', return_value=dagbag) @mock.patch('airflow.models.DagBag.collect_dags') def do_schedule(function, function2): # Use a empty file since the above mock will return the # expected DAGs. Also specify only a single file so that it doesn't # try to schedule the above DAG repeatedly. scheduler = SchedulerJob(num_runs=1, executor=executor, subdir=os.path.join(models.DAGS_FOLDER, "no_dags.py")) scheduler.heartrate = 0 scheduler.run() do_schedule() self.assertEquals(1, len(executor.queued_tasks)) executor.queued_tasks.clear() do_schedule() self.assertEquals(2, len(executor.queued_tasks)) def test_scheduler_run_duration(self): """ Verifies that the scheduler run duration limit is followed. """ dag_id = 'test_start_date_scheduling' dag = self.dagbag.get_dag(dag_id) dag.clear() self.assertTrue(dag.start_date > DEFAULT_DATE) expected_run_duration = 5 start_time = datetime.datetime.now() scheduler = SchedulerJob(dag_id, run_duration=expected_run_duration, **self.default_scheduler_args) scheduler.run() end_time = datetime.datetime.now() run_duration = (end_time - start_time).total_seconds() _log.info("Test ran in %.2fs, expected %.2fs", run_duration, expected_run_duration) assert run_duration - expected_run_duration < 5.0 def test_dag_with_system_exit(self): """ Test to check that a DAG with a system.exit() doesn't break the scheduler. """ dag_id = 'exit_test_dag' dag_ids = [dag_id] dag_directory = os.path.join(models.DAGS_FOLDER, "..", "dags_with_system_exit") dag_file = os.path.join(dag_directory, 'b_test_scheduler_dags.py') dagbag = DagBag(dag_folder=dag_file) for dag_id in dag_ids: dag = dagbag.get_dag(dag_id) dag.clear() scheduler = SchedulerJob(dag_ids=dag_ids, subdir= dag_directory, num_runs=1, **self.default_scheduler_args) scheduler.run() session = settings.Session() self.assertEqual( len(session.query(TI).filter(TI.dag_id == dag_id).all()), 1) def test_dag_get_active_runs(self): """ Test to check that a DAG returns it's active runs """ now = datetime.datetime.now() six_hours_ago_to_the_hour = (now - datetime.timedelta(hours=6)).replace(minute=0, second=0, microsecond=0) START_DATE = six_hours_ago_to_the_hour DAG_NAME1 = 'get_active_runs_test' default_args = { 'owner': 'airflow', 'depends_on_past': False, 'start_date': START_DATE } dag1 = DAG(DAG_NAME1, schedule_interval='* * * * *', max_active_runs=1, default_args=default_args ) run_this_1 = DummyOperator(task_id='run_this_1', dag=dag1) run_this_2 = DummyOperator(task_id='run_this_2', dag=dag1) run_this_2.set_upstream(run_this_1) run_this_3 = DummyOperator(task_id='run_this_3', dag=dag1) run_this_3.set_upstream(run_this_2) session = settings.Session() orm_dag = DagModel(dag_id=dag1.dag_id) session.merge(orm_dag) session.commit() session.close() scheduler = SchedulerJob() dag1.clear() dr = scheduler.create_dag_run(dag1) # We had better get a dag run self.assertIsNotNone(dr) execution_date = dr.execution_date running_dates = dag1.get_active_runs() try: running_date = running_dates[0] except: running_date = 'Except' self.assertEqual(execution_date, running_date, 'Running Date must match Execution Date')
def setUpClass(cls): cls.dagbag = DagBag(include_examples=True)
class BackfillJobTest(unittest.TestCase): def setUp(self): self.parser = cli.CLIFactory.get_parser() self.dagbag = DagBag(include_examples=True) @unittest.skipIf('sqlite' in configuration.get('core', 'sql_alchemy_conn'), "concurrent access not supported in sqlite") 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') dag.clear() target_dag.clear() scheduler = SchedulerJob() queue = mock.Mock() scheduler._process_task_instances(target_dag, queue=queue) self.assertFalse(queue.append.called) job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True ) job.run() scheduler = SchedulerJob() queue = mock.Mock() scheduler._process_task_instances(target_dag, queue=queue) self.assertTrue(queue.append.called) target_dag.clear() dag.clear() @unittest.skipIf('sqlite' in configuration.get('core', 'sql_alchemy_conn'), "concurrent access not supported in sqlite") def test_backfill_multi_dates(self): dag = self.dagbag.get_dag('example_bash_operator') dag.clear() job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE+datetime.timedelta(days=1), ignore_first_depends_on_past=True ) job.run() session = settings.Session() drs = session.query(DagRun).filter( DagRun.dag_id=='example_bash_operator' ).order_by(DagRun.execution_date).all() self.assertTrue(drs[0].execution_date == DEFAULT_DATE) self.assertTrue(drs[0].state == State.SUCCESS) self.assertTrue(drs[1].execution_date == DEFAULT_DATE+datetime.timedelta(days=1)) self.assertTrue(drs[1].state == State.SUCCESS) dag.clear() session.close() @unittest.skipIf('sqlite' in configuration.get('core', 'sql_alchemy_conn'), "concurrent access not supported in sqlite") def test_backfill_examples(self): """ Test backfilling example dags """ # some DAGs really are just examples... but try to make them work! skip_dags = [ 'example_http_operator', 'example_twitter_dag', 'example_trigger_target_dag', 'example_trigger_controller_dag', # tested above 'test_utils', # sleeps forever ] logger = logging.getLogger('BackfillJobTest.test_backfill_examples') dags = [ dag for dag in self.dagbag.dags.values() if 'example_dags' in dag.full_filepath and dag.dag_id not in skip_dags ] for dag in dags: dag.clear( start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) for i, dag in enumerate(sorted(dags, key=lambda d: d.dag_id)): logger.info('*** Running example DAG #{}: {}'.format(i, dag.dag_id)) job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True) job.run() def test_backfill_pooled_tasks(self): """ Test that queued tasks are executed by BackfillJob Test for https://github.com/airbnb/airflow/pull/1225 """ session = settings.Session() pool = Pool(pool='test_backfill_pooled_task_pool', slots=1) session.add(pool) session.commit() dag = self.dagbag.get_dag('test_backfill_pooled_task_dag') dag.clear() job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # run with timeout because this creates an infinite loop if not # caught with timeout(seconds=30): job.run() ti = TI( task=dag.get_task('test_backfill_pooled_task'), execution_date=DEFAULT_DATE) ti.refresh_from_db() self.assertEqual(ti.state, State.SUCCESS) 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_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) dag.clear()
class BackfillJobTest(unittest.TestCase): def setUp(self): self.parser = cli.CLIFactory.get_parser() self.dagbag = DagBag(include_examples=True) def test_backfill_examples(self): """ Test backfilling example dags """ # some DAGs really are just examples... but try to make them work! skip_dags = [ 'example_http_operator', ] logger = logging.getLogger('BackfillJobTest.test_backfill_examples') dags = [ dag for dag in self.dagbag.dags.values() if 'example_dags' in dag.full_filepath and dag.dag_id not in skip_dags ] for dag in dags: dag.clear( start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) for i, dag in enumerate(sorted(dags, key=lambda d: d.dag_id)): logger.info('*** Running example DAG #{}: {}'.format(i, dag.dag_id)) job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_first_depends_on_past=True) job.run() def test_trap_executor_error(self): """ Test that errors setting up tasks (before tasks run) are caught. Executors run tasks with the `airflow run` command. If a task runs, its state (success, failure, or other) is stored in the database and `airflow run` exits without error. However, if an error is raised before the task runs, then the task won't be able to update its status. This can put the executor into an infinite loop of trying to run the task. To counteract that, the executor traps errors coming from `airflow run` (essentially looking for returncode != 0). This unit test creates such an error by trying to run a subdag whose dag_id has changed and therefore can't be found. If the trap is working properly, the error will be caught and the Backfill will report failures. If the trap is not working, the job will run infinitely (the unit test uses a timeout to protect against that case). Test for https://github.com/airbnb/airflow/pull/1220 """ dag = self.dagbag.get_dag('test_raise_executor_error') dag.clear() job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # run with timeout because this creates an infinite loop if not # caught def run_with_timeout(): with timeout(seconds=30): job.run() self.assertRaises(AirflowException, run_with_timeout) def test_backfill_pooled_tasks(self): """ Test that queued tasks are executed by BackfillJob Test for https://github.com/airbnb/airflow/pull/1225 """ session = settings.Session() pool = Pool(pool='test_backfill_pooled_task_pool', slots=1) session.add(pool) session.commit() dag = self.dagbag.get_dag('test_backfill_pooled_task_dag') dag.clear() job = BackfillJob( dag=dag, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE) # run with timeout because this creates an infinite loop if not # caught with timeout(seconds=30): job.run() ti = TI( task=dag.get_task('test_backfill_pooled_task'), execution_date=DEFAULT_DATE) ti.refresh_from_db() self.assertEqual(ti.state, State.SUCCESS) def test_backfill_depends_on_past(self): """ Test that backfill resects 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_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 _trigger_dag( dag_id: str, dag_bag: DagBag, dag_run: DagModel, run_id: Optional[str], conf: Optional[Union[dict, str]], execution_date: Optional[datetime], replace_microseconds: bool, ) -> List[DagRun]: # pylint: disable=too-many-arguments """Triggers DAG run. :param dag_id: DAG ID :param dag_bag: DAG Bag model :param dag_run: DAG Run model :param run_id: ID of the dag_run :param conf: configuration :param execution_date: date of execution :param replace_microseconds: whether microseconds should be zeroed :return: list of triggered dags """ dag = dag_bag.get_dag(dag_id) # prefetch dag if it is stored serialized if dag_id not in dag_bag.dags: raise DagNotFound("Dag id {} not found".format(dag_id)) execution_date = execution_date if execution_date else timezone.utcnow() if not timezone.is_localized(execution_date): raise ValueError("The execution_date should be localized") if replace_microseconds: execution_date = execution_date.replace(microsecond=0) if dag.default_args and 'start_date' in dag.default_args: min_dag_start_date = dag.default_args["start_date"] if min_dag_start_date and execution_date < min_dag_start_date: raise ValueError( "The execution_date [{0}] should be >= start_date [{1}] from DAG's default_args".format( execution_date.isoformat(), min_dag_start_date.isoformat())) if not run_id: run_id = "manual__{0}".format(execution_date.isoformat()) dag_run_id = dag_run.find(dag_id=dag_id, run_id=run_id) if dag_run_id: raise DagRunAlreadyExists("Run id {} already exists for dag id {}".format( run_id, dag_id )) run_conf = None if conf: if isinstance(conf, dict): run_conf = conf else: run_conf = json.loads(conf) triggers = [] dags_to_trigger = [dag] while dags_to_trigger: dag = dags_to_trigger.pop() trigger = dag.create_dagrun( run_id=run_id, execution_date=execution_date, state=State.RUNNING, conf=run_conf, external_trigger=True, ) triggers.append(trigger) if dag.subdags: dags_to_trigger.extend(dag.subdags) return triggers
def setUp(self): self.dagbag = DagBag(include_examples=True) self.cluster = LocalCluster()
def setUpClass(cls): cls.dagbag = DagBag(include_examples=True) cls.parser = cli_parser.get_parser()
class KeywordDifficultyPipelineTest(TestCase): def setUp(self): self.dagbag = DagBag() self.dag_id = self.dagbag.get_dag( 'update_keywords_difficulty_pipeline') self.dag_login_aws = self.dag_id.tasks[0] self.dag_update_keywords_difficulty = self.dag_id.tasks[1] def test_dag_login_aws_should_see_task_id(self): actual = self.dag_login_aws.task_id expected = 'login_aws' self.assertEqual(actual, expected) def test_dag_login_aws_should_see_bash_command(self): actual = self.dag_login_aws.bash_command expected = '$(aws ecr get-login --region eu-west-1 --no-include-email)' self.assertEqual(actual, expected) def test_dag_update_keywords_difficulty_pipeline_should_see_task_id(self): actual = self.dag_update_keywords_difficulty.task_id expected = 'sent_ketwords_to_authoritylabs_pipeline' self.assertEqual(actual, expected) @patch.dict( 'os.environ', { 'DATABASE_HOST': '192.68.33.61', 'ELASTICSEARCH_URL': 'http://192.68.33.61:9200', 'DYNAMODB_HOST': 'http://192.68.33.61:4567', }) def test_dag_update_keywords_difficulty_pipeline_should_see_environment( self): expected = { 'DATABASE_HOST': '192.68.33.61', 'ELASTICSEARCH_URL': 'http://192.68.33.61:9200', 'DYNAMODB_HOST': 'http://192.68.33.61:4567', } self.assertEqual(os.environ['DATABASE_HOST'], expected.get('DATABASE_HOST')) self.assertEqual(os.environ['ELASTICSEARCH_URL'], expected.get('ELASTICSEARCH_URL')) self.assertEqual(os.environ['DYNAMODB_HOST'], expected.get('DYNAMODB_HOST')) def test_dag_update_keywords_difficulty_pipeline_should_see_auto_remove( self): actual = self.dag_update_keywords_difficulty.auto_remove self.assertTrue(actual) def test_dag_update_keywords_difficulty_pipeline_should_see_image(self): actual = self.dag_update_keywords_difficulty.image expected = '133506877714.dkr.ecr.eu-west-1.amazonaws.com/pronto-dashboard' self.assertIn(expected, actual) def test_dag_update_keywords_difficulty_pipeline_should_see_command(self): actual = self.dag_update_keywords_difficulty.command expected = 'python pronto_dashboard/manage.py update_keywords_difficulty --settings=pronto_dashboard.settings' self.assertIn(expected, actual)