dtf-system/ops/backup.py

"""Local PostgreSQL and clean-object backup with isolated restore verification."""
import argparse
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import subprocess
from uuid import uuid4

ROOT = Path(__file__).resolve().parent.parent
BACKUPS = ROOT / 'backups'

def docker(script, *args, **kwargs):
    return subprocess.run(['docker','compose','exec','-T','db','sh','-c',script,'sh',*args],
                          cwd=ROOT,check=True,**kwargs)

def checksum(path):
    digest=hashlib.sha256()
    with path.open('rb') as stream:
        for block in iter(lambda:stream.read(1048576),b''):digest.update(block)
    return digest.hexdigest()

def compose_exec(service, *command, **kwargs):
    return subprocess.run(['docker','compose','exec','-T',service,*command],
                          cwd=ROOT,check=True,**kwargs)

def create():
    BACKUPS.mkdir(mode=0o700,exist_ok=True)
    prefix='dtf-'+datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')+'-'+uuid4().hex[:8]
    database=BACKUPS/(prefix+'.dump')
    objects=BACKUPS/(prefix+'.objects.tar.gz')
    manifest_path=BACKUPS/(prefix+'.manifest.json')
    sidecar=BACKUPS/(prefix+'.manifest.sha256')
    created=[]
    try:
        with database.open('xb') as stream:
            created.append(database);database.chmod(0o600)
            docker('pg_dump -U "$POSTGRES_USER" -d "$POSTGRES_DB" --format=custom',stdout=stream)
        with objects.open('xb') as stream:
            created.append(objects);objects.chmod(0o600)
            result=compose_exec('api','python','-m','local.storage_backup','export',
                                stdout=stream,stderr=subprocess.PIPE)
        summaries=[line.removeprefix('DTF_BACKUP_SUMMARY ') for line in
                   result.stderr.decode().splitlines() if line.startswith('DTF_BACKUP_SUMMARY ')]
        if len(summaries)!=1:
            raise RuntimeError('Object backup did not return a valid summary')
        manifest={'format':'dtf-local-backup-v2',
                  'created_at':datetime.now(timezone.utc).isoformat(),
                  'database':{'file':database.name,'sha256':checksum(database)},
                  'objects':{'file':objects.name,'sha256':checksum(objects),
                             **json.loads(summaries[0])}}
        manifest_path.write_text(json.dumps(manifest,indent=2,sort_keys=True)+'\n')
        created.append(manifest_path);manifest_path.chmod(0o600)
        sidecar.write_text(checksum(manifest_path)+'\n')
        created.append(sidecar);sidecar.chmod(0o600)
    except Exception:
        for path in reversed(created):path.unlink(missing_ok=True)
        raise
    print(manifest_path.relative_to(ROOT))
    return manifest_path

def verify_database(path):
    if checksum(path)!=path.with_suffix('.sha256').read_text().strip():
        raise ValueError('Backup checksum mismatch')
    database='dtf_verify_'+uuid4().hex
    docker('createdb -U "$POSTGRES_USER" "$1"',database)
    try:
        with path.open('rb') as stream:
            docker('pg_restore -U "$POSTGRES_USER" -d "$1" --exit-on-error --no-owner --no-privileges',database,stdin=stream)
        result=docker('psql -U "$POSTGRES_USER" -d "$1" -At -v ON_ERROR_STOP=1 -c "SELECT json_build_object(\'orders\',(SELECT count(*) FROM dtf_local.orders),\'uploads\',(SELECT count(*) FROM dtf_local.uploads),\'accounts\',(SELECT count(*) FROM dtf_local.accounts),\'history\',(SELECT count(*) FROM dtf_local.movements));"',database,capture_output=True,text=True)
        print('PASS: backup restored to isolated database; restored counts '+result.stdout.strip())
    finally:
        # Only the freshly created UUID-named verification database is removed.
        docker('dropdb -U "$POSTGRES_USER" "$1"',database)
    print('Removed temporary verification database. Active database was untouched.')

def bundle_file(manifest_path, name):
    if not isinstance(name,str) or Path(name).name!=name:
        raise ValueError('Backup manifest contains an invalid filename')
    path=(manifest_path.parent/name).resolve()
    if path.parent!=BACKUPS.resolve():
        raise ValueError('Backup manifest references a file outside backups/')
    return path

def verify(path):
    path=path.resolve()
    if path.parent!=BACKUPS.resolve():
        raise ValueError('Choose a backup generated in this repository backups/ directory')
    if path.suffix=='.dump':
        return verify_database(path)
    if not path.name.endswith('.manifest.json'):
        raise ValueError('Choose a v2 .manifest.json or legacy .dump backup')
    sidecar=path.with_name(path.name.removesuffix('.json')+'.sha256')
    if checksum(path)!=sidecar.read_text().strip():
        raise ValueError('Backup manifest checksum mismatch')
    manifest=json.loads(path.read_text())
    if manifest.get('format')!='dtf-local-backup-v2':
        raise ValueError('Unsupported backup format')
    database=bundle_file(path,manifest.get('database',{}).get('file'))
    objects=bundle_file(path,manifest.get('objects',{}).get('file'))
    if checksum(database)!=manifest['database'].get('sha256') or checksum(objects)!=manifest['objects'].get('sha256'):
        raise ValueError('Backup bundle checksum mismatch')
    # Reuse the legacy database verifier with a temporary matching sidecar.
    legacy_sidecar=database.with_suffix('.sha256')
    had_legacy=legacy_sidecar.exists()
    previous=legacy_sidecar.read_bytes() if had_legacy else None
    legacy_sidecar.write_text(manifest['database']['sha256']+'\n');legacy_sidecar.chmod(0o600)
    try:
        verify_database(database)
    finally:
        if had_legacy:legacy_sidecar.write_bytes(previous)
        else:legacy_sidecar.unlink(missing_ok=True)
    with objects.open('rb') as stream:
        compose_exec('api','python','-m','local.storage_backup','verify',stdin=stream)
    print('PASS: combined database and clean-object backup verified. Active data was untouched.')

if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('command',choices=['create','verify','create-and-verify'])
    parser.add_argument('path',nargs='?',type=Path,
                        help='v2 manifest printed by create, or a legacy .dump')
    args=parser.parse_args()
    if args.command=='verify':
        if not args.path:parser.error('verify requires the path printed by create')
        verify(args.path)
    else:
        path=create()
        if args.command=='create-and-verify':verify(path)