utils/hf_storage.py

"""
Hugging Face Dataset 持久化后端（免费档可用）
=============================================

HF 免费 CPU Space 的文件系统是**临时的**：重启 / 休眠唤醒 / 重建后本地
``data/*.db`` 会被清空。为在不花钱（不挂载付费持久磁盘）的前提下保留用户账号与
管理员配置，本模块把本地 SQLite 文件与一个**私有 HF Dataset 仓库**做最小同步：

- 进程启动（首次建连）时：从 Dataset 拉取最新 DB 覆盖本地（``pull_db``）。
- 关键写操作后（注册新用户、保存管理员配置）：把本地 DB 推回 Dataset（``push_db``）。

设计原则：

- **可选 & 优雅降级**：未安装 ``huggingface_hub``、或未配置环境变量
  (``HF_TOKEN`` + ``PHARMAK_DATASET_REPO``) 时，``is_enabled()`` 返回 ``False``，
  所有同步调用静默 no-op，应用退回纯本地 SQLite（本地开发即此模式）。
- **不泄密**：日志只记录仓库 id 与文件名，绝不记录 token。
- **小规模适用**：Dataset 本质是 git 仓，适合低并发 / 低频写的小规模使用
  （竞赛 Demo、内部试用）。不适合高并发生产写入。

环境变量：

- ``PHARMAK_DATASET_REPO``：私有 Dataset 仓库 id，如 ``"your-name/pharma-k-data"``。
- ``HF_TOKEN``：具有该 Dataset **写权限**的访问令牌（放 HF Space Secrets）。
- ``PHARMAK_DB_FILENAME``（可选）：仓库内文件名，默认 ``"pharma_k.db"``。
"""

from __future__ import annotations

import logging
import os
import shutil
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

#: 环境变量名。
ENV_REPO = "PHARMAK_DATASET_REPO"
ENV_TOKEN = "HF_TOKEN"
ENV_DB_FILENAME = "PHARMAK_DB_FILENAME"

#: 仓库内默认文件名。
DEFAULT_DB_FILENAME = "pharma_k.db"


def _repo() -> str:
    return (os.environ.get(ENV_REPO) or "").strip()


def _token() -> str:
    return (os.environ.get(ENV_TOKEN) or "").strip()


def _db_filename() -> str:
    return (os.environ.get(ENV_DB_FILENAME) or "").strip() or DEFAULT_DB_FILENAME


def is_enabled() -> bool:
    """是否启用 Dataset 同步：依赖可导入且 token + repo 均已配置。"""
    if not _repo() or not _token():
        return False
    try:
        import huggingface_hub  # noqa: F401
    except Exception:
        logger.info("未安装 huggingface_hub，跳过 Dataset 同步（退回纯本地存储）。")
        return False
    return True


def ensure_repo() -> bool:
    """确保私有 Dataset 仓库存在（幂等）。失败不抛出，返回是否成功。"""
    if not is_enabled():
        return False
    try:
        from huggingface_hub import HfApi

        HfApi().create_repo(
            repo_id=_repo(),
            repo_type="dataset",
            private=True,
            exist_ok=True,
            token=_token(),
        )
        return True
    except Exception as exc:  # noqa: BLE001 - 同步失败不得影响应用主流程
        logger.warning("创建 / 校验 Dataset 仓库失败（仓库=%s）：%s", _repo(), exc)
        return False


def pull_db(local_path: Path) -> bool:
    """从 Dataset 拉取 DB 覆盖到 ``local_path``。

    - 同步未启用：no-op，返回 ``False``。
    - 仓库中尚无该文件（首次部署）：视为正常，返回 ``False`` 并让本地新建空库。
    - 成功覆盖：返回 ``True``。
    """
    if not is_enabled():
        return False
    try:
        from huggingface_hub import hf_hub_download
        from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError

        try:
            cached = hf_hub_download(
                repo_id=_repo(),
                filename=_db_filename(),
                repo_type="dataset",
                token=_token(),
            )
        except (EntryNotFoundError, RepositoryNotFoundError):
            logger.info("Dataset 中暂无 %s，将以空库初始化（首次部署属正常）。", _db_filename())
            return False

        local_path.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(cached, local_path)
        logger.info("已从 Dataset 拉取数据库（仓库=%s，文件=%s）。", _repo(), _db_filename())
        return True
    except Exception as exc:  # noqa: BLE001
        logger.warning("从 Dataset 拉取数据库失败（仓库=%s）：%s", _repo(), exc)
        return False


def push_db(local_path: Path) -> bool:
    """把本地 DB 推回 Dataset。失败不抛出，返回是否成功。"""
    if not is_enabled():
        return False
    if not Path(local_path).exists():
        return False
    try:
        from huggingface_hub import HfApi

        ensure_repo()
        HfApi().upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=_db_filename(),
            repo_id=_repo(),
            repo_type="dataset",
            token=_token(),
            commit_message="chore: sync pharma_k.db",
        )
        logger.info("已推送数据库到 Dataset（仓库=%s，文件=%s）。", _repo(), _db_filename())
        return True
    except Exception as exc:  # noqa: BLE001
        logger.warning("推送数据库到 Dataset 失败（仓库=%s）：%s", _repo(), exc)
        return False


__all__ = [
    "is_enabled",
    "ensure_repo",
    "pull_db",
    "push_db",
    "ENV_REPO",
    "ENV_TOKEN",
    "ENV_DB_FILENAME",
    "DEFAULT_DB_FILENAME",
]