Python + Subprocesses = ⚡

$ ./verify_postgres_restore --help

usage: verify_postgres_restore [-h] -r RCLONE_SOURCE
    [--max-age MAX_AGE] -q QUERIES_DIR
    [--container-name CONTAINER_NAME] [--postgres-version POSTGRES_VERSION]
    [--postgres-password POSTGRES_PASSWORD] [--port PORT] [--max-wait MAX_WAIT]
    [-v]

PostgreSQL Backup Testing Tool

This script validates PostgreSQL backups by:
1. Taking an existing logical backup file (.sql.gz). This is pulled from a local
   or remote file or directory/prefix (newest file by modified time) using
   rclone. Therefore all of the rclone environment variables and configuration
   options apply to control how the backup file is retrieved. See the rclone
   documentation. E.x. https://rclone.org/docs/#environment-variables
2. Spinning up a temporary PostgreSQL container using podman
3. Restoring the backup
4. Running validation queries from SQL scripts organized by database
5. Ensuring all queries return 0 (success)

Query directory structure:
  queries/
  ├── giteadb/
  │   ├── test1.sql
  │   └── test2.sql
  ├── shioridb/
  │   └── test3.sql
  └── anotherdb/
      └── test4.sql

Each SQL script should return a single integer:
  - 0 indicates success/test passed
  - Any other value indicates failure

options:
  -h, --help            show this help message and exit
  -q, --queries-dir QUERIES_DIR
                        Directory containing database subdirectories with SQL
                        test scripts
  -v, --verbose         Enable verbose debug logging

Rclone options:
  Rclone options for retrieving the backup file

  -r, --rclone-source RCLONE_SOURCE
                        Can be a path to gzipped PostgreSQL backup file
                        (.sql.gz), a directory with backups (uses the most
                        recent by modified time), or
                        `rclone_profile:bucket/prefix/pg_backup.sql.gz` or a
                        prefix with many such backups like
                        `rclone_profile:bucket/prefix/`. This can also use an
                        rclone connection string. See
                        https://rclone.org/docs/#connection-strings for more
                        information. E.x. `:s3,env_auth:bucket/path` or
                        `:s3,env_auth,profile=my_profile_name_in_aws_config_file,region=us-east-2:mybucket/prefix/`
  --max-age MAX_AGE     Rclone max age. This is helpful because if no backups
                        were taken in, for example, the last 2 days, this
                        command will fail. This can help detect broken backup
                        procedures. Though if you don't have good alerting for
                        broken backup procedures, you probably don't have good
                        alerting for broken restore testing, let alone any
                        restore testing :) . Examples: `2d`, `1w`, `2025-01-01`.
                        See https://rclone.org/docs/#time-options for more
                        information.

container options:
  Options for configuring the PostgreSQL test container

  --container-name CONTAINER_NAME
                        Name for the test container (default: pg-backup-test)
  --postgres-version POSTGRES_VERSION
                        PostgreSQL version tag for container (default: 16)
  --postgres-password POSTGRES_PASSWORD
                        Password for PostgreSQL superuser (default: testpass123)
  --port PORT           Host port to map to container port 5432 (default: 5432)
  --max-wait MAX_WAIT   Maximum seconds to wait for PostgreSQL to be
                        ready (default: 30)

import logging
from pathlib import Path
import subprocess
from typing import Final

_RCLONE_MAX_TRANSFER_SIZE_REACHED_RC: Final[int] = 8
_RCLONE_NO_FILES_TRANSFERED_RC: Final[int] = 9
_RCLONE_SUCCESS_RC: Final[int] = 0

