def clear_dag_runs(dag_id, start_date, end_date):
"""
Clears all the DagRuns and corrects the DagStats for an interval passed in the clear command because the
clear command only clears the TaskInstances.
:param dag_id: Dag id name. String.
:param start_date: Start date. String of form %Y-%m-%d %H:%M:%S.
:param end_date: End date. String of form %Y-%m-%d %H:%M:%S.
:return: None.
"""
start_date = datetime.strptime(start_date, '%Y-%m-%d %H:%M:%S')
end_date = datetime.strptime(end_date, '%Y-%m-%d %H:%M:%S')
session = Session()
try:
dagrun_query = session.query(DagRun) \
.filter(DagRun.dag_id == dag_id) \
.filter(DagRun.execution_date >= start_date) \
.filter(DagRun.execution_date < end_date)
dagrun_query_result = dagrun_query.all()
# remove dagruns with this state for clear command
for result in dagrun_query_result:
session.delete(result)
# fix DagStats
for state in State.dag_states:
removed_state_counts = dagrun_query.filter(
DagRun.state == state).count()
dagstat_query = session.query(DagStat) \
.filter(DagStat.dag_id == dag_id) \
.filter(DagStat.state == state)
dagstat_query_result = dagstat_query.first(
) # only one row every time
dagstat_query_result.count = max(
dagstat_query_result.count - removed_state_counts, 0)
session.commit()
except:
session.rollback()
finally:
session.close()
