def test_dagrun_find(self): session = settings.Session() now = timezone.utcnow() dag_id1 = "test_dagrun_find_externally_triggered" dag_run = models.DagRun( dag_id=dag_id1, run_type=DagRunType.MANUAL, execution_date=now, start_date=now, state=State.RUNNING, external_trigger=True, ) session.add(dag_run) dag_id2 = "test_dagrun_find_not_externally_triggered" dag_run = models.DagRun( dag_id=dag_id2, run_type=DagRunType.MANUAL, execution_date=now, start_date=now, state=State.RUNNING, external_trigger=False, ) session.add(dag_run) session.commit() assert 1 == len(models.DagRun.find(dag_id=dag_id1, external_trigger=True)) assert 0 == len(models.DagRun.find(dag_id=dag_id1, external_trigger=False)) assert 0 == len(models.DagRun.find(dag_id=dag_id2, external_trigger=True)) assert 1 == len(models.DagRun.find(dag_id=dag_id2, external_trigger=False))
def test_dagrun_find(self): session = settings.Session() now = datetime.datetime.now() dag_id1 = "test_dagrun_find_externally_triggered" dag_run = models.DagRun( dag_id=dag_id1, run_id='manual__' + now.isoformat(), execution_date=now, start_date=now, state=State.RUNNING, external_trigger=True, ) session.add(dag_run) dag_id2 = "test_dagrun_find_not_externally_triggered" dag_run = models.DagRun( dag_id=dag_id2, run_id='manual__' + now.isoformat(), execution_date=now, start_date=now, state=State.RUNNING, external_trigger=False, ) session.add(dag_run) session.commit() self.assertEqual(1, len(models.DagRun.find(dag_id=dag_id1, external_trigger=True))) self.assertEqual(0, len(models.DagRun.find(dag_id=dag_id1, external_trigger=False))) self.assertEqual(0, len(models.DagRun.find(dag_id=dag_id2, external_trigger=True))) self.assertEqual(1, len(models.DagRun.find(dag_id=dag_id2, external_trigger=False)))
def test_get_task_instance_on_empty_dagrun(self): """ Make sure that a proper value is returned when a dagrun has no task instances """ dag = DAG(dag_id='test_get_task_instance_on_empty_dagrun', start_date=timezone.datetime(2017, 1, 1)) ShortCircuitOperator(task_id='test_short_circuit_false', dag=dag, python_callable=lambda: False) session = settings.Session() now = timezone.utcnow() # Don't use create_dagrun since it will create the task instances too which we # don't want dag_run = models.DagRun( dag_id=dag.dag_id, run_id='manual__' + now.isoformat(), execution_date=now, start_date=now, state=State.RUNNING, external_trigger=False, ) session.add(dag_run) session.commit() ti = dag_run.get_task_instance('test_short_circuit_false') self.assertEqual(None, ti)
def test_schedule_dag_fake_scheduled_previous(self): """ Test scheduling a dag where there is a prior DagRun which has the same run_id as the next run should have """ delta = timedelta(hours=1) dag = DAG(TEST_DAG_ID + 'test_schedule_dag_fake_scheduled_previous', schedule_interval=delta, start_date=DEFAULT_DATE) dag.tasks = [ models.BaseOperator(task_id="faketastic", owner='Also fake', start_date=DEFAULT_DATE) ] scheduler = jobs.SchedulerJob(test_mode=True) trigger = models.DagRun(dag_id=dag.dag_id, run_id=models.DagRun.id_for_date(DEFAULT_DATE), execution_date=DEFAULT_DATE, state=utils.State.SUCCESS, external_trigger=True) settings.Session().add(trigger) settings.Session().commit() dag_run = scheduler.schedule_dag(dag) assert dag_run is not None assert dag_run.dag_id == dag.dag_id assert dag_run.run_id is not None assert dag_run.run_id != '' assert dag_run.execution_date == DEFAULT_DATE + delta, ( 'dag_run.execution_date did not match expectation: {0}'.format( dag_run.execution_date)) assert dag_run.state == models.State.RUNNING assert dag_run.external_trigger == False
def test_get_task_instance_on_empty_dagrun(self): """ Make sure that a proper value is returned when a dagrun has no task instances """ session = settings.Session() # Any dag will work for this dag = self.dagbag.get_dag('test_dagrun_short_circuit_false') now = datetime.datetime.now() # Don't use create_dagrun since it will create the task instances too which we # don't want dag_run = models.DagRun( dag_id=dag.dag_id, run_id='manual__' + now.isoformat(), execution_date=now, start_date=now, state=State.RUNNING, external_trigger=False, ) session.add(dag_run) session.commit() ti = dag_run.get_task_instance('test_short_circuit_false') self.assertEqual(None, ti)
def update_last_run(self): last_dag_run = self.last_run() if last_dag_run: dag_task_execution_date = self.previous_schedule( last_dag_run.execution_date) print dag_task_execution_date if dag_task_execution_date.date( ) != last_dag_run.execution_date.date(): dag_task_execution_date = datetime.combine( last_dag_run.execution_date.date(), dag_task_execution_date.time()) print dag_task_execution_date print last_dag_run.execution_date if dag_task_execution_date != last_dag_run.execution_date: session = settings.Session dag_re_schedule_run = models.DagRun( dag_id=self.dag_id, run_id='scheduled__' + dag_task_execution_date.isoformat(), execution_date=dag_task_execution_date, start_date=datetime.now(), end_date=datetime.now(), state=State.SUCCESS, external_trigger=False) session.add(dag_re_schedule_run) session.commit() for dag_task_id in self.task_ids: task_instance = models.TaskInstance( self.get_task(dag_task_id), execution_date=dag_task_execution_date, state=State.SUCCESS) task_instance.start_date = datetime.now() task_instance.end_date = datetime.now() session.add(task_instance) session.commit()