Point-in-Time Targeting for PostgreSQL WAL

#!/usr/bin/env python3
"""Compute a PostgreSQL point-in-time recovery target for a restore drill.

Given a target timestamp and a backup catalog, this tool selects the base
backup to restore from, maps the WAL segment range required to replay up to
the target, validates that the WAL timeline is continuous across that range,
and emits the recovery.conf-style settings the sandbox restore will apply.

Exit codes (consumed by the DR drill orchestrator):
    0  target resolved and the WAL timeline is continuous -> proceed
    1  target unreachable: no base backup before it, target beyond the
       archive, or a gap in the WAL timeline                -> halt, escalate
    2  usage / configuration / catalog-parse error          -> abort
"""
from __future__ import annotations

import argparse
import datetime
import json
import sys
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

WAL_SEGMENT_SIZE = 1 << 24            # 16 MiB, the default segment size
SEGMENTS_PER_LOGID = 1 << 8          # 2**32 / WAL_SEGMENT_SIZE = 256


class Unreachable(Exception):
    """The requested instant cannot be recovered from the catalog."""


@dataclass(frozen=True)
class BaseBackup:
    label: str
    timeline: int
    stop_time: int      # epoch seconds when the backup finished
    stop_lsn: str       # "X/Y" LSN at backup stop


@dataclass(frozen=True)
class WalSegment:
    name: str           # 24 hex chars: timeline + high + low
    end_time: int       # epoch seconds of the last record in the segment


def parse_epoch(iso: str) -> int:
    dt = datetime.datetime.fromisoformat(iso.replace("Z", "+00:00"))
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=datetime.timezone.utc)
    return int(dt.timestamp())


def lsn_to_segno(lsn: str) -> int:
    """Map an 'X/Y' LSN to its logical WAL segment number."""
    hi, _, lo = lsn.partition("/")
    value = (int(hi, 16) << 32) | int(lo, 16)
    return value >> 24


def segno_to_name(segno: int, timeline: int) -> str:
    """Render a segment number + timeline as a 24-hex WAL file name."""
    return "%08X%08X%08X" % (timeline,
                             segno // SEGMENTS_PER_LOGID,
                             segno % SEGMENTS_PER_LOGID)


def name_to_tli_segno(name: str) -> Tuple[int, int]:
    tli = int(name[0:8], 16)
    seg = int(name[8:16], 16) * SEGMENTS_PER_LOGID + int(name[16:24], 16)
    return tli, seg


def select_base_backup(backups: List[BaseBackup], target: int) -> BaseBackup:
    """Latest base backup that finished at or before the target."""
    candidates = [b for b in backups if b.stop_time <= target]
    if not candidates:
        raise Unreachable(
            f"no base backup finished at or before {target}; "
            "target predates the earliest full backup")
    return max(candidates, key=lambda b: b.stop_time)


def resolve_target_segment(segments: List[WalSegment],
                           target: int) -> Tuple[int, int]:
    """Return (timeline, segno) of the earliest WAL segment covering the target.

    Segments are ordered by end time across every timeline, so a target that
    falls on a branched timeline still resolves to the segment that must be
    replayed to reach it.
    """
    for seg in sorted(segments, key=lambda s: s.end_time):
        if seg.end_time >= target:
            return name_to_tli_segno(seg.name)
    raise Unreachable(
        f"target {target} is beyond the newest archived WAL; the archive "
        "does not yet cover the requested instant")


def validate_continuity(start_segno: int, end_segno: int, timeline: int,
                        present: set, switches: Dict[int, Tuple[int, int]]
                        ) -> Tuple[int, List[str]]:
    """Walk the segment range, following timeline switches, list any gaps."""
    missing: List[str] = []
    current_tli = timeline
    for segno in range(start_segno, end_segno + 1):
        if (current_tli, segno) not in present:
            missing.append(segno_to_name(segno, current_tli))
        switch = switches.get(current_tli)
        if switch and segno == switch[0]:
            current_tli = switch[1]      # branch onto the child timeline
    return current_tli, missing


def build_recovery_config(target_iso: str, timeline: int,
                          restore_command: str,
                          target_lsn: Optional[str]) -> str:
    lines = [
        f"recovery_target_time = '{target_iso}'",
        "recovery_target_inclusive = true",
        f"recovery_target_timeline = '{timeline}'",
        f"restore_command = '{restore_command}'",
    ]
    if target_lsn:
        lines.insert(1, f"recovery_target_lsn = '{target_lsn}'")
    return "\n".join(lines)


def resolve(catalog: dict, target_iso: str, restore_command: str,
            target_lsn: Optional[str]) -> str:
    target = parse_epoch(target_iso)
    backups = [BaseBackup(b["label"], b["timeline"],
                          parse_epoch(b["stop_time"]), b["stop_lsn"])
               for b in catalog["base_backups"]]
    segments = [WalSegment(s["name"], parse_epoch(s["end_time"]))
                for s in catalog["wal_segments"]]
    switches = {h["parent_timeline"]:
                (lsn_to_segno(h["switch_lsn"]), h["timeline"])
                for h in catalog.get("timeline_history", [])}

    base = select_base_backup(backups, target)
    start_segno = lsn_to_segno(base.stop_lsn)
    _end_tli, end_segno = resolve_target_segment(segments, target)
    present = {name_to_tli_segno(s.name) for s in segments}

    final_tli, missing = validate_continuity(
        start_segno, end_segno, base.timeline, present, switches)
    if missing:
        raise Unreachable(
            "WAL timeline is not continuous to the target; missing "
            f"segments: {', '.join(missing[:5])}")

    print(f"BASE_BACKUP {base.label} timeline={base.timeline} "
          f"start_segment={segno_to_name(start_segno, base.timeline)}")
    print(f"WAL_RANGE {segno_to_name(start_segno, base.timeline)}.."
          f"{segno_to_name(end_segno, final_tli)} "
          f"segments={end_segno - start_segno + 1}")
    return build_recovery_config(target_iso, final_tli, restore_command,
                                 target_lsn)