def _download_most_recent_postgres_backup_sqlgz_archive(
    rclone_source: str,
    workdir: Path,
    rclone_config: Path | None = None,
    max_age: str | None = None,
) -> Path:
    """Download the most recently modified '.sql.gz' file from the rclone source
    into `workdir` and return the Path, which is _guaranteed_ to be in
    `workdir`. `max_age` is used to error if there were no sufficiently recent
    backups (e.x. the backup process is broken). `rclone_source` _can_ refer to
    a local path, `profile_from_rclone_config:path` or an inline rclone
    connection string. Raises `FileNotFoundError` if not files matching the
    criterion were found, or `RuntimeError` if there are any other rclone errors
    (in which case, inspect the logs in stderr)"""
    _cmd = (
        "rclone",
        *(("--config", str(rclone_config)) if rclone_config else ()),
        "copy",
        *(("--max-age", str(max_age)) if max_age else ()),
        "--include",
        "**.sql.gz",
        "--check-first",
        "--order-by",
        "modtime,descending",
        # Max transfer of 1 byte after sorting and 'completing the sort before
        # downloading anything' collectively have the effect of downloading only
        # the 'most recently modified file
        "--max-transfer",
        "1B",
        "--cutoff-mode",
        "soft",
        "--transfers",
        "1",
        "--inplace",
        "--error-on-no-transfer",
        "--use-json-log",
        "-vv",
        rclone_source,
        str(workdir),
    )
    _LOGGER.info("Running cmd %s", shlex.join(_cmd))

    proc = subprocess.run(_cmd)
    rc = proc.returncode
    if rc == _RCLONE_NO_FILES_TRANSFERED_RC:
        _msg = f"No files found in {rclone_source}"
        raise FileNotFoundError(_msg)
    if rc not in (_RCLONE_MAX_TRANSFER_SIZE_REACHED_RC, _RCLONE_SUCCESS_RC):
        _msg = f"Rclone failed and exited with the unexpected code {rc}"
        raise RuntimeError(_msg)
    downloaded_file = next((x for x in workdir.rglob("**/*") if x != workdir), None)
    if not downloaded_file:
        _msg = f"No files downloaded to {workdir}??"
        raise FileNotFoundError(_msg)
    return downloaded_file

import shutil

if not shutil.which("rclone"):
    raise FileNotFoundError("Rclone _must_ be installed to run this script")

find /path/to/parquet/dataset -print0 -type f -name '*.parquet' \
    xargs -0 -I{} -P8 sh -c \
    'python3 my_parquet_transformation_script.py -o /path/to/parquet/dataset {}'

import asyncio
import asyncio.subprocess

class _CommandResult(NamedTuple):
    """Result of a command execution."""

    returncode: int
    stdout: bytes
    stderr: bytes

async def do_stuff(cmd: list[str], timeout_seconds: float) -> _CommandResult
    process = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    try:
        stdout, stderr = await asyncio.wait_for(
            process.communicate(),
            timeout=timeout_seconds,
        )
    except asyncio.TimeoutError:
        process.kill()
        await process.wait()
        raise
    if process.returncode is None:
        _msg = f"Why is returncode not set? {_cmd_str=}"
        raise RuntimeError(_msg)
    return _CommandResult(
        returncode=process.returncode,
        stdout=stdout if capture_output else b"",
        stderr=stderr if capture_output else b"",
    )

import asyncio
import sys
from collections.abc import Sequence
import logging

_LOGGER = logging.getLogger(__name__)

async def async_main(argv: Sequence[str] | None = None) -> int:
    """
    Async main entry point for the CLI.

    Returns:
        Exit code (0 for success, 1 for failure)
    """
    _opts = _parse_arguments(argv if argv is not None else sys.argv[1:])

    logging.basicConfig(
        level=logging.DEBUG if _opts.verbose else logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    try:
        return await _actual_core_logic(opts=_opts)
    except Exception as e:
        _LOGGER.exception("Test execution failed")
        return 1


def main(argv: Sequence[str] | None = None) -> int:
    """
    Main entry point for the CLI.

    Returns:
        Exit code (0 for success, 1 for failure)
    """
    return asyncio.run(async_main(argv))


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

import time

async def poll_long_running_task_and_log_progress(
    cmd_args: tuple[str, ...],
    poll_interval_seconds: float,
) -> None:
    cmd = ("long-running-cli", *cmd_args)
    process = await asyncio.create_subprocess_exec(
        *cmd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    while True:
        try:
            stdout, stderr = await asyncio.wait_for(
                process.communicate(),
                timeout=poll_interval_seconds,
            )
        except asyncio.TimeoutError:
            await _report_on_progress(["cmd", "args"])