""" 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", ]