baiyue-logistics-chat/logistics_chat/database_backup.py

"""Verified SQLite snapshots and offline restore, independent of Flask lifecycle."""

from contextlib import closing
from datetime import datetime, timedelta, timezone
import os
from pathlib import Path
import re
import sqlite3
import tempfile
from threading import Event, Thread
import time
import uuid

CHINA_TIME = timezone(timedelta(hours=8))


def read_connection(path):
    return sqlite3.connect(Path(path).resolve().as_uri() + '?mode=ro', uri=True, timeout=30)


def check_database(path):
    """Read-only full integrity check; never create a missing source."""
    with closing(read_connection(path)) as connection:
        rows = connection.execute('PRAGMA integrity_check').fetchall()
    if rows != [('ok',)]:
        raise sqlite3.DatabaseError('Database integrity check failed')
    return 'ok'


def snapshot_database(source, destination):
    """Include committed WAL data, verify, then publish without overwriting a file."""
    source, destination = Path(source).resolve(), Path(destination).resolve()
    if source == destination:
        raise ValueError('Backup must differ from source database')
    if destination.exists():
        raise FileExistsError(destination)
    destination.parent.mkdir(parents=True, exist_ok=True)
    descriptor, temporary = tempfile.mkstemp(prefix='.sqlite-backup-', suffix='.tmp', dir=destination.parent)
    os.close(descriptor)
    temporary = Path(temporary)
    deadline = time.monotonic() + 120

    def progress(status, remaining, total):
        if time.monotonic() > deadline:
            raise TimeoutError('SQLite backup exceeded 120 seconds')

    try:
        with closing(read_connection(source)) as original, closing(sqlite3.connect(temporary)) as target:
            original.backup(target, pages=256, progress=progress, sleep=0.05)
            # Publish a standalone file even when the source database uses WAL.
            target.execute('PRAGMA journal_mode=DELETE')
        check_database(temporary)
        with temporary.open('r+b') as handle:
            os.fsync(handle.fileno())
        # Hard-link publication is atomic and refuses an existing destination.
        os.link(temporary, destination)
        return destination
    finally:
        temporary.unlink(missing_ok=True)
        for suffix in ('-wal', '-shm', '-journal'):
            Path(str(temporary) + suffix).unlink(missing_ok=True)


def unique_backup_path(directory, prefix):
    timestamp = datetime.now(CHINA_TIME).strftime('%Y%m%d-%H%M%S-%f')
    return Path(directory) / f'{prefix}-{timestamp}-{uuid.uuid4().hex[:8]}.db'


def prune_migration_snapshots(directory, keep=5):
    """Call after publishing a verified snapshot, under the startup write lock."""
    if keep < 1:
        raise ValueError('At least one migration snapshot must be retained')
    snapshots = []
    for candidate in Path(directory).iterdir():
        if (re.fullmatch(r'pre-migration-\d{8}-\d{6}-\d{6}-[0-9a-f]{8}\.db', candidate.name)
                and not candidate.is_symlink() and candidate.is_file()):
            snapshots.append(candidate)
    # Generated names sort by their creation timestamps, independently of copied mtimes.
    for candidate in sorted(snapshots, key=lambda path: path.name, reverse=True)[keep:]:
        candidate.unlink(missing_ok=True)


def daily_backup(database, directory, now=None):
    now = now or datetime.now(CHINA_TIME)
    directory = Path(directory)
    destination = directory / ('history_' + now.strftime('%Y%m%d') + '.db')
    if destination.exists():
        check_database(destination)
    else:
        try:
            snapshot_database(database, destination)
        except FileExistsError:
            # Another worker published the same day's complete snapshot.
            check_database(destination)
    cutoff = now.date() - timedelta(days=6)
    for candidate in directory.iterdir():
        match = re.fullmatch(r'history_(\d{8})\.db', candidate.name)
        if not match or candidate.is_symlink() or not candidate.is_file():
            continue
        try:
            old = datetime.strptime(match[1], '%Y%m%d').date()
        except ValueError:
            continue
        if old < cutoff:
            candidate.unlink(missing_ok=True)
    return destination


def restore_database(database, backup, directory):
    """Offline only: keep current state, stage a verified copy, and replace atomically."""
    database, backup, directory = Path(database).resolve(), Path(backup).resolve(), Path(directory).resolve()
    if backup.parent != directory or backup == database or backup.suffix != '.db':
        raise ValueError('Restore input must be a .db file inside the configured backup directory')
    check_database(backup)
    # Avoid silently detaching a live WAL/journal from a replacement database.
    for suffix in ('-wal', '-shm', '-journal'):
        if Path(str(database) + suffix).exists():
            raise RuntimeError('Database has SQLite sidecar files; stop all users and checkpoint before restoring')
    safety = None
    if database.exists():
        safety = snapshot_database(database, unique_backup_path(directory, 'pre-restore'))
        # Reading a clean WAL database may create coordination files. Let SQLite
        # checkpoint/remove these under the declared offline maintenance window.
        with closing(sqlite3.connect(database, timeout=30)) as current:
            current.execute('PRAGMA wal_checkpoint(TRUNCATE)')
            current.execute('PRAGMA journal_mode=DELETE')
        for suffix in ('-wal', '-shm', '-journal'):
            if Path(str(database) + suffix).exists():
                raise RuntimeError('SQLite sidecars remain; database is not ready for offline restore')
    staged = database.with_name('.restore-' + uuid.uuid4().hex + '.db')
    try:
        snapshot_database(backup, staged)
        os.replace(staged, database)
    finally:
        staged.unlink(missing_ok=True)
    return safety


class DailyBackupWorker:
    """One stoppable task per app; daily files also coordinate independent processes."""

    def __init__(self, app):
        self.app = app
        self.stop_event = Event()
        # First tick rechecks startup's backup, including a startup crossing midnight.
        self.last_date = None
        self.thread = Thread(target=self.run, name='sqlite-daily-backup', daemon=True)

    def tick(self, now=None):
        now = now or datetime.now(CHINA_TIME)
        if now.date() == self.last_date:
            return False
        daily_backup(self.app.config['DATABASE'], self.app.config['BACKUP_DIR'], now)
        self.last_date = now.date()
        return True

    def run(self):
        while not self.stop_event.wait(60):
            try:
                self.tick()
            except Exception:
                self.app.logger.exception('Daily SQLite backup failed; will retry in 60 seconds')

    def start(self):
        self.thread.start()

    def stop(self):
        self.stop_event.set()
        self.thread.join(timeout=2)