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()
