Backup, Recovery & Disaster Planning

# PostgreSQL Full Backup with pg_dump
import subprocess
from datetime import datetime
from pathlib import Path

# pg_dump is a client binary, not part of psycopg, and it must exist on the
# machine running this script. With PostgreSQL in Docker it lives inside the
# container, so a bare pg_dump call fails with
# "FileNotFoundError: [Errno 2] No such file or directory: 'pg_dump'".
# Two ways out: install the client (apt install postgresql-client-16), or
# run the binary through docker exec, which is what this does so the
# example needs no host install.
CONTAINER = "pg-demo"


def full_backup_postgres(db_name, backup_dir):
    backup_dir = Path(backup_dir)
    backup_dir.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    # -F c is a compressed binary archive, not SQL text, so ".dump" is the
    # honest extension. Only pg_restore reads it; psql cannot.
    backup_file = backup_dir / f"{db_name}_full_{timestamp}.dump"

    cmd = [
        'docker', 'exec',
        '-e', 'PGPASSWORD=demo',   # pg_dump never prompts under subprocess
        CONTAINER,
        'pg_dump',
        '-U', 'demo',
        '-d', db_name,
        '-F', 'c',                 # Custom format (compressed)
    ]

    # No text=True: the custom format is binary and decoding it corrupts it.
    result = subprocess.run(cmd, capture_output=True)

    if result.returncode != 0:
        print(f"Backup failed: {result.stderr.decode().strip()}")
        return None

    backup_file.write_bytes(result.stdout)
    size_mb = backup_file.stat().st_size / (1024 * 1024)
    print(f"Full backup completed: {backup_file}")
    print(f"   Size: {size_mb:.2f} MB")
    return backup_file


# Usage
backup_file = full_backup_postgres('demo', './backups')

# WAL-based incremental backup, start to finish. This block turns on
# archiving, generates traffic, and measures the result on its own: it
# needs only the pg-demo container and the seed data from the top of the
# lesson.
import subprocess
import time

CONTAINER = "pg-demo"
WAL_ARCHIVE = "/var/lib/postgresql/wal_archive"


def psql(sql):
    result = subprocess.run(
        ["docker", "exec", "-e", "PGPASSWORD=demo", CONTAINER,
         "psql", "-U", "demo", "-d", "demo", "-tAc", sql],
        capture_output=True, text=True, check=True)
    return result.stdout.strip()


def in_container(*args):
    # -u postgres matters: docker exec defaults to root, and a root-owned
    # archive directory makes every archive_command fail with
    # "cp: can't create ...: Permission denied".
    result = subprocess.run(["docker", "exec", "-u", "postgres", CONTAINER, *args],
                            capture_output=True, text=True, check=True)
    return result.stdout


def switch_and_wait():
    """Close the current WAL segment and block until it reaches the archive."""
    closed = psql("SELECT pg_walfile_name(pg_switch_wal())")
    for _ in range(30):
        archived = psql("SELECT coalesce(last_archived_wal, '') FROM pg_stat_archiver")
        if archived >= closed:
            return closed
        time.sleep(1)
    raise RuntimeError(f"segment {closed} was never archived")


def archive_snapshot():
    """Map every archived segment to its size in bytes."""
    out = in_container("sh", "-c", f"stat -c '%s %n' {WAL_ARCHIVE}/* 2>/dev/null || true")
    snapshot = {}
    for line in out.splitlines():
        size, path = line.split(" ", 1)
        snapshot[path.rsplit("/", 1)[-1]] = int(size)
    return snapshot


def incremental_backup_size(before, after):
    new_files = sorted(set(after) - set(before))
    total_bytes = sum(after[f] for f in new_files)
    print(f"WAL segments archived: {len(new_files)}")
    print(f"Raw size: {total_bytes / (1024 * 1024):.2f} MB")
    return new_files, total_bytes


# 1. Turn on continuous WAL archiving. archive_mode only takes effect on
#    restart; archive_command on its own would just need a reload.
in_container("mkdir", "-p", WAL_ARCHIVE)
psql("ALTER SYSTEM SET wal_level = replica")
psql("ALTER SYSTEM SET archive_mode = on")
psql("ALTER SYSTEM SET archive_command = "
     f"'test ! -f {WAL_ARCHIVE}/%f && cp %p {WAL_ARCHIVE}/%f'")

subprocess.run(["docker", "restart", CONTAINER], capture_output=True, check=True)
for _ in range(30):
    if subprocess.run(["docker", "exec", CONTAINER, "pg_isready", "-U", "demo"],
                      capture_output=True).returncode == 0:
        break
    time.sleep(1)
