#!/usr/bin/env bash
#
# postgresql-basebackup-ubuntu.sh
#
# Physical backup for PostgreSQL on Ubuntu using pg_basebackup.
#
#   - full cluster copy in tar format, one file per tablespace, gzip-compressed
#   - a replication slot keeps required WAL on the server for the duration
#   - backup_manifest verified with pg_verifybackup when available
#   - retention by number of days, applied only after a successful run
#   - a single lock so overlapping cron runs cannot corrupt the set
#
# A base backup alone restores to the moment the backup ended. For
# point-in-time recovery you also need continuous WAL archiving configured on
# the server (archive_mode = on, archive_command writing to durable storage).
#
# Usage:
#   sudo -u postgres ./postgresql-basebackup-ubuntu.sh
#   sudo -u postgres BACKUP_ROOT=/backup/base RETENTION_DAYS=14 ./postgresql-basebackup-ubuntu.sh
#
# Environment:
#   BACKUP_ROOT     destination directory                (default /var/backups/postgresql-base)
#   RETENTION_DAYS  delete backup sets older than N days  (default 7)
#   SLOT_NAME       temporary replication slot name       (default pg_basebackup_script)
#   COMPRESS_LEVEL  gzip level for the tar output 1-9     (default 6)
#   MAX_RATE        throttle, e.g. 100M, empty = no limit (default empty)
#   PGHOST/PGPORT/PGUSER  standard libpq variables; the role needs REPLICATION
#
set -Eeuo pipefail

BACKUP_ROOT="${BACKUP_ROOT:-/var/backups/postgresql-base}"
RETENTION_DAYS="${RETENTION_DAYS:-7}"
SLOT_NAME="${SLOT_NAME:-pg_basebackup_script}"
COMPRESS_LEVEL="${COMPRESS_LEVEL:-6}"
MAX_RATE="${MAX_RATE:-}"
LOCK_FILE="${LOCK_FILE:-/var/lock/postgresql-basebackup.lock}"

log() {
  echo "[$(date --iso-8601=seconds)] $*"
}

fail() {
  echo "ERROR: $*" >&2
  exit 1
}

[[ "$RETENTION_DAYS" =~ ^[0-9]+$ ]] || fail "RETENTION_DAYS must be an integer."
[[ "$COMPRESS_LEVEL" =~ ^[1-9]$ ]] || fail "COMPRESS_LEVEL must be 1-9."
[[ "$SLOT_NAME" =~ ^[a-z0-9_]+$ ]] || fail "SLOT_NAME must match ^[a-z0-9_]+$."

for command_name in psql pg_basebackup flock; do
  command -v "$command_name" >/dev/null || fail "Missing command: $command_name"
done

exec 9>"$LOCK_FILE" || fail "Cannot open lock file $LOCK_FILE"
flock -n 9 || fail "Another base backup is already in progress."

SERVER_VERSION_NUM="$(psql -XAtqc 'SHOW server_version_num' postgres)" ||
  fail "Cannot connect to PostgreSQL. The role needs the REPLICATION attribute."
SERVER_VERSION="$(psql -XAtqc 'SHOW server_version' postgres)"

# pg_basebackup opens a second connection for WAL streaming, so the server needs
# at least one spare walsender slot beyond any replicas already connected.
WAL_SENDERS="$(psql -XAtqc 'SHOW max_wal_senders' postgres)"
(( WAL_SENDERS >= 1 )) || fail "max_wal_senders is 0; pg_basebackup cannot stream WAL."

ARCHIVE_MODE="$(psql -XAtqc 'SHOW archive_mode' postgres)"
if [[ "$ARCHIVE_MODE" != "on" ]]; then
  log "WARNING: archive_mode is '$ARCHIVE_MODE'. This backup restores only to its"
  log "WARNING: own end point. Point-in-time recovery requires WAL archiving."
fi

TIMESTAMP="$(date +%Y%m%dT%H%M%S)"
STAGING_DIR="${BACKUP_ROOT}/.in-progress-${TIMESTAMP}"
FINAL_DIR="${BACKUP_ROOT}/${TIMESTAMP}"

install -d -m 700 "$BACKUP_ROOT"
install -d -m 700 "$STAGING_DIR"

