hf_backup.py

"""Резервное копирование данных SalGram в приватный репозиторий Hugging Face.

Зачем: на бесплатном HF Space диск эфемерный — при пересборке/перезапуске/
«засыпании» Space всё в контейнере пропадает. Этот модуль раз в
HF_BACKUP_INTERVAL секунд снимает консистентную копию БД и выгружает её
(вместе с ключом шифрования и медиафайлами) в ПРИВАТНЫЙ dataset-репозиторий,
а при старте контейнера восстанавливает оттуда. Так код остаётся в публичном
Space, а данные лежат только в приватном репо — их видишь и качаешь только ты.

Включается только если заданы переменные окружения:
    HF_BACKUP_REPO      — id приватного репо, напр. "username/salgram-data"
    HF_TOKEN            — токен с правом записи в этот репо
    HF_BACKUP_INTERVAL  — (необяз.) период в секундах, по умолчанию 1200 (20 мин)

Без них модуль ничего не делает — локальная разработка и Amvera не затронуты.
"""
import json
import os
import shutil
import sqlite3
import tempfile
import threading
import time
import traceback

REPO_ENV = "HF_BACKUP_REPO"
TOKEN_ENV = "HF_TOKEN"
INTERVAL_ENV = "HF_BACKUP_INTERVAL"
DEFAULT_INTERVAL = 1200  # 20 минут

REPO_TYPE = "dataset"
DB_NAME = "salgram.db"
EXTRA_FILES = ("secret.key", "vapid.json")  # без них восстановленную БД не расшифровать
MEDIA_DIRS = ("avatars", "photos", "files", "wallpapers", "gifs", "audio")
_INDEX = ".hf_uploaded.json"  # локальный список уже выгруженных файлов (чтобы не слать повторно)


def _config():
    """(repo, token, interval) если бэкап настроен, иначе None."""
    repo = os.environ.get(REPO_ENV)
    token = os.environ.get(TOKEN_ENV)
    if not repo or not token:
        return None
    try:
        interval = int(os.environ.get(INTERVAL_ENV, DEFAULT_INTERVAL))
    except ValueError:
        interval = DEFAULT_INTERVAL
    return repo, token, max(60, interval)


def _log(msg):
    print(f"[hf_backup] {msg}", flush=True)


def _snapshot_db(data_dir, dst_path):
    """Консистентная копия БД через online-backup API SQLite (безопасно при WAL)."""
    src = sqlite3.connect(os.path.join(data_dir, DB_NAME), timeout=30)
    try:
        dst = sqlite3.connect(dst_path)
        try:
            with dst:
                src.backup(dst)
        finally:
            dst.close()
    finally:
        src.close()


def restore_on_boot(data_dir):
    """Если локальной БД нет — скачать последнюю копию из приватного репо.

    Вызывать ДО _load_key() и init_db(): восстановленные ключ и БД нужны им.
    """
    cfg = _config()
    if not cfg:
        return
    repo, token, _ = cfg
    if os.path.exists(os.path.join(data_dir, DB_NAME)):
        _log("локальная БД уже на месте — восстановление пропущено")
        return
    try:
        from huggingface_hub import snapshot_download
        snap = snapshot_download(repo_id=repo, repo_type=REPO_TYPE, token=token)
    except Exception as e:  # репо ещё нет / пустой / нет сети — стартуем с чистой БД
        _log(f"копий пока нет или репо недоступен ({e!r}) — старт с чистой БД")
        return
    count = 0
    for root, _dirs, files in os.walk(snap):
        for name in files:
            if name.startswith("."):
                continue  # служебные .gitattributes / .cache
            rel = os.path.relpath(os.path.join(root, name), snap)
            target = os.path.join(data_dir, rel)
            os.makedirs(os.path.dirname(target) or ".", exist_ok=True)
            # snapshot_download кладёт симлинки на кеш — копируем именно содержимое.
            shutil.copyfile(os.path.realpath(os.path.join(root, name)), target)
            count += 1
    _log(f"восстановлено файлов из бэкапа: {count}")


def _remote_files(repo, token):
    """Список файлов, уже лежащих в репо (для первичного заполнения индекса)."""
    try:
        from huggingface_hub import HfApi
        return set(HfApi(token=token).list_repo_files(repo, repo_type=REPO_TYPE))
    except Exception:
        return set()


def backup_now(data_dir):
    """Один цикл бэкапа: снимок БД + новые медиа/ключи -> приватный репо."""
    cfg = _config()
    if not cfg:
        return
    repo, token, _ = cfg
    from huggingface_hub import CommitOperationAdd, HfApi
    api = HfApi(token=token)
    api.create_repo(repo, repo_type=REPO_TYPE, private=True, exist_ok=True)

    index_path = os.path.join(data_dir, _INDEX)
    if os.path.exists(index_path):
        with open(index_path) as f:
            uploaded = set(json.load(f))
    else:  # первый бэкап в этом контейнере — берём, что уже есть в репо
        uploaded = _remote_files(repo, token)

    with tempfile.TemporaryDirectory() as tmp:
        snap = os.path.join(tmp, DB_NAME)
        _snapshot_db(data_dir, snap)
        ops = [CommitOperationAdd(DB_NAME, snap)]  # БД меняется всегда — шлём каждый раз
        added = []

        for name in EXTRA_FILES:
            path = os.path.join(data_dir, name)
            if os.path.exists(path) and name not in uploaded:
                ops.append(CommitOperationAdd(name, path))
                added.append(name)

        for d in MEDIA_DIRS:
            ddir = os.path.join(data_dir, d)
            if not os.path.isdir(ddir):
                continue
            for name in sorted(os.listdir(ddir)):
                rel = f"{d}/{name}"
                path = os.path.join(ddir, name)
                if os.path.isfile(path) and rel not in uploaded:
                    ops.append(CommitOperationAdd(rel, path))
                    added.append(rel)

        api.create_commit(repo_id=repo, repo_type=REPO_TYPE, operations=ops,
                          commit_message="SalGram backup")

    uploaded.update(added)
    with open(index_path, "w") as f:
        json.dump(sorted(uploaded), f)
    return len(added)


def start(data_dir):
    """Запустить фоновый поток периодического бэкапа (no-op, если не настроено)."""
    cfg = _config()
    if not cfg:
        return
    _, _, interval = cfg

    def loop():
        while True:
            time.sleep(interval)
            try:
                n = backup_now(data_dir)
                _log(f"бэкап выполнен (+{n} новых файлов)")
            except Exception:
                _log("ошибка бэкапа:\n" + traceback.format_exc())

    threading.Thread(target=loop, name="hf-backup", daemon=True).start()
    _log(f"периодический бэкап включён: каждые {interval} с -> {os.environ[REPO_ENV]}")