Automating Partition Creation with Apache Airflow

This walkthrough builds an Apache Airflow DAG that provisions next month’s RANGE partitions across a fleet of PostgreSQL shards — the orchestrator-driven variant of the automated partition creation workflows described in partitioning implementation patterns & routing. Database-local schedulers stop scaling once the same partition must exist on eight shards at once: one missed cron on one shard and next month’s inserts start failing there while the other seven look healthy. Centralizing the DDL in Airflow gives you one schedule, one audit trail, and a validation gate that treats the fleet as a unit.

Prerequisites

• Airflow 2.7+ with apache-airflow-providers-postgres (which pulls in the common.sql provider) installed on scheduler and workers
• Network reachability from Airflow workers to every shard’s PostgreSQL port
• A dedicated partition_bot role on each shard with CREATE on the schema that owns the parent table
• The parent table (e.g. orders) already declared PARTITION BY RANGE (created_at) on every shard
• A secrets backend or Fernet-encrypted metadata database for connection passwords
• An incident webhook or SMTP relay reachable from workers for failure callbacks
• Agreement that all partition boundaries are defined in UTC

Step 1 — Register one connection per shard

Each shard becomes an Airflow connection with a predictable id (pg_shard_0 … pg_shard_3). The connection id is the unit of task mapping later, so the naming convention matters: the DAG derives its shard list from these ids, and per-shard task logs are labeled by them. How rows were routed onto those shards in the first place — typically one of the hash routing algorithms — is irrelevant here; this DAG only cares that every shard carries the same parent table.

for i in 0 1 2 3; do
  airflow connections add "pg_shard_${i}" \
    --conn-type postgres \
    --conn-host "shard${i}.db.internal" \
    --conn-port 5432 \
    --conn-schema orders_db \
    --conn-login partition_bot \
    --conn-password "${SHARD_DB_PASSWORD}"
done

Operational note: Grant partition_bot only CONNECT on the database and CREATE on the schema — it never needs INSERT or SELECT on data. A leaked orchestrator credential that can create empty tables is an annoyance; one that can read the orders table is an incident.

SRE tip: Prefer a secrets backend (Vault, AWS Secrets Manager, GCP Secret Manager) over airflow connections add in production so password rotation never requires touching Airflow. The connection ids stay stable either way, which is all the DAG depends on.

Step 2 — Write the idempotent partition DDL

The DDL must be safe to run twice on the same shard and deterministic for a given schedule interval. CREATE TABLE IF NOT EXISTS ... PARTITION OF covers re-runs and retries; deriving the boundary dates from data_interval_end (rather than “now”) covers determinism. The DAG runs on the 25th of each month, so data_interval_end falls on the 25th and the following month is the partition being provisioned — a five-to-six-day buffer before those rows arrive.

-- Rendered by Airflow's Jinja templating; runs on the 25th, creates NEXT month.
CREATE TABLE IF NOT EXISTS orders_p{{ data_interval_end.add(months=1).format('YYYYMM') }}
  PARTITION OF orders
  FOR VALUES FROM ('{{ data_interval_end.add(months=1).format('YYYY-MM-01') }}')
              TO ('{{ data_interval_end.add(months=2).format('YYYY-MM-01') }}');

Operational note: Never compute boundaries with datetime.now() inside the DAG file or the SQL. Parse-time timestamps make the rendered DDL depend on when the scheduler parsed the file, which breaks retries that cross midnight and makes airflow dags test unreproducible. data_interval_end is frozen per run.

SRE tip: Keep the suffix format (YYYYMM) and the boundary format (YYYY-MM-01) derived from the same expression. Teams that compute the name and the range in two places eventually ship a partition named orders_p202608 whose bounds cover September — and every query that prunes by name-based assumptions lies from then on.

Step 3 — Build the dynamically mapped DAG

Dynamic task mapping expands one operator definition into one task instance per shard at runtime. SQLExecuteQueryOperator.partial(...).expand(conn_id=...) gives each shard its own retries, its own log, and its own failure state — a network blip on shard 2 retries shard 2 only, instead of re-running the fleet.

# dags/create_shard_partitions.py
import pendulum
from airflow.decorators import dag
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator

SHARD_CONN_IDS = ["pg_shard_0", "pg_shard_1", "pg_shard_2", "pg_shard_3"]

CREATE_PARTITION_SQL = """
CREATE TABLE IF NOT EXISTS orders_p{{ data_interval_end.add(months=1).format('YYYYMM') }}
  PARTITION OF orders
  FOR VALUES FROM ('{{ data_interval_end.add(months=1).format('YYYY-MM-01') }}')
              TO ('{{ data_interval_end.add(months=2).format('YYYY-MM-01') }}');
"""

@dag(
    dag_id="create_shard_partitions",
    schedule="0 3 25 * *",          # 03:00 UTC on the 25th of every month
    start_date=pendulum.datetime(2026, 6, 1, tz="UTC"),
    catchup=False,
    max_active_runs=1,
    default_args={"retries": 2, "retry_delay": pendulum.duration(minutes=10)},
    tags=["partitioning", "ddl"],
)
def create_shard_partitions():
    SQLExecuteQueryOperator.partial(
        task_id="create_partition",
        sql=CREATE_PARTITION_SQL,
        autocommit=True,
        hook_params={"options": "-c lock_timeout=5s"},
    ).expand(conn_id=SHARD_CONN_IDS)

create_shard_partitions()