else:
    raise RuntimeError("PostgreSQL did not come back up")

print("archive_mode:", psql("SHOW archive_mode"))

# 2. Baseline: flush anything already pending, then snapshot the archive.
switch_and_wait()
before = archive_snapshot()

# 3. Do some work: 3,000 new transactions on top of the 200,000-row seed.
psql("""INSERT INTO transactions (user_id, amount, description)
        SELECT (i % 5000) + 1, ((i * 91) % 50000)::numeric / 100,
               'Incremental demo row ' || i
        FROM generate_series(1, 3000) AS i""")
switch_and_wait()

# 4. Whatever is new in the archive IS the incremental backup.
after = archive_snapshot()
files, total_bytes = incremental_backup_size(before, after)

# PITR base backup. Like the pg_dump script, this runs pg_basebackup
# through docker exec, so nothing has to be installed on the host. Here
# there is a second reason to go through the container: see the note below.
import subprocess
from datetime import datetime

CONTAINER = "pg-demo"
# A path inside the container, owned by postgres. pg_basebackup writes a
# directory, not a single file, so it cannot be streamed to stdout the way
# the pg_dump archive was.
BACKUP_DIR = "/var/lib/postgresql/backups"


def create_pitr_base_backup():
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_path = f"{BACKUP_DIR}/base_{timestamp}"

    cmd = [
        "docker", "exec", "-u", "postgres", "-e", "PGPASSWORD=demo", CONTAINER,
        "pg_basebackup",
        "-h", "localhost",
        "-U", "demo",
        "-D", backup_path,
        "-Ft",                   # Tar format
        "-z",                    # Compress
        "-P",                    # Show progress
        "--wal-method=stream",   # Include the WAL generated during the copy
        "--checkpoint=fast",     # Force an immediate checkpoint
    ]

    print("Creating PITR base backup...")
    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        print(f"Backup failed: {result.stderr.strip()}")
        return None

    print(f"Base backup created: {backup_path}")
    print("   Includes WAL files for PITR")
    return backup_path


backup_path = create_pitr_base_backup()

# Simulate the disaster PITR will undo, and capture the timestamp to
# rewind to. Run this after the base backup, before the restore.
import subprocess
import time
from pathlib import Path

CONTAINER = "pg-demo"
TARGET_FILE = Path("pitr_target.txt")


def psql(sql):
    result = subprocess.run(
        ["docker", "exec", "-e", "PGPASSWORD=demo", CONTAINER,
         "psql", "-U", "demo", "-d", "demo", "-tAc", sql],
        capture_output=True, text=True, check=True)
    return result.stdout.strip()


# 1. A row worth saving, inserted just before things go wrong.
psql("""INSERT INTO transactions (user_id, amount, description)
        VALUES (1, 99.99, 'SAFE MARKER - before disaster')""")
print("Rows before disaster:", psql("SELECT count(*) FROM transactions"))

# 2. The moment to rewind to. Read it from the server, not the host clock:
#    recovery_target_time is compared against timestamps the server wrote.
#    Saved to a file so the restore step picks up YOUR timestamp instead of
#    one copied out of this page.
time.sleep(2)
target_time = psql("SELECT now()")
TARGET_FILE.write_text(target_time)
print("Recovery target:", target_time)
print(f"   saved to {TARGET_FILE}")
time.sleep(2)

# 3. The disaster: someone forgets the WHERE clause.
psql("DELETE FROM transactions")
print("Rows after disaster:", psql("SELECT count(*) FROM transactions"))

# 4. Recovery can only replay what reached the archive, and the segment
#    holding these changes is still being written. Force it out.
psql("SELECT pg_switch_wal()")
time.sleep(2)
print("Last archived segment:", psql("SELECT last_archived_wal FROM pg_stat_archiver"))

# Perform PITR. The recovered data directory is built inside the
# container, next to the one still running, so every filesystem step goes
# through docker exec.
import subprocess
from pathlib import Path

CONTAINER = "pg-demo"
BACKUP_DIR = "/var/lib/postgresql/backups"


def in_container(*args):
    result = subprocess.run(["docker", "exec", "-u", "postgres", CONTAINER, *args],
                            capture_output=True, text=True)
    if result.returncode != 0:
        # Worth doing by hand rather than with check=True: that raises a bare
        # CalledProcessError and throws away the message that explains it,
        # such as "tar: can't open ...: No such file or directory".
        raise RuntimeError(f"{args[0]} failed: {result.stderr.strip()}")
    return result.stdout


