Apache Airflow and Target Database: Ensuring Write Idempotency During Task Retries

Answer the question directly

To make a write operation idempotent when Airflow retries a task, map the logical_date (execution date) to a date‑based partition in the target table and then use either a replace‑partition pattern or a unique‑key‑based UPSERT. Both approaches guarantee that a retry overwrites any partial write from the previous attempt.

Confirmed facts

• Airflow does not provide distributed transactions; idempotency must be implemented inside the task logic or relied upon from the database.
• The logical_date value is constant across all retry attempts for a given task instance, assuming the DAG run is not re‑triggered with a different execution date.
• Most modern relational databases (PostgreSQL, MySQL 8+, SQL Server, Oracle) support either range/list partitioning on a timestamp column or UPSERT syntax (INSERT … ON CONFLICT DO UPDATE or MERGE).

Likely explanation / steps for this case

1. Design the table with a date partition column (e.g., event_date DATE). Create the table partitioned by RANGE on that column so each logical_date maps to its own partition.
2. In the Airflow task, use the logical_date value (available via Jinja templating or the context) as the filter for the write operation:

# Example using PostgresHook
from airflow.providers.postgres.hooks.postgres import postgres

def write_data(**context):
    logical_date = context['logical_date'].date()  # YYYY‑MM‑DD
    hook = PostgresHook(postgres_conn_id='my_db')
    # Replace‑partition strategy
    hook.run(f"DELETE FROM target_table WHERE event_date = '{logical_date}';")
    hook.run("""
        INSERT INTO target_table (event_date, col1, col2)
        VALUES (%s, %s, %s)
    """, parameters=(logical_date, val1, val2))

1. Alternative: unique‑key UPSERT – add a unique constraint that includes the partition column plus any business key (e.g., UNIQUE (event_date, business_key)). Then the task can issue a single UPSERT statement:

hook.run("""
    INSERT INTO target_table (event_date, business_key, col1, col2)
    VALUES (%s, %s, %s, %s)
    ON CONFLICT (event_date, business_key) DO UPDATE
    SET col1 = EXCLUDED.col1,
        col2 = EXCLUDED.col2;
""", parameters=(logical_date, biz_key, val1, val2))

Trade‑offs under high‑frequency retries

• Replace‑partition strategy
  • Pros: Simple to reason about; guarantees that the partition contains only the data from the latest successful run.
  • Cons: Requires two statements (DELETE then INSERT). If the DELETE succeeds but the INSERT fails, the partition is left empty (data loss). Higher latency and more WAL pressure under frequent retries.
• Unique‑key UPSERT strategy
  • Pros: Single statement, atomic within the database; lower overhead and no risk of leaving an empty partition.
  • Cons: Requires a suitable unique key; if the business logic expects only inserts (never updates), the UPDATE side may need to be a no‑update or raise an error. Update conflicts can occur if two retries try to modify the same row concurrently (rare with Airflow’s serial retry but possible with parallelism).

Verification steps (recommended)

1. Create a test table with the chosen partitioning or unique key.
2. Run an Airflow task that simulates a partial failure (e.g., raise an exception after the INSERT but before commit).
3. Observe the retry: the final row count should match a successful single run, and the data should be identical.
4. Check Airflow logs to confirm that logical_date remains unchanged across attempts.

If you are unsure whether your target database supports UPSERT or efficient partition‑level DELETE, please confirm that capability before choosing a strategy.