Operational note: hook_params={"options": "-c lock_timeout=5s"} makes the session fail fast if the ACCESS EXCLUSIVE lock on orders queues behind a long-running query, handing control to Airflow’s retry logic instead of letting DDL sit in the lock queue blocking every writer behind it.

SRE tip: Pin max_active_runs=1. If a stuck run from last month is cleared and re-fired while this month’s run is live, two sessions can race on the same parent table’s lock queue. Serializing runs costs nothing here — the DAG takes seconds — and removes an entire class of overlap incidents.

Step 4 — Validate the partition exists on every shard

Success of the DDL task proves the statement ran; it does not prove the fleet is consistent — IF NOT EXISTS happily no-ops against a partition someone hand-created with the wrong bounds. A mapped validation task queries each shard’s catalog and fails loudly if the expected partition is missing, turning silent drift into a red task instance naming the exact shard.

from airflow.decorators import task
from airflow.providers.postgres.hooks.postgres import PostgresHook

@task
def assert_partition_exists(conn_id: str, data_interval_end=None):
    suffix = data_interval_end.add(months=1).format("YYYYMM")
    partition = f"public.orders_p{suffix}"
    hook = PostgresHook(postgres_conn_id=conn_id)
    row = hook.get_first("SELECT to_regclass(%s);", parameters=(partition,))
    if row is None or row[0] is None:
        raise ValueError(f"{partition} missing on {conn_id}")
    return f"{conn_id}: {partition} present"

# inside create_shard_partitions(), after the create task:
#   create = SQLExecuteQueryOperator.partial(...).expand(conn_id=SHARD_CONN_IDS)
#   create >> assert_partition_exists.expand(conn_id=SHARD_CONN_IDS)

Operational note: to_regclass() returns NULL instead of raising when the relation is absent, which keeps the check to a single round-trip and makes the failure path an explicit Python exception with the shard id in the message — exactly what lands in the alert.

SRE tip: Extend the check beyond existence when you can afford one more query: compare pg_get_expr(relpartbound, oid) against the expected FROM/TO strings. Existence catches the missed shard; a bounds comparison also catches the mis-created one.

Step 5 — Alert on failure and stay backfill-safe

Partition DDL failures are urgent — every day of delay eats the buffer before month-end — so wire on_failure_callback at the DAG level to page with the shard-level detail (map_index identifies which shard). Backfill safety comes from three properties already in place: catchup=False stops the scheduler from firing months of historical runs after a pause, IF NOT EXISTS makes any deliberate airflow dags backfill of a past interval a harmless no-op, and interval-derived boundaries mean a re-run of July’s run always emits July’s DDL.

import json, os, urllib.request

def notify_failure(context):
    ti = context["task_instance"]
    payload = {
        "text": (
            f"Partition DDL failure: dag={ti.dag_id} task={ti.task_id} "
            f"shard_index={ti.map_index} run={context['run_id']} "
            f"try={ti.try_number}"
        )
    }
    req = urllib.request.Request(
        os.environ["ALERT_WEBHOOK_URL"],
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    urllib.request.urlopen(req, timeout=10)

# in the @dag(...) decorator:
#   default_args={"retries": 2, "retry_delay": pendulum.duration(minutes=10),
#                 "on_failure_callback": notify_failure},

Operational note: The callback fires only after retries are exhausted, so the page means “a shard is really failing”, not “a TCP reset happened once”. Keep retries low (2) for DDL — if two attempts ten minutes apart both fail, the cause is a lock, a permission, or a down shard, none of which a third retry fixes.

SRE tip: Add a second, independent alarm outside Airflow: a monthly check on each shard that the next month’s partition exists by the 28th. If the Airflow deployment itself is down, no callback will ever fire — the dead-man switch on the shard side is what catches that.

Verification

Exercise the whole DAG against real shards without registering a run in the scheduler:

airflow dags test create_shard_partitions 2026-07-25

[2026-07-25T03:00:04] create_partition (map_index=0) SUCCESS
[2026-07-25T03:00:04] create_partition (map_index=1) SUCCESS
[2026-07-25T03:00:05] create_partition (map_index=2) SUCCESS
[2026-07-25T03:00:05] create_partition (map_index=3) SUCCESS
[2026-07-25T03:00:07] assert_partition_exists (map_index=0..3) SUCCESS
Dag run create_shard_partitions ... state=success

Then confirm the partition and its bounds directly on a shard — repeat per shard or loop with psql over the connection list:

SELECT c.relname AS partition,
       pg_get_expr(c.relpartbound, c.oid) AS bounds
FROM pg_partition_tree('public.orders') t
JOIN pg_class c ON c.oid = t.relid
WHERE t.isleaf AND c.relname = 'orders_p202608';

   partition    |                            bounds
----------------+---------------------------------------------------------------
 orders_p202608 | FOR VALUES FROM ('2026-08-01 00:00:00+00') TO ('2026-09-01 00:00:00+00')

An empty result on any shard means the validation task should have failed — if it did not, the suffix expressions in Step 2 and Step 4 have drifted apart and must be reunified before the next run.

Failure mode table

FAQ

Related

• Automated Partition Creation Workflows — parent guide comparing scheduler-driven, extension-driven, and orchestrator-driven partition automation
• Automating Quarterly Partition Creation in PostgreSQL — the single-instance, pg_partman-based approach this DAG generalizes across a fleet
• Automating Monthly Partition Creation in MySQL 8.0 — the same monthly cadence implemented with MySQL’s native EVENT scheduler and stored procedures