upstream-zulip/puppet/zulip/files/postgresql/pg_backup_and_purge

#!/usr/bin/env python3
import configparser
import glob
import json
import logging
import os
import subprocess
import sys
import time
from datetime import datetime, timedelta, timezone

logging.Formatter.converter = time.gmtime
logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s")
logger = logging.getLogger(__name__)


config_file = configparser.RawConfigParser()
config_file.read("/etc/zulip/zulip.conf")


def get_config(
    section: str,
    key: str,
    default_value: str = "",
) -> str:
    if config_file.has_option(section, key):
        return config_file.get(section, key)
    return default_value


replicas = subprocess.check_output(
    ["psql", "-v", "ON_ERROR_STOP=1", "-t", "-c", "SELECT COUNT(*) FROM pg_stat_replication"],
    stdin=subprocess.DEVNULL,
    text=True,
).strip()
if int(replicas) > 0:
    # We are the primary and we have replicas; we expect that backups
    # will be taken on one of them.
    sys.exit(0)

skip_backups = get_config("postgresql", "skip_backups", "")
if skip_backups.lower() in [1, "y", "t", "yes", "true", "enable", "enabled"]:
    sys.exit(0)

is_rhel_based = os.path.exists("/etc/redhat-release")
if is_rhel_based:
    pg_data_paths = glob.glob("/var/lib/pgsql/*/data")
else:
    pg_data_paths = glob.glob("/var/lib/postgresql/*/main")
if len(pg_data_paths) != 1:
    print(f"PostgreSQL installation is not unique: {pg_data_paths}")
    sys.exit(1)
pg_data_path = pg_data_paths[0]

disk_concurrency = get_config("postgresql", "backups_disk_concurrency", "1")
env = os.environ.copy()
env["WALG_UPLOAD_DISK_CONCURRENCY"] = disk_concurrency
subprocess.check_call(["env-wal-g", "backup-push", pg_data_path], env=env)

data = json.loads(subprocess.check_output(["env-wal-g", "backup-list", "--json"]))
# Use strptime for Python < 3.11 compatibility; we might not have
# backports-datetime-fromisoformat here
backups = {
    datetime.strptime(backup["time"], "%Y-%m-%dT%H:%M:%S%z"): backup["backup_name"]
    for backup in data
}

one_month_ago = datetime.now(tz=timezone.utc) - timedelta(days=30)
for date in sorted(backups.keys(), reverse=True):
    if date < one_month_ago:
        # We pass `FIND_FULL` such that if delta backups are being
        # used, we keep the prior FULL backup and all of the deltas
        # that we need.  In practice, this means that if we're doing
        # weekly full backups (`backups_incremental = 6`), we only
        # delete any data once a week.
        subprocess.check_call(
            ["env-wal-g", "delete", "--confirm", "before", "FIND_FULL", backups[date]]
        )
        # Because we're going from most recent to least recent, we
        # only have to do one delete operation
        break