Fetch master wal-file via replication protocol

pg-receivewal-follow.sh

#!/usr/bin/env bash
#
# pg-receivewal-follow.sh — get one WAL file from the PostgreSQL primary.
#
# Run with:  timeout -k 1s 30s bash pg-receivewal-follow.sh %f %p
#
# =====================================================================
# 1. WHAT THIS SCRIPT DOES
# =====================================================================
#
# PostgreSQL writes all changes to WAL files (Write-Ahead Log).
# A standby server replays these files to stay equal to the primary.
# The standby asks for the files one by one. It runs `restore_command`
# for each file.
#
# This script is a `restore_command`. PostgreSQL calls it with two
# arguments:
#
#   %f  the name of the file that PostgreSQL wants
#   %p  the path where PostgreSQL wants the file
#
# The script gets the file from the primary and puts it at that path.
# It uses only the PostgreSQL replication protocol.
#
# Put this line in postgresql.conf of the standby:
#
#   restore_command = 'timeout -k 1s 30s bash /path/pg-receivewal-follow.sh %f %p'
#
# The connection settings come from the usual libpq variables:
# PGHOST, PGPORT, PGUSER, PGPASSWORD (or ~/.pgpass), PGSSLMODE.
# The user must have the REPLICATION privilege.
#
# =====================================================================
# 2. THE TWO KINDS OF FILES
# =====================================================================
#
# A) A WAL segment. The name has 24 hex digits: 000000010000000A00000003
#
#      digits  1..8  = timeline            (here: 1)
#      digits  9..16 = high part of the LSN (here: 0xA)
#      digits 17..24 = segment number       (here: 3)
#
#    One segment is 16 MB by default. If your cluster uses another
#    size, set WAL_SEGMENT_MB. Look at `pg_controldata` to find it.
#
# B) A timeline history file: 00000002.history
#
#    This file tells where timeline 2 branched from timeline 1.
#    Timeline 1 has no history file.
#
# An LSN is a position in the WAL. PostgreSQL shows it as two hex
# numbers "high/low". The real position is: high * 2^32 + low.
#
# =====================================================================
# 3. HOW THE SCRIPT WORKS
# =====================================================================
#
#   Check the two arguments.
#   Make a private work directory next to the destination.
#   │
#   ├─ Does the name end with ".history"?
#   │    Ask the primary: TIMELINE_HISTORY <n>.
#   │    Write the answer to a work file.
#   │
#   └─ Is the name a WAL segment?
#        Ask the primary: IDENTIFY_SYSTEM.
#        The answer has the flush LSN. All bytes before it are safe
#        on the disk of the primary.
#        Compute the first and the last byte of the segment.
#        │
#        ├─ Is the segment fully after the flush LSN?
#        │    Exit 1. PostgreSQL will ask again later.
#        │
#        └─ Is at least one byte of the segment before the flush LSN?
#             Run pg_receivewal. Stop at the smaller of:
#             the segment end, or the flush LSN.
#             Check the result (see section 4).
#             Make the file exactly one segment long.
#
#   Write the work file to the disk (sync).
#   Move it to the destination.
#
# =====================================================================
# 4. WHY THE CHECKS ARE STRICT
# =====================================================================
#
# PostgreSQL trusts this script. If the script gives exit 0 and a bad
# file, the standby can replay wrong data or stop forever. So the
# script has one rule: if in doubt, exit 1. Exit 1 is always safe.
# PostgreSQL waits and calls the script again.
#
# pg_receivewal has some behaviours that the script must handle:
#
# - It names the file ".partial" until it has the full segment.
#   The script accepts a ".partial" file only if the file has the
#   full size, and the log confirms the received position.
#
# - Exit code 0 is not proof. The script reads the log line
#   "stopped log streaming at X/Y (timeline N)" and compares the
#   numbers with the request.
#
# - When the segment is on an old timeline, PostgreSQL 16 writes an
#   error "server stopped streaming timeline ...". The file is still
#   complete. The script accepts this one error only if all numbers
#   in it prove that the segment is complete. All other errors are
#   fatal.
#
# - pg_receivewal follows timeline switches and can make more files.
#   The script takes only the requested file.
#
# Partial segments: if the primary has not written the full segment,
# the script takes the bytes up to the flush LSN. The rest of the
# file is zero bytes. PostgreSQL stops at the first zero record and
# asks for the segment again later.
#
# =====================================================================
# 5. EXIT CODES
# =====================================================================
#
#   0  the file is at the destination and is complete
#   1  all other cases; nothing is left at the destination
#
# =====================================================================