def latest_base_backup():
    """Newest base_* directory. The timestamped names sort chronologically."""
    listing = in_container("sh", "-c", f"ls -d {BACKUP_DIR}/base_* 2>/dev/null || true")
    backups = sorted(listing.split())
    if not backups:
        raise RuntimeError(f"no base backup in {BACKUP_DIR}: run the previous step first")
    return backups[-1]


def perform_pitr(base_backup_path, wal_archive_dir, target_time, restore_dir):
    """Restore the database to a specific point in time.

    base_backup_path: directory holding base.tar.gz, from the previous step
    wal_archive_dir:  directory the archive_command has been filling
    target_time:      timestamp to stop recovery at
    restore_dir:      where the recovered data directory is built
    """
    print(f"Starting PITR to {target_time}...")
    print(f"Using base backup: {base_backup_path}")

    print("Step 1: Extracting base backup...")
    # Recovery has to start from an empty directory. Leftovers from an
    # earlier attempt mix old and new files, and PostgreSQL either refuses
    # to start or recovers something you did not intend.
    if in_container("sh", "-c", f"ls -A {restore_dir} 2>/dev/null || true").strip():
        raise RuntimeError(
            f"{restore_dir} is not empty. Remove it first:\n"
            f"  docker exec -u postgres {CONTAINER} rm -rf {restore_dir}")
    in_container("mkdir", "-p", restore_dir)
    in_container("tar", "-xzf", f"{base_backup_path}/base.tar.gz", "-C", restore_dir)

    # PostgreSQL refuses to start on a data directory with the wrong
    # permissions. tar preserves the permissions of files INSIDE the
    # archive, but not of the extraction target itself, so skipping this
    # gives "FATAL: data directory ... has invalid permissions".
    in_container("chmod", "700", restore_dir)

    print("Step 2: Writing recovery configuration...")
    recovery_conf = (
        f"restore_command = 'cp {wal_archive_dir}/%f %p'\n"
        f"recovery_target_time = '{target_time}'\n"
        f"recovery_target_action = 'promote'\n"
    )
    # recovery.signal is what puts the instance into recovery mode at all.
    in_container("touch", f"{restore_dir}/recovery.signal")
    in_container("sh", "-c",
                 f"cat >> {restore_dir}/postgresql.auto.conf <<'CONF'\n{recovery_conf}CONF")

    print("PITR configuration created")
    print(f"   Target time: {target_time}")
    print(f"   Restore directory: {restore_dir}")
    print("\n   Start PostgreSQL on this data directory to complete recovery")
    return restore_dir


# Usage: restore to just before the disaster. Neither value is hardcoded:
# the base backup is discovered, and the target comes from the file the
# previous step wrote. In a real incident you would type the timestamp of
# the moment before the bad query instead.
restore_dir = perform_pitr(
    base_backup_path=latest_base_backup(),
    wal_archive_dir="/var/lib/postgresql/wal_archive",
    target_time=Path("pitr_target.txt").read_text().strip(),
    restore_dir="/var/lib/postgresql/restored",
)

# The script only stages the recovery. Starting a second postmaster on
# the restored directory, on a spare port, is what actually replays WAL.
docker exec -u postgres -d pg-demo \
  sh -c 'postgres -D /var/lib/postgresql/restored -p 5433 > /tmp/restore.log 2>&1'

docker exec pg-demo cat /tmp/restore.log

# Then query the recovered instance on port 5433:
docker exec -e PGPASSWORD=demo pg-demo psql -U demo -d demo -p 5433 \
  -c "SELECT count(*) FROM transactions;" \
  -c "SELECT id, description FROM transactions WHERE description LIKE 'SAFE MARKER%';"

# Automated backup verification. Install: pip install "psycopg[binary]"
import subprocess
from datetime import datetime
from pathlib import Path

import psycopg

CONTAINER = "pg-demo"
BACKUP_DIR = Path("./backups")
ADMIN_DSN = "host=localhost port=5432 dbname=postgres user=demo password=demo"
EXPECTED_TABLES = ("users", "transactions")


def latest_backup():
    """Newest archive in the backup directory: the names sort chronologically."""
    dumps = sorted(BACKUP_DIR.glob("*.dump"))
    if not dumps:
        raise RuntimeError(f"no .dump file in {BACKUP_DIR}: run the full backup script first")
    return dumps[-1]


def test_backup_restore(backup_file, test_db_name="test_restore_db"):
    """Test a backup by actually restoring it and verifying the data."""
    print(f"Testing backup: {backup_file}")
    start_time = datetime.now()

    # Step 1: Recreate the scratch database. CREATE/DROP DATABASE cannot run
    # inside a transaction, hence autocommit.
    print("Step 1: Cleaning up old test database...")
    with psycopg.connect(ADMIN_DSN, autocommit=True) as conn:
        conn.execute(f"DROP DATABASE IF EXISTS {test_db_name}")
        conn.execute(f"CREATE DATABASE {test_db_name}")

    # Step 2: Restore into it. The archive sits on the host and pg_restore
    # lives in the container, so feed it over stdin: "docker exec -i" plus
    # no filename argument makes pg_restore read from standard input.
    print("Step 2: Restoring backup...")
    cmd = [
        "docker", "exec", "-i", "-e", "PGPASSWORD=demo", CONTAINER,
        "pg_restore", "-U", "demo", "-d", test_db_name,
    ]
    # Stream the archive rather than read_bytes() it: production dumps are
    # measured in gigabytes and do not belong in memory.
    with open(backup_file, "rb") as archive:
        result = subprocess.run(cmd, stdin=archive, capture_output=True)

    if result.returncode != 0:
        print(f"Restore failed: {result.stderr.decode().strip()}")
        return False

    # Step 3: Verify the data actually arrived, not just that the command
    # exited 0. A backup that restores an empty schema still "succeeds".
    print("Step 3: Verifying data...")
    test_dsn = ADMIN_DSN.replace("dbname=postgres", f"dbname={test_db_name}")
    with psycopg.connect(test_dsn) as conn:
        found = {row[0] for row in conn.execute("""
            SELECT tablename FROM pg_tables
            WHERE schemaname NOT IN ('pg_catalog', 'information_schema')
        """).fetchall()}
        print(f"   Found {len(found)} tables: {', '.join(sorted(found))}")

        # Check the tables exist before counting them. Querying a missing
        # table would raise UndefinedTable and crash the checker instead of
        # reporting the backup as bad, which is the one job it has.
        missing = [t for t in EXPECTED_TABLES if t not in found]
        if missing:
            print(f"Verification FAILED: missing tables: {', '.join(missing)}")
            return False

        for table in EXPECTED_TABLES:
            count = conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
            print(f"   {table} table: {count} rows")
            if count == 0:
                print(f"Verification FAILED: {table} restored empty")
                return False

    # Step 4: Cleanup
    print("Step 4: Cleaning up test database...")
    with psycopg.connect(ADMIN_DSN, autocommit=True) as conn:
        conn.execute(f"DROP DATABASE {test_db_name}")

    elapsed = (datetime.now() - start_time).total_seconds()
    print(f"Backup test PASSED in {elapsed:.2f}s")
    print("   Backup is valid and restorable!")
    return True


# Usage: run weekly backup verification against the newest archive the full
# backup script produced.
test_backup_restore(latest_backup())

-- What the primary needs. On PostgreSQL 16 all three are already the
-- defaults, so a stock server needs no configuration change at all:
--
--   wal_level       = replica   -- enough WAL to feed a replica
--   max_wal_senders = 10        -- concurrent replication connections
--
-- hot_standby = on is a STANDBY setting, not a primary one: it controls
-- whether the standby answers read queries during recovery. It is also on
-- by default. Setting it on the primary does nothing.
--
-- Confirm rather than assume:
SHOW wal_level;
SHOW max_wal_senders;
SHOW hot_standby;

# Set up streaming replication from pg-primary to pg-replica.
import subprocess
import time

PRIMARY = "pg-primary"
REPLICA = "pg-replica"
PGDATA = "/var/lib/postgresql/data"
SLOT = "replica1_slot"


def psql(container, sql, password="demo"):
    result = subprocess.run(
        ["docker", "exec", "-e", f"PGPASSWORD={password}", container,
         "psql", "-U", "demo", "-d", "demo", "-tAc", sql],
        capture_output=True, text=True, check=True)
    return result.stdout.strip()


def in_container(container, *args):
    result = subprocess.run(["docker", "exec", "-u", "postgres", container, *args],
                            capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"{args[0]} failed: {result.stderr.strip()}")
    return result.stdout


