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.

Practice

1. What is the main benefit of optimizing the Airflow database backend?

easy

A. Increases the number of DAGs you can create

B. Improves task scheduling and UI speed

C. Automatically fixes all DAG errors

D. Enables Airflow to run without a scheduler

Solution

1. Step 1: Understand Airflow database role

   The Airflow database stores metadata about tasks and DAGs, which affects scheduling and UI responsiveness.

2. Step 2: Connect optimization to performance

   Optimizing the database backend reduces delays in task scheduling and speeds up the UI loading times.

3. Final Answer:

   Improves task scheduling and UI speed -> Option B

4. Quick Check:

   Database optimization = better scheduling and UI [OK]

Hint: Think about what slows down task scheduling and UI [OK]

Common Mistakes:

• Confusing database optimization with adding more DAGs
• Assuming it fixes DAG code errors automatically
• Thinking it removes the need for a scheduler

2. Which command correctly upgrades the Airflow database schema?

easy

A. airflow db refresh

B. airflow db migrate

C. airflow db update

D. airflow db upgrade

Solution

1. Step 1: Recall Airflow database commands

   The official command to upgrade the Airflow database schema is airflow db upgrade.

2. Step 2: Verify command correctness

   Other options like migrate, update, or refresh are not valid Airflow CLI commands for database upgrade.

3. Final Answer:

   airflow db upgrade -> Option D

4. Quick Check:

   Upgrade command = airflow db upgrade [OK]

Hint: Remember 'upgrade' is the official schema update command [OK]

Common Mistakes:

• Using 'migrate' instead of 'upgrade'
• Confusing with generic database commands
• Typing 'update' or 'refresh' which are invalid

3. Given this airflow.cfg snippet, what is the effect of setting sql_alchemy_pool_size = 10?

[database]
sql_alchemy_pool_size = 10

medium

A. Limits the number of simultaneous DB connections to 10

B. Sets the max retries for DB connections to 10

C. Defines the timeout in seconds for DB queries

D. Enables 10 read replicas for the database

Solution

1. Step 1: Understand sql_alchemy_pool_size meaning

   This setting controls the maximum number of open connections in the SQLAlchemy connection pool.

2. Step 2: Interpret the value 10

   Setting it to 10 means Airflow can have up to 10 simultaneous database connections open.

3. Final Answer:

   Limits the number of simultaneous DB connections to 10 -> Option A

4. Quick Check:

   sql_alchemy_pool_size = max DB connections [OK]

Hint: Pool size means max open DB connections [OK]

Common Mistakes:

• Confusing pool size with retry count
• Thinking it sets query timeout
• Assuming it controls DB replicas

4. You run airflow db check and get a connection timeout error. What is the most likely cause?

medium

A. Database server is unreachable or down

B. The airflow.cfg file is missing

C. The DAG files have syntax errors

D. The Airflow scheduler is not running

Solution

1. Step 1: Analyze the error type

   A connection timeout error means Airflow cannot reach the database server.

2. Step 2: Rule out other options

   Missing config or DAG syntax errors do not cause connection timeouts. Scheduler status is unrelated to DB connection.

3. Final Answer:

   Database server is unreachable or down -> Option A

4. Quick Check:

   Connection timeout = DB unreachable [OK]

Hint: Timeout means no DB response, check server status [OK]

Common Mistakes:

• Blaming missing config for connection timeout
• Thinking DAG errors cause DB connection issues
• Assuming scheduler status affects DB connection

5. To improve Airflow database performance under heavy load, which combination of settings is best to tune in airflow.cfg?

hard

A. Increase sql_alchemy_pool_size and reduce sql_alchemy_max_overflow

B. Disable connection pooling by setting sql_alchemy_pool_size to 1

C. Decrease sql_alchemy_pool_size and increase sql_alchemy_max_overflow

D. Set sql_alchemy_pool_size and sql_alchemy_max_overflow both to zero

Solution

1. Step 1: Understand pool size and max overflow

   sql_alchemy_pool_size sets base DB connections; max_overflow allows extra temporary connections beyond pool size.

2. Step 2: Tune for heavy load

   Decreasing pool size but increasing max_overflow lets Airflow handle bursts by opening extra connections temporarily.

3. Step 3: Evaluate other options

   Setting both to zero disables pooling (bad), increasing pool size but reducing overflow limits burst handling, disabling pooling hurts performance.

4. Final Answer:

   Decrease sql_alchemy_pool_size and increase sql_alchemy_max_overflow -> Option C

5. Quick Check:

   Pool size low + max overflow high = better burst handling [OK]

Hint: Use small pool + high overflow for burst DB load [OK]

Common Mistakes:

• Setting both pool size and overflow to zero disables pooling
• Increasing pool size but limiting overflow reduces flexibility
• Disabling pooling hurts DB performance