set -Eeuo pipefail
export LC_ALL=C PG_COLOR=never PGTARGETSESSIONATTRS=any
unset PGOPTIONS PGSERVICE PGSERVICEFILE
umask 077

# Constants that come from the WAL format.
readonly LSN_HIGH_STEP=4294967296 # 2^32: one step of the "high" part of an LSN
readonly LSN_LOW_MASK=4294967295  # 2^32 - 1: the "low" part of an LSN
readonly MEGABYTE=1048576

# --- Step 1: read and check the arguments ----------------------------

[[ $# == 2 ]] || exit 1
wal_name=$1
destination=$2

# Accept only the two name formats from section 2.
[[ $wal_name =~ ^[0-9A-F]{24}$ || $wal_name =~ ^[0-9A-F]{8}\.history$ ]] || exit 1

timeline=$((16#${wal_name:0:8}))
((timeline > 0)) || exit 1

# --- Step 2: make a private work directory ---------------------------

# The directory is next to the destination. So the final `mv` stays on
# one file system and is atomic. PostgreSQL never sees a half file.
work_dir=$(mktemp -d -- "$(dirname -- "$destination")/.wal-fetch.XXXXXXXX")
trap 'rm -rf -- "$work_dir"' EXIT
trap 'exit 1' ERR HUP INT TERM

result_file="$work_dir/$wal_name"

# psql in replication mode. It speaks the replication protocol, not SQL.
#   -X  do not read ~/.psqlrc      -w  never ask for a password
#   -q  quiet                      -A  no column alignment
#   -t  print only the data rows
psql=(psql -XwqAt -v ON_ERROR_STOP=1 -d 'replication=yes')

# --- Helper: change a text LSN "high/low" into one integer -----------

# One integer is easy to compare and to use in arithmetic.
# Exit 1 if the text is not an LSN.
lsn_to_number() {
    [[ $1 =~ ^([0-9A-F]{1,8})/([0-9A-F]{1,8})$ ]] || exit 1
    local high=$((16#${BASH_REMATCH[1]}))
    local low=$((16#${BASH_REMATCH[2]}))
    ((high < 2147483648)) || exit 1 # Keep the result inside a signed 64-bit integer.
    echo "$((high * LSN_HIGH_STEP + low))"
}

if [[ $wal_name == *.history ]]; then
    # ===== Branch A: timeline history file ===========================

    # Timeline 1 has no history file.
    ((timeline > 1)) || exit 1

    # The server answers with two fields: the file name and the content.
    # The content can have newlines. So -z separates the fields with a
    # NUL byte, and -0 ends the row with a NUL byte.
    "${psql[@]}" -z -0 -c "TIMELINE_HISTORY $timeline" >"$work_dir/response"
    {
        IFS= read -r -d '' answer_name
        IFS= read -r -d '' answer_content
    } <"$work_dir/response"

    [[ $answer_name == "$wal_name" && -n $answer_content ]] || exit 1
    printf '%s' "$answer_content" >"$result_file"

else
    # ===== Branch B: WAL segment =====================================

    # B1. Segment size. It must be a power of two from 1 MB to 1024 MB.
    segment_mb=${WAL_SEGMENT_MB:-16}
    [[ $segment_mb =~ ^[1-9][0-9]{0,3}$ ]] || exit 1
    ((segment_mb <= 1024 && (segment_mb & (segment_mb - 1)) == 0)) || exit 1
    segment_size=$((segment_mb * MEGABYTE))

    # B2. Ask the primary who it is and how far it has written.
    # IDENTIFY_SYSTEM answers: system id, timeline, flush LSN, db name.
    identity=$("${psql[@]}" -F ' ' -c IDENTIFY_SYSTEM)
    read -r _system_id primary_timeline flush_lsn _db_name <<<"$identity"
    flush_pos=$(lsn_to_number "$flush_lsn")

    # B3. Compute the first and the last byte of the segment from its name.
    segment_number=$((16#${wal_name:16:8}))
    ((segment_number < LSN_HIGH_STEP / segment_size)) || exit 1
    segment_start=$((16#${wal_name:8:8} * LSN_HIGH_STEP + segment_number * segment_size))
    segment_end=$((segment_start + segment_size))
    ((segment_start > 0 && segment_end > segment_start)) || exit 1
    ((timeline <= primary_timeline)) || exit 1 # The primary cannot know a future timeline.

    # B4. The primary can only give bytes that it has flushed.
    stop_pos=$((segment_end < flush_pos ? segment_end : flush_pos))
    ((stop_pos > segment_start)) || exit 1 # Nothing of this segment is flushed yet.

    # pg_receivewal -E stops *after* the given LSN. So give it stop - 1.
    last_pos=$((stop_pos - 1))
    printf -v end_lsn '%X/%X' "$((last_pos >> 32))" "$((last_pos & LSN_LOW_MASK))"

    # B5. Stream the segment.
    # An empty ".partial" file tells pg_receivewal to start this segment
    # from its first byte. Without it, pg_receivewal starts at the flush
    # LSN and the start of the segment is lost.
    #   -w  never ask for a password   -n  do not loop after a stop
    #   -v  verbose: writes the "stopped log streaming" line that B7 needs
    : >"$result_file.partial"
    pg_receivewal -w -n -v -D "$work_dir" -E "$end_lsn" \
        -d 'application_name=wal-fetch' 2>"$work_dir/receive.log" || exit 1

    # B6. Find the file. pg_receivewal follows timeline switches and can
    # make more files. Take only the requested one. It can still have the
    # ".partial" name when the stream stopped inside the segment.
    [[ -f $result_file ]] || result_file="$result_file.partial"
    [[ -f $result_file && $(stat -c %s -- "$result_file") == "$segment_size" ]] || exit 1

    # B7. Exit code 0 from pg_receivewal is not proof. Read the real
    # received LSN and timeline from the log. Compare them with the request.
    stopped_line='^pg_receivewal: stopped log streaming at ([0-9A-F]+/[0-9A-F]+) \(timeline ([0-9]+)\)$'
    read -r received_lsn received_timeline < <(sed -nE "s|$stopped_line|\1 \2|p" "$work_dir/receive.log")
    received_pos=$(lsn_to_number "$received_lsn")
    ((received_pos >= stop_pos)) || exit 1
    ((received_timeline >= timeline && received_timeline <= primary_timeline)) || exit 1

    # B8. Look for errors in the log. Every error is fatal, with one
    # exception. PostgreSQL 16 writes this error when the segment is on an
    # old timeline and the stream reached the end of that timeline:
    #   "server stopped streaming timeline T at X/Y, but reported next
    #    timeline N to begin at X/Z"
    # The file is complete in this case. Accept this error only if every
    # number in it proves that the segment is complete.
    early_stop_line='^pg_receivewal: error: server stopped streaming timeline ([0-9]+) at ([0-9A-F]+/[0-9A-F]+), but reported next timeline ([0-9]+) to begin at ([0-9A-F]+/[0-9A-F]+)$'
    errors=$(sed -nE '/: (error|fatal):|could not finish writing WAL files/p' "$work_dir/receive.log")
    if [[ -n $errors ]]; then
        [[ $errors =~ $early_stop_line ]] || exit 1 # Any other error, or more than one line.
        old_timeline=${BASH_REMATCH[1]}
        old_end_pos=$(lsn_to_number "${BASH_REMATCH[2]}")
        next_timeline=${BASH_REMATCH[3]}
        fork_pos=$(lsn_to_number "${BASH_REMATCH[4]}")

        ((received_timeline == timeline)) || exit 1 # We stayed on the requested timeline.
        ((old_timeline == timeline)) || exit 1      # The message is about that timeline.
        ((next_timeline > timeline)) || exit 1      # The switch goes forward in time.
        ((stop_pos == segment_end)) || exit 1       # We asked for the full segment.
        [[ -f $work_dir/$wal_name ]] || exit 1      # pg_receivewal gave it the final name.
        ((old_end_pos == stop_pos)) || exit 1       # The old timeline ended exactly at our end.
        ((fork_pos > stop_pos)) || exit 1           # The new timeline starts after our segment.
    fi

    # B9. Make the file exactly one segment long.
    # Bytes after the stop position must be zero. Cut the file at the stop
    # position. Then grow it back to the full size. The new bytes are zero.
    truncate -s "$((stop_pos - segment_start))" -- "$result_file"
    truncate -s "$segment_size" -- "$result_file"
fi

# --- Step 3: put the file in place -----------------------------------

# sync writes the data to the disk. The move is atomic: PostgreSQL sees
# no file, or a complete file. Never a half file.
sync -d -- "$result_file"
mv -fT -- "$result_file" "$destination"