Database backend optimization in Apache Airflow - Mini Project: Build & Apply

Database backend optimization

📖 Scenario: You are managing an Airflow setup that uses a database backend to store task metadata. Over time, the database has grown large and slow. You want to optimize the database backend by cleaning up old task instances and setting a retention policy.

🎯 Goal: Build an Airflow DAG that performs database cleanup by deleting task instances older than a certain number of days, using a configuration variable for the retention period.

📋 What You'll Learn

Create an Airflow DAG with a cleanup task

Use a configuration variable for retention days

Implement the cleanup logic using Airflow's ORM

Print the number of deleted records after cleanup

💡 Why This Matters

🌍 Real World

Airflow uses a database backend to store metadata about tasks and DAG runs. Over time, this data grows and can slow down the system. Cleaning up old task instances helps keep the database fast and manageable.

💼 Career

DevOps engineers and data engineers often need to maintain Airflow environments. Knowing how to optimize the database backend by cleaning old data is a practical skill to improve system performance and reliability.

1

Create the Airflow DAG skeleton

Create an Airflow DAG called db_cleanup_dag with start_date set to 2024-01-01 and schedule_interval set to @daily. Import DAG from airflow and datetime from datetime.

Apache Airflow

# Your code here

Hint

Use DAG constructor with dag_id, start_date, and schedule_interval.

2

Add a retention period configuration variable

Add a variable called retention_days and set it to 30 to represent the number of days to keep task instances.

Apache Airflow

from airflow import DAG
from datetime import datetime

# Add retention_days variable below

Hint

Just create a variable retention_days and assign 30.

3

Implement the cleanup task logic

Import PythonOperator from airflow.operators.python, TaskInstance from airflow.models, and timedelta from datetime. Define a function cleanup_old_task_instances that deletes TaskInstance records older than retention_days. Use Airflow's ORM session to query and delete these records. Then create a PythonOperator named cleanup_task that runs this function inside the db_cleanup_dag.

Apache Airflow

from airflow import DAG
from datetime import datetime

retention_days = 30

# Define cleanup_old_task_instances function and cleanup_task operator below

Hint

Use @provide_session decorator for the function. Query TaskInstance with execution_date less than cutoff date. Delete and commit. Create PythonOperator with python_callable=cleanup_old_task_instances.

4

Print the cleanup result

Add a print statement inside the cleanup_old_task_instances function that outputs the number of deleted task instances in the format: Deleted X old task instances, where X is the count of deleted records.

Apache Airflow

from airflow import DAG
from datetime import datetime, timedelta
from airflow.operators.python import PythonOperator
from airflow.models import TaskInstance
from airflow.utils.session import provide_session

retention_days = 30

db_cleanup_dag = DAG(
    dag_id='db_cleanup_dag',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@daily'
)

@provide_session
def cleanup_old_task_instances(session=None):
    cutoff_date = datetime.utcnow() - timedelta(days=retention_days)
    old_tasks = session.query(TaskInstance).filter(TaskInstance.execution_date < cutoff_date)
    count = old_tasks.count()
    old_tasks.delete(synchronize_session=False)
    session.commit()
    # Add print statement below

cleanup_task = PythonOperator(
    task_id='cleanup_old_tasks',
    python_callable=cleanup_old_task_instances,
    dag=db_cleanup_dag
)

Hint

Use print(f"Deleted {count} old task instances") inside the function.