Clean up the Airflow metadata database using DAGs

In addition to storing configurations about your Airflow environment, the Airflow metadata database stores data about past and present task runs. Airflow never automatically removes metadata, so the longer you use it, the more task run data is stored in your metadata DB. Over a long enough time, this can result in a bloated metadata DB, which can affect performance across your Airflow environment.

When a table in the metadata DB is larger than 50GiB, you might start to experience degraded scheduler performance. This can result in:

• Slow task scheduling
• Slow DAG parsing
• Gunicorn timing out when using the Celery executor
• Slower Airflow UI load times

The following tables in the database are at risk of becoming too large over time:

• dag_run
• job
• log
• rendered_task_instance_fields
• task_instance
• xcom

To keep your Airflow environment running at optimal performance, you can clean the metadata DB using the Airflow CLI airflow db clean command. This command was created as a way to safely clean up your metadata DB without querying it directly. This tutorial describes how to implement a cleanup DAG in Airflow so that you can clean your database using the command directly from the Airflow UI.

Warnings

Deleting data from the metadata database can be an extremely destructive action. If you delete data that future task runs depend on, it’s difficult to restore the database to its previous state without interrupting your data pipelines. Before implementing the DAG in this tutorial, consider the following:

• When specifying the clean_before_timestamp value, use as old a date as possible. The older the deleted data, the less likely it is to affect your currently running DAGs.
• The DAG drops the archived tables it created in the cleanup process by default using the --skip-archive option and does not maintain any history. If the task fails for example if it runs for longer than five minutes, the archive tables are not cleared. By running airflow db drop-archived in the second task of the DAG, we ensure all archive tables are dropped even in the event of the first task failing.

Prerequisites

• An Airflow project project
• The Astro CLI

Step 1: Create your DAG

1. In your dags folder, create a file called db_cleanup.py.

