pg_refclone.py

#!/usr/bin/env python3
import argparse
import logging
import os
import random
import string
import sys
import shutil
import subprocess
from pathlib import Path


def generate_slot_name():
    return f"pg_refclone_{''.join(random.choices(string.ascii_lowercase + string.digits, k=6))}"


import psycopg
from psycopg import connect
from psycopg.rows import namedtuple_row

# to ensure that we are actually connecting to the right database,
# get socket dir and port from the postmaster.pid file in original_dir
def conn_str_from_postmaster_pid(data_dir):
    postmaster_pid = data_dir / "postmaster.pid"
    if not postmaster_pid.exists():
        raise ValueError(f"No postmaster.pid found in {data_dir}")
    lines = postmaster_pid.read_text().strip().splitlines()
    if len(lines) < 5:
        raise ValueError(f"Invalid postmaster.pid in {data_dir}")
    port = lines[3].strip()
    socket_dir = lines[4].strip()
    return f"host={socket_dir} port={port} dbname=postgres"


def get_pg_bin_dir(data_dir):
    postmaster_opts = data_dir / "postmaster.opts"
    if not postmaster_opts.exists():
        raise ValueError(f"No postmaster.opts found in {data_dir}")
    lines = postmaster_opts.read_text().strip().splitlines()
    if not lines:
        raise ValueError(f"Invalid postmaster.opts in {data_dir}")
    postgres_path = lines[0].split()[0]
    return str(Path(postgres_path).parent)


def get_pg_receivewal_path(data_dir):
    bin_dir = get_pg_bin_dir(data_dir)
    return str(Path(bin_dir) / "pg_receivewal")

INCOMPLETE_BACKUP_LABEL = """pg_refclone.py is not finished!
Removing this file and starting PostgreSQL
will implicitly accept data corruption,
as the recovery will not run through
enough WAL to make things consistent.
"""

logging.basicConfig(level=logging.DEBUG, format='%(message)s')
logger = logging.getLogger(__name__)