cleanup_staging() {
  local exit_code=$?
  if (( exit_code != 0 )); then
    [[ -d "$STAGING_DIR" ]] && { log "Run failed; removing $STAGING_DIR"; rm -rf -- "$STAGING_DIR"; }
    # A slot left behind after a failure makes the server retain WAL forever,
    # which eventually fills pg_wal and stops the primary.
    psql -XAtqc "SELECT pg_drop_replication_slot('${SLOT_NAME}')
                   WHERE EXISTS (SELECT 1 FROM pg_replication_slots
                                  WHERE slot_name = '${SLOT_NAME}')" postgres >/dev/null 2>&1 || true
  fi
  exit "$exit_code"
}
trap cleanup_staging EXIT

# A stale slot from a killed previous run would make --create-slot fail.
psql -XAtqc "SELECT pg_drop_replication_slot('${SLOT_NAME}')
               WHERE EXISTS (SELECT 1 FROM pg_replication_slots
                              WHERE slot_name = '${SLOT_NAME}' AND active IS NOT TRUE)" \
  postgres >/dev/null

log "PostgreSQL $SERVER_VERSION"
log "Destination: $FINAL_DIR"
log "Starting pg_basebackup (this reads the whole cluster)"

BASEBACKUP_ARGS=(
  --pgdata="$STAGING_DIR"
  --format=tar
  --gzip
  --compress="$COMPRESS_LEVEL"
  --wal-method=stream          # ships the WAL generated during the copy
  --checkpoint=fast            # do not wait for the next scheduled checkpoint
  --create-slot
  --slot="$SLOT_NAME"
  --progress
  --verbose
  --no-password
)
[[ -n "$MAX_RATE" ]] && BASEBACKUP_ARGS+=( --max-rate="$MAX_RATE" )

pg_basebackup "${BASEBACKUP_ARGS[@]}"

# pg_basebackup drops the slot itself on success; this only covers the gap where
# the server kept it after an unclean disconnect.
psql -XAtqc "SELECT pg_drop_replication_slot('${SLOT_NAME}')
               WHERE EXISTS (SELECT 1 FROM pg_replication_slots
                              WHERE slot_name = '${SLOT_NAME}' AND active IS NOT TRUE)" \
  postgres >/dev/null

# backup_manifest exists from PostgreSQL 13 onward; pg_verifybackup reads the
# tar set directly from PostgreSQL 17, so older versions only get a file check.
if [[ -f "${STAGING_DIR}/backup_manifest" ]]; then
  if command -v pg_verifybackup >/dev/null && (( SERVER_VERSION_NUM >= 170000 )); then
    log "Verifying backup manifest with pg_verifybackup"
    pg_verifybackup --format=tar "$STAGING_DIR" || fail "pg_verifybackup reported a corrupt backup."
  else
    log "backup_manifest present; verify after extraction with pg_verifybackup"
  fi
else
  log "No backup_manifest (PostgreSQL < 13); skipping manifest verification."
fi

[[ -f "${STAGING_DIR}/base.tar.gz" ]] || fail "base.tar.gz is missing from the backup."
gzip --test "${STAGING_DIR}"/*.tar.gz || fail "A compressed archive failed the gzip integrity test."

(cd "$STAGING_DIR" && sha256sum ./*.tar.gz > SHA256SUMS)

cat > "${STAGING_DIR}/MANIFEST.txt" <<MANIFEST
backup_type      physical (pg_basebackup, tar + gzip)
started_at       ${TIMESTAMP}
finished_at      $(date --iso-8601=seconds)
server_version   ${SERVER_VERSION}
archive_mode     ${ARCHIVE_MODE}
host             $(hostname -f 2>/dev/null || hostname)
restore          stop postgres, empty PGDATA, extract base.tar.gz into it,
                 extract pg_wal.tar.gz into PGDATA/pg_wal, then start the server
pitr             add restore_command and recovery_target_time to postgresql.conf
                 and create the file PGDATA/recovery.signal before starting
MANIFEST

chmod -R go-rwx "$STAGING_DIR"
mv -- "$STAGING_DIR" "$FINAL_DIR"
trap - EXIT

if (( RETENTION_DAYS > 0 )); then
  log "Applying retention: removing sets older than ${RETENTION_DAYS} days"
  find "$BACKUP_ROOT" -mindepth 1 -maxdepth 1 -type d \
    -name '20*T*' -mtime "+${RETENTION_DAYS}" \
    -print -exec rm -rf -- {} +
fi

log "Base backup completed: $FINAL_DIR"
log "Total size: $(du -sh "$FINAL_DIR" | cut -f1)"
log "Reminder: WAL older than the oldest retained base backup can be recycled;"
log "keep archived WAL for at least as long as the base backups it covers."