def test_delete_old_records(self, rtif_num, num_to_keep, remaining_rtifs, expected_query_count):
"""
Test that old records are deleted from rendered_task_instance_fields table
for a given task_id and dag_id.
"""
session = settings.Session()
dag = DAG("test_delete_old_records", start_date=START_DATE)
with dag:
task = BashOperator(task_id="test", bash_command="echo {{ ds }}")
rtif_list = [
RTIF(TI(task=task, execution_date=EXECUTION_DATE + timedelta(days=num)))
for num in range(rtif_num)
]
session.add_all(rtif_list)
session.commit()
result = session.query(RTIF).filter(RTIF.dag_id == dag.dag_id, RTIF.task_id == task.task_id).all()
for rtif in rtif_list:
self.assertIn(rtif, result)
self.assertEqual(rtif_num, len(result))
# Verify old records are deleted and only 'num_to_keep' records are kept
with assert_queries_count(expected_query_count):
RTIF.delete_old_records(task_id=task.task_id, dag_id=task.dag_id, num_to_keep=num_to_keep)
result = session.query(RTIF).filter(RTIF.dag_id == dag.dag_id, RTIF.task_id == task.task_id).all()
self.assertEqual(remaining_rtifs, len(result))
def test_delete_old_records(self):
"""
Test that old records are deleted from rendered_task_instance_fields table
for a given task_id and dag_id.
"""
session = settings.Session()
dag = DAG("test_delete_old_records", start_date=START_DATE)
with dag:
task = BashOperator(task_id="test", bash_command="echo {{ ds }}")
rtif_1 = RTIF(TI(task=task, execution_date=EXECUTION_DATE))
rtif_2 = RTIF(
TI(task=task, execution_date=EXECUTION_DATE + timedelta(days=1)))
rtif_3 = RTIF(
TI(task=task, execution_date=EXECUTION_DATE + timedelta(days=2)))
session.add(rtif_1)
session.add(rtif_2)
session.add(rtif_3)
session.commit()
result = session.query(RTIF)\
.filter(RTIF.dag_id == dag.dag_id, RTIF.task_id == task.task_id).all()
self.assertIn(rtif_1, result)
self.assertIn(rtif_2, result)
self.assertIn(rtif_3, result)
self.assertEqual(3, len(result))
# Verify old records are deleted and only 1 record is kept
RTIF.delete_old_records(task_id=task.task_id,
dag_id=task.dag_id,
num_to_keep=1)
result = session.query(RTIF) \
.filter(RTIF.dag_id == dag.dag_id, RTIF.task_id == task.task_id).all()
self.assertEqual(1, len(result))
self.assertEqual(rtif_3.execution_date, result[0].execution_date)
def test_delete_old_records(self, rtif_num, num_to_keep, remaining_rtifs, expected_query_count):
"""
Test that old records are deleted from rendered_task_instance_fields table
for a given task_id and dag_id.
"""
session = settings.Session()
dag = DAG("test_delete_old_records", start_date=START_DATE)
with dag:
task = BashOperator(task_id="test", bash_command="echo {{ ds }}")
rtif_list = [
RTIF(TI(task=task, execution_date=EXECUTION_DATE + timedelta(days=num)))
for num in range(rtif_num)
]
session.add_all(rtif_list)
session.commit()
result = session.query(RTIF).filter(RTIF.dag_id == dag.dag_id, RTIF.task_id == task.task_id).all()
for rtif in rtif_list:
assert rtif in result
assert rtif_num == len(result)
# Verify old records are deleted and only 'num_to_keep' records are kept
# For other DBs,an extra query is fired in RenderedTaskInstanceFields.delete_old_records
expected_query_count_based_on_db = (
expected_query_count + 1
if session.bind.dialect.name == "mssql" and expected_query_count != 0
else expected_query_count
)
with assert_queries_count(expected_query_count_based_on_db):
RTIF.delete_old_records(task_id=task.task_id, dag_id=task.dag_id, num_to_keep=num_to_keep)
result = session.query(RTIF).filter(RTIF.dag_id == dag.dag_id, RTIF.task_id == task.task_id).all()
assert remaining_rtifs == len(result)