def main():
    parser = argparse.ArgumentParser(description="Clone a PostgreSQL data directory")
    parser.add_argument("original_dir", help="Path to the original PostgreSQL data directory")
    parser.add_argument("clone_dir", help="Path for the cloned database")
    parser.add_argument(
        "--start-fast",
        action="store_true",
        default=False,
        help="Force immediate checkpoint when starting backup"
    )
    args = parser.parse_args()

    original_dir = Path(args.original_dir)
    clone_dir = Path(args.clone_dir)

    pg_receivewal_path = get_pg_receivewal_path(original_dir)

    if clone_dir.exists() and any(clone_dir.iterdir()):
        logger.error(f"Error: clone directory {clone_dir} already exists and is not empty")
        sys.exit(1)

    # test if the original_dir and clone_dir are on the same device,
    # otherwise reflink cloning will definitely not work.
    # we can check clone_dir's parent (as clone_dir does not exist yet).
    if os.stat(original_dir).st_dev != os.stat(clone_dir.parent).st_dev:
        logger.error(f"Error: {original_dir} and {clone_dir} must be on the same filesystem")
        sys.exit(1)

    # check for tablespaces (symlinks in pg_tblspc/ subdirectory)
    pg_tblspc_dir = original_dir / "pg_tblspc"
    if pg_tblspc_dir.exists():
        for entry in pg_tblspc_dir.iterdir():
            if entry.is_symlink():
                logger.error(f"Error: symlink found in {pg_tblspc_dir}, indicating tablespace usage. Tablespaces are not supported.")
                sys.exit(1)

    # create directory for clone
    logger.info(f"Creating clone directory: {clone_dir}")
    os.makedirs(clone_dir, exist_ok=True)

    # make sure data directory cannot be started before we finish properly
    logger.info(f"FILE: write {(clone_dir / 'backup_label')}")
    (clone_dir / "backup_label").write_text(INCOMPLETE_BACKUP_LABEL)

    # open a connection to the original database, we'll use this to create a
    # replication slot, and start and stop the backup, and drop the slot later.
    conn_str = conn_str_from_postmaster_pid(original_dir)
    try:
        conn = connect(conn_str)
    except psycopg.OperationalError as e:
        logger.error(f"Error: could not connect to database: {e}")
        sys.exit(1)

    conn.autocommit = True

    slot_name = generate_slot_name()

    try:
        # create replication slot so that we can later pull all the required WAL
        # (will be cleaned up on error or clean exit by the finally statement below)
        with conn.cursor() as cur:
            logger.info(f"DB: SELECT pg_create_physical_replication_slot('{slot_name}', TRUE)")
            cur.execute(f"SELECT pg_create_physical_replication_slot('{slot_name}', TRUE)")
            cur.fetchone()

            # start backup
            logger.info(f"DB: SELECT pg_backup_start('pg_refclone', {args.start_fast})")
            cur.execute("SELECT pg_backup_start('pg_refclone', %s)", (args.start_fast,))
            start_lsn = cur.fetchone()[0]
            logger.info(f"BACKUP START LSN: {start_lsn}")

        # do the reflink copy
        # NOTE: shutil.copytree in Python 3.14+ will use optimizations like copy_file_range(2) for
        # local file copies when copy_function=shutil.copy2 is used.
        # See https://docs.python.org/3/library/shutil.html#platform-dependent-efficient-copy-operations
        # We will not recreate symlinks, but instead copy the symlinked files directly.
        # This will most likely not be blazing fast, e.g. if you have whole tablespaces symlinked,
        # but it will be safe that way. Recreating the symlinks would result in two database directories
        # linking to the same exact files on disk, leading to corruption.
        logger.info(f"FILE: Copying {original_dir} to {clone_dir}")
        shutil.copytree(
            original_dir, clone_dir,
            symlinks=False,
            copy_function=shutil.copy2,
            dirs_exist_ok=True,
            ignore=shutil.ignore_patterns("pg_wal", "postmaster.pid", "postmaster.opts", "pg_control")
        )

        # make sure the clone is started in standby mode
        standby_signal = clone_dir / "standby.signal"
        logger.info(f"FILE: touch {standby_signal}")
        standby_signal.touch()

        # stop backup, saving return values for later use
        # the return values are lsn, labelfile, spcmapfile
        with conn.cursor(row_factory=namedtuple_row) as cur:
            logger.info("DB: SELECT * FROM pg_backup_stop(FALSE)")
            cur.execute("SELECT * FROM pg_backup_stop(FALSE)")
            result = cur.fetchone()
            end_lsn = result.lsn
            labelfile = result.labelfile
            spcmapfile = result.spcmapfile
            logger.info(f"BACKUP STOP LSN: {end_lsn}")

        # create backup_label file
        logger.info(f"FILE: write {(clone_dir / 'backup_label')}")
        (clone_dir / "backup_label").write_text(labelfile)

        # copy pg_control (was excluded from copytree)
        shutil.copy2(original_dir / "global" / "pg_control", clone_dir / "global" / "pg_control")

        # create tablespace_map file
        tablespace_map_path = clone_dir / "tablespace_map"
        if spcmapfile is not None:
            logger.info(f"FILE: write {tablespace_map_path}")
            tablespace_map_path.write_bytes(spcmapfile.encode() if isinstance(spcmapfile, str) else spcmapfile)

        # advance slot to start of backup (otherwise we get all the WAL produced before pg_backup_start was called)
        with conn.cursor() as cur:
            logger.info(f"DB: SELECT pg_replication_slot_advance('{slot_name}', {start_lsn!r})")
            cur.execute(f"SELECT pg_replication_slot_advance('{slot_name}', %s)", (start_lsn,))
            cur.fetchone()

        # create pg_wal directory (not copied via ignore_patterns)
        pg_wal_path = clone_dir / "pg_wal"
        logger.info(f"FILE: mkdir {pg_wal_path}")
        pg_wal_path.mkdir()

        # pull all the WAL that we need, should end automatically when we reach backup end LSN
        logger.info(f"PG_RECEIVEWAL: launch pg_receivewal to stream WAL to {pg_wal_path}")
        result = subprocess.run(
            [
                pg_receivewal_path,
                "-v",
                "-d", conn_str,
                "-D", str(pg_wal_path) + "/",
                "-S", slot_name,
                "-E", str(end_lsn),
                "-s", "1"
            ],
            capture_output=True,
            text=True
        )
        logger.info(f"PG_RECEIVEWAL output begin:\n{result.stdout}\nPG_RECEIVEWAL output end:")
        if result.returncode != 0:
            logger.error(f"PG_RECEIVEWAL error begin:\n: {result.stderr}\nPG_RECEIVEWAL error end:")
            sys.exit(1)

        # get the WAL file name matching the backup end LSN, ensuring the latter is properly quoted
        with conn.cursor() as cur:
            logger.info(f"DB: SELECT pg_walfile_name({end_lsn!r})")
            cur.execute("SELECT pg_walfile_name(%s)", (end_lsn,))
            stop_walfile = cur.fetchone()[0]
            logger.info(f"BACKUP STOP WAL FILE: {stop_walfile}")

        # rename the WAL file containing the backup stop LSN,
        # as pg_receivewal keeps it as a ".partial" file,
        # even though pg_backup_stop() switches to a new file.
        partial_wal = pg_wal_path / f"{stop_walfile}.partial"
        if partial_wal.exists():
            final_wal = pg_wal_path / stop_walfile
            logger.info(f"FILE: rename {partial_wal} -> {final_wal}")
            partial_wal.rename(final_wal)

    finally:
        # drop slot later
        with conn.cursor() as cur:
            logger.info(f"DB: SELECT pg_drop_replication_slot('{slot_name}')")
            cur.execute(f"SELECT pg_drop_replication_slot('{slot_name}')")
        conn.close()


if __name__ == "__main__":
    main()