SELFTEST_CATALOG = {
    "base_backups": [
        {"label": "full-0009", "timeline": 1,
         "stop_time": "2024-11-15T12:00:00Z", "stop_lsn": "0/3000000"}
    ],
    "wal_segments": [
        {"name": "000000010000000000000003", "end_time": "2024-11-15T12:40:00Z"},
        {"name": "000000010000000000000004", "end_time": "2024-11-15T13:20:00Z"},
        {"name": "000000010000000000000005", "end_time": "2024-11-15T14:10:00Z"},
        {"name": "000000010000000000000006", "end_time": "2024-11-15T15:00:00Z"}
    ],
    "timeline_history": []
}


def selftest() -> int:
    assert lsn_to_segno("0/3000000") == 3
    assert segno_to_name(3, 1) == "000000010000000000000003"
    assert name_to_tli_segno("000000010000000000000006") == (1, 6)
    cfg = resolve(SELFTEST_CATALOG, "2024-11-15T14:30:00Z",
                  "cp /wal/%f %p", None)
    assert "recovery_target_time = '2024-11-15T14:30:00Z'" in cfg
    assert "recovery_target_timeline = '1'" in cfg
    # A gap makes the same target unreachable.
    holed = json.loads(json.dumps(SELFTEST_CATALOG))
    holed["wal_segments"] = [s for s in holed["wal_segments"]
                             if s["name"] != "000000010000000000000005"]
    try:
        resolve(holed, "2024-11-15T14:30:00Z", "cp /wal/%f %p", None)
    except Unreachable:
        print("SELFTEST ok: resolves a target and detects a WAL gap")
        return 0
    print("SELFTEST failed: gap not detected", file=sys.stderr)
    return 1


def current_timeline(dsn: str) -> int:
    """Read the source cluster's live timeline as a cross-check.

    Lazily imports psycopg so the catalog-only path needs no driver. The
    live timeline must match the resolved recovery_target_timeline, or the
    catalog is stale relative to the running cluster.
    """
    import psycopg
    with psycopg.connect(dsn) as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT timeline_id FROM pg_control_checkpoint()")
            return int(cur.fetchone()[0])


def main() -> int:
    ap = argparse.ArgumentParser(description="Resolve a PostgreSQL PITR target.")
    ap.add_argument("--catalog", help="path to the backup catalog JSON")
    ap.add_argument("--target", help="ISO-8601 recovery target time")
    ap.add_argument("--restore-command", default="cp /wal_archive/%f %p",
                    help="restore_command the sandbox will use")
    ap.add_argument("--target-lsn", default=None,
                    help="optional explicit recovery_target_lsn to emit")
    ap.add_argument("--verify-dsn", default=None,
                    help="optional source-cluster DSN for a live timeline check")
    ap.add_argument("--selftest", action="store_true",
                    help="run offline checks with an embedded catalog")
    args = ap.parse_args()

    if args.selftest:
        return selftest()
    if not args.catalog or not args.target:
        print("usage: pg_pitr_target.py --catalog C --target ISO "
              "[--restore-command CMD] [--target-lsn X/Y] [--verify-dsn DSN]",
              file=sys.stderr)
        return 2

    try:
        catalog = json.loads(open(args.catalog, encoding="utf-8").read())
        config = resolve(catalog, args.target, args.restore_command,
                         args.target_lsn)
    except Unreachable as exc:
        print(f"UNREACHABLE: {exc}", file=sys.stderr)
        return 1
    except (OSError, KeyError, ValueError, json.JSONDecodeError) as exc:
        print(f"config error: {exc}", file=sys.stderr)
        return 2

    if args.verify_dsn:
        try:
            import psycopg  # noqa: F401  (import checked before connecting)
        except ImportError:
            print("psycopg not installed: pip install 'psycopg[binary]'",
                  file=sys.stderr)
            return 2
        try:
            live = current_timeline(args.verify_dsn)
            print(f"LIVE_TIMELINE {live}")
        except psycopg.Error as exc:
            print(f"cluster cross-check failed: {exc}", file=sys.stderr)
            return 2

    print("--- recovery settings ---")
    print(config)
    return 0


if __name__ == "__main__":
    sys.exit(main())

BASE_BACKUP full-0009 timeline=1 start_segment=000000010000000000000003
WAL_RANGE 000000010000000000000003..000000010000000000000006 segments=4
--- recovery settings ---
recovery_target_time = '2024-11-15T14:30:00Z'
recovery_target_inclusive = true
recovery_target_timeline = '1'
restore_command = 'cp /wal_archive/%f %p'

UNREACHABLE: WAL timeline is not continuous to the target; missing segments: 000000010000000000000005

#!/usr/bin/env bash
set -euo pipefail

OUT=$(python3 pg_pitr_target.py --catalog "$CATALOG" --target "$TARGET_ISO")
case $? in
  0) echo "$OUT" | sed -n '/--- recovery settings ---/,$p' | tail -n +2 \
       > "$SANDBOX_PGDATA/postgresql.auto.conf.d/recovery.conf" ;;
  1) curl -s -X POST "$PAGERDUTY_WEBHOOK" -d '{"event":"pitr_target_unreachable"}'; exit 1 ;;
  *) echo "[$(date -u)] resolver misconfigured; aborting"; exit 2 ;;
esac