2. Copy the following code into the file.

   1
   """A DB Cleanup DAG maintained by Astronomer."""
   2
   3
   import logging
   4
   from datetime import UTC, datetime, timedelta
   5
   from typing import Optional, List
   6
   7
   from sqlalchemy import func
   8
   from sqlalchemy.orm import Session
   9
   10
   from airflow.cli.commands.db_command import all_tables
   11
   from airflow.decorators import dag, task
   12
   from airflow.models.param import Param
   13
   from airflow.operators.bash import BashOperator
   14
   from airflow.utils.db import reflect_tables
   15
   from airflow.utils.db_cleanup import _effective_table_names
   16
   from airflow.utils.session import provide_session, NEW_SESSION
   17
   18
   19
   @dag(
   20
       dag_id="astronomer_db_cleanup_dag",
   21
       schedule_interval=None,
   22
       start_date=datetime(2024, 1, 1),
   23
       catchup=False,
   24
       is_paused_upon_creation=False,
   25
       description=__doc__,
   26
       doc_md=__doc__,
   27
       render_template_as_native_obj=True,
   28
       max_active_tasks=1,
   29
       tags=["cleanup"],
   30
       params={
   31
           "clean_before_timestamp": Param(
   32
               default=(datetime.now(tz=UTC) - timedelta(days=90)).isoformat(),
   33
               type="string",
   34
               format="date-time",
   35
               description="Delete records older than this timestamp. Default is 90 days ago.",
   36
           ),
   37
           "tables": Param(
   38
               default=[],
   39
               type=["null", "array"],
   40
               examples=all_tables,
   41
               description="List of tables to clean. Default is all tables.",
   42
           ),
   43
           "dry_run": Param(
   44
               default=False,
   45
               type="boolean",
   46
               description="Print the SQL queries that would be run, but do not execute them. Default is False.",
   47
           ),
   48
           "batch_size_days": Param(
   49
               default=7,
   50
               type="integer",
   51
               description="Number of days in each batch for the cleanup. Default is 7 days.",
   52
           ),
   53
       },
   54
   )
   55
   def astronomer_db_cleanup_dag():
   56
   57
       @provide_session
   58
       def get_oldest_timestamp(
   59
           tables,
   60
           session: Session = NEW_SESSION,
   61
       ) -> Optional[str]:
   62
           oldest_timestamp_list = []
   63
           existing_tables = reflect_tables(tables=None, session=session).tables
   64
           _, effective_config_dict = _effective_table_names(table_names=tables)
   65
           for table_name, table_config in effective_config_dict.items():
   66
               if table_name in existing_tables:
   67
                   orm_model = table_config.orm_model
   68
                   recency_column = table_config.recency_column
   69
                   oldest_execution_date = (
   70
                       session.query(func.min(recency_column))
   71
                       .select_from(orm_model)
   72
                       .scalar()
   73
                   )
   74
                   if oldest_execution_date:
   75
                       oldest_timestamp_list.append(oldest_execution_date.isoformat())
   76
                   else:
   77
                       logging.info(f"No data found for {table_name}, skipping...")
   78
               else:
   79
                   logging.warning(f"Table {table_name} not found. Skipping.")
   80
   81
           if oldest_timestamp_list:
   82
               return min(oldest_timestamp_list)
   83
   84
       @task
   85
       def get_chunked_timestamps(**context) -> List:
   86
           batches = []
   87
           start_chunk_time = get_oldest_timestamp(context["params"]["tables"])
   88
           if start_chunk_time:
   89
               start_ts = datetime.fromisoformat(start_chunk_time)
   90
               end_ts = datetime.fromisoformat(context["params"]["clean_before_timestamp"])
   91
               batch_size_days = context["params"]["batch_size_days"]
   92
   93
               while start_ts < end_ts:
   94
                   batch_end = min(start_ts + timedelta(days=batch_size_days), end_ts)
   95
                   batches.append({"BATCH_TS": batch_end.isoformat()})
   96
                   start_ts += timedelta(days=batch_size_days)
   97
           return batches
   98
   99
       # The "clean_archive_tables" task drops archived tables created by the previous "clean_db" task, in case that task fails due to an error or timeout.
   100
       db_archive_cleanup = BashOperator(
   101
           task_id="clean_archive_tables",
   102
           bash_command="""\
   103
               airflow db drop-archived \
   104
           {% if params.tables -%}
   105
                --tables {{ params.tables|join(',') }} \
   106
           {% endif -%}
   107
                --yes \
   108
           """,
   109
           do_xcom_push=False,
   110
           trigger_rule="all_done",
   111
       )
   112
   113
       chunked_timestamps = get_chunked_timestamps()
   114
   115
       (
   116
           BashOperator.partial(
   117
               task_id="db_cleanup",
   118
               bash_command="""\
   119
               airflow db clean \
   120
                --clean-before-timestamp $BATCH_TS \
   121
           {% if params.dry_run -%}
   122
                --dry-run \
   123
           {% endif -%}
   124
                --skip-archive \
   125
           {% if params.tables -%}
   126
                --tables '{{ params.tables|join(',') }}' \
   127
           {% endif -%}
   128
                --verbose \
   129
                --yes \
   130
           """,
   131
               append_env=True,
   132
               do_xcom_push=False,
   133
           ).expand(env=chunked_timestamps)
   134
           >> db_archive_cleanup
   135
       )
   136
   137
   138
   astronomer_db_cleanup_dag()

   Rather than running on a schedule, this DAG is triggered manually by default and includes params so that you’re in full control over how you clean the metadata DB.

   It includes two tasks:

   • clean_db, that runs airflow db clean
   • clean_archive_tables, that runs airflow db drop-archived

   These two tasks run with params you specify at runtime. The params let you specify:

   • What age of data to delete. Any data that was created before the specified time will be deleted. The default is to delete all data older than 90 days.
   • Whether to run the cleanup as a dry run, meaning that no data is deleted. The DAG will instead return the SQL that would be executed based on other parameters you have specified. The default is to run the deletion without a dry run.
   • Which tables to delete data from. The default is all tables.

Step 2: Run the DAG

In this step, run the DAG in a local Airflow environment to practice the workflow for cleaning metadata DB data. When completing this process in a production environment, you would complete this process only after your Airflow environment has been running for a while.

1. Run astro dev start in your Astro project to start Airflow, then open the Airflow UI at localhost:8080.

2. In the Airflow UI, run the db_cleanup DAG by clicking the play button, then click Trigger DAG w/ Config. Configure the following params:

   • dry_run: true
   • tables: all_tables
   • clean_before_timestamp: (datetime.now(tz=UTC) - timedelta(days=90)).isoformat()

3. Click Trigger.

4. After the task completes, click on the Graph tab.

5. Click on the clean_db task.

6. Click on the Logs tab.

7. Check that the airflow db cleanup command completed successfully. Note that if you created a new Astro project for this tutorial, the run will not show much data to be deleted.

You can now use this DAG to periodically clean data from the Airflow metadata DB as needed.