# --- On the PRIMARY -------------------------------------------------
# A dedicated role with REPLICATION. It needs no rights on any table:
# physical replication ships WAL, not rows.
psql(PRIMARY, "CREATE USER replicator WITH REPLICATION LOGIN PASSWORD 'replpass'")

# A replication slot makes the primary retain WAL until this specific
# replica has consumed it. wal_keep_size only keeps a fixed amount and then
# drops the rest, which silently breaks a replica that falls behind.
psql(PRIMARY, f"SELECT pg_create_physical_replication_slot('{SLOT}')")

# The image's pg_hba.conf trusts replication only from 127.0.0.1, and its
# catch-all "host all all all" rule does NOT cover the replication
# pseudo-database. Without this line pg_basebackup is rejected.
in_container(PRIMARY, "sh", "-c",
             f"echo 'host replication replicator all scram-sha-256' >> {PGDATA}/pg_hba.conf")
psql(PRIMARY, "SELECT pg_reload_conf()")
print("Primary ready: replication role, slot, and pg_hba rule in place")

# --- On the REPLICA -------------------------------------------------
# -R writes the connection settings AND creates standby.signal, the file
# that actually makes this a standby rather than an independent server.
# -S ties the replica to the slot created above.
subprocess.run(
    ["docker", "exec", "-u", "postgres", "-e", "PGPASSWORD=replpass", REPLICA,
     "pg_basebackup", "-h", PRIMARY, "-U", "replicator", "-D", PGDATA,
     "-Fp",          # Plain format: this IS the replica's data directory
     "-Xs",          # Stream WAL alongside the copy
     "-R",           # Write standby.signal + primary_conninfo
     "-S", SLOT],
    capture_output=True, text=True, check=True)

# Same trap as the PITR restore: PostgreSQL refuses to start unless the
# data directory is 0700, and the one the image created is not.
in_container(REPLICA, "chmod", "700", PGDATA)
print("Replica seeded from base backup")

signal_file = in_container(REPLICA, "ls", f"{PGDATA}/standby.signal").strip()
print(f"   {signal_file} exists -> starts in standby mode")

subprocess.run(["docker", "exec", "-u", "postgres", "-d", REPLICA,
                "sh", "-c", f"postgres -D {PGDATA} > /tmp/replica.log 2>&1"],
               capture_output=True, check=True)

# Wait for the primary to actually report the replica, rather than assuming.
for _ in range(30):
    state = psql(PRIMARY, "SELECT coalesce(state, '') FROM pg_stat_replication")
    if state == "streaming":
        break
    time.sleep(1)
else:
    raise RuntimeError("replica never reached streaming state")

print(f"Replication state: {state}")

# Prove the replica is streaming, read-only, and promotable.
import subprocess
import time

PRIMARY = "pg-primary"
REPLICA = "pg-replica"
PGDATA = "/var/lib/postgresql/data"


def psql(container, sql, table=False):
    # -c prints a formatted table, -tAc a bare value.
    flag = "-c" if table else "-tAc"
    result = subprocess.run(
        ["docker", "exec", "-e", "PGPASSWORD=demo", container,
         "psql", "-U", "demo", "-d", "demo", flag, sql],
        capture_output=True, text=True)
    return (result.stdout + result.stderr).strip()


psql(PRIMARY, "CREATE TABLE dr_test (id serial primary key, note text)")
psql(PRIMARY, "INSERT INTO dr_test (note) VALUES ('written on primary')")
time.sleep(1)

print("Replica reads:", psql(REPLICA, "SELECT note FROM dr_test ORDER BY id"))
print("Replica writes:", psql(REPLICA, "INSERT INTO dr_test (note) VALUES ('nope')"))
print()
print("Primary view of the replica:")
print(psql(PRIMARY, "SELECT client_addr, state, sync_state, replay_lag "
                    "FROM pg_stat_replication", table=True))
print(psql(PRIMARY, "SELECT slot_name, active, wal_status "
                    "FROM pg_replication_slots", table=True))

# Failover: promotion ends recovery and the standby becomes a primary.
print("In recovery before promotion:", psql(REPLICA, "SELECT pg_is_in_recovery()"))
subprocess.run(["docker", "exec", "-u", "postgres", REPLICA,
                "pg_ctl", "promote", "-D", PGDATA], capture_output=True, check=True)
time.sleep(3)
print("In recovery after promotion: ", psql(REPLICA, "SELECT pg_is_in_recovery()"))
print("Write to promoted replica:",
      psql(REPLICA, "INSERT INTO dr_test (note) "
                    "VALUES ('written after promotion') RETURNING note"))