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

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.

One connection per shard, discovered by prefix at parse time Four Airflow connections named shard_pg_0 through shard_pg_3 are stored in the metadata database. A helper lists connection ids matching the prefix and returns them to the DAG, which builds one PostgresHook per shard. Adding shard four means adding a connection, not editing the DAG file. Airflow connections shard_pg_0 → db-0.internal:5432 shard_pg_1 → db-1.internal:5432 shard_pg_2 → db-2.internal:5432 shard_pg_3 → db-3.internal:5432 list_shard_conn_ids() filters on the shard_pg_ prefix PostgresHook(shard_pg_0) PostgresHook(shard_pg_1) PostgresHook(shard_pg_2) PostgresHook(shard_pg_3) Keep credentials in the connection store, never in the DAG. The DAG file then contains no hostnames, which is what lets the same code run against staging and production without a branch — and lets a new shard join by configuration alone. 
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.

One mapped task per shard, with per-shard retries and a single verify gate The DAG starts by listing shards from a connection prefix. A dynamically mapped ensure_partitions task then runs once per shard, each with its own retry budget, so one unreachable shard fails alone instead of aborting the run. A verify task queries every shard for the expected bound and only succeeds when all report it; otherwise it routes to the alert task. list_shards() from conn prefix ensure_partitions[shard_0] retries=3, backoff 5 min ensure_partitions[shard_1] retries=3, backoff 5 min ensure_partitions[shard_2] retries=3, backoff 5 min ensure_partitions[shard_3] unreachable → fails alone verify_all_shards expects the same max bound run succeeds metric published alert_on_gap pages the on-call DBA Mapping over shards rather than looping inside one task is what makes the failure legible: the Airflow grid shows exactly which shard is behind, and a retry re-runs only that shard's DDL — which is safe because the DDL is idempotent. 
# 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.

Why partition DAGs must read the wall clock, not the logical date During a backfill of three historical runs, a DAG that derives the target month from the logical execution date re-creates partitions that already exist and never advances the horizon. A DAG that derives it from the current date creates the same correct future partitions on every run, so a backfill is harmless. Backfill of 3 past runs — target month from logical date Same backfill — target month from CURRENT_DATE run for 2026-05-01creates 2026_06…2026_08already exist — no-op run for 2026-06-01creates 2026_07…2026_09already exist — no-op run for 2026-07-01creates 2026_08…2026_10horizon never reaches 2026_11 run for 2026-05-01creates 2026_09…2026_11from today's date run for 2026-06-01creates 2026_09…2026_11identical, idempotent run for 2026-07-01creates 2026_09…2026_11horizon correct after any run Partition maintenance is not a per-interval data task — it is a convergence task. Its output depends on when it runs, not on which interval it represents, so clearing and re-running old DAG runs must never move the horizon backwards. 
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

Failure mode Root cause SRE mitigation
DDL succeeds on three shards, fails on one Network partition, shard failover in progress, or rotated partition_bot credential on that shard alone Per-map-index retries re-run only the failed shard; the callback includes map_index, so the page names the shard — clear and re-run just that task instance once the shard is reachable
Partition exists but with wrong bounds (validation passes, inserts fail Aug 1) Suffix and boundary computed from different expressions, or a hand-created partition pre-empting the IF NOT EXISTS Extend validation to compare pg_get_expr(relpartbound, ...) against expected bounds; treat manual partition DDL on shards as a change-control violation
Scheduler pause causes a pile-up of runs firing at once on resume catchup enabled or max_active_runs unset, letting queued intervals run concurrently Keep catchup=False and max_active_runs=1; after a long outage, run airflow dags test for the missed interval manually and verify before re-enabling the schedule

FAQ

Why orchestrate partition DDL from Airflow instead of running pg_partman on each shard?

Per-shard extensions work well until the fleet grows. Airflow gives you one schedule, one audit trail, and one alert channel for the whole fleet, and the validation task catches the shard where DDL silently failed — exactly the drift that per-instance cron makes invisible. Many teams run both: pg_partman locally as a safety net (as in automating quarterly partition creation in PostgreSQL), with Airflow as the source of truth and the fleet-wide monitor.

How do I add a new shard without editing the DAG file?

Store the shard list in an Airflow Variable (a JSON array of connection ids) and read it at parse time with Variable.get("shard_conn_ids", deserialize_json=True), or derive it from a naming convention over registered connections. Registering the connection and updating the Variable brings a new shard into the next run with no code deploy. Put the Variable under change control — it now defines your fleet, and a typo in it is a missing partition.

Does CREATE TABLE ... PARTITION OF block writes on the shard?

It takes a brief ACCESS EXCLUSIVE lock on the parent while the child is attached — milliseconds, since no data is scanned — but it queues behind long-running queries, and every writer then queues behind it. Schedule the DAG off-peak, set lock_timeout in the session (Step 3 does this via hook_params), and let Airflow’s retry handle the fail-fast rather than allowing the DDL to camp in the lock queue.

Back to Automated Partition Creation Workflows