def test_emit_scheduling_delay(self, stats_mock):
"""
Tests that dag scheduling delay stat is set properly once running scheduled dag.
dag_run.update_state() invokes the _emit_true_scheduling_delay_stats_for_finished_state method.
"""
dag = DAG(dag_id='test_emit_dag_stats', start_date=days_ago(1))
dag_task = DummyOperator(task_id='dummy', dag=dag, owner='airflow')
session = settings.Session()
orm_dag = DagModel(
dag_id=dag.dag_id,
has_task_concurrency_limits=False,
next_dagrun=dag.start_date,
next_dagrun_create_after=dag.following_schedule(dag.start_date),
is_active=True,
)
session.add(orm_dag)
session.flush()
dag_run = dag.create_dagrun(
run_type=DagRunType.SCHEDULED,
state=State.SUCCESS,
execution_date=dag.start_date,
start_date=dag.start_date,
session=session,
)
ti = dag_run.get_task_instance(dag_task.task_id)
ti.set_state(State.SUCCESS, session)
session.commit()
session.close()
dag_run.update_state()
true_delay = (ti.start_date - dag.following_schedule(dag_run.execution_date)).total_seconds()
stats_mock.assert_called()
sched_delay_stat_call = call(f'dagrun.{dag.dag_id}.first_task_scheduling_delay', true_delay)
self.assertIn(sched_delay_stat_call, stats_mock.mock_calls)
def test_emit_scheduling_delay(self, schedule_interval, expected):
"""
Tests that dag scheduling delay stat is set properly once running scheduled dag.
dag_run.update_state() invokes the _emit_true_scheduling_delay_stats_for_finished_state method.
"""
dag = DAG(dag_id='test_emit_dag_stats',
start_date=days_ago(1),
schedule_interval=schedule_interval)
dag_task = DummyOperator(task_id='dummy', dag=dag, owner='airflow')
session = settings.Session()
try:
orm_dag = DagModel(
dag_id=dag.dag_id,
has_task_concurrency_limits=False,
next_dagrun=dag.start_date,
next_dagrun_create_after=dag.following_schedule(
dag.start_date),
is_active=True,
)
session.add(orm_dag)
session.flush()
dag_run = dag.create_dagrun(
run_type=DagRunType.SCHEDULED,
state=State.SUCCESS,
execution_date=dag.start_date,
start_date=dag.start_date,
session=session,
)
ti = dag_run.get_task_instance(dag_task.task_id, session)
ti.set_state(State.SUCCESS, session)
session.flush()
with mock.patch.object(Stats, 'timing') as stats_mock:
dag_run.update_state(session)
metric_name = f'dagrun.{dag.dag_id}.first_task_scheduling_delay'
if expected:
true_delay = ti.start_date - dag.following_schedule(
dag_run.execution_date)
sched_delay_stat_call = call(metric_name, true_delay)
assert sched_delay_stat_call in stats_mock.mock_calls
else:
# Assert that we never passed the metric
sched_delay_stat_call = call(metric_name, mock.ANY)
assert sched_delay_stat_call not in stats_mock.mock_calls
finally:
# Don't write anything to the DB
session.rollback()
session.close()
def test_following_previous_schedule_daily_dag_CET_to_CEST(self):
"""
Make sure DST transitions are properly observed
"""
local_tz = pendulum.timezone('Europe/Zurich')
start = local_tz.convert(datetime.datetime(2018, 3, 25, 2),
dst_rule=pendulum.PRE_TRANSITION)
utc = timezone.convert_to_utc(start)
dag = DAG('tz_dag', start_date=start, schedule_interval='0 3 * * *')
prev = dag.previous_schedule(utc)
prev_local = local_tz.convert(prev)
self.assertEqual(prev_local.isoformat(), "2018-03-24T03:00:00+01:00")
self.assertEqual(prev.isoformat(), "2018-03-24T02:00:00+00:00")
_next = dag.following_schedule(utc)
next_local = local_tz.convert(_next)
self.assertEqual(next_local.isoformat(), "2018-03-25T03:00:00+02:00")
self.assertEqual(_next.isoformat(), "2018-03-25T01:00:00+00:00")
prev = dag.previous_schedule(_next)
prev_local = local_tz.convert(prev)
self.assertEqual(prev_local.isoformat(), "2018-03-24T03:00:00+01:00")
self.assertEqual(prev.isoformat(), "2018-03-24T02:00:00+00:00")
def test_following_previous_schedule(self):
"""
Make sure DST transitions are properly observed
"""
local_tz = pendulum.timezone('Europe/Zurich')
start = local_tz.convert(datetime.datetime(2018, 10, 28, 2, 55),
dst_rule=pendulum.PRE_TRANSITION)
self.assertEqual(start.isoformat(), "2018-10-28T02:55:00+02:00",
"Pre-condition: start date is in DST")
utc = timezone.convert_to_utc(start)
dag = DAG('tz_dag', start_date=start, schedule_interval='*/5 * * * *')
_next = dag.following_schedule(utc)
next_local = local_tz.convert(_next)
self.assertEqual(_next.isoformat(), "2018-10-28T01:00:00+00:00")
self.assertEqual(next_local.isoformat(), "2018-10-28T02:00:00+01:00")
prev = dag.previous_schedule(utc)
prev_local = local_tz.convert(prev)
self.assertEqual(prev_local.isoformat(), "2018-10-28T02:50:00+02:00")
prev = dag.previous_schedule(_next)
prev_local = local_tz.convert(prev)
self.assertEqual(prev_local.isoformat(), "2018-10-28T02:55:00+02:00")
self.assertEqual(prev, utc)
def test_following_previous_schedule_daily_dag_CET_to_CEST(self):
"""
Make sure DST transitions are properly observed
"""
local_tz = pendulum.timezone('Europe/Zurich')
start = local_tz.convert(datetime.datetime(2018, 3, 25, 2),
dst_rule=pendulum.PRE_TRANSITION)
utc = timezone.convert_to_utc(start)
dag = DAG('tz_dag', start_date=start, schedule_interval='0 3 * * *')
prev = dag.previous_schedule(utc)
prev_local = local_tz.convert(prev)
self.assertEqual(prev_local.isoformat(), "2018-03-24T03:00:00+01:00")
self.assertEqual(prev.isoformat(), "2018-03-24T02:00:00+00:00")
_next = dag.following_schedule(utc)
next_local = local_tz.convert(_next)
self.assertEqual(next_local.isoformat(), "2018-03-25T03:00:00+02:00")
self.assertEqual(_next.isoformat(), "2018-03-25T01:00:00+00:00")
prev = dag.previous_schedule(_next)
prev_local = local_tz.convert(prev)
self.assertEqual(prev_local.isoformat(), "2018-03-24T03:00:00+01:00")
self.assertEqual(prev.isoformat(), "2018-03-24T02:00:00+00:00")
def test_following_previous_schedule(self):
"""
Make sure DST transitions are properly observed
"""
local_tz = pendulum.timezone('Europe/Zurich')
start = local_tz.convert(datetime.datetime(2018, 10, 28, 2, 55),
dst_rule=pendulum.PRE_TRANSITION)
self.assertEqual(start.isoformat(), "2018-10-28T02:55:00+02:00",
"Pre-condition: start date is in DST")
utc = timezone.convert_to_utc(start)
dag = DAG('tz_dag', start_date=start, schedule_interval='*/5 * * * *')
_next = dag.following_schedule(utc)
next_local = local_tz.convert(_next)
self.assertEqual(_next.isoformat(), "2018-10-28T01:00:00+00:00")
self.assertEqual(next_local.isoformat(), "2018-10-28T02:00:00+01:00")
prev = dag.previous_schedule(utc)
prev_local = local_tz.convert(prev)
self.assertEqual(prev_local.isoformat(), "2018-10-28T02:50:00+02:00")
prev = dag.previous_schedule(_next)
prev_local = local_tz.convert(prev)
self.assertEqual(prev_local.isoformat(), "2018-10-28T02:55:00+02:00")
self.assertEqual(prev, utc)
def test_next_dagruns_to_examine_only_unpaused(self):
"""
Check that "next_dagruns_to_examine" ignores runs from paused/inactive DAGs
"""
dag = DAG(dag_id='test_dags', start_date=DEFAULT_DATE)
DummyOperator(task_id='dummy', dag=dag, owner='airflow')
session = settings.Session()
orm_dag = DagModel(
dag_id=dag.dag_id,
has_task_concurrency_limits=False,
next_dagrun=dag.start_date,
next_dagrun_create_after=dag.following_schedule(DEFAULT_DATE),
is_active=True,
)
session.add(orm_dag)
session.flush()
dr = dag.create_dagrun(
run_type=DagRunType.SCHEDULED,
state=State.RUNNING,
execution_date=DEFAULT_DATE,
start_date=DEFAULT_DATE,
session=session,
)
runs = DagRun.next_dagruns_to_examine(session).all()
assert runs == [dr]
orm_dag.is_paused = True
session.flush()
runs = DagRun.next_dagruns_to_examine(session).all()
assert runs == []
session.rollback()
session.close()
def manage_slas(self, dag: DAG, session: Session = None) -> None:
"""
Finding all tasks that have SLAs defined, and sending alert emails
where needed. New SLA misses are also recorded in the database.
We are assuming that the scheduler runs often, so we only check for
tasks that should have succeeded in the past hour.
"""
self.log.info("Running SLA Checks for %s", dag.dag_id)
if not any(isinstance(ti.sla, timedelta) for ti in dag.tasks):
self.log.info(
"Skipping SLA check for %s because no tasks in DAG have SLAs",
dag)
return
qry = (session.query(
TI.task_id,
func.max(TI.execution_date).label('max_ti')).with_hint(
TI, 'USE INDEX (PRIMARY)',
dialect_name='mysql').filter(TI.dag_id == dag.dag_id).filter(
or_(TI.state == State.SUCCESS,
TI.state == State.SKIPPED)).filter(
TI.task_id.in_(dag.task_ids)).group_by(
TI.task_id).subquery('sq'))
max_tis: List[TI] = (session.query(TI).filter(
TI.dag_id == dag.dag_id,
TI.task_id == qry.c.task_id,
TI.execution_date == qry.c.max_ti,
).all())
ts = timezone.utcnow()
for ti in max_tis:
task = dag.get_task(ti.task_id)
if task.sla and not isinstance(task.sla, timedelta):
raise TypeError(
f"SLA is expected to be timedelta object, got "
f"{type(task.sla)} in {task.dag_id}:{task.task_id}")
dttm = dag.following_schedule(ti.execution_date)
while dttm < timezone.utcnow():
following_schedule = dag.following_schedule(dttm)
if following_schedule + task.sla < timezone.utcnow():
session.merge(
SlaMiss(task_id=ti.task_id,
dag_id=ti.dag_id,
execution_date=dttm,
timestamp=ts))
dttm = dag.following_schedule(dttm)
session.commit()
# pylint: disable=singleton-comparison
slas: List[SlaMiss] = (
session.query(SlaMiss).filter(SlaMiss.notification_sent == False,
SlaMiss.dag_id == dag.dag_id) # noqa
.all())
# pylint: enable=singleton-comparison
if slas: # pylint: disable=too-many-nested-blocks
sla_dates: List[datetime.datetime] = [
sla.execution_date for sla in slas
]
fetched_tis: List[TI] = (session.query(TI).filter(
TI.state != State.SUCCESS, TI.execution_date.in_(sla_dates),
TI.dag_id == dag.dag_id).all())
blocking_tis: List[TI] = []
for ti in fetched_tis:
if ti.task_id in dag.task_ids:
ti.task = dag.get_task(ti.task_id)
blocking_tis.append(ti)
else:
session.delete(ti)
session.commit()
task_list = "\n".join(sla.task_id + ' on ' +
sla.execution_date.isoformat()
for sla in slas)
blocking_task_list = "\n".join(ti.task_id + ' on ' +
ti.execution_date.isoformat()
for ti in blocking_tis)
# Track whether email or any alert notification sent
# We consider email or the alert callback as notifications
email_sent = False
notification_sent = False
if dag.sla_miss_callback:
# Execute the alert callback
self.log.info('Calling SLA miss callback')
try:
dag.sla_miss_callback(dag, task_list, blocking_task_list,
slas, blocking_tis)
notification_sent = True
except Exception: # pylint: disable=broad-except
self.log.exception(
"Could not call sla_miss_callback for DAG %s",
dag.dag_id)
email_content = f"""\
Here's a list of tasks that missed their SLAs:
<pre><code>{task_list}\n<code></pre>
Blocking tasks:
<pre><code>{blocking_task_list}<code></pre>
Airflow Webserver URL: {conf.get(section='webserver', key='base_url')}
"""
tasks_missed_sla = []
for sla in slas:
try:
task = dag.get_task(sla.task_id)
except TaskNotFound:
# task already deleted from DAG, skip it
self.log.warning(
"Task %s doesn't exist in DAG anymore, skipping SLA miss notification.",
sla.task_id)
continue
tasks_missed_sla.append(task)
emails: Set[str] = set()
for task in tasks_missed_sla:
if task.email:
if isinstance(task.email, str):
emails |= set(get_email_address_list(task.email))
elif isinstance(task.email, (list, tuple)):
emails |= set(task.email)
if emails:
try:
send_email(emails,
f"[airflow] SLA miss on DAG={dag.dag_id}",
email_content)
email_sent = True
notification_sent = True
except Exception: # pylint: disable=broad-except
Stats.incr('sla_email_notification_failure')
self.log.exception(
"Could not send SLA Miss email notification for DAG %s",
dag.dag_id)
# If we sent any notification, update the sla_miss table
if notification_sent:
for sla in slas:
sla.email_sent = email_sent
sla.notification_sent = True
session.merge(sla)
session.commit()
