diff --git a/astrbot/core/backup/__init__.py b/astrbot/core/backup/__init__.py index 8e33ef9705..565bcd0fd7 100644 --- a/astrbot/core/backup/__init__.py +++ b/astrbot/core/backup/__init__.py @@ -6,8 +6,13 @@ # 从 constants 模块导入共享常量 from .constants import ( BACKUP_MANIFEST_VERSION, + HARD_FAIL_COMPONENTS, KB_METADATA_MODELS, MAIN_DB_MODELS, + SPECIAL_COMPONENTS, + component_of_entry, + derive_component_states, + get_backup_components, get_backup_directories, ) @@ -21,6 +26,11 @@ "ImportPreCheckResult", "MAIN_DB_MODELS", "KB_METADATA_MODELS", + "SPECIAL_COMPONENTS", + "HARD_FAIL_COMPONENTS", + "get_backup_components", "get_backup_directories", + "component_of_entry", + "derive_component_states", "BACKUP_MANIFEST_VERSION", ] diff --git a/astrbot/core/backup/constants.py b/astrbot/core/backup/constants.py index 041fa407bc..ff15bd8adb 100644 --- a/astrbot/core/backup/constants.py +++ b/astrbot/core/backup/constants.py @@ -77,11 +77,136 @@ def get_backup_directories() -> dict[str, str]: "plugin_data": get_astrbot_plugin_data_path(), # 插件数据 "config": get_astrbot_config_path(), # 配置目录 "t2i_templates": get_astrbot_t2i_templates_path(), # T2I 模板 - "webchat": get_astrbot_webchat_path(), # WebChat 数据 + "webchat": get_astrbot_webchat_path(), # Legacy images within attachments. "temp": get_astrbot_temp_path(), # 临时文件 "skills": get_astrbot_skills_path(), # Skills } # 备份清单版本号 -BACKUP_MANIFEST_VERSION = "1.1" +# 1.2: additive — adds "components" and "component_checksums" fields and extends +# checksums to every archive entry. Older importers ignore unknown fields, and +# directory import requires >= 1.1, so 1.2 backups remain readable by old versions. +BACKUP_MANIFEST_VERSION = "1.2" + +# Non-directory backup components that can be selected individually. +SPECIAL_COMPONENTS: list[str] = [ + "database", + "knowledge_base", + "cmd_config", + "attachments", +] + +# Components whose verification failure aborts the whole import before any +# modification. All other components are soft-fail: corrupted entries are +# skipped with warnings and the rest of the import continues. +HARD_FAIL_COMPONENTS: frozenset[str] = frozenset( + {"database", "knowledge_base", "cmd_config"} +) + +# Fixed archive entries that must exist for a declared component to be usable. +# Components not listed here are validated by prefix scan (see _component_has_entries). +_REQUIRED_ENTRIES: dict[str, tuple[str, ...]] = { + "database": ("databases/main_db.json",), + "knowledge_base": ("databases/kb_metadata.json",), + "cmd_config": ("config/cmd_config.json",), +} + + +def get_backup_components() -> list[str]: + """Return selectable component ids, including legacy images under attachments. + + Returns: + Special components followed by independently selectable data directories. + """ + return SPECIAL_COMPONENTS + [ + name for name in get_backup_directories() if name != "webchat" + ] + + +def component_of_entry(name: str) -> str | None: + """Map an archive entry path to its owning backup component id. + + Args: + name: Entry path inside the backup ZIP (e.g. "directories/plugins/x.py"). + + Returns: + The component id, or None for entries not owned by any component + (e.g. "manifest.json" or unknown paths). + """ + if name == "databases/main_db.json": + return "database" + if name == "config/cmd_config.json": + return "cmd_config" + if name.startswith("databases/kb_") or name.startswith("files/kb_media/"): + return "knowledge_base" + if name.startswith(("files/attachments/", "directories/webchat/imgs/")): + return "attachments" + if name.startswith("directories/"): + parts = name.split("/") + if ( + len(parts) >= 3 + and parts[1] != "webchat" + and parts[1] in get_backup_directories() + ): + return parts[1] + return None + + +def _component_has_entries(component: str, names: set[str], manifest: dict) -> bool: + """Check whether a component has its required entries in the archive.""" + required = _REQUIRED_ENTRIES.get(component) + if required is not None: + return all(entry in names for entry in required) + if component == "attachments": + prefixes = ("files/attachments/",) + if "webchat" in manifest.get("directories", []): + prefixes += ("directories/webchat/imgs/",) + return any(n.startswith(prefixes) and not n.endswith("/") for n in names) + # Directory component: declared in the manifest and has at least one entry. + return component in manifest.get("directories", []) and any( + n.startswith(f"directories/{component}/") for n in names + ) + + +def derive_component_states( + manifest: dict, namelist: list[str] +) -> tuple[list[str], list[str]]: + """Derive (available, broken) components from the manifest and real entries. + + Derivation trusts actual ZIP entries, never self-reported manifest booleans. + + - Manifests with a "components" field (v1.2+): a declared component is + available when its required entries exist, otherwise it is broken + (declared but corrupted — callers must not silently drop it). + - Legacy manifests (no "components" field): the exporter always attempted + a full backup, so every known component is probed against the actual + entries; there is no "declared" concept and nothing is broken. + + Args: + manifest: Parsed manifest.json content. + namelist: Entry names of the backup ZIP. + + Returns: + A (available, broken) tuple of component id lists. + """ + names = set(namelist) + known = get_backup_components() + declared = manifest.get("components") + + if declared is None: + available = [ + comp for comp in known if _component_has_entries(comp, names, manifest) + ] + return available, [] + + available: list[str] = [] + broken: list[str] = [] + for comp in declared: + if comp not in known: + continue + if _component_has_entries(comp, names, manifest): + available.append(comp) + else: + broken.append(comp) + return available, broken diff --git a/astrbot/core/backup/exporter.py b/astrbot/core/backup/exporter.py index a922375998..d16fecd0cc 100644 --- a/astrbot/core/backup/exporter.py +++ b/astrbot/core/backup/exporter.py @@ -4,19 +4,26 @@ 导出格式为 JSON，这是数据库无关的方案，支持未来向 MySQL/PostgreSQL 迁移。 """ +import asyncio import hashlib +import io import json import os import zipfile +from collections import deque +from collections.abc import AsyncIterator from datetime import datetime, timezone +from itertools import chain from pathlib import Path from typing import TYPE_CHECKING, Any -from sqlalchemy import select +from sqlalchemy import JSON, LargeBinary, String, Text, cast, func, inspect, select +from sqlmodel.sql.sqltypes import AutoString from astrbot.core import logger from astrbot.core.config.default import VERSION from astrbot.core.db import BaseDatabase +from astrbot.core.db.po import Attachment from astrbot.core.utils.astrbot_path import ( get_astrbot_backups_path, get_astrbot_data_path, @@ -27,8 +34,17 @@ BACKUP_MANIFEST_VERSION, KB_METADATA_MODELS, MAIN_DB_MODELS, + component_of_entry, + get_backup_components, get_backup_directories, ) +from .importer import AstrBotImporter +from .resources import ( + MAX_JSON_RECORD_BYTES, + BackupTableStream, + backup_json_limit, + check_backup_json_field, +) if TYPE_CHECKING: from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager @@ -37,21 +53,10 @@ class AstrBotExporter: - """AstrBot 数据导出器 - - 导出内容： - - 主数据库所有表（data/data_v4.db） - - 知识库元数据（data/knowledge_base/kb.db） - - 每个知识库的向量文档数据 - - 配置文件（data/cmd_config.json） - - 附件文件 - - 知识库多媒体文件 - - 插件目录（data/plugins） - - 插件数据目录（data/plugin_data） - - 配置目录（data/config） - - T2I 模板目录（data/t2i_templates） - - WebChat 数据目录（data/webchat） - - 临时文件目录（data/temp） + """Export selected database, configuration, attachment, and extension data. + + Attachments include legacy WebChat images in data/webchat/imgs; upload + fragments in data/webchat/.chunks are excluded. """ def __init__( @@ -64,21 +69,36 @@ def __init__( self.kb_manager = kb_manager self.config_path = config_path self._checksums: dict[str, str] = {} + # Export report: components actually written and entries skipped + # (with reasons). Surfaced by the service layer in the task result. + self.exported_components: list[str] = [] + self.skipped_entries: list[dict[str, str]] = [] async def export_all( self, output_dir: str | None = None, progress_callback: Any | None = None, + components: list[str] | None = None, ) -> str: - """导出所有数据到 ZIP 文件 + """导出选定组件到 ZIP 文件 Args: output_dir: 输出目录 progress_callback: 进度回调函数，接收参数 (stage, current, total, message) + components: 要导出的组件 id 列表。None 表示全量（向后兼容）。 Returns: str: 生成的 ZIP 文件路径 + + Raises: + ValueError: components 不含任何有效组件 id。 + RuntimeError: ZIP 条目写入中途失败；半成品 ZIP 会被清理， + 不会留下无法通过完整性校验的备份产物。 """ + selected = self._normalize_components(components) + self._checksums.clear() + self.exported_components.clear() + self.skipped_entries.clear() if output_dir is None: output_dir = get_astrbot_backups_path() @@ -89,21 +109,25 @@ async def export_all( zip_filename = f"astrbot_backup_{timestamp}.zip" zip_path = os.path.join(output_dir, zip_filename) - logger.info(f"开始导出备份到 {zip_path}") + logger.info(f"Starting backup export to {zip_path}") try: + included: list[str] = [] with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: # 1. 导出主数据库 - if progress_callback: - await progress_callback("main_db", 0, 100, "正在导出主数据库...") - main_data = await self._export_main_database() - main_db_json = json.dumps( - main_data, ensure_ascii=False, indent=2, default=str - ) - zf.writestr("databases/main_db.json", main_db_json) - self._add_checksum("databases/main_db.json", main_db_json) - if progress_callback: - await progress_callback("main_db", 100, 100, "主数据库导出完成") + main_data: dict[str, Any] = {} + if "database" in selected: + if progress_callback: + await progress_callback( + "main_db", 0, 100, "正在导出主数据库..." + ) + main_data = await self._export_main_database() + await self._write_table_dump( + zf, "databases/main_db.json", main_data + ) + included.append("database") + if progress_callback: + await progress_callback("main_db", 100, 100, "主数据库导出完成") # 2. 导出知识库数据 kb_meta_data: dict[str, Any] = { @@ -111,17 +135,15 @@ async def export_all( "kb_documents": [], "kb_media": [], } - if self.kb_manager: + if "knowledge_base" in selected and self.kb_manager: if progress_callback: await progress_callback( "kb_metadata", 0, 100, "正在导出知识库元数据..." ) kb_meta_data = await self._export_kb_metadata() - kb_meta_json = json.dumps( - kb_meta_data, ensure_ascii=False, indent=2, default=str + await self._write_table_dump( + zf, "databases/kb_metadata.json", kb_meta_data ) - zf.writestr("databases/kb_metadata.json", kb_meta_json) - self._add_checksum("databases/kb_metadata.json", kb_meta_json) if progress_callback: await progress_callback( "kb_metadata", 100, 100, "知识库元数据导出完成" @@ -139,175 +161,586 @@ async def export_all( f"正在导出知识库 {kb_helper.kb.kb_name} 的文档数据...", ) doc_data = await self._export_kb_documents(kb_helper) - doc_json = json.dumps( - doc_data, ensure_ascii=False, indent=2, default=str - ) doc_path = f"databases/kb_{kb_id}/documents.json" - zf.writestr(doc_path, doc_json) - self._add_checksum(doc_path, doc_json) + await self._write_table_dump(zf, doc_path, doc_data) + del doc_data # 导出 FAISS 索引文件 - await self._export_faiss_index(zf, kb_helper, kb_id) + await self._run_io( + self._export_faiss_index, zf, kb_helper, kb_id + ) # 导出知识库多媒体文件 - await self._export_kb_media_files(zf, kb_helper, kb_id) + await self._run_io( + self._export_kb_media_files, zf, kb_helper, kb_id + ) if progress_callback: await progress_callback( "kb_documents", total_kbs, total_kbs, "知识库文档导出完成" ) + included.append("knowledge_base") # 3. 导出配置文件 - if progress_callback: - await progress_callback("config", 0, 100, "正在导出配置文件...") - if os.path.exists(self.config_path): - with open(self.config_path, encoding="utf-8") as f: - config_content = f.read() - zf.writestr("config/cmd_config.json", config_content) - self._add_checksum("config/cmd_config.json", config_content) - if progress_callback: - await progress_callback("config", 100, 100, "配置文件导出完成") + if "cmd_config" in selected: + if progress_callback: + await progress_callback("config", 0, 100, "正在导出配置文件...") + if os.path.exists(self.config_path): + await self._run_io( + self._write_entry, + zf, + "config/cmd_config.json", + Path(self.config_path), + ) + included.append("cmd_config") + else: + self._record_skip( + "config/cmd_config.json", "config file does not exist" + ) + if progress_callback: + await progress_callback("config", 100, 100, "配置文件导出完成") # 4. 导出附件文件 - if progress_callback: - await progress_callback("attachments", 0, 100, "正在导出附件...") - await self._export_attachments(zf, main_data.get("attachments", [])) - if progress_callback: - await progress_callback("attachments", 100, 100, "附件导出完成") + if "attachments" in selected: + if progress_callback: + await progress_callback( + "attachments", 0, 100, "正在导出附件..." + ) + attachment_rows = await self._export_attachment_records() + attachment_count = 0 + if isinstance(attachment_rows, list): + attachment_count = await self._run_io( + self._export_attachments, zf, attachment_rows + ) + else: + try: + async for batch in attachment_rows: + attachment_count += await self._run_io( + self._export_attachments, zf, batch + ) + finally: + await attachment_rows.aclose() + if attachment_count > 0: + included.append("attachments") + if progress_callback: + await progress_callback("attachments", 100, 100, "附件导出完成") # 5. 导出插件和其他目录 - if progress_callback: - await progress_callback( - "directories", 0, 100, "正在导出插件和数据目录..." + dir_names = [ + d + for d in get_backup_directories() + if d in selected or (d == "webchat" and "attachments" in selected) + ] + dir_stats: dict[str, dict[str, int]] = {} + if dir_names: + if progress_callback: + await progress_callback( + "directories", 0, 100, "正在导出插件和数据目录..." + ) + dir_stats = await self._run_io( + self._export_directories, zf, dir_names + ) + for directory, stats in dir_stats.items(): + component = ( + "attachments" if directory == "webchat" else directory + ) + if stats["files"] > 0 and component not in included: + included.append(component) + if progress_callback: + await progress_callback("directories", 100, 100, "目录导出完成") + + if "attachments" in selected and "attachments" not in included: + self._record_skip( + "files/attachments/", "no attachment files to export" ) - dir_stats = await self._export_directories(zf) - if progress_callback: - await progress_callback("directories", 100, 100, "目录导出完成") # 6. 生成 manifest if progress_callback: await progress_callback("manifest", 0, 100, "正在生成清单...") - manifest = self._generate_manifest(main_data, kb_meta_data, dir_stats) - manifest_json = json.dumps(manifest, ensure_ascii=False, indent=2) - zf.writestr("manifest.json", manifest_json) + manifest = await self._run_io( + self._generate_manifest, + main_data, + kb_meta_data, + dir_stats, + included, + ) + await self._run_io( + self._write_json_entry, zf, "manifest.json", manifest + ) + await self._run_io(zf.close) + # Apply the same ZIP metadata limits before publishing the file. + check = await self._run_io( + AstrBotImporter(self.main_db).pre_check, zip_path + ) + if check.error or not check.valid: + raise ValueError(check.error or "Backup archive validation failed") if progress_callback: await progress_callback("manifest", 100, 100, "清单生成完成") - logger.info(f"备份导出完成: {zip_path}") + self.exported_components = included + logger.info(f"Backup export completed: {zip_path}") return zip_path - except Exception as e: - logger.error(f"备份导出失败: {e}") + except (Exception, asyncio.CancelledError) as e: + logger.error(f"Backup export failed: {e}") # 清理失败的文件 if os.path.exists(zip_path): os.remove(zip_path) raise - async def _export_main_database(self) -> dict[str, list[dict]]: - """导出主数据库所有表""" - export_data: dict[str, list[dict]] = {} + async def _run_io(self, operation: Any, *args: Any) -> Any: + """Run blocking work and wait for it before closing a cancelled export. - async with self.main_db.get_db() as session: - for table_name, model_class in MAIN_DB_MODELS.items(): - try: - result = await session.execute(select(model_class)) - records = result.scalars().all() - export_data[table_name] = [ - self._model_to_dict(record) for record in records - ] - logger.debug( - f"导出表 {table_name}: {len(export_data[table_name])} 条记录" - ) - except Exception as e: - logger.warning(f"导出表 {table_name} 失败: {e}") - export_data[table_name] = [] + Args: + operation: Synchronous operation using the archive. + *args: Positional arguments for the operation. - return export_data + Returns: + The operation result. + """ + task = asyncio.create_task(asyncio.to_thread(operation, *args)) + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + # A worker cannot be cancelled; it must finish before ZipFile closes. + try: + await task + finally: + raise - async def _export_kb_metadata(self) -> dict[str, list[dict]]: - """导出知识库元数据库""" - if not self.kb_manager: - return {"knowledge_bases": [], "kb_documents": [], "kb_media": []} + async def _write_table_dump( + self, zf: zipfile.ZipFile, name: str, data: dict + ) -> None: + """Write table batches directly to ZIP and retain only manifest row counts. - export_data: dict[str, list[dict]] = {} + Args: + zf: Destination archive. + name: Database JSON entry name. + data: Table names mapped to row lists or async batch iterators. Values + are replaced with lightweight ranges holding the final row counts. - async with self.kb_manager.kb_db.get_db() as session: - for table_name, model_class in KB_METADATA_MODELS.items(): - try: - result = await session.execute(select(model_class)) - records = result.scalars().all() - export_data[table_name] = [ - self._model_to_dict(record) for record in records - ] - logger.debug( - f"导出知识库表 {table_name}: {len(export_data[table_name])} 条记录" - ) - except Exception as e: - logger.warning(f"导出知识库表 {table_name} 失败: {e}") - export_data[table_name] = [] + Raises: + ValueError: A record or the entire entry exceeds its resource budget. + """ + limit = backup_json_limit(name) + hasher = hashlib.sha256() + written = 0 + first_row = True + encoder = json.JSONEncoder( + ensure_ascii=False, separators=(",", ":"), default=str, allow_nan=False + ) + dest = io.BufferedWriter( + zf.open(name, "w", force_zip64=True), buffer_size=1 << 20 + ) + + def write(value: str | list[dict]) -> None: + """Encode one marker or batch in the worker thread. + + Args: + value: A raw JSON delimiter or a batch of table rows. + + Raises: + ValueError: The entry or an individual record is oversized. + """ + nonlocal written, first_row + for row in [None] if isinstance(value, str) else value: + if row is None: + parts = (value,) + else: + parts = chain(("" if first_row else ",",), encoder.iterencode(row)) + first_row = False + row_size = 0 + for part in parts: + chunk = part.encode("utf-8") + written += len(chunk) + row_size += len(chunk) + if written > limit: + raise ValueError(f"Backup JSON {name} exceeds the size limit") + if row_size > MAX_JSON_RECORD_BYTES: + raise ValueError("Backup JSON record exceeds the size limit") + hasher.update(chunk) + dest.write(chunk) - return export_data + try: + await self._run_io(write, "{") + for index, (table, rows) in enumerate(data.items()): + first_row = True + prefix = ("," if index else "") + json.dumps(table) + ":[" + await self._run_io(write, prefix) + count = 0 + if isinstance(rows, list): + await self._run_io(write, rows) + count = len(rows) + else: + try: + async for batch in rows: + await self._run_io(write, batch) + count += len(batch) + finally: + await rows.aclose() + data[table] = range(count) + await self._run_io(write, "]") + await self._run_io(write, "}") + finally: + await self._run_io(dest.close) + # Consume without retaining rows, using the importer's exact parser and + # resource limits before publishing an archive as successfully exported. + await self._run_io( + deque, + chain.from_iterable( + rows for _, rows in BackupTableStream(zf, name).items() + ), + 0, + ) + self._checksums[name] = f"sha256:{hasher.hexdigest()}" + + def _write_json_entry(self, zf: zipfile.ZipFile, name: str, data: Any) -> None: + """Encode, hash and compress JSON incrementally in the worker thread. - async def _export_kb_documents(self, kb_helper: Any) -> dict[str, Any]: - """导出知识库的文档块数据""" + Args: + zf: Destination archive. + name: JSON entry name. + data: JSON-serializable data. + + Raises: + ValueError: The JSON exceeds the importer's resource budget. + """ + limit = backup_json_limit(name) + hasher = hashlib.sha256() + size = 0 + buffer = bytearray() + encoder = json.JSONEncoder(ensure_ascii=False, indent=2, default=str) + with zf.open(name, "w", force_zip64=True) as dest: + for text in encoder.iterencode(data): + chunk = text.encode("utf-8") + size += len(chunk) + if size > limit: + raise ValueError(f"Backup JSON {name} exceeds the size limit") + buffer.extend(chunk) + if len(buffer) >= 1 << 20: + hasher.update(buffer) + dest.write(buffer) + buffer.clear() + if buffer: + hasher.update(buffer) + dest.write(buffer) + if name != "manifest.json": + self._checksums[name] = f"sha256:{hasher.hexdigest()}" + + def _normalize_components(self, components: list[str] | None) -> set[str]: + """Validate requested component ids against the known set. + + Args: + components: Requested ids, or None for a full backup. + + Returns: + The effective set of component ids to export. + + Raises: + ValueError: If a list was given but contains no valid ids. + """ + known = set(get_backup_components()) + if components is None: + return known + selected = set(components) & known + unknown = set(components) - known + if unknown: + logger.warning(f"Ignoring unknown backup components: {sorted(unknown)}") + if not selected: + raise ValueError("No valid backup components selected") + return selected + + def _record_skip(self, arcname: str, reason: str) -> None: + """Record a skipped entry with its reason (never silent).""" + self.skipped_entries.append({"entry": arcname, "reason": reason}) + logger.warning(f"Export entry skipped: {arcname} ({reason})") + + def _write_entry(self, zf: zipfile.ZipFile, arcname: str, src_path: Path) -> None: + """Stream a disk file into the ZIP while computing its sha256. + + Single pass, constant memory. The source is opened before the ZIP + entry is created, so a missing/unreadable source raises OSError + *before* any entry bytes exist and callers may treat it as a + skippable pre-write failure. Any error after the entry write started + leaves a partial entry in the archive, so it is re-raised as + RuntimeError and the whole export aborts (the half-written ZIP is + removed by export_all). + + Args: + zf: Open ZIP file being written. + arcname: Entry path inside the archive. + src_path: Source file on disk. + + Raises: + OSError: Source cannot be opened (pre-write, skippable). + RuntimeError: Entry failed mid-write (must abort the export). + """ + fsrc = open(src_path, "rb") try: - from astrbot.core.db.vec_db.faiss_impl.vec_db import FaissVecDB + hasher = hashlib.sha256() + try: + size = 0 + with zf.open(arcname, mode="w", force_zip64=True) as fdst: + while chunk := fsrc.read(1 << 20): + if arcname == "config/cmd_config.json": + size += len(chunk) + if size > backup_json_limit(arcname): + raise ValueError("Backup JSON exceeds the size limit") + hasher.update(chunk) + fdst.write(chunk) + except Exception as e: + raise RuntimeError(f"Entry {arcname} failed mid-write: {e}") from e + finally: + fsrc.close() + self._checksums[arcname] = f"sha256:{hasher.hexdigest()}" + + async def _export_attachment_records(self) -> AsyncIterator[list[dict]]: + """Return attachment batches without retaining the whole table. + + Returns: + An async iterator of attachment record batches. + """ + return self._export_records( + self.main_db.get_db, Attachment, self._model_to_dict + ) - vec_db: FaissVecDB = kb_helper.vec_db - if not vec_db or not vec_db.document_storage: - return {"documents": []} + async def _export_records( + self, get_session: Any, model_class: type, convert: Any + ) -> AsyncIterator[list[dict]]: + """Fetch bounded raw rows and decode JSON in the export worker. + + Args: + get_session: Factory returning an async database session. + model_class: Model to export. + convert: Function converting a transient model to its dump schema. + JSON fields are placeholders here and filled after validation. - # 获取所有文档 - docs = await vec_db.document_storage.get_documents( - metadata_filters={}, - offset=0, - limit=None, # 获取全部 + Yields: + A batch of rows serialized to dictionaries. + """ + columns = [] + json_fields = set() + text_fields = set() + for attr in inspect(model_class).column_attrs: + column = attr.columns[0] + expression = column + if isinstance(column.type, JSON): + json_fields.add(attr.key) + expression = cast(column, Text) + if isinstance(column.type, (JSON, String, AutoString)): + text_fields.add(attr.key) + # Fetch one extra byte to detect oversize values without + # transferring an arbitrarily large field into Python first. + # SQLite TEXT substr stops at NUL; BLOB substr preserves it. + # Include SQLModel's AutoString decorator explicitly. + # Stored JSON may contain ASCII escapes and separator spaces; + # allow their expansion relative to our compact UTF-8 dump. + raw_limit = MAX_JSON_RECORD_BYTES * ( + 6 if attr.key in json_fields else 1 + ) + expression = func.substr( + cast(expression, LargeBinary), 1, raw_limit + 1 + ) + columns.append(expression.label(attr.key)) + + def convert_batch(batch: list[dict]) -> list[dict]: + """Decode raw fields and convert transient models in the worker. + + Args: + batch: Bounded raw database rows. + + Returns: + Records in the existing backup schema. + """ + converted = [] + for values in batch: + # Convert only scalar fields first. This preserves the existing + # schema without making model_dump copy decoded JSON containers. + scalar_values = { + key: None + if key in json_fields + else ( + value.decode("utf-8") + if key in text_fields and value is not None + else value + ) + for key, value in values.items() + } + row = convert(model_class(**scalar_values)) + size = len( + json.dumps( + row, + ensure_ascii=False, + separators=(",", ":"), + default=str, + allow_nan=False, + ).encode("utf-8") + ) + size -= 4 * sum(values[key] is not None for key in json_fields) + if size > MAX_JSON_RECORD_BYTES: + raise ValueError("Backup JSON record exceeds the size limit") + # Validate every JSON field before materializing any of them. + # The placeholder null already accounts for four bytes per field. + for key in json_fields: + if values[key] is not None: + size += check_backup_json_field( + values[key], MAX_JSON_RECORD_BYTES - size + ) + for key in json_fields: + if values[key] is not None: + row[key] = json.loads(values[key]) + converted.append(row) + return converted + + async with get_session() as session: + result = await session.stream( + select(*columns).execution_options(yield_per=1) + ) + try: + batch = [] + batch_size = 0 + async for row in result.mappings(): + values = dict(row) + if any( + isinstance(value, (str, bytes)) + and len(value) + > MAX_JSON_RECORD_BYTES * (6 if key in json_fields else 1) + for key, value in values.items() + ): + raise ValueError("Backup JSON record exceeds the size limit") + # Four bytes per character conservatively bounds Unicode + # storage; count small scalar fields too. Never prefetch 500 + # potentially multi-megabyte ORM objects on the event loop. + size = sum( + len(value) * 4 if isinstance(value, (str, bytes)) else 32 + for value in values.values() + ) + if batch and ( + len(batch) >= 500 or batch_size + size > MAX_JSON_RECORD_BYTES + ): + yield await self._run_io(convert_batch, batch) + batch = [] + batch_size = 0 + batch.append(values) + batch_size += size + if batch: + yield await self._run_io(convert_batch, batch) + finally: + await result.close() + + def _component_checksums(self, included: list[str]) -> dict[str, str]: + """Derive a deterministic per-component digest from entry checksums. + + The digest is sha256 over the sorted "path:entry_hash" lines of every + entry owned by the component, giving import a component-level + verification anchor without re-reading raw bytes. + """ + result: dict[str, str] = {} + for comp in included: + lines = sorted( + f"{path}:{checksum}" + for path, checksum in self._checksums.items() + if component_of_entry(path) == comp ) + digest = hashlib.sha256("\n".join(lines).encode("utf-8")).hexdigest() + result[comp] = f"sha256:{digest}" + return result + + async def _export_main_database(self) -> dict[str, Any]: + """Prepare replay-free streams for the main database tables. - return {"documents": docs} - except Exception as e: - logger.warning(f"导出知识库文档失败: {e}") + Returns: + Table names mapped to async row-batch iterators. + """ + return { + table: self._export_records(self.main_db.get_db, model, self._model_to_dict) + for table, model in MAIN_DB_MODELS.items() + } + + async def _export_kb_metadata(self) -> dict[str, Any]: + """Prepare streams for the knowledge base metadata tables. + + Returns: + Table names mapped to async row-batch iterators. + """ + if not self.kb_manager: + return {} + return { + table: self._export_records( + self.kb_manager.kb_db.get_db, model, self._model_to_dict + ) + for table, model in KB_METADATA_MODELS.items() + } + + async def _export_kb_documents(self, kb_helper: Any) -> dict[str, Any]: + """Prepare the knowledge base document stream using its existing database. + + Args: + kb_helper: Knowledge base being exported. + + Returns: + Document table mapped to an async row-batch iterator. + """ + from astrbot.core.db.vec_db.faiss_impl.document_storage import Document + + vec_db = kb_helper.vec_db + if not vec_db or not vec_db.document_storage: return {"documents": []} + # DocumentStorage exposes get_session rather than get_db. Its own + # conversion also preserves the exported field names used by import. + storage = vec_db.document_storage + return { + "documents": self._export_records( + storage.get_session, Document, storage._document_to_dict + ) + } - async def _export_faiss_index( + def _export_faiss_index( self, zf: zipfile.ZipFile, kb_helper: Any, kb_id: str, ) -> None: """导出 FAISS 索引文件""" + index_path = kb_helper.kb_dir / "index.faiss" + if not index_path.exists(): + return + archive_path = f"databases/kb_{kb_id}/index.faiss" try: - index_path = kb_helper.kb_dir / "index.faiss" - if index_path.exists(): - archive_path = f"databases/kb_{kb_id}/index.faiss" - zf.write(str(index_path), archive_path) - logger.debug(f"导出 FAISS 索引: {archive_path}") - except Exception as e: - logger.warning(f"导出 FAISS 索引失败: {e}") - - async def _export_kb_media_files( + self._write_entry(zf, archive_path, index_path) + logger.debug(f"Exported FAISS index: {archive_path}") + except OSError as e: + # Source unreadable before the entry write started: skippable. + self._record_skip(archive_path, str(e)) + + def _export_kb_media_files( self, zf: zipfile.ZipFile, kb_helper: Any, kb_id: str ) -> None: """导出知识库的多媒体文件""" - try: - media_dir = kb_helper.kb_medias_dir - if not media_dir.exists(): - return - - for root, _, files in os.walk(media_dir): - for file in files: - file_path = Path(root) / file - # 计算相对路径 - rel_path = file_path.relative_to(kb_helper.kb_dir) - archive_path = f"files/kb_media/{kb_id}/{rel_path}" - zf.write(str(file_path), archive_path) - except Exception as e: - logger.warning(f"导出知识库媒体文件失败: {e}") - - async def _export_directories( - self, zf: zipfile.ZipFile + media_dir = kb_helper.kb_medias_dir + if not media_dir.exists(): + return + + for root, _, files in os.walk(media_dir): + for file in files: + file_path = Path(root) / file + # 计算相对路径 + rel_path = file_path.relative_to(kb_helper.kb_dir).as_posix() + archive_path = f"files/kb_media/{kb_id}/{rel_path}" + try: + self._write_entry(zf, archive_path, file_path) + except OSError as e: + # Source unreadable before the entry write started. + self._record_skip(archive_path, str(e)) + + def _export_directories( + self, zf: zipfile.ZipFile, dir_names: list[str] ) -> dict[str, dict[str, int]]: - """导出插件和其他数据目录 + """导出选定的插件和其他数据目录 + + Args: + zf: 打开的 ZIP 文件对象 + dir_names: 要导出的目录键列表（get_backup_directories 的子集） Returns: dict: 每个目录的统计信息 {dir_name: {"files": count, "size": bytes}} @@ -315,17 +748,18 @@ async def _export_directories( stats: dict[str, dict[str, int]] = {} backup_directories = get_backup_directories() - for dir_name, dir_path in backup_directories.items(): - full_path = Path(dir_path) - if not full_path.exists(): - logger.debug(f"目录不存在，跳过: {full_path}") + for dir_name in dir_names: + full_path = Path(backup_directories[dir_name]) + scan_path = full_path / "imgs" if dir_name == "webchat" else full_path + if not scan_path.exists(): + logger.debug(f"Skipping missing directory: {scan_path}") continue file_count = 0 total_size = 0 try: - for root, dirs, files in os.walk(full_path): + for root, dirs, files in os.walk(scan_path): # 跳过 __pycache__ 目录 dirs[:] = [d for d in dirs if d != "__pycache__"] @@ -335,41 +769,54 @@ async def _export_directories( continue file_path = Path(root) / file + # 计算相对路径 + rel_path = file_path.relative_to(full_path).as_posix() + archive_path = f"directories/{dir_name}/{rel_path}" try: - # 计算相对路径 - rel_path = file_path.relative_to(full_path) - archive_path = f"directories/{dir_name}/{rel_path}" - zf.write(str(file_path), archive_path) + self._write_entry(zf, archive_path, file_path) file_count += 1 total_size += file_path.stat().st_size - except Exception as e: - logger.warning(f"导出文件 {file_path} 失败: {e}") + except OSError as e: + # Source vanished or unreadable before the entry + # write started: skip and record, never silent. + self._record_skip(archive_path, str(e)) stats[dir_name] = {"files": file_count, "size": total_size} logger.debug( - f"导出目录 {dir_name}: {file_count} 个文件, {total_size} 字节" + f"Exported directory {dir_name}: {file_count} files, {total_size} bytes" ) + except RuntimeError: + # Mid-write entry failure: the archive now holds a partial + # entry — the whole export must fail. + raise except Exception as e: - logger.warning(f"导出目录 {dir_path} 失败: {e}") + logger.warning(f"Failed to export directory {full_path}: {e}") stats[dir_name] = {"files": 0, "size": 0} return stats - async def _export_attachments( - self, zf: zipfile.ZipFile, attachments: list[dict] - ) -> None: - """导出附件文件""" + def _export_attachments(self, zf: zipfile.ZipFile, attachments: list[dict]) -> int: + """导出附件文件 + + Returns: + int: 实际写入的附件文件数量 + """ + count = 0 for attachment in attachments: + file_path = attachment.get("path", "") + if not file_path or not os.path.exists(file_path): + continue + # 使用 attachment_id 作为文件名 + attachment_id = attachment.get("attachment_id", "") + ext = os.path.splitext(file_path)[1] + archive_path = f"files/attachments/{attachment_id}{ext}" try: - file_path = attachment.get("path", "") - if file_path and os.path.exists(file_path): - # 使用 attachment_id 作为文件名 - attachment_id = attachment.get("attachment_id", "") - ext = os.path.splitext(file_path)[1] - archive_path = f"files/attachments/{attachment_id}{ext}" - zf.write(file_path, archive_path) - except Exception as e: - logger.warning(f"导出附件失败: {e}") + self._write_entry(zf, archive_path, Path(file_path)) + count += 1 + except OSError as e: + # Source unreadable before the entry write started. + self._record_skip(archive_path, str(e)) + return count def _model_to_dict(self, record: Any) -> dict: """将 SQLModel 实例转换为字典 @@ -408,40 +855,44 @@ def _add_checksum(self, path: str, content: str | bytes) -> None: def _generate_manifest( self, - main_data: dict[str, list[dict]], - kb_meta_data: dict[str, list[dict]], + main_data: dict[str, Any], + kb_meta_data: dict[str, Any], dir_stats: dict[str, dict[str, int]] | None = None, + included_components: list[str] | None = None, ) -> dict: - """生成备份清单""" + """生成备份清单 + + Args: + main_data: 主数据库导出数据 + kb_meta_data: 知识库元数据 + dir_stats: 目录导出统计 + included_components: 实际写入 ZIP 的组件 id 列表 + """ if dir_stats is None: dir_stats = {} - # 收集知识库 ID - kb_document_tables = {} - if self.kb_manager: - for kb_id in self.kb_manager.kb_insts.keys(): - kb_document_tables[kb_id] = "documents" - - # 收集附件文件列表 - attachment_files = [] - for attachment in main_data.get("attachments", []): - attachment_id = attachment.get("attachment_id", "") - path = attachment.get("path", "") - if attachment_id and path: - ext = os.path.splitext(path)[1] - attachment_files.append(f"{attachment_id}{ext}") + if included_components is None: + included_components = [] + kb_document_tables = { + path.removeprefix("databases/kb_").removesuffix( + "/documents.json" + ): "documents" + for path in self._checksums + if path.startswith("databases/kb_") and path.endswith("/documents.json") + } + + attachment_files = [ + path.removeprefix("files/attachments/") + for path in self._checksums + if path.startswith("files/attachments/") + ] - # 收集知识库媒体文件 kb_media_files: dict[str, list[str]] = {} - if self.kb_manager: - for kb_id, kb_helper in self.kb_manager.kb_insts.items(): - media_files: list[str] = [] - media_dir = kb_helper.kb_medias_dir - if media_dir.exists(): - for root, _, files in os.walk(media_dir): - for file in files: - media_files.append(file) - if media_files: - kb_media_files[kb_id] = media_files + for path in self._checksums: + if path.startswith("files/kb_media/"): + kb_id, _, relative_path = path.removeprefix( + "files/kb_media/" + ).partition("/") + kb_media_files.setdefault(kb_id, []).append(Path(relative_path).name) manifest = { "version": BACKUP_MANIFEST_VERSION, @@ -452,6 +903,8 @@ def _generate_manifest( "main_db": "v4", "kb_db": "v1", }, + "components": sorted(included_components), + "component_checksums": self._component_checksums(included_components), "tables": { "main_db": list(main_data.keys()), "kb_metadata": list(kb_meta_data.keys()), diff --git a/astrbot/core/backup/importer.py b/astrbot/core/backup/importer.py index 7d30d27c39..dec38fac22 100644 --- a/astrbot/core/backup/importer.py +++ b/astrbot/core/backup/importer.py @@ -7,10 +7,16 @@ - 版本匹配时也需要用户确认 """ +import hashlib import json import os import shutil +import sqlite3 +import tempfile import zipfile +import zlib +from collections.abc import Iterable, Iterator +from contextlib import closing from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path @@ -30,10 +36,22 @@ # 从共享常量模块导入 from .constants import ( + HARD_FAIL_COMPONENTS, KB_METADATA_MODELS, MAIN_DB_MODELS, + component_of_entry, + derive_component_states, + get_backup_components, get_backup_directories, ) +from .resources import ( + MAX_JSON_RECORD_BYTES, + BackupTableStream, + backup_json_limit, + backup_row_batches, + open_backup, + read_backup_json, +) if TYPE_CHECKING: from astrbot.core.knowledge_base.kb_mgr import KnowledgeBaseManager @@ -119,14 +137,14 @@ def warn_invalid_count(self, value: Any, key_for_log: tuple[Any, ...]) -> None: if self.limit > 0: if self._count < self.limit: logger.warning( - "platform_stats count 非法，已按 0 处理: value=%r, key=%s", + "Invalid platform_stats count; using 0: value=%r, key=%s", value, key_for_log, ) self._count += 1 if self._count == self.limit and not self._suppression_logged: logger.warning( - "platform_stats 非法 count 告警已达到上限 (%d)，后续将抑制", + "Invalid platform_stats count warning limit reached (%d); suppressing further warnings", self.limit, ) self._suppression_logged = True @@ -135,7 +153,7 @@ def warn_invalid_count(self, value: Any, key_for_log: tuple[Any, ...]) -> None: if not self._suppression_logged: # limit <= 0: emit only one suppression warning. logger.warning( - "platform_stats 非法 count 告警已达到上限 (%d)，后续将抑制", + "Invalid platform_stats count warning limit reached (%d); suppressing further warnings", self.limit, ) self._suppression_logged = True @@ -169,6 +187,10 @@ class ImportPreCheckResult: error: str = "" # 备份包含的内容摘要 backup_summary: dict = field(default_factory=dict) + # 可恢复的组件（声明且条目完整），按 ZIP 实际条目推导 + available_components: list[str] = field(default_factory=list) + # 已声明但条目缺失的损坏组件（不可恢复，默认恢复遇之中止） + broken_components: list[str] = field(default_factory=list) def to_dict(self) -> dict: return { @@ -182,6 +204,8 @@ def to_dict(self) -> dict: "warnings": self.warnings, "error": self.error, "backup_summary": self.backup_summary, + "available_components": self.available_components, + "broken_components": self.broken_components, } @@ -221,20 +245,10 @@ class DatabaseClearError(RuntimeError): class AstrBotImporter: - """AstrBot 数据导入器 - - 导入备份文件中的所有数据，包括： - - 主数据库所有表 - - 知识库元数据和文档 - - 配置文件 - - 附件文件 - - 知识库多媒体文件 - - 插件目录（data/plugins） - - 插件数据目录（data/plugin_data） - - 配置目录（data/config） - - T2I 模板目录（data/t2i_templates） - - WebChat 数据目录（data/webchat） - - 临时文件目录（data/temp） + """Restore selected database, configuration, attachment, and extension data. + + Attachments include legacy WebChat images in data/webchat/imgs; upload + fragments in data/webchat/.chunks are excluded. """ def __init__( @@ -265,20 +279,19 @@ def pre_check(self, zip_path: str) -> ImportPreCheckResult: result.current_version = VERSION if not os.path.exists(zip_path): - result.error = f"备份文件不存在: {zip_path}" + result.error = f"Backup file does not exist: {zip_path}" return result try: - with zipfile.ZipFile(zip_path, "r") as zf: + with open_backup(zip_path) as zf: # 读取 manifest try: - manifest_data = zf.read("manifest.json") - manifest = json.loads(manifest_data) + manifest = read_backup_json(zf, "manifest.json") except KeyError: - result.error = "备份文件缺少 manifest.json，不是有效的 AstrBot 备份" + result.error = "Invalid AstrBot backup: manifest.json is missing" return result except json.JSONDecodeError as e: - result.error = f"manifest.json 格式错误: {e}" + result.error = f"Invalid manifest.json: {e}" return result # 提取基本信息 @@ -286,14 +299,26 @@ def pre_check(self, zip_path: str) -> ImportPreCheckResult: result.backup_time = manifest.get("exported_at", "未知") result.valid = True - # 构建备份摘要 + # 构建备份摘要：只信 ZIP 实际条目，不信 manifest 自报字段 + namelist = zf.namelist() result.backup_summary = { "tables": list(manifest.get("tables", {}).keys()), - "has_knowledge_bases": manifest.get("has_knowledge_bases", False), - "has_config": manifest.get("has_config", False), + "has_knowledge_bases": "databases/kb_metadata.json" in namelist, + "has_config": "config/cmd_config.json" in namelist, "directories": manifest.get("directories", []), } + # 三态推导可用/损坏组件 + available, broken = derive_component_states(manifest, namelist) + result.available_components = available + result.broken_components = broken + if any(n.startswith("directories/webchat/.chunks/") for n in namelist): + result.warnings.append( + "WebChat upload fragments cannot restore upload sessions and will be ignored." + ) + if not available and not broken: + result.warnings.append("This backup contains no restorable data.") + # 检查版本兼容性 version_check = self._check_version_compatibility(result.backup_version) result.version_status = version_check["status"] @@ -306,10 +331,10 @@ def pre_check(self, zip_path: str) -> ImportPreCheckResult: return result except zipfile.BadZipFile: - result.error = "无效的 ZIP 文件" + result.error = "Invalid ZIP file" return result except Exception as e: - result.error = f"检查备份文件失败: {e}" + result.error = f"Failed to inspect backup file: {e}" return result def _check_version_compatibility(self, backup_version: str) -> dict: @@ -326,7 +351,7 @@ def _check_version_compatibility(self, backup_version: str) -> dict: return { "status": "major_diff", "can_import": False, - "message": "备份文件缺少版本信息", + "message": "Backup is missing version information", } # 提取主版本（前两位）进行比较 @@ -339,8 +364,8 @@ def _check_version_compatibility(self, backup_version: str) -> dict: "status": "major_diff", "can_import": False, "message": ( - f"主版本不兼容: 备份版本 {backup_version}, 当前版本 {VERSION}。" - f"跨主版本导入可能导致数据损坏，请使用相同主版本的 AstrBot。" + f"Incompatible major version: backup={backup_version}, current={VERSION}. " + f"Importing across major versions may corrupt data; use the same AstrBot major version." ), } @@ -366,13 +391,19 @@ async def import_all( zip_path: str, mode: str = "replace", # "replace" 清空后导入 progress_callback: Any | None = None, + components: list[str] | None = None, ) -> ImportResult: - """从 ZIP 文件导入所有数据 + """从 ZIP 文件导入选定组件的数据 + + 两阶段执行：阶段一全量预检（条目 hash、聚合 hash、JSON 结构干跑， + 零修改），阶段二落盘（临时文件原子替换、主库单事务）。 Args: zip_path: ZIP 备份文件路径 mode: 导入模式，目前仅支持 "replace"（清空后导入） progress_callback: 进度回调函数，接收参数 (stage, current, total, message) + components: 要恢复的组件 id 列表。None 恢复全部可用组件 + （遇 broken 组件中止）；空列表直接拒绝。 Returns: ImportResult: 导入结果 @@ -380,25 +411,24 @@ async def import_all( result = ImportResult() if not os.path.exists(zip_path): - result.add_error(f"备份文件不存在: {zip_path}") + result.add_error(f"Backup file does not exist: {zip_path}") return result - logger.info(f"开始从 {zip_path} 导入备份") + logger.info(f"Starting backup import from {zip_path}") try: - with zipfile.ZipFile(zip_path, "r") as zf: + with open_backup(zip_path) as zf: # 1. 读取并验证 manifest if progress_callback: await progress_callback("validate", 0, 100, "正在验证备份文件...") try: - manifest_data = zf.read("manifest.json") - manifest = json.loads(manifest_data) + manifest = read_backup_json(zf, "manifest.json") except KeyError: - result.add_error("备份文件缺少 manifest.json") + result.add_error("Backup is missing manifest.json") return result except json.JSONDecodeError as e: - result.add_error(f"manifest.json 格式错误: {e}") + result.add_error(f"Invalid manifest.json: {e}") return result # 版本校验 @@ -408,106 +438,658 @@ async def import_all( result.add_error(str(e)) return result + namelist = zf.namelist() + available, broken = derive_component_states(manifest, namelist) + + selected = self._resolve_selection( + components, available, broken, result + ) + if selected is None: + return result + + # Validate every selected component before making any changes. + precheck = await self._run_pre_verify( + zf, manifest, namelist, selected, result, progress_callback + ) + if precheck is None: + return result + if progress_callback: await progress_callback("validate", 100, 100, "验证完成") - # 2. 导入主数据库 - if progress_callback: - await progress_callback("main_db", 0, 100, "正在导入主数据库...") + skip_components = precheck["skip_components"] + bad_entries = precheck["bad_entries"] + + attachment_rows = [] + attachment_paths = ( + { + Path(name).stem: None + for name in namelist + if name.startswith("files/attachments/") + and not name.endswith("/") + } + if "attachments" in selected + and "attachments" not in skip_components + else None + ) - try: - main_data_content = zf.read("databases/main_db.json") - main_data = json.loads(main_data_content) + # ========== 阶段二：落盘 ========== - if mode == "replace": - await self._clear_main_db() + # 2. 导入主数据库 + if "database" in selected: + if progress_callback: + await progress_callback( + "main_db", 0, 100, "正在导入主数据库..." + ) - imported = await self._import_main_database(main_data) - result.imported_tables.update(imported) - except DatabaseClearError as e: - result.add_error(f"清空主数据库失败: {e}") - return result - except Exception as e: - result.add_error(f"导入主数据库失败: {e}") - return result + try: + main_data = precheck["json"].pop("main_db") + imported = await self._import_main_database( + main_data, + clear=(mode == "replace"), + attachment_paths=attachment_paths, + ) + attachment_rows = ( + {"attachment_id": key, "path": value} + for key, value in (attachment_paths or {}).items() + if value is not None + ) + del main_data + result.imported_tables.update(imported) + except DatabaseClearError as e: + result.add_error(f"Failed to clear main database: {e}") + return result + except Exception as e: + result.add_error( + f"Main database import failed (transaction rolled back): {e}" + ) + return result - if progress_callback: - await progress_callback("main_db", 100, 100, "主数据库导入完成") + if progress_callback: + await progress_callback("main_db", 100, 100, "主数据库导入完成") # 3. 导入知识库 - if self.kb_manager and "databases/kb_metadata.json" in zf.namelist(): + if "knowledge_base" in selected and self.kb_manager: if progress_callback: await progress_callback("kb", 0, 100, "正在导入知识库...") try: - kb_meta_content = zf.read("databases/kb_metadata.json") - kb_meta_data = json.loads(kb_meta_content) - - if mode == "replace": - await self._clear_kb_data() - - await self._import_knowledge_bases(zf, kb_meta_data, result) + await self._import_knowledge_bases( + zf, + precheck["json"].pop("kb_metadata"), + result, + clear=(mode == "replace"), + json_ctx=precheck["json"], + bad_entries=bad_entries.get("knowledge_base", []), + ) except Exception as e: - result.add_warning(f"导入知识库失败: {e}") + result.add_warning(f"Knowledge base import failed: {e}") if progress_callback: await progress_callback("kb", 100, 100, "知识库导入完成") # 4. 导入配置文件 - if progress_callback: - await progress_callback("config", 0, 100, "正在导入配置文件...") + if "cmd_config" in selected: + if progress_callback: + await progress_callback("config", 0, 100, "正在导入配置文件...") - if "config/cmd_config.json" in zf.namelist(): try: - config_content = zf.read("config/cmd_config.json") # 备份现有配置 if os.path.exists(self.config_path): backup_path = f"{self.config_path}.bak" shutil.copy2(self.config_path, backup_path) - with open(self.config_path, "wb") as f: - f.write(config_content) + # 阶段一已验证 hash 与结构，经临时文件原子替换 + self._write_entry_safe( + zf, "config/cmd_config.json", Path(self.config_path) + ) result.imported_files["config"] = 1 except Exception as e: - result.add_warning(f"导入配置文件失败: {e}") + result.add_warning(f"Configuration import failed: {e}") - if progress_callback: - await progress_callback("config", 100, 100, "配置文件导入完成") + if progress_callback: + await progress_callback("config", 100, 100, "配置文件导入完成") # 5. 导入附件文件 - if progress_callback: - await progress_callback("attachments", 0, 100, "正在导入附件...") - - attachment_count = await self._import_attachments( - zf, main_data.get("attachments", []) - ) - result.imported_files["attachments"] = attachment_count + if "attachments" in selected and "attachments" not in skip_components: + if progress_callback: + await progress_callback( + "attachments", 0, 100, "正在导入附件..." + ) + + # 附件原始路径来自主库 attachments 表；未恢复主库时仅作 + # 路径提示读取（路径仍强制校验在附件目录内）。 + if ( + "database" not in selected + and attachment_paths + and "databases/main_db.json" in namelist + ): + try: + hint_data = read_backup_json(zf, "databases/main_db.json") + attachment_rows = hint_data.get("attachments", []) + del hint_data + except Exception as exc: + result.add_warning(f"Attachment path hints skipped: {exc}") + attachment_rows = [] + + attachment_count = await self._import_attachments( + zf, attachment_rows, bad_entries.get("attachments", []) + ) + result.imported_files["attachments"] = attachment_count - if progress_callback: - await progress_callback("attachments", 100, 100, "附件导入完成") + if progress_callback: + await progress_callback("attachments", 100, 100, "附件导入完成") # 6. 导入插件和其他目录 - if progress_callback: - await progress_callback( - "directories", 0, 100, "正在导入插件和数据目录..." + selected_dirs = [ + d + for d in get_backup_directories() + if ("attachments" if d == "webchat" else d) in selected + and ("attachments" if d == "webchat" else d) not in skip_components + ] + if selected_dirs: + if progress_callback: + await progress_callback( + "directories", 0, 100, "正在导入插件和数据目录..." + ) + + dir_stats = await self._import_directories( + zf, + manifest, + result, + selected_dirs=selected_dirs, + bad_entries=bad_entries, ) + result.imported_directories = dir_stats - dir_stats = await self._import_directories(zf, manifest, result) - result.imported_directories = dir_stats - - if progress_callback: - await progress_callback("directories", 100, 100, "目录导入完成") + if progress_callback: + await progress_callback("directories", 100, 100, "目录导入完成") - logger.info(f"备份导入完成: {result.to_dict()}") + logger.info( + "Backup import finished: success=%s, warnings=%d, errors=%d", + result.success, + len(result.warnings), + len(result.errors), + ) return result except zipfile.BadZipFile: - result.add_error("无效的 ZIP 文件") + result.add_error("Invalid ZIP file") return result except Exception as e: - result.add_error(f"导入失败: {e}") + result.add_error(f"Backup import failed: {e}") return result + def _resolve_selection( + self, + components: list[str] | None, + available: list[str], + broken: list[str], + result: ImportResult, + ) -> set[str] | None: + """Resolve the requested component selection. + + Broken components are checked FIRST so an explicitly requested broken + component is never downgraded to a plain "unavailable" skip. + + Args: + components: Requested ids, or None for the default restore. + available: Components whose required entries exist. + broken: Declared-but-missing components. + result: Import result collecting warnings/errors. + + Returns: + The effective component set, or None when the import must not + start (error already recorded, nothing modified). + """ + if components is not None and len(components) == 0: + result.add_error( + "No components selected for restoration (components is empty)" + ) + return None + + if components is None: + # Default restore must not silently drop corrupted components. + if broken: + result.add_error( + f"Declared components have missing entries: {sorted(broken)}. " + "Default restoration cannot skip broken components; explicitly select available components to restore." + ) + return None + if not available: + result.add_error("This backup contains no restorable data") + return None + return set(available) + + requested = set(components) + unknown = requested - set(get_backup_components()) + if unknown: + result.add_warning(f"Ignoring unknown component ids: {sorted(unknown)}") + requested -= unknown + + # Explicit selection legitimately excludes broken components, but + # the exclusion is always noted — never silent. + excluded_broken = set(broken) - requested + if excluded_broken: + result.add_warning( + f"Declared components have missing entries; excluded from this restore: {sorted(excluded_broken)}" + ) + + broken_req = requested & set(broken) + hard_broken = broken_req & HARD_FAIL_COMPONENTS + if hard_broken: + result.add_error( + f"Requested declared components have missing entries: {sorted(hard_broken)}; import aborted" + ) + return None + for comp in broken_req - HARD_FAIL_COMPONENTS: + result.add_error( + f"Requested declared component {comp} has missing entries; skipped" + ) + requested -= broken_req + + unavailable = requested - set(available) + if unavailable: + result.add_warning( + f"Skipping components absent from the backup: {sorted(unavailable)}" + ) + requested -= unavailable + + if not requested: + result.add_error("None of the requested components can be restored") + return None + return requested + + def _hash_entry(self, zf: zipfile.ZipFile, name: str) -> str: + """Stream an entry and return its hex sha256 (constant memory).""" + hasher = hashlib.sha256() + with zf.open(name) as src: + while chunk := src.read(1 << 20): + hasher.update(chunk) + return hasher.hexdigest() + + def _load_json_entry( + self, zf: zipfile.ZipFile, name: str, result: ImportResult + ) -> Any | None: + """Read and parse a JSON entry, recording an error on failure. + + A literal "null" payload is structurally invalid for every JSON + entry consumed here and is rejected explicitly — never confused + with a valid empty dump or a parse failure. + """ + try: + value = read_backup_json(zf, name) + except (KeyError, ValueError) as e: + result.add_error(f"Failed to parse JSON entry {name}: {e}") + return None + if value is None: + result.add_error(f"JSON entry is null: {name}") + return None + return value + + def _dry_run_table_dump( + self, + data: Any, + models: dict[str, type], + label: str, + result: ImportResult, + preprocess: bool = False, + ) -> bool: + """Validate a table dump: root dict[str, list[dict]] plus per-row checks. + + Every row goes through the same normalization as the real import + (strict datetime conversion — conversion failures are errors here, + not swallowed) followed by explicit model_validate(). Note that the + plain SQLModel constructor bypasses pydantic validation for table + models, so only model_validate can reject bad values. + """ + if not isinstance(data, (dict, BackupTableStream)): + result.add_error(f"Invalid {label} structure: root must be an object") + return False + for table_name, rows in data.items(): + model_class = models.get(table_name) + if model_class is None: + continue # unknown tables are skipped at import with a warning + if not isinstance(data, BackupTableStream) and not isinstance(rows, list): + result.add_error( + f"Invalid {label} structure: table {table_name} must be an array" + ) + return False + if preprocess: + rows = self._preprocess_main_table_rows(table_name, rows) + # Apply the same normalized byte budget as the write phase before + # any component can clear existing data. + for row in (row for batch in backup_row_batches(rows) for row in batch): + if not isinstance(row, dict): + result.add_error( + f"Invalid {label} structure: table {table_name} contains a non-object record" + ) + return False + try: + row = self._convert_datetime_fields(row, model_class, strict=True) + model_class.model_validate(row) + except Exception as e: + result.add_error( + f"Record validation failed in {label} table {table_name}: {e}" + ) + return False + return True + + def _dry_run_documents(self, data: Any, name: str, result: ImportResult) -> bool: + """Validate a KB documents.json payload structure.""" + if isinstance(data, BackupTableStream): + found = False + for table, rows in data.items(): + if table == "documents": + found = True + for row in ( + row for batch in backup_row_batches(rows) for row in batch + ): + if not self._dry_run_documents( + {"documents": [row]}, name, result + ): + return False + if not found: + result.add_error( + f"Invalid {name} structure: documents array is missing" + ) + return found + if not isinstance(data, dict) or not isinstance(data.get("documents"), list): + result.add_error(f"Invalid {name} structure: documents array is missing") + return False + for doc in data["documents"]: + if not isinstance(doc, dict) or "doc_id" not in doc or "text" not in doc: + result.add_error( + f"Invalid {name} structure: document is missing doc_id/text" + ) + return False + try: + json.loads(doc.get("metadata", "{}")) + except (TypeError, json.JSONDecodeError): + result.add_error( + f"Invalid {name} structure: metadata is not valid JSON" + ) + return False + return True + + async def _run_pre_verify( + self, + zf: zipfile.ZipFile, + manifest: dict, + namelist: list[str], + selected: set[str], + result: ImportResult, + progress_callback: Any | None = None, + ) -> dict | None: + """Verify all selected components before restoring any data. + + Import is a maintenance operation and may block the event loop. Keeping + verification here also ensures cancellation cannot close the archive + while a worker is still reading it. + + Returns: + A context dict for phase two, or None when a hard-fail component + failed verification — nothing has been modified. + """ + # Check type-specific limits before hashing or making any changes. + # Database dumps are replayable streams, so no total JSON cache is needed. + json_entries = [ + entry + for entry in zf.infolist() + if component_of_entry(entry.filename) in selected + and ( + entry.filename + in { + "databases/main_db.json", + "databases/kb_metadata.json", + "config/cmd_config.json", + } + or ( + entry.filename.startswith("databases/kb_") + and entry.filename.endswith("/documents.json") + ) + ) + ] + if any( + entry.file_size > backup_json_limit(entry.filename) + for entry in json_entries + ): + result.add_error("Backup JSON exceeds the size limit for its entry type") + return None + + checksums: dict[str, str] = manifest.get("checksums", {}) + comp_checksums: dict[str, str] = manifest.get("component_checksums", {}) + # v1.2 format = declared via the components field OR the manifest + # version itself; both enforce full checksum coverage. A manifest + # claiming version 1.2 without the components field is malformed + # and must not fall into the legacy degradation path. + is_v12 = ( + "components" in manifest + or VersionComparator.compare_version( + str(manifest.get("version", "1.0")), "1.2" + ) + >= 0 + ) + + ctx: dict[str, Any] = { + "json": {}, + "bad_entries": {}, + "skip_components": set(), + "unverified_counts": {}, + } + + comps = sorted(selected) + for i, comp in enumerate(comps): + if progress_callback: + await progress_callback( + "validate", i, len(comps), f"正在校验组件 {comp}..." + ) + ok = self._pre_verify_component( + zf, + comp, + checksums=checksums, + comp_checksums=comp_checksums, + is_v12=is_v12, + namelist=namelist, + result=result, + ctx=ctx, + ) + if not ok: + return None + + for comp, count in ctx["unverified_counts"].items(): + if count: + result.add_warning( + f"Component {comp}: {count} entries have no checksum; verification skipped (legacy backup)" + ) + + return ctx + + def _pre_verify_component( + self, + zf: zipfile.ZipFile, + comp: str, + *, + checksums: dict[str, str], + comp_checksums: dict[str, str], + is_v12: bool, + namelist: list[str], + result: ImportResult, + ctx: dict[str, Any], + ) -> bool: + """Verify one component's entries before restoration. + + Entry hash checks, component aggregate check and JSON structural + dry-runs, before anything is modified. Hard-fail components abort + the import on any failure; soft-fail components record corrupted + entries (skipped individually) or manifest-level errors (whole + component skipped). + + Returns: + False when a hard failure aborts the import (nothing has been + modified). Soft-fail issues are recorded in ctx/result. + """ + hard = comp in HARD_FAIL_COMPONENTS + entries = [ + n for n in namelist if component_of_entry(n) == comp and not n.endswith("/") + ] + # Missing entries are structural errors, even if other files are corrupt. + missing = sorted( + {name for name in checksums if component_of_entry(name) == comp} + - set(namelist) + ) + if missing: + result.add_error( + f"Declared component {comp} has missing entries: {missing}" + ) + if hard: + return False + ctx["skip_components"].add(comp) + return True + + bad: list[str] = [] + hashes: dict[str, str] = {} + unverified = 0 + + for name in entries: + expected = checksums.get(name) + if expected is None: + if is_v12: + # v1.2 requires full checksum coverage: this is a + # manifest-level (structural) format violation. + if hard: + result.add_error( + f"v1.2 backup entry is missing a checksum: {name}" + ) + return False + result.add_error( + f"v1.2 backup entry is missing a checksum: {name}; skipping component {comp}" + ) + ctx["skip_components"].add(comp) + return True + unverified += 1 + continue + try: + actual = f"sha256:{self._hash_entry(zf, name)}" + except (zipfile.BadZipFile, OSError, zlib.error) as e: + # Entry-level read failure (e.g. CRC error): classify like a + # checksum mismatch instead of failing the whole task. + if hard: + result.add_error(f"Failed to read entry: {name} ({e})") + return False + bad.append(name) + continue + if actual != expected: + if hard: + result.add_error(f"Entry checksum verification failed: {name}") + return False + bad.append(name) + continue + hashes[name] = actual + + # Component aggregate check (entry hashes already computed). + expected_agg = comp_checksums.get(comp) + if expected_agg is None: + if is_v12: + msg = f"v1.2 backup is missing the aggregate checksum for component {comp}" + if hard: + result.add_error(msg) + return False + result.add_error(f"{msg}; skipping component {comp}") + ctx["skip_components"].add(comp) + return True + else: + lines = sorted(f"{p}:{h}" for p, h in hashes.items()) + actual_agg = ( + "sha256:" + hashlib.sha256("\n".join(lines).encode("utf-8")).hexdigest() + ) + if actual_agg != expected_agg: + if hard: + result.add_error( + f"Component {comp} aggregate verification failed (manifest does not match contents)" + ) + return False + if not bad: + # Entries all match but the aggregate does not: + # manifest structural error, not known corruption. + result.add_error( + f"Component {comp} aggregate verification failed (invalid manifest structure); skipping component" + ) + ctx["skip_components"].add(comp) + return True + # Known corrupted entries already explain the mismatch; + # the entry-level warnings suffice, do not escalate. + + ctx["unverified_counts"][comp] = unverified + if bad: + ctx["bad_entries"][comp] = bad + # Never silent: corrupted entries surface in the import result, + # not only in per-file phase-two logs. + result.add_warning( + f"Component {comp}: verification failed or unreadable data for {len(bad)} entries; skipping these entries during import" + ) + + # JSON structural dry-run for hard-fail components. + if comp == "database": + data = self._load_json_entry(zf, "databases/main_db.json", result) + if data is None or not self._dry_run_table_dump( + data, MAIN_DB_MODELS, "main_db", result, preprocess=True + ): + return False + ctx["json"]["main_db"] = data + elif comp == "knowledge_base": + meta = self._load_json_entry(zf, "databases/kb_metadata.json", result) + if meta is None or not self._dry_run_table_dump( + meta, KB_METADATA_MODELS, "kb_metadata", result + ): + return False + ctx["json"]["kb_metadata"] = meta + for name in entries: + if name.startswith("databases/kb_") and name.endswith( + "/documents.json" + ): + doc = self._load_json_entry(zf, name, result) + if doc is None or not self._dry_run_documents(doc, name, result): + return False + ctx["json"][name] = doc + elif comp == "cmd_config": + cfg = self._load_json_entry(zf, "config/cmd_config.json", result) + if cfg is None: + return False + if not isinstance(cfg, dict): + result.add_error( + "Invalid cmd_config.json structure: root must be a JSON object" + ) + return False + + return True + + def _write_entry_safe( + self, zf: zipfile.ZipFile, name: str, target_path: Path + ) -> None: + """Write a ZIP entry to target_path via a same-dir temp file. + + The entry was already verified in phase one, so no re-hashing here. + The temp file is atomically moved into place with os.replace; on + failure only the temp file is removed and the pre-existing target + (if any) stays intact. + """ + tmp_fd, tmp_name = tempfile.mkstemp( + dir=target_path.parent, + prefix=f".{target_path.name}.", + suffix=".restore-tmp", + ) + try: + with os.fdopen(tmp_fd, "wb") as dst, zf.open(name) as src: + shutil.copyfileobj(src, dst) + os.replace(tmp_name, target_path) + except Exception: + try: + os.unlink(tmp_name) + except OSError: + pass + raise + def _validate_version(self, manifest: dict) -> None: """验证版本兼容性 - 仅允许相同主版本导入 @@ -516,7 +1098,7 @@ def _validate_version(self, manifest: dict) -> None: """ backup_version = manifest.get("astrbot_version") if not backup_version: - raise ValueError("备份文件缺少版本信息") + raise ValueError("Backup is missing version information") # 使用新的版本兼容性检查 version_check = self._check_version_compatibility(backup_version) @@ -526,7 +1108,9 @@ def _validate_version(self, manifest: dict) -> None: # minor_diff 和 match 都允许导入 if version_check["status"] == "minor_diff": - logger.warning(f"版本差异警告: {version_check['message']}") + logger.warning( + "Backup version differs: backup=%s, current=%s", backup_version, VERSION + ) async def _clear_main_db(self) -> None: """清空主数据库所有表""" @@ -535,27 +1119,21 @@ async def _clear_main_db(self) -> None: for table_name, model_class in MAIN_DB_MODELS.items(): try: await session.execute(delete(model_class)) - logger.debug(f"已清空表 {table_name}") + logger.debug(f"Cleared table {table_name}") except Exception as e: raise DatabaseClearError( - f"清空表 {table_name} 失败: {e}" + f"Failed to clear table {table_name}: {e}" ) from e async def _clear_kb_data(self) -> None: - """清空知识库数据""" + """清空知识库文件目录与实例 + + 注意：元数据表的清理由 _import_knowledge_bases(clear=True) 在导入 + 事务内完成，保证清表与插入原子性；此处只处理文件目录与内存实例。 + """ if not self.kb_manager: return - # 清空知识库元数据表 - async with self.kb_manager.kb_db.get_db() as session: - async with session.begin(): - for table_name, model_class in KB_METADATA_MODELS.items(): - try: - await session.execute(delete(model_class)) - logger.debug(f"已清空知识库表 {table_name}") - except Exception as e: - logger.warning(f"清空知识库表 {table_name} 失败: {e}") - # 删除知识库文件目录 for kb_id in list(self.kb_manager.kb_insts.keys()): try: @@ -564,56 +1142,176 @@ async def _clear_kb_data(self) -> None: if kb_helper.kb_dir.exists(): shutil.rmtree(kb_helper.kb_dir) except Exception as e: - logger.warning(f"清理知识库 {kb_id} 失败: {e}") + logger.warning(f"Failed to clean up knowledge base {kb_id}: {e}") self.kb_manager.kb_insts.clear() async def _import_main_database( - self, data: dict[str, list[dict]] + self, + data: dict[str, list[dict]] | BackupTableStream, + clear: bool = False, + attachment_paths: dict[str, str | None] | None = None, ) -> dict[str, int]: - """导入主数据库数据""" + """导入主数据库数据 + + 清表与插入在同一事务内执行：任何失败（包括逐条校验无法发现的 + 记录间约束冲突）整体回滚，旧数据不丢失。 + + Args: + data: 表名到记录列表的映射 + clear: 是否在插入前清空所有主库表 + attachment_paths: Optional archive attachment ids mapped to path hints, + filled during this pass with a bounded total path size. + + Returns: + dict: 每个表导入的记录数 + + Raises: + DatabaseClearError: 清空表失败（事务回滚，旧数据保留）。 + """ imported: dict[str, int] = {} + hint_bytes = 0 async with self.main_db.get_db() as session: async with session.begin(): + if clear: + for table_name, model_class in MAIN_DB_MODELS.items(): + try: + await session.execute(delete(model_class)) + logger.debug(f"Cleared table {table_name}") + except Exception as e: + raise DatabaseClearError( + f"Failed to clear table {table_name}: {e}" + ) from e + for table_name, rows in data.items(): model_class = MAIN_DB_MODELS.get(table_name) if not model_class: - logger.warning(f"未知的表: {table_name}") + logger.warning(f"Skipping unknown table: {table_name}") continue normalized_rows = self._preprocess_main_table_rows(table_name, rows) count = 0 - for row in normalized_rows: - try: - # 转换 datetime 字符串为 datetime 对象 - row = self._convert_datetime_fields(row, model_class) - obj = model_class(**row) - session.add(obj) - count += 1 - except Exception as e: - logger.warning(f"导入记录到 {table_name} 失败: {e}") + for batch in backup_row_batches(normalized_rows): + for row in batch: + try: + # 转换 datetime 字符串为 datetime 对象 + row = self._convert_datetime_fields(row, model_class) + # 与预检一致：显式 model_validate，普通构造会 + # 绕过 pydantic 校验 + obj = model_class.model_validate(row) + session.add(obj) + count += 1 + if table_name == "attachments" and attachment_paths: + key, path = ( + row.get("attachment_id"), + row.get("path"), + ) + if ( + key in attachment_paths + and attachment_paths[key] is None + and isinstance(path, str) + ): + hint_bytes += len(path.encode("utf-8")) + if hint_bytes <= MAX_JSON_RECORD_BYTES: + attachment_paths[key] = path + else: + logger.warning( + "Remaining attachment path hints exceed " + "the memory budget and will be skipped" + ) + # Keep the hints already collected, but + # stop collecting through this local reference. + attachment_paths = None + except Exception as e: + logger.warning( + f"Failed to import record into {table_name}: {e}" + ) + continue + # Flush without committing so failures roll back all rows. + await session.flush() imported[table_name] = count - logger.debug(f"导入表 {table_name}: {count} 条记录") + logger.debug(f"Imported table {table_name}: {count} records") return imported def _preprocess_main_table_rows( - self, table_name: str, rows: list[dict[str, Any]] - ) -> list[dict[str, Any]]: + self, table_name: str, rows: Iterable[dict[str, Any]] + ) -> Iterable[dict[str, Any]]: if table_name == "platform_stats": + if not isinstance(rows, list): + return self._merge_streamed_platform_stats(rows) normalized_rows = self._merge_platform_stats_rows(rows) duplicate_count = len(rows) - len(normalized_rows) if duplicate_count > 0: logger.warning( - "检测到 %s 重复键 %d 条，已在导入前聚合", + "Merged duplicate rows before import: table=%s, duplicates=%d", table_name, duplicate_count, ) return normalized_rows return rows + def _merge_streamed_platform_stats( + self, rows: Iterable[dict[str, Any]] + ) -> Iterator[dict[str, Any]]: + """Merge legacy statistics on disk without retaining all unique keys. + + Args: + rows: Streamed statistics records in their original order. + + Yields: + Normalized records in first-occurrence order, with duplicate counts summed. + """ + limiter = _InvalidCountWarnLimiter(PLATFORM_STATS_INVALID_COUNT_WARN_LIMIT) + # An empty SQLite filename creates a private temporary on-disk database + # that is removed on close. It never touches the application's database. + with closing(sqlite3.connect("")) as cache: + cache.execute("PRAGMA cache_size = -2048") + cache.execute( + "CREATE TABLE merged (position INTEGER PRIMARY KEY, " + "merge_key TEXT UNIQUE, payload TEXT NOT NULL)" + ) + duplicates = 0 + for row in rows: + normalized, timestamp, count = self._normalize_platform_stats_entry( + row, limiter + ) + platform_id = normalized.get("platform_id") + platform_type = normalized.get("platform_type") + key = None + if ( + timestamp is not None + and isinstance(platform_id, str) + and isinstance(platform_type, str) + ): + key = json.dumps((timestamp, platform_id, platform_type)) + existing = cache.execute( + "SELECT payload FROM merged WHERE merge_key = ?", (key,) + ).fetchone() + if existing is not None: + first = json.loads(existing[0]) + first["count"] += count + cache.execute( + "UPDATE merged SET payload = ? WHERE merge_key = ?", + (json.dumps(first), key), + ) + duplicates += 1 + continue + cache.execute( + "INSERT INTO merged (merge_key, payload) VALUES (?, ?)", + (key, json.dumps(normalized)), + ) + if duplicates: + logger.warning( + f"Merged {duplicates} duplicate platform statistics rows" + ) + for (payload,) in cache.execute( + "SELECT payload FROM merged ORDER BY position" + ): + yield json.loads(payload) + def _merge_platform_stats_rows( self, rows: list[dict[str, Any]] ) -> list[dict[str, Any]]: @@ -714,32 +1412,69 @@ def _normalize_platform_stats_timestamp(self, value: Any) -> str | None: async def _import_knowledge_bases( self, zf: zipfile.ZipFile, - kb_meta_data: dict[str, list[dict]], + kb_meta_data: dict[str, list[dict]] | BackupTableStream, result: ImportResult, + clear: bool = False, + json_ctx: dict[str, Any] | None = None, + bad_entries: list[str] | None = None, ) -> None: - """导入知识库数据""" + """导入知识库数据 + + Args: + zf: 打开的 ZIP 文件对象 + kb_meta_data: Validated KB metadata, replayed from the archive on demand. + result: 导入结果对象 + clear: 是否在导入事务内先清空元数据表（清表与插入原子） + json_ctx: Replayable JSON streams validated during phase one. + bad_entries: 阶段一检出的损坏条目，逐条跳过 + """ if not self.kb_manager: return - # 1. 导入知识库元数据 + if json_ctx is None: + json_ctx = {} + if bad_entries is None: + bad_entries = [] + + # 1. 导入知识库元数据（清表与插入在同一事务） + imported = {} async with self.kb_manager.kb_db.get_db() as session: async with session.begin(): + if clear: + for table_name, model_class in KB_METADATA_MODELS.items(): + try: + await session.execute(delete(model_class)) + logger.debug(f"Cleared knowledge base table {table_name}") + except Exception as e: + raise DatabaseClearError( + f"Failed to clear knowledge base table {table_name}: {e}" + ) from e + for table_name, rows in kb_meta_data.items(): model_class = KB_METADATA_MODELS.get(table_name) if not model_class: continue count = 0 - for row in rows: - try: - row = self._convert_datetime_fields(row, model_class) - obj = model_class(**row) - session.add(obj) - count += 1 - except Exception as e: - logger.warning(f"导入知识库记录到 {table_name} 失败: {e}") - - result.imported_tables[f"kb_{table_name}"] = count + for batch in backup_row_batches(rows): + for row in batch: + try: + row = self._convert_datetime_fields(row, model_class) + obj = model_class.model_validate(row) + session.add(obj) + count += 1 + except Exception as e: + logger.warning( + f"Failed to import knowledge base record into {table_name}: {e}" + ) + continue + await session.flush() + + imported[f"kb_{table_name}"] = count + + result.imported_tables.update(imported) + if clear: + await self._clear_kb_data() # 2. 导入每个知识库的文档和文件 for kb_data in kb_meta_data.get("knowledge_bases", []): @@ -753,47 +1488,55 @@ async def _import_knowledge_bases( # 导入文档数据 doc_path = f"databases/kb_{kb_id}/documents.json" - if doc_path in zf.namelist(): + if doc_path in zf.namelist() and doc_path not in bad_entries: try: - doc_content = zf.read(doc_path) - doc_data = json.loads(doc_content) + # Replay the validated entry without retaining its document list. + doc_data = json_ctx.pop(doc_path, None) + if doc_data is None: + doc_data = read_backup_json(zf, doc_path) # 导入到文档存储数据库 await self._import_kb_documents(kb_id, doc_data) + del doc_data except Exception as e: - result.add_warning(f"导入知识库 {kb_id} 的文档失败: {e}") + result.add_warning( + f"Failed to import documents for knowledge base {kb_id}: {e}" + ) # 导入 FAISS 索引 faiss_path = f"databases/kb_{kb_id}/index.faiss" - if faiss_path in zf.namelist(): + if faiss_path in zf.namelist() and faiss_path not in bad_entries: try: - target_path = kb_dir / "index.faiss" - with zf.open(faiss_path) as src, open(target_path, "wb") as dst: - dst.write(src.read()) + self._write_entry_safe(zf, faiss_path, kb_dir / "index.faiss") except Exception as e: - result.add_warning(f"导入知识库 {kb_id} 的 FAISS 索引失败: {e}") + result.add_warning( + f"Failed to import FAISS index for knowledge base {kb_id}: {e}" + ) # 导入媒体文件 media_prefix = f"files/kb_media/{kb_id}/" for name in zf.namelist(): - if name.startswith(media_prefix): + if name.startswith(media_prefix) and name not in bad_entries: try: rel_path = name[len(media_prefix) :] target_path = kb_dir / rel_path # Validate path is within kb directory (CWE-22) if not _validate_path_within(target_path, kb_dir): - logger.warning(f"媒体文件路径越界，已跳过: {target_path}") + logger.warning( + f"Skipping media file outside the knowledge base directory: {target_path}" + ) continue target_path.parent.mkdir(parents=True, exist_ok=True) - with zf.open(name) as src, open(target_path, "wb") as dst: - dst.write(src.read()) + self._write_entry_safe(zf, name, target_path) except Exception as e: - result.add_warning(f"导入媒体文件 {name} 失败: {e}") + result.add_warning(f"Failed to import media file {name}: {e}") # 3. 重新加载知识库实例 await self.kb_manager.load_kbs() - async def _import_kb_documents(self, kb_id: str, doc_data: dict) -> None: + async def _import_kb_documents( + self, kb_id: str, doc_data: dict | BackupTableStream + ) -> None: """导入知识库文档到向量数据库""" from astrbot.core.db.vec_db.faiss_impl.document_storage import DocumentStorage @@ -805,41 +1548,79 @@ async def _import_kb_documents(self, kb_id: str, doc_data: dict) -> None: await doc_storage.initialize() try: - documents = doc_data.get("documents", []) - for doc in documents: - try: - await doc_storage.insert_document( - doc_id=doc.get("doc_id", ""), - text=doc.get("text", ""), - metadata=json.loads(doc.get("metadata", "{}")), - ) - except Exception as e: - logger.warning(f"导入文档块失败: {e}") + for batch in backup_row_batches(doc_data.get("documents", [])): + await doc_storage.insert_documents_batch( + doc_ids=[doc.get("doc_id", "") for doc in batch], + texts=[doc.get("text", "") for doc in batch], + metadatas=[json.loads(doc.get("metadata", "{}")) for doc in batch], + ) finally: await doc_storage.close() async def _import_attachments( self, zf: zipfile.ZipFile, - attachments: list[dict], + attachments: Iterable[dict], + bad_entries: list[str] | None = None, ) -> int: - """导入附件文件""" + """导入附件文件 + + Args: + zf: ZIP 文件对象 + attachments: 附件记录（用于恢复原始路径） + bad_entries: 阶段一检出的损坏条目，逐条跳过 + + Returns: + int: 导入的附件数量 + """ + if bad_entries is None: + bad_entries = [] count = 0 attachments_dir = Path(self.config_path).parent / "attachments" - attachments_dir.mkdir(parents=True, exist_ok=True) attachment_prefix = "files/attachments/" + attachment_ids = { + Path(name).stem + for name in zf.namelist() + if name.startswith(attachment_prefix) and not name.endswith("/") + } + if not attachment_ids: + return 0 + attachments_dir.mkdir(parents=True, exist_ok=True) + attachment_paths = {} + hint_bytes = 0 + try: + for attachment in attachments: + attachment_id = attachment.get("attachment_id") + path = attachment.get("path") + if ( + attachment_id in attachment_ids + and attachment_id not in attachment_paths + and isinstance(path, str) + ): + hint_bytes += len(path.encode("utf-8")) + if hint_bytes > MAX_JSON_RECORD_BYTES: + raise ValueError( + "Attachment path hints exceed the memory budget" + ) + attachment_paths[attachment_id] = path + except Exception as exc: + # Hints are optional; decoding failures must not discard attachments. + logger.warning(f"Attachment path hints skipped: {exc}") for name in zf.namelist(): if name.startswith(attachment_prefix) and name != attachment_prefix: + if name in bad_entries: + # Phase one reported this entry as corrupted; skip it. + # The pre-existing target file (if any) is untouched. + logger.warning( + f"Skipping attachment with failed verification: {name}" + ) + continue try: # 从附件记录中找到原始路径 attachment_id = os.path.splitext(os.path.basename(name))[0] - original_path = None - for att in attachments: - if att.get("attachment_id") == attachment_id: - original_path = att.get("path") - break + original_path = attachment_paths.get(attachment_id) if original_path: target_path = Path(original_path) @@ -848,15 +1629,16 @@ async def _import_attachments( # Validate path is within attachments directory (CWE-22) if not _validate_path_within(target_path, attachments_dir): - logger.warning(f"附件路径越界，已跳过: {target_path}") + logger.warning( + f"Skipping attachment outside the attachments directory: {target_path}" + ) continue target_path.parent.mkdir(parents=True, exist_ok=True) - with zf.open(name) as src, open(target_path, "wb") as dst: - dst.write(src.read()) + self._write_entry_safe(zf, name, target_path) count += 1 except Exception as e: - logger.warning(f"导入附件 {name} 失败: {e}") + logger.warning(f"Failed to import attachment {name}: {e}") return count @@ -865,6 +1647,8 @@ async def _import_directories( zf: zipfile.ZipFile, manifest: dict, result: ImportResult, + selected_dirs: list[str] | None = None, + bad_entries: dict[str, list[str]] | None = None, ) -> dict[str, int]: """导入插件和其他数据目录 @@ -872,16 +1656,22 @@ async def _import_directories( zf: ZIP 文件对象 manifest: 备份清单 result: 导入结果对象 + selected_dirs: 只导入这些目录键；None 导入清单中的全部目录 + bad_entries: 阶段一检出的损坏条目（按组件分组），逐文件跳过 Returns: dict: 每个目录导入的文件数量 """ + if bad_entries is None: + bad_entries = {} dir_stats: dict[str, int] = {} # 检查备份版本是否支持目录备份（需要版本 >= 1.1） backup_version = manifest.get("version", "1.0") if VersionComparator.compare_version(backup_version, "1.1") < 0: - logger.info("备份版本不支持目录备份，跳过目录导入") + logger.info( + "Skipping directory import: backup version does not support directory backups" + ) return dir_stats backed_up_dirs = manifest.get("directories", []) @@ -889,11 +1679,20 @@ async def _import_directories( for dir_name in backed_up_dirs: if dir_name not in backup_directories: - result.add_warning(f"未知的目录类型: {dir_name}") + result.add_warning(f"Unknown directory type: {dir_name}") + continue + if selected_dirs is not None and dir_name not in selected_dirs: continue target_dir = Path(backup_directories[dir_name]) archive_prefix = f"directories/{dir_name}/" + if dir_name == "webchat": + # Preserve legacy filenames without replacing active upload data. + target_dir = target_dir / "imgs" + archive_prefix += "imgs/" + bad_files = bad_entries.get( + "attachments" if dir_name == "webchat" else dir_name, [] + ) file_count = 0 @@ -907,6 +1706,14 @@ async def _import_directories( if not dir_files: continue + if dir_name == "webchat" and not any( + not name.endswith("/") and name not in bad_files + for name in dir_files + ): + result.add_warning( + "No valid legacy WebChat images to restore; existing images were preserved." + ) + continue # 备份现有目录（如果存在） if target_dir.exists(): @@ -914,7 +1721,9 @@ async def _import_directories( if backup_path.exists(): shutil.rmtree(backup_path) shutil.move(str(target_dir), str(backup_path)) - logger.debug(f"已备份现有目录 {target_dir} 到 {backup_path}") + logger.debug( + f"Backed up existing directory {target_dir} to {backup_path}" + ) # 创建目标目录 target_dir.mkdir(parents=True, exist_ok=True) @@ -922,6 +1731,15 @@ async def _import_directories( # 解压文件 for name in dir_files: try: + if name in bad_files: + # Phase one reported this entry as corrupted; + # the old version stays recoverable in the .bak dir. + result.add_warning( + f"File {name} verification failed; skipped. " + f"The previous version can be recovered manually from {target_dir}.bak" + ) + continue + # 计算相对路径 rel_path = name[len(archive_prefix) :] if not rel_path: # 跳过目录条目 @@ -930,7 +1748,9 @@ async def _import_directories( target_path = target_dir / rel_path # Validate path is within target directory (CWE-22) if not _validate_path_within(target_path, target_dir): - result.add_warning(f"文件路径越界，已跳过: {name}") + result.add_warning( + f"Skipping file outside the target directory: {name}" + ) continue if zf.getinfo(name).is_dir(): @@ -939,23 +1759,38 @@ async def _import_directories( target_path.parent.mkdir(parents=True, exist_ok=True) - with zf.open(name) as src, open(target_path, "wb") as dst: - dst.write(src.read()) + self._write_entry_safe(zf, name, target_path) file_count += 1 except Exception as e: - result.add_warning(f"导入文件 {name} 失败: {e}") + result.add_warning(f"Failed to import file {name}: {e}") dir_stats[dir_name] = file_count - logger.debug(f"导入目录 {dir_name}: {file_count} 个文件") + logger.debug(f"Imported directory {dir_name}: {file_count} files") except Exception as e: - result.add_warning(f"导入目录 {dir_name} 失败: {e}") + result.add_warning(f"Failed to import directory {dir_name}: {e}") dir_stats[dir_name] = 0 return dir_stats - def _convert_datetime_fields(self, row: dict, model_class: type) -> dict: - """转换 datetime 字符串字段为 datetime 对象""" + def _convert_datetime_fields( + self, row: dict, model_class: type, strict: bool = False + ) -> dict: + """转换 datetime 字符串字段为 datetime 对象 + + Args: + row: 记录字典 + model_class: 目标模型类 + strict: 严格模式（预检使用）：转换失败抛出异常而不是静默 + 保留原值。导入路径保持宽松（默认 False），单条失败按 + warning 跳过。 + + Returns: + 转换后的记录字典 + + Raises: + ValueError: strict 模式下 datetime 解析失败。 + """ result = row.copy() # 获取模型的 datetime 字段 @@ -972,8 +1807,16 @@ def _convert_datetime_fields(self, row: dict, model_class: type) -> dict: value = result[column.name] if isinstance(value, str): # 解析 ISO 格式的日期时间字符串 - result[column.name] = datetime.fromisoformat(value) + try: + result[column.name] = datetime.fromisoformat(value) + except ValueError: + if strict: + raise + except ValueError: + if strict: + raise except Exception: - pass + if strict: + raise return result diff --git a/astrbot/core/backup/resources.py b/astrbot/core/backup/resources.py new file mode 100644 index 0000000000..15697181dc --- /dev/null +++ b/astrbot/core/backup/resources.py @@ -0,0 +1,414 @@ +"""Resource limits shared by backup inspection and restoration.""" + +import io +import json +import re +import zipfile +from collections.abc import Iterable, Iterator +from contextlib import closing, contextmanager, suppress +from decimal import Decimal +from itertools import groupby +from pathlib import Path +from typing import Any + +import ijson + +MAX_DIRECTORY_BYTES = 8 * 1024 * 1024 +MAX_ENTRIES = 40_000 +MAX_MANIFEST_BYTES = 8 * 1024 * 1024 +MAX_JSON_BYTES = 32 * 1024 * 1024 +MAX_DATABASE_JSON_BYTES = 2 * 1024 * 1024 * 1024 +MAX_KB_DOCUMENT_JSON_BYTES = 256 * 1024 * 1024 +MAX_JSON_RECORD_BYTES = 8 * 1024 * 1024 +MAX_JSON_DEPTH = 64 +MAX_EXTRACTED_BYTES = 32 * 1024 * 1024 * 1024 +_JSON_PUNCTUATION = re.compile(rb'["\\{}\[\],]') + + +def check_backup_json_field(data: bytes, remaining: int) -> int: + """Check a raw database JSON field before constructing its containers. + + Count its compact UTF-8 representation without materializing arrays or maps. + The enclosing dump, table array and row already consume three nesting levels. + + Args: + data: Bounded raw JSON bytes from the database. + remaining: Remaining serialized byte budget for the complete row. + + Returns: + The field's compact serialized byte count. + + Raises: + ValueError: The field exceeds the row size or nesting budget. + ijson.JSONError: The field contains malformed JSON. + """ + size = 0 + # Each pair records whether a container is a map and its member count. + containers: list[list[int]] = [] + with io.BytesIO(data) as source: + for event, value in ijson.basic_parse(source, use_float=False): + if event in ("end_map", "end_array"): + size += 1 + containers.pop() + else: + if containers and (not containers[-1][0] or event == "map_key"): + size += int(containers[-1][1] > 0) # Comma between members. + containers[-1][1] += 1 + if event in ("start_map", "start_array"): + size += 1 + containers.append([int(event == "start_map"), 0]) + if len(containers) + 3 > MAX_JSON_DEPTH: + raise ValueError("Backup JSON nesting exceeds the depth limit") + else: + if event == "map_key": + size += 1 # Colon after a map key. + if isinstance(value, str) and len(value) > remaining - size: + raise ValueError("Backup JSON record exceeds the size limit") + if isinstance(value, Decimal): + value = float(value) + size += len( + json.dumps(value, ensure_ascii=False, allow_nan=False).encode( + "utf-8" + ) + ) + if size > remaining: + raise ValueError("Backup JSON record exceeds the size limit") + return size + + +def backup_row_batches(rows: Iterable[dict]) -> Iterator[list[dict]]: + """Bound pending database objects by row count and serialized bytes. + + Args: + rows: Records from a validated table dump. + + Yields: + Batches of at most 500 rows and approximately one record budget of bytes. + + Raises: + ValueError: An individual row exceeds the record budget. + """ + batch: list[dict] = [] + size = 0 + for row in rows: + row_size = len( + json.dumps( + row, ensure_ascii=False, separators=(",", ":"), default=str + ).encode("utf-8") + ) + if row_size > MAX_JSON_RECORD_BYTES: + raise ValueError("Backup JSON record exceeds the size limit") + if batch and (len(batch) >= 500 or size + row_size > MAX_JSON_RECORD_BYTES): + yield batch + batch = [] + size = 0 + batch.append(row) + size += row_size + if batch: + yield batch + + +def backup_json_limit(name: str) -> int: + """Return the expanded size limit for a JSON entry. + + Args: + name: Archive entry name. + + Returns: + Maximum expanded bytes allowed for this entry type. + """ + if name == "manifest.json": + return MAX_MANIFEST_BYTES + if name in {"databases/main_db.json", "databases/kb_metadata.json"}: + return MAX_DATABASE_JSON_BYTES + if name.startswith("databases/kb_") and name.endswith("/documents.json"): + return MAX_KB_DOCUMENT_JSON_BYTES + return MAX_JSON_BYTES + + +class _RecordLimitedReader: + """Bound individual records and nesting before feeding the JSON parser. + + Streaming parsers still buffer a complete string token. Scanning punctuation + in bounded chunks prevents a single huge string or record exhausting memory. + """ + + def __init__(self, source: Any) -> None: + self.source = source + self.offset = 0 + self.depth = 0 + self.string_start: int | None = None + self.escaped_at = -1 + self.record_start: int | None = None + self.segment_start = 0 + self.safe_end = 0 + self.pending = b"" + + def read(self, size: int = -1) -> bytes: + """Read a chunk and check string, record and nesting budgets. + + Args: + size: Requested number of bytes, capped to the parser's chunk size. + + Returns: + At most 64 KiB of validated input bytes. + + Raises: + ValueError: A token, record or nesting depth exceeds its budget. + """ + size = min(size, 65536) if size >= 0 else 65536 + if size == 0: + return b"" + if len(self.pending) >= size: + chunk, self.pending = self.pending[:size], self.pending[size:] + return chunk + chunk = self.source.read(size - len(self.pending)) + for match in _JSON_PUNCTUATION.finditer(chunk): + pos = self.offset + match.start() + # Also bound scalars outside a valid table (e.g. a huge root number). + if pos - self.segment_start > MAX_JSON_RECORD_BYTES: + raise ValueError("Backup JSON token exceeds the size limit") + self.segment_start = pos + 1 + if ( + self.record_start is not None + and pos - self.record_start > MAX_JSON_RECORD_BYTES + ): + raise ValueError("Backup JSON record exceeds the size limit") + symbol = match[0] + if self.string_start is not None: + if pos - self.string_start > MAX_JSON_RECORD_BYTES: + raise ValueError("Backup JSON string exceeds the size limit") + if pos == self.escaped_at: + continue + if symbol == b"\\": + self.escaped_at = pos + 1 + elif symbol == b'"': + self.string_start = None + self.safe_end = pos + 1 + continue + if symbol == b'"': + self.string_start = pos + elif symbol in (b"{", b"["): + self.safe_end = pos + 1 + self.depth += 1 + if self.depth > MAX_JSON_DEPTH: + raise ValueError("Backup JSON nesting exceeds the depth limit") + if self.depth == 2 and symbol == b"[": + self.record_start = pos + 1 + elif symbol in (b"}", b"]"): + self.safe_end = pos + 1 + if self.depth == 2: + self.record_start = None + self.depth -= 1 + elif symbol == b",": + self.safe_end = pos + 1 + if self.depth == 2: + self.record_start = pos + 1 + self.offset += len(chunk) + if any( + start is not None and self.offset - start > MAX_JSON_RECORD_BYTES + for start in (self.record_start, self.string_start, self.segment_start) + ): + raise ValueError("Backup JSON token or record exceeds the size limit") + # End on a complete token whenever possible. The pure-Python ijson + # lexer otherwise retains preceding input when chunks keep ending inside + # strings, causing its buffer to grow with the entire database dump. + combined = self.pending + chunk + boundary = self.safe_end - (self.offset - len(combined)) + if 0 < boundary < len(combined): + self.pending = combined[boundary:] + return combined[:boundary] + self.pending = b"" + return combined + + +class BackupTableStream: + """Replay a table dump one row at a time from the still-open backup. + + The importer validates one pass before modifying anything, then replays the + same archive for restoration. No full table or cross-component cache is held. + """ + + def __init__(self, archive: zipfile.ZipFile, name: str) -> None: + self.archive = archive + self.name = name + + def _records(self) -> Iterator[tuple[str, dict | None]]: + """Parse table headers and bounded rows, enforcing the dump structure. + + Yields: + A table name and row, or None for a table header (including empty tables). + + Raises: + ValueError: The dump is malformed or exceeds a resource limit. + """ + + def parse_events(source: Any) -> Iterator[tuple[str, Any]]: + """Feed bounded chunks and explicitly close the parser on every exit. + + Args: + source: Uncompressed archive entry. + + Yields: + JSON parser events. + """ + pending = ijson.sendable_list() + # Preserve arbitrary-sized JSON integers; use_float=True overflows + # on valid integers with the native backend. + parser = ijson.basic_parse_coro(pending, use_float=False) + reader = _RecordLimitedReader(source) + try: + while chunk := reader.read(65536): + parser.send(chunk) + for event, value in pending: + yield ( + event, + float(value) if isinstance(value, Decimal) else value, + ) + pending.clear() + parser.close() + for event, value in pending: + yield event, float(value) if isinstance(value, Decimal) else value + finally: + # A resource-limit error can stop in the middle of a token. + # Preserve that error rather than a secondary incomplete-JSON error. + with suppress(ijson.JSONError): + parser.close() + + with ( + self.archive.open(self.name) as source, + closing(parse_events(source)) as events, + ): + if next(events, None) != ("start_map", None): + raise ValueError(f"{self.name}: table dump must be an object") + seen = set() + for event, table in events: + if event == "end_map": + if next(events, None) is not None: + raise ValueError(f"{self.name}: trailing JSON data") + return + if event != "map_key" or len(table) > 256 or table in seen: + raise ValueError(f"{self.name}: invalid or duplicate table name") + seen.add(table) + if len(seen) > MAX_ENTRIES: + raise ValueError("Backup table count exceeds the resource limit") + if next(events, None) != ("start_array", None): + raise ValueError(f"{self.name}: table {table} must be an array") + yield table, None + for event, value in events: + if event == "end_array": + break + if event != "start_map": + raise ValueError( + f"{self.name}: table {table} contains a non-object row" + ) + builder = ijson.common.ObjectBuilder() + builder.event(event, value) + depth = 1 + for event, value in events: + builder.event(event, value) + if event in ("start_map", "start_array"): + depth += 1 + elif event in ("end_map", "end_array"): + depth -= 1 + if depth == 0: + break + if depth: + raise ValueError(f"{self.name}: incomplete row") + yield table, builder.value + del builder + else: + raise ValueError(f"{self.name}: incomplete table") + raise ValueError(f"{self.name}: incomplete table dump") + + def items(self) -> Iterator[tuple[str, Iterator[dict]]]: + """Iterate tables in archive order; consume each row iterator before advancing. + + Yields: + A table name and its row iterator. + """ + with closing(self._records()) as records: + for table, group in groupby(records, key=lambda item: item[0]): + yield table, (row for _, row in group if row is not None) + + def get(self, key: str, default: Any = ()) -> Iterator[dict]: + """Replay one table without retaining the other tables in memory. + + Args: + key: Requested table name. + default: Rows to yield if the table does not exist. + + Yields: + Rows from the requested table, or the default iterable. + """ + with closing(self.items()) as tables: + for table, rows in tables: + if table == key: + yield from rows + return + yield from default + + +@contextmanager +def open_backup(path: str | Path, mode: str = "r") -> Iterator[zipfile.ZipFile]: + """Open a backup after bounding ZIP metadata allocation. + + Args: + path: Backup file on disk. + mode: Read mode, or append mode for updating the manifest. + + Yields: + The open archive, with bounded metadata and expanded size. + + Raises: + ValueError: The archive exceeds a resource limit. + zipfile.BadZipFile: The archive has no valid end record. + """ + with Path(path).open("r+b" if mode == "a" else "rb") as stream: + # ZipFile eagerly allocates the entire central directory. Use its own + # ZIP64-aware end-record reader to reject oversized metadata first. + end = zipfile._EndRecData(stream) + if end is None: + raise zipfile.BadZipFile("Missing ZIP end record") + if ( + end[zipfile._ECD_SIZE] > MAX_DIRECTORY_BYTES + or end[zipfile._ECD_ENTRIES_TOTAL] > MAX_ENTRIES + ): + raise ValueError("Backup ZIP directory exceeds the resource limit") + stream.seek(0) + with zipfile.ZipFile(stream, mode) as archive: + entries = archive.infolist() + if len(entries) > MAX_ENTRIES: + raise ValueError("Backup ZIP entry count exceeds the resource limit") + if sum(entry.file_size for entry in entries) > MAX_EXTRACTED_BYTES: + raise ValueError("Backup expanded size exceeds the resource limit") + yield archive + + +def read_backup_json(archive: zipfile.ZipFile, name: str) -> Any: + """Read small JSON eagerly, or expose a replayable stream for database dumps. + + Args: + archive: An open backup archive. + name: JSON entry to read. + + Returns: + The decoded JSON value, or a lazy stream of database rows. + + Raises: + ValueError: The entry is oversized or the manifest is not an object. + KeyError: The entry is missing. + """ + limit = backup_json_limit(name) + if archive.getinfo(name).file_size > limit: + raise ValueError(f"Backup JSON {name} exceeds the size limit ({limit} bytes)") + if name.startswith("databases/"): + return BackupTableStream(archive, name) + with archive.open(name) as source: + content = source.read(limit + 1) + if len(content) > limit: + raise ValueError(f"Backup JSON {name} exceeds the size limit ({limit} bytes)") + value = json.loads(content) + if name == "manifest.json" and not isinstance(value, dict): + raise ValueError("Backup manifest must be a JSON object") + return value diff --git a/astrbot/dashboard/api/backups.py b/astrbot/dashboard/api/backups.py index bd1ec026cc..091dd9d1e6 100644 --- a/astrbot/dashboard/api/backups.py +++ b/astrbot/dashboard/api/backups.py @@ -1,12 +1,13 @@ from __future__ import annotations -from fastapi import APIRouter, Depends, File, Form, Query, Request, UploadFile +from fastapi import APIRouter, Body, Depends, File, Form, Query, Request, UploadFile from fastapi.responses import FileResponse from astrbot.core import logger from astrbot.dashboard.async_utils import run_maybe_async from astrbot.dashboard.responses import error, ok from astrbot.dashboard.schemas import ( + BackupExportRequest, BackupImportRequest, BackupRenameRequest, BackupUploadInitRequest, @@ -70,7 +71,7 @@ async def _run(operation, *, prefix: str): except BackupServiceError as exc: return error(str(exc)) except Exception as exc: - logger.error("%s: %s", prefix, exc, exc_info=True) + logger.error("Backup API operation failed: %s", exc, exc_info=True) return error(f"{prefix}: {exc!s}") @@ -100,7 +101,7 @@ def _download_backup( except BackupServiceError as exc: return error(str(exc)) except Exception as exc: - logger.error("下载备份失败: %s", exc, exc_info=True) + logger.error("Failed to download backup: %s", exc, exc_info=True) return error(f"下载备份失败: {exc!s}") @@ -132,10 +133,15 @@ async def list_dashboard_backups( @router.post("/backups") async def create_backup( + payload: BackupExportRequest | None = Body(default=None), _auth: AuthContext = Depends(require_system_scope), service: BackupService = Depends(get_service), ): - return await _run(service.export_backup, prefix="创建备份失败") + # Body is optional: a bare POST keeps the legacy full-backup behavior. + return await _run( + lambda: service.export_backup(_model_dict(payload) if payload else {}), + prefix="创建备份失败", + ) @legacy_router.post("/export") diff --git a/astrbot/dashboard/schemas.py b/astrbot/dashboard/schemas.py index d54a4649e4..fef5dac3bd 100644 --- a/astrbot/dashboard/schemas.py +++ b/astrbot/dashboard/schemas.py @@ -80,8 +80,13 @@ class ChatUploadSessionRequest(OpenModel): upload_id: str | None = None +class BackupExportRequest(OpenModel): + components: list[str] | None = None + + class BackupImportRequest(OpenModel): confirmed: bool | None = None + components: list[str] | None = None class BackupRenameRequest(OpenModel): diff --git a/astrbot/dashboard/services/backup_service.py b/astrbot/dashboard/services/backup_service.py index 291e5d408c..c3181b4669 100644 --- a/astrbot/dashboard/services/backup_service.py +++ b/astrbot/dashboard/services/backup_service.py @@ -6,7 +6,6 @@ import re import traceback import uuid -import zipfile from dataclasses import dataclass from datetime import datetime from pathlib import Path @@ -17,6 +16,11 @@ from astrbot.core import logger from astrbot.core.backup.exporter import AstrBotExporter from astrbot.core.backup.importer import AstrBotImporter +from astrbot.core.backup.resources import ( + MAX_MANIFEST_BYTES, + open_backup, + read_backup_json, +) from astrbot.core.core_lifecycle import AstrBotCoreLifecycle from astrbot.core.db import BaseDatabase from astrbot.core.utils.astrbot_path import ( @@ -93,6 +97,11 @@ def _validate_backup_filename(filename: str | None, *, missing: str) -> str: return filename def _init_task(self, task_id: str, task_type: str, status: str = "pending") -> None: + if any( + task["status"] in ("pending", "processing") + for task in self.backup_tasks.values() + ): + raise BackupServiceError("已有备份任务正在运行，请等待完成后重试") self.backup_tasks[task_id] = { "type": task_type, "status": status, @@ -105,6 +114,9 @@ def _init_task(self, task_id: str, task_type: str, status: str = "pending") -> N "current": 0, "total": 100, "message": "", + # Accumulated per-stage completion entries (append-only) so the + # UI can show every component's outcome, not just the latest. + "stages": [], } def _set_task_result( @@ -160,6 +172,11 @@ async def _callback( total=total, message=message, ) + progress = self.backup_progress.get(task_id) + if progress is not None and total > 0 and current >= total: + stages = progress.setdefault("stages", []) + if not stages or stages[-1]["component"] != stage: + stages.append({"component": stage, "message": message}) return _callback @@ -171,13 +188,12 @@ async def cleanup_upload_session(self, upload_id: str) -> None: def get_backup_manifest(self, zip_path: str) -> dict | None: try: - with zipfile.ZipFile(zip_path, "r") as zf: + with open_backup(zip_path) as zf: if "manifest.json" in zf.namelist(): - manifest_data = zf.read("manifest.json") - return json.loads(manifest_data.decode("utf-8")) + return read_backup_json(zf, "manifest.json") return None except Exception as exc: - logger.debug(f"读取备份 manifest 失败: {exc}") + logger.debug(f"Failed to read backup manifest: {exc}") return None def list_backups(self, *, page: int, page_size: int) -> dict: @@ -195,7 +211,7 @@ def list_backups(self, *, page: int, page_size: int) -> dict: manifest = self.get_backup_manifest(file_path) if manifest is None: - logger.debug(f"跳过无效备份文件: {filename}") + logger.debug(f"Skipping invalid backup: {filename}") continue stat = os.stat(file_path) @@ -221,16 +237,25 @@ def list_backups(self, *, page: int, page_size: int) -> dict: "page_size": page_size, } - def export_backup(self) -> dict: + def export_backup(self, data: object = None) -> dict: + payload = self._payload(data) + components = payload.get("components") + if components is not None and not ( + isinstance(components, list) and all(isinstance(c, str) for c in components) + ): + raise BackupServiceError("components 必须是字符串数组") + task_id = str(uuid.uuid4()) self._init_task(task_id, "export", "pending") - asyncio.create_task(self.background_export_task(task_id)) + asyncio.create_task(self.background_export_task(task_id, components)) return { "task_id": task_id, "message": "export task created, processing in background", } - async def background_export_task(self, task_id: str) -> None: + async def background_export_task( + self, task_id: str, components: list[str] | None = None + ) -> None: try: self._update_progress(task_id, status="processing", message="正在初始化...") kb_manager = getattr(self.core_lifecycle, "kb_manager", None) @@ -242,6 +267,7 @@ async def background_export_task(self, task_id: str) -> None: zip_path = await exporter.export_all( output_dir=self.backup_dir, progress_callback=self._make_progress_callback(task_id), + components=components, ) self._set_task_result( task_id, @@ -250,10 +276,15 @@ async def background_export_task(self, task_id: str) -> None: "filename": os.path.basename(zip_path), "path": zip_path, "size": os.path.getsize(zip_path), + "components": exporter.exported_components, + "skipped": exporter.skipped_entries, }, ) + except asyncio.CancelledError: + self._set_task_result(task_id, "failed", error="Backup task cancelled") + raise except Exception as exc: - logger.error(f"后台导出任务 {task_id} 失败: {exc}") + logger.error(f"Background export task {task_id} failed: {exc}") logger.error(traceback.format_exc()) self._set_task_result(task_id, "failed", error=str(exc)) @@ -277,7 +308,7 @@ async def upload_backup(self, file: Any | None) -> dict: ) from exc logger.info( - f"上传的备份文件已保存: {unique_filename} (原始名称: {file.filename})" + f"Saved uploaded backup: {unique_filename} (original name: {file.filename})" ) return { "filename": unique_filename, @@ -352,20 +383,21 @@ async def upload_chunk( def mark_backup_as_uploaded(self, zip_path: str) -> None: try: manifest = {"origin": "uploaded", "uploaded_at": datetime.now().isoformat()} - with zipfile.ZipFile(zip_path, "r") as zf: + with open_backup(zip_path) as zf: if "manifest.json" in zf.namelist(): - manifest_data = zf.read("manifest.json") - manifest = json.loads(manifest_data.decode("utf-8")) + manifest = read_backup_json(zf, "manifest.json") manifest["origin"] = "uploaded" manifest["uploaded_at"] = datetime.now().isoformat() - with zipfile.ZipFile(zip_path, "a") as zf: + with open_backup(zip_path, "a") as zf: new_manifest = json.dumps(manifest, ensure_ascii=False, indent=2) + if len(new_manifest.encode("utf-8")) > MAX_MANIFEST_BYTES: + raise ValueError("Updated backup manifest exceeds the size limit") zf.writestr("manifest.json", new_manifest) - logger.debug(f"已标记备份为上传来源: {zip_path}") + logger.debug(f"Marked backup as uploaded: {zip_path}") except Exception as exc: - logger.warning(f"标记备份来源失败: {exc}") + logger.warning(f"Failed to mark backup origin: {exc}") async def upload_complete(self, data: object, *, owner: str = "") -> dict: payload = self._payload(data) @@ -386,7 +418,7 @@ async def upload_complete(self, data: object, *, owner: str = "") -> dict: self.mark_backup_as_uploaded(output_path) logger.info( - f"分片上传完成: {session.filename}, size={file_size}, " + f"Chunked upload completed: {session.filename}, size={file_size}, " f"chunks={session.total_chunks}" ) @@ -406,7 +438,7 @@ async def upload_abort( try: if await self.chunked_uploads.abort(upload_id, owner=owner): - logger.info(f"取消分片上传: {upload_id}") + logger.info(f"Aborted chunked upload: {upload_id}") except ChunkedUploadError as exc: raise BackupServiceError(str(exc)) from exc @@ -457,16 +489,24 @@ def import_backup(self, data: object) -> dict: if not os.path.exists(zip_path): raise BackupServiceError(f"备份文件不存在: {filename}") + components = payload.get("components") + if components is not None and not ( + isinstance(components, list) and all(isinstance(c, str) for c in components) + ): + raise BackupServiceError("components 必须是字符串数组") + task_id = str(uuid.uuid4()) self._init_task(task_id, "import", "pending") - asyncio.create_task(self.background_import_task(task_id, zip_path)) + asyncio.create_task(self.background_import_task(task_id, zip_path, components)) return { "task_id": task_id, "message": "import task created, processing in background", } - async def background_import_task(self, task_id: str, zip_path: str) -> None: + async def background_import_task( + self, task_id: str, zip_path: str, components: list[str] | None = None + ) -> None: try: self._update_progress(task_id, status="processing", message="正在初始化...") kb_manager = getattr(self.core_lifecycle, "kb_manager", None) @@ -479,18 +519,25 @@ async def background_import_task(self, task_id: str, zip_path: str) -> None: zip_path=zip_path, mode="replace", progress_callback=self._make_progress_callback(task_id), + components=components, ) if result.success: self._set_task_result(task_id, "completed", result=result.to_dict()) else: + # Keep the full result on failure: warnings and the already + # restored scope must survive, not just the error string. self._set_task_result( task_id, "failed", + result=result.to_dict(), error="; ".join(result.errors), ) + except asyncio.CancelledError: + self._set_task_result(task_id, "failed", error="Backup task cancelled") + raise except Exception as exc: - logger.error(f"后台导入任务 {task_id} 失败: {exc}") + logger.error(f"Background import task {task_id} failed: {exc}") logger.error(traceback.format_exc()) self._set_task_result(task_id, "failed", error=str(exc)) @@ -510,7 +557,7 @@ def get_progress(self, task_id: str | None) -> dict: if status == "processing" and task_id in self.backup_progress: response_data["progress"] = self.backup_progress[task_id] - if status == "completed": + if status in ("completed", "failed") and task_info.get("result") is not None: response_data["result"] = task_info["result"] if status == "failed": response_data["error"] = task_info["error"] @@ -592,7 +639,7 @@ def rename_backup(self, data: object) -> dict: raise BackupServiceError(f"文件名 '{new_filename}' 已存在") os.rename(old_path, new_path) - logger.info(f"备份文件重命名: {filename} -> {new_filename}") + logger.info(f"Renamed backup: {filename} -> {new_filename}") return { "old_filename": filename, "new_filename": new_filename, diff --git a/dashboard/src/api/generated/openapi-v1/types.gen.ts b/dashboard/src/api/generated/openapi-v1/types.gen.ts index 9b9332e447..bca0c16251 100644 --- a/dashboard/src/api/generated/openapi-v1/types.gen.ts +++ b/dashboard/src/api/generated/openapi-v1/types.gen.ts @@ -7,12 +7,18 @@ export type BackupChunkUploadRequest = { }; export type BackupExportRequest = { - include?: Array<(string)>; - exclude?: Array<(string)>; + /** + * Component ids to export. Omitted exports everything. + */ + components?: Array<(string)>; }; export type BackupImportRequest = { confirmed?: boolean; + /** + * Component ids to restore. Omitted restores all available. + */ + components?: Array<(string)>; }; export type BackupRenameRequest = { diff --git a/dashboard/src/api/v1.ts b/dashboard/src/api/v1.ts index 9c6e11324c..83e56a15a3 100644 --- a/dashboard/src/api/v1.ts +++ b/dashboard/src/api/v1.ts @@ -782,11 +782,11 @@ export const backupApi = { openApiV1.checkBackup({ path: { filename } }), ); }, - import(filename: string, confirmed = true) { + import(filename: string, confirmed = true, components?: string[]) { return typed( openApiV1.importBackup({ path: { filename }, - body: { confirmed } as any, + body: { confirmed, components } as any, }), ); }, diff --git a/dashboard/src/components/shared/BackupDialog.vue b/dashboard/src/components/shared/BackupDialog.vue index 8fe2172003..773e279708 100644 --- a/dashboard/src/components/shared/BackupDialog.vue +++ b/dashboard/src/components/shared/BackupDialog.vue @@ -6,9 +6,9 @@ {{ t('features.settings.backup.dialog.title') }} - + - + mdi-export {{ t('features.settings.backup.tabs.export') }} @@ -26,20 +26,43 @@ -
- mdi-cloud-upload -

{{ t('features.settings.backup.export.title') }}

-

{{ t('features.settings.backup.export.description') }}

- - - {{ t('features.settings.backup.export.includes') }} +
+
+ + mdi-cloud-upload + +
+

{{ t('features.settings.backup.export.title') }}

+

{{ t('features.settings.backup.export.description') }}

+
+
+
+ {{ t('features.settings.backup.export.selectComponents') }} + +
+ + {{ t('features.settings.backup.export.selectAll') }} + + + {{ t('features.settings.backup.export.clearAll') }} + +
+
+ + +
{{ w }}
 - - mdi-export - {{ t('features.settings.backup.export.button') }} - +
+ + mdi-export + {{ t('features.settings.backup.export.button') }} + +
@@ -53,6 +76,14 @@ mdi-check-circle

{{ t('features.settings.backup.export.completed') }}

{{ exportResult?.filename }}

+
+ + {{ t(`features.settings.backup.components.${comp}`) }} + +
+ +
{{ s.entry }}: {{ s.reason }}
+ mdi-download {{ t('features.settings.backup.export.download') }} @@ -153,29 +184,26 @@
- - - - mdi-package-variant - {{ t('features.settings.backup.import.backupContents') }} - - -
- - {{ checkResult.backup_summary.tables.length }} {{ t('features.settings.backup.import.tables') }} - - - {{ t('features.settings.backup.import.knowledgeBases') }} - - - {{ t('features.settings.backup.import.configFiles') }} - - - {{ dir }} - -
- - +

{{ t('features.settings.backup.import.restoreScope') }}

+ + + + + + + +
{{ w }}
+ @@ -197,6 +225,7 @@ color="error" size="large" variant="tonal" + :disabled="importComponents.length === 0" @click="confirmImport" > mdi-alert @@ -211,11 +240,20 @@

{{ t('features.settings.backup.import.processing') }}

{{ importProgress.message || t('features.settings.backup.import.wait') }}

+
+
+ mdi-check-circle + {{ s.message }} +
+
mdi-check-circle

{{ t('features.settings.backup.import.completed') }}

+ +
{{ w }}
+ {{ t('features.settings.backup.import.restartRequired') }} @@ -234,6 +272,18 @@ {{ importError }} + +
{{ w }}
+ + + + mdi-database-refresh + {{ t('features.settings.backup.import.restoredBeforeFailure') }} + + +
{{ s }}
+ + group.components) +const exportComponents = ref(BACKUP_COMPONENTS.filter(component => component !== 'temp')) +// Restore all available components by default after checking the backup. +const importComponents = ref([]) +// 导入结果（完成页展示 warnings） +const importResult = ref(null) + // 分片上传状态（调度由 useChunkedUpload 管理） const uploader = useChunkedUpload(backupApi) const { canResume } = uploader @@ -472,6 +536,49 @@ const versionAlertMessage = computed(() => { return t('features.settings.backup.import.version.matchMessage') }) +// 导出侧连锁警告：附件依赖主库；非全量备份旧版不可恢复 +const exportLinkWarnings = computed(() => { + const warnings = [] + const selected = exportComponents.value + if (!selected.includes('database') && selected.includes('attachments')) { + warnings.push(t('features.settings.backup.export.warningAttachmentsWithoutDb')) + } + if (selected.length > 0 && selected.length < BACKUP_COMPONENTS.length) { + warnings.push(t('features.settings.backup.export.warningSelectiveOldVersion')) + } + return warnings +}) + +// 恢复侧连锁警告：单独恢复附件可能成为孤儿文件 +const importLinkWarnings = computed(() => { + if (!importComponents.value.includes('database') && importComponents.value.includes('attachments')) { + return [t('features.settings.backup.import.warningAttachmentsWithoutDb')] + } + return [] +}) + +// 失败前已恢复的内容统计（失败页展示，让用户判断哪些数据已被修改） +// 零计数项同样展示并标注"已清空"——恢复空表意味着旧数据已被清除 +const importRestoredStats = computed(() => { + const r = importResult.value + if (!r) return [] + const annotate = (key, count) => + count > 0 + ? `${key}: ${count}` + : `${key}: ${t('features.settings.backup.import.restoredEmpty')}` + const stats = [] + for (const [table, count] of Object.entries(r.imported_tables || {})) { + stats.push(annotate(table, count)) + } + for (const [key, count] of Object.entries(r.imported_files || {})) { + stats.push(annotate(key, count)) + } + for (const [key, count] of Object.entries(r.imported_directories || {})) { + stats.push(annotate(key, count)) + } + return stats +}) + // 监听对话框打开 watch(isOpen, (newVal) => { if (newVal) { @@ -509,7 +616,7 @@ const startExport = async () => { exportProgress.value = { current: 0, total: 100, message: '' } try { - const response = await backupApi.create() + const response = await backupApi.create({ components: [...exportComponents.value] }) if (response.data.status === 'ok') { exportTaskId.value = response.data.data.task_id pollExportProgress() @@ -624,6 +731,9 @@ const checkUploadedBackup = async () => { return } + // 恢复范围默认全选可用组件（broken 组件不可勾选） + importComponents.value = [...(checkResult.value.available_components || [])] + // 显示确认对话框 importStatus.value = 'confirm' @@ -635,13 +745,17 @@ const checkUploadedBackup = async () => { // 确认导入 const confirmImport = async () => { - if (!uploadedFilename.value) return + if (!uploadedFilename.value || !checkResult.value?.can_import || !importComponents.value.length) return importStatus.value = 'processing' importProgress.value = { current: 0, total: 100, message: '' } try { - const response = await backupApi.import(uploadedFilename.value, true) + const response = await backupApi.import( + uploadedFilename.value, + true, + [...importComponents.value] + ) if (response.data.status === 'ok') { importTaskId.value = response.data.data.task_id @@ -669,13 +783,16 @@ const pollImportProgress = async () => { importProgress.value = { current: data.progress.current || 0, total: data.progress.total || 100, - message: data.progress.message || '' + message: data.progress.message || '', + stages: data.progress.stages || [] } setTimeout(pollImportProgress, 1000) } else if (data.status === 'completed') { importStatus.value = 'completed' + importResult.value = data.result } else if (data.status === 'failed') { importStatus.value = 'failed' + importResult.value = data.result importError.value = data.error || 'Import failed' } else { setTimeout(pollImportProgress, 1000) @@ -699,6 +816,8 @@ const resetImport = async () => { importError.value = '' uploadedFilename.value = '' checkResult.value = null + importComponents.value = [] + importResult.value = null uploadMessageOverride.value = '' } @@ -738,12 +857,15 @@ const restoreFromList = async (filename) => { } checkResult.value = checkResponse.data.data - + if (!checkResult.value.valid) { alert(checkResult.value.error || t('features.settings.backup.import.invalidBackup')) return } + // 恢复范围默认全选可用组件（broken 组件不可勾选） + importComponents.value = [...(checkResult.value.available_components || [])] + // 切换到导入标签页并显示确认 activeTab.value = 'import' importStatus.value = 'confirm' @@ -866,6 +988,7 @@ const restartAstrBot = async () => { // 重置所有状态 const resetAll = async () => { resetExport() + exportComponents.value = BACKUP_COMPONENTS.filter(component => component !== 'temp') await resetImport() activeTab.value = 'export' } @@ -885,6 +1008,30 @@ defineExpose({ open }) diff --git a/dashboard/src/i18n/locales/en-US/features/settings.json b/dashboard/src/i18n/locales/en-US/features/settings.json index 5c1d821ef9..baa680839e 100644 --- a/dashboard/src/i18n/locales/en-US/features/settings.json +++ b/dashboard/src/i18n/locales/en-US/features/settings.json @@ -171,10 +171,26 @@ "import": "Import Backup", "list": "Backup List" }, + "components": { + "database": "Main Database", + "knowledge_base": "Knowledge Base", + "cmd_config": "Main Config File", + "attachments": "Attachments", + "plugins": "Plugins", + "plugin_data": "Plugin Data", + "config": "Config Directory", + "t2i_templates": "T2I Templates", + "temp": "Temporary Files", + "skills": "Skills" + }, "export": { "title": "Create Backup", - "description": "Export all data as a ZIP backup file, including database, knowledge base, config, attachments and skills.", - "includes": "Backup includes: Main database, Knowledge bases (metadata + vector index + documents), Config files, Attachment files, Skills", + "description": "Export selected data as a ZIP backup. Main Data and Plugins & Extensions are selected by default; include Temporary Files if needed.", + "selectComponents": "Select what to back up", + "selectAll": "Select All", + "clearAll": "Clear", + "warningAttachmentsWithoutDb": "Attachments need matching main database records to be accessible. Back up the main database with them.", + "warningSelectiveOldVersion": "Selective backups cannot be restored on older AstrBot versions (full backups only).", "button": "Start Export", "processing": "Exporting...", "wait": "Please wait, packaging data...", @@ -186,7 +202,7 @@ }, "import": { "title": "Import Backup", - "warning": "⚠️ Import will clear and overwrite existing data! Please make sure you have backed up your current data.", + "warning": "Restoration clears and replaces existing data within the selected scope. Back up your current data first.", "selectFile": "Select backup file (.zip)", "uploadAndCheck": "Upload & Check", "uploading": "Uploading...", @@ -196,10 +212,11 @@ "uploadComplete": "Upload complete, merging file...", "checking": "Checking backup file...", "invalidBackup": "Invalid backup file", - "backupContents": "Backup Contents", - "tables": "tables", - "knowledgeBases": "Knowledge Bases", - "configFiles": "Config Files", + "restoreScope": "Select what to restore", + "brokenHint": "Declared in this backup, but required entries are missing. Cannot be restored.", + "warningAttachmentsWithoutDb": "Restoring attachments without the main database leaves them unreferenced and possibly inaccessible.", + "restoredBeforeFailure": "Restored before failure", + "restoredEmpty": "0 (emptied)", "confirmImport": "Confirm Import", "button": "Start Import", "processing": "Importing...", @@ -215,12 +232,13 @@ "currentVersion": "Current Version", "backupTime": "Backup Time", "matchTitle": "✅ Version Match", - "matchMessage": "Import will clear and overwrite all existing data, including:\n• Main database (conversations, settings, etc.)\n• Knowledge bases\n• Plugins and plugin data\n• Configuration files\n• Skills\n\nThis action cannot be undone! Do you want to continue?", + "matchMessage": "Versions match. Review the restoration scope below.", "minorDiffTitle": "⚠️ Version Difference Warning", - "minorDiffMessage": "Minor version differences can usually be imported, but there may be some data structure changes.\nImport will clear and overwrite all existing data!\n\nDo you want to continue?", + "minorDiffMessage": "Minor version differences are usually compatible, but data structures may differ. Review the restoration scope below.", "majorDiffTitle": "⛔ Cannot Import", "majorDiffMessage": "Major version numbers are different. Cross-major-version import may cause data corruption.\nPlease use the same major version of AstrBot for import." - } + }, + "replacementSummary": "Existing data for the following items will be cleared and replaced. This cannot be undone:" }, "list": { "empty": "No backup files", @@ -236,6 +254,37 @@ "renameInvalidChars": "Filename contains invalid characters", "renameFailed": "Rename failed", "ftpHint": "For large backup files, you can also upload directly to the data/backups directory via FTP/SFTP" + }, + "groups": { + "main": { + "title": "Main Data", + "description": "Conversations, knowledge bases, settings and attachments." + }, + "extensions": { + "title": "Plugins & Extensions", + "description": "Plugins, plugin data, skills and text-to-image templates." + }, + "temporary": { + "title": "Temporary Files", + "description": "Temporary files from message downloads, media processing and tools; usually unnecessary to back up." + } + }, + "componentDescriptions": { + "database": "Conversations, personas, statistics and other data stored in the database.", + "knowledge_base": "Knowledge base documents, metadata, vector indexes and media files.", + "cmd_config": "The main bot configuration file.", + "config": "Configuration profiles, plugin settings and other configuration files.", + "attachments": "Images, audio and other chat files, including legacy WebChat attachments. Back up and restore together with the main database.", + "plugins": "Program files of installed plugins.", + "plugin_data": "Persistent data saved by plugins.", + "skills": "Installed skill files.", + "t2i_templates": "Custom text-to-image templates.", + "temp": "Temporary files from message downloads, media processing and tools; usually unnecessary to back up." + }, + "scope": { + "selectedCount": "{selected} / {total} selected", + "notIncluded": "Not included in this backup.", + "selectAtLeastOne": "Select at least one item to restore." } }, "apiKey": { diff --git a/dashboard/src/i18n/locales/ja-JP/features/settings.json b/dashboard/src/i18n/locales/ja-JP/features/settings.json index 19ac7f7dd2..1ce8aff745 100644 --- a/dashboard/src/i18n/locales/ja-JP/features/settings.json +++ b/dashboard/src/i18n/locales/ja-JP/features/settings.json @@ -171,10 +171,26 @@ "import": "バックアップをインポート", "list": "バックアップ一覧" }, + "components": { + "database": "メインデータベース", + "knowledge_base": "ナレッジベース", + "cmd_config": "メイン設定ファイル", + "attachments": "添付ファイル", + "plugins": "プラグイン", + "plugin_data": "プラグインデータ", + "config": "設定ディレクトリ", + "t2i_templates": "T2Iテンプレート", + "temp": "一時ファイル", + "skills": "スキル" + }, "export": { "title": "バックアップを作成", - "description": "データベース、ナレッジベース、設定、添付ファイル、スキルを含むすべてのデータを ZIP バックアップファイルとしてエクスポートします。", - "includes": "バックアップ内容：メインデータベース、ナレッジベース（メタデータ＋ベクトルインデックス＋ドキュメント）、設定ファイル、添付ファイル、スキル", + "description": "選択したデータを ZIP バックアップとして出力します。主要データとプラグイン・拡張機能は初期選択され、一時ファイルは必要に応じて選択できます。", + "selectComponents": "バックアップする内容を選択", + "selectAll": "すべて選択", + "clearAll": "クリア", + "warningAttachmentsWithoutDb": "添付ファイルへのアクセスには対応するメインデータベースのレコードが必要です。メインデータベースも一緒にバックアップしてください。", + "warningSelectiveOldVersion": "選択的バックアップは旧バージョンの AstrBot では復元できません（旧版はフルバックアップのみ対応）。", "button": "エクスポートを開始", "processing": "エクスポート中…", "wait": "データをパッケージ化しています。しばらくお待ちください…", @@ -186,7 +202,7 @@ }, "import": { "title": "バックアップをインポート", - "warning": "⚠️ インポートすると既存のデータが削除され、バックアップの内容で上書きされます。現在のデータを必ずバックアップしてください。", + "warning": "復元すると、選択した範囲の既存データが消去・置換されます。事前に現在のデータをバックアップしてください。", "selectFile": "バックアップファイルを選択（.zip）", "uploadAndCheck": "アップロードして確認", "uploading": "アップロード中…", @@ -196,10 +212,11 @@ "uploadComplete": "アップロードが完了しました。ファイルを結合中…", "checking": "バックアップファイルを確認中…", "invalidBackup": "無効なバックアップファイル", - "backupContents": "バックアップ内容", - "tables": "テーブル", - "knowledgeBases": "ナレッジベース", - "configFiles": "設定ファイル", + "restoreScope": "復元する内容を選択", + "brokenHint": "バックアップに記載されていますが、必要なエントリが欠落しているため復元できません。", + "warningAttachmentsWithoutDb": "メインデータベースなしで添付ファイルを復元すると、レコード参照がなくアクセスできない孤立ファイルになる可能性があります。", + "restoredBeforeFailure": "失敗前に復元された内容", + "restoredEmpty": "0（クリア済み）", "confirmImport": "インポートを確認", "button": "インポートを開始", "processing": "インポート中…", @@ -215,12 +232,13 @@ "currentVersion": "現在のバージョン", "backupTime": "バックアップ日時", "matchTitle": "✅ バージョン一致", - "matchMessage": "インポートすると、以下を含む既存のすべてのデータが削除され、バックアップの内容で上書きされます：\n• メインデータベース（会話履歴、設定など）\n• ナレッジベース\n• プラグインとプラグインデータ\n• 設定ファイル\n• スキル\n\nこの操作は元に戻せません。続行しますか？", + "matchMessage": "バージョンが一致しています。以下の復元範囲を確認してください。", "minorDiffTitle": "⚠️ バージョン差異の警告", - "minorDiffMessage": "マイナーバージョンが異なっていても通常はインポートできますが、データ構造が一部異なる可能性があります。\nインポートすると既存のすべてのデータが削除され、バックアップの内容で上書きされます。\n\n続行しますか？", + "minorDiffMessage": "マイナーバージョンの違いは通常互換性がありますが、データ構造が異なる場合があります。以下の復元範囲を確認してください。", "majorDiffTitle": "⛔ インポートできません", "majorDiffMessage": "メジャーバージョンが異なります。メジャーバージョンをまたぐインポートはデータ破損の原因となる可能性があります。\n同じメジャーバージョンの AstrBot を使用してください。" - } + }, + "replacementSummary": "以下の項目の既存データが消去・置換されます。この操作は取り消せません：" }, "list": { "empty": "バックアップファイルがありません", @@ -236,6 +254,37 @@ "renameInvalidChars": "ファイル名に使用できない文字が含まれています", "renameFailed": "名前の変更に失敗しました", "ftpHint": "大きなバックアップファイルは、FTP/SFTP などを使用して data/backups ディレクトリへ直接アップロードすることもできます" + }, + "groups": { + "main": { + "title": "主要データ", + "description": "会話、ナレッジベース、設定、添付ファイル。" + }, + "extensions": { + "title": "プラグインと拡張機能", + "description": "プラグイン、プラグインデータ、スキル、画像生成テンプレート。" + }, + "temporary": { + "title": "一時ファイル", + "description": "メッセージのダウンロード、メディア処理、ツールの実行で生成される一時ファイル。通常はバックアップ不要です。" + } + }, + "componentDescriptions": { + "database": "会話履歴、ペルソナ、統計などのデータベース内のデータ。", + "knowledge_base": "ナレッジベースの文書、メタデータ、ベクトルインデックス、メディアファイル。", + "cmd_config": "ボットのメイン設定ファイル。", + "config": "設定プロファイル、プラグイン設定などの設定ファイル。", + "attachments": "旧版 WebChat の添付ファイルを含む、チャットの画像・音声・その他のファイル。メインデータベースと一緒にバックアップ・復元してください。", + "plugins": "インストール済みプラグインのプログラムファイル。", + "plugin_data": "プラグインが保存した永続データ。", + "skills": "インストール済みスキルのファイル。", + "t2i_templates": "カスタムのテキスト画像変換テンプレート。", + "temp": "メッセージのダウンロード、メディア処理、ツールの実行で生成される一時ファイル。通常はバックアップ不要です。" + }, + "scope": { + "selectedCount": "{total} 項目中 {selected} 項目を選択", + "notIncluded": "このバックアップには含まれていません。", + "selectAtLeastOne": "復元する項目を少なくとも1つ選択してください。" } }, "apiKey": { diff --git a/dashboard/src/i18n/locales/ru-RU/features/settings.json b/dashboard/src/i18n/locales/ru-RU/features/settings.json index 3a89471ea8..a1b3e2ac5e 100644 --- a/dashboard/src/i18n/locales/ru-RU/features/settings.json +++ b/dashboard/src/i18n/locales/ru-RU/features/settings.json @@ -171,10 +171,26 @@ "import": "Импорт", "list": "Список копий" }, + "components": { + "database": "Основная база данных", + "knowledge_base": "База знаний", + "cmd_config": "Основной файл конфигурации", + "attachments": "Вложения", + "plugins": "Плагины", + "plugin_data": "Данные плагинов", + "config": "Каталог конфигурации", + "t2i_templates": "Шаблоны T2I", + "temp": "Временные файлы", + "skills": "Навыки" + }, "export": { "title": "Создать резервную копию", - "description": "Экспорт всех данных в ZIP-архив, включая базы данных, базу знаний, конфигурации, вложения и Skills.", - "includes": "Включает: основную БД, векторные индексы знаний, файлы конфигурации, медиа-вложения, Skills.", + "description": "Экспорт выбранных данных в ZIP. Основные данные, плагины и расширения выбраны по умолчанию; временные файлы можно добавить при необходимости.", + "selectComponents": "Выберите данные для резервного копирования", + "selectAll": "Выбрать все", + "clearAll": "Очистить", + "warningAttachmentsWithoutDb": "Для доступа к вложениям нужны соответствующие записи основной базы данных. Сохраните базу данных вместе с вложениями.", + "warningSelectiveOldVersion": "Выборочная резервная копия не может быть восстановлена в старых версиях AstrBot (только полные копии).", "button": "Начать экспорт", "processing": "Экспорт...", "wait": "Пожалуйста, подождите, мы упаковываем данные...", @@ -186,7 +202,7 @@ }, "import": { "title": "Восстановление из копии", - "warning": "⚠️ Внимание! Импорт полностью удалит и перезапишет текущие данные! Убедитесь, что у вас есть копия текущего состояния.", + "warning": "Восстановление удалит и заменит существующие данные в выбранных категориях. Сначала сохраните текущие данные.", "selectFile": "Выберите ZIP-архив", "uploadAndCheck": "Загрузить и проверить", "uploading": "Загрузка...", @@ -196,10 +212,11 @@ "uploadComplete": "Загружено, идет сборка...", "checking": "Проверка структуры...", "invalidBackup": "Некорректный файл резервной копии", - "backupContents": "Состав архива", - "tables": "таблиц БД", - "knowledgeBases": "баз знаний", - "configFiles": "конфигов", + "restoreScope": "Выберите данные для восстановления", + "brokenHint": "Указано в резервной копии, но необходимые записи отсутствуют. Восстановление невозможно.", + "warningAttachmentsWithoutDb": "Восстановление вложений без основной базы данных оставит их без связанных записей и, возможно, недоступными.", + "restoredBeforeFailure": "Восстановлено до сбоя", + "restoredEmpty": "0 (очищено)", "confirmImport": "Подтвердите импорт", "button": "Начать восстановление", "processing": "Восстановление...", @@ -215,12 +232,13 @@ "currentVersion": "Текущая версия", "backupTime": "Дата создания", "matchTitle": "✅ Версии совпадают", - "matchMessage": "Импорт перезапишет все текущие данные, включая:\n• Основную БД (чаты, настройки)\n• Базы знаний\n• Плагины и их данные\n• Файлы конфигурации\n• Skills\n\nЭто действие необратимо! Продолжить?", + "matchMessage": "Версии совпадают. Проверьте выбранные данные для восстановления ниже.", "minorDiffTitle": "⚠️ Разница в минорной версии", - "minorDiffMessage": "Разница в минорных версиях обычно допустима, но структура данных могла немного измениться. Все текущие данные будут удалены!\n\nПродолжить импорт?", + "minorDiffMessage": "Различия минорных версий обычно совместимы, но структуры данных могут отличаться. Проверьте выбранные данные для восстановления ниже.", "majorDiffTitle": "⛔ Импорт невозможен", "majorDiffMessage": "Версии основного выпуска различаются. Импорт между мажорными версиями может привести к фатальному повреждению данных.\nИспользуйте AstrBot той же основной версии." - } + }, + "replacementSummary": "Существующие данные следующих элементов будут удалены и заменены. Это действие нельзя отменить:" }, "list": { "empty": "Резервные копии не найдены", @@ -236,6 +254,37 @@ "renameInvalidChars": "Имя содержит недопустимые символы", "renameFailed": "Ошибка переименования", "ftpHint": "Для больших архивов вы можете загружать их напрямую в папку data/backups через FTP/SFTP." + }, + "groups": { + "main": { + "title": "Основные данные", + "description": "Диалоги, базы знаний, настройки и вложения." + }, + "extensions": { + "title": "Плагины и расширения", + "description": "Плагины, их данные, навыки и шаблоны преобразования текста в изображения." + }, + "temporary": { + "title": "Временные файлы", + "description": "Файлы загрузок сообщений, обработки медиа и работы инструментов; обычно не требуют резервного копирования." + } + }, + "componentDescriptions": { + "database": "История диалогов, персонажи, статистика и другие данные в базе данных.", + "knowledge_base": "Документы базы знаний, метаданные, векторные индексы и медиафайлы.", + "cmd_config": "Основной файл конфигурации бота.", + "config": "Профили конфигурации, настройки плагинов и другие файлы настроек.", + "attachments": "Изображения, аудио и другие файлы чата, включая вложения старых версий WebChat. Сохраняйте и восстанавливайте вместе с основной базой данных.", + "plugins": "Программные файлы установленных плагинов.", + "plugin_data": "Постоянные данные, сохранённые плагинами.", + "skills": "Файлы установленных навыков.", + "t2i_templates": "Пользовательские шаблоны преобразования текста в изображения.", + "temp": "Файлы загрузок сообщений, обработки медиа и работы инструментов; обычно не требуют резервного копирования." + }, + "scope": { + "selectedCount": "Выбрано {selected} из {total}", + "notIncluded": "Не включено в эту резервную копию.", + "selectAtLeastOne": "Выберите хотя бы один элемент для восстановления." } }, "apiKey": { diff --git a/dashboard/src/i18n/locales/zh-CN/features/settings.json b/dashboard/src/i18n/locales/zh-CN/features/settings.json index fceb9e25a1..9007789080 100644 --- a/dashboard/src/i18n/locales/zh-CN/features/settings.json +++ b/dashboard/src/i18n/locales/zh-CN/features/settings.json @@ -171,10 +171,26 @@ "import": "导入备份", "list": "备份列表" }, + "components": { + "database": "主数据库", + "knowledge_base": "知识库", + "cmd_config": "主配置文件", + "attachments": "附件文件", + "plugins": "插件", + "plugin_data": "插件数据", + "config": "配置目录", + "t2i_templates": "T2I 模板", + "temp": "临时文件", + "skills": "Skills" + }, "export": { "title": "创建备份", - "description": "将所有数据导出为 ZIP 备份文件，包括数据库、知识库、配置、附件和技能。", - "includes": "备份包含：主数据库、知识库（元数据+向量索引+文档）、配置文件、附件文件、技能", + "description": "将所选数据导出为 ZIP 备份。默认选择主要数据和插件与扩展，临时文件按需选择。", + "selectComponents": "选择要备份的内容", + "selectAll": "全选", + "clearAll": "清空", + "warningAttachmentsWithoutDb": "附件需要对应的主数据库记录才能正常访问，建议同时备份主数据库。", + "warningSelectiveOldVersion": "选择性备份无法在旧版本 AstrBot 上恢复（旧版仅支持全量备份）。", "button": "开始导出", "processing": "正在导出...", "wait": "请稍候，正在打包数据...", @@ -186,7 +202,7 @@ }, "import": { "title": "导入备份", - "warning": "⚠️ 导入将会清空并覆盖现有数据！请确保已备份当前数据。", + "warning": "恢复会清空并替换所选范围内的现有数据，请先备份当前数据。", "selectFile": "选择备份文件 (.zip)", "uploadAndCheck": "上传并检查", "uploading": "正在上传...", @@ -196,10 +212,11 @@ "uploadComplete": "上传完成，正在合并文件...", "checking": "正在检查备份文件...", "invalidBackup": "无效的备份文件", - "backupContents": "备份内容", - "tables": "个数据表", - "knowledgeBases": "知识库", - "configFiles": "配置文件", + "restoreScope": "选择要恢复的内容", + "brokenHint": "备份声明了该项，但所需条目缺失，无法恢复。", + "warningAttachmentsWithoutDb": "单独恢复附件将没有数据库记录引用，可能成为无法访问的孤儿文件。", + "restoredBeforeFailure": "失败前已恢复的内容", + "restoredEmpty": "0（已清空）", "confirmImport": "确认导入", "button": "开始导入", "processing": "正在导入...", @@ -215,12 +232,13 @@ "currentVersion": "当前版本", "backupTime": "备份时间", "matchTitle": "✅ 版本匹配", - "matchMessage": "导入将会清空并覆盖现有的所有数据，包括：\n• 主数据库（对话记录、配置等）\n• 知识库数据\n• 插件及插件数据\n• 配置文件\n• 技能\n\n此操作不可撤销！是否确认继续？", + "matchMessage": "版本匹配，请在下方确认要恢复的范围。", "minorDiffTitle": "⚠️ 版本差异警告", - "minorDiffMessage": "小版本差异通常是兼容的，但可能存在少量数据结构变化。\n导入将会清空并覆盖现有的所有数据！\n\n是否确认继续导入？", + "minorDiffMessage": "小版本差异通常兼容，但可能存在数据结构变化，请在下方确认要恢复的范围。", "majorDiffTitle": "⛔ 无法导入", "majorDiffMessage": "主版本号不同，跨主版本导入可能导致数据损坏。\n请使用相同主版本的 AstrBot 进行导入。" - } + }, + "replacementSummary": "以下内容的现有数据将被清空并替换，此操作不可撤销：" }, "list": { "empty": "暂无备份文件", @@ -236,6 +254,37 @@ "renameInvalidChars": "文件名包含非法字符", "renameFailed": "重命名失败", "ftpHint": "对于较大的备份文件，也可以通过 FTP/SFTP 等方式直接上传到 data/backups 目录" + }, + "groups": { + "main": { + "title": "主要数据", + "description": "对话、知识库、配置和附件等主要数据。" + }, + "extensions": { + "title": "插件与扩展", + "description": "插件、插件数据、技能和文转图模板。" + }, + "temporary": { + "title": "临时文件", + "description": "消息下载、媒体处理和工具运行产生的临时文件，通常无需备份。" + } + }, + "componentDescriptions": { + "database": "对话记录、人格、统计和其他存储在数据库中的数据。", + "knowledge_base": "知识库文档、元数据、向量索引和媒体文件。", + "cmd_config": "机器人的主配置文件。", + "config": "配置档案、插件配置等配置文件。", + "attachments": "聊天中使用的图片、音频和其他文件，包含旧版 WebChat 附件。建议与主数据库一起备份和恢复。", + "plugins": "已安装插件的程序文件。", + "plugin_data": "插件保存的持久化数据。", + "skills": "已安装的技能文件。", + "t2i_templates": "自定义的文字转图片模板。", + "temp": "消息下载、媒体处理和工具运行产生的临时文件，通常无需备份。" + }, + "scope": { + "selectedCount": "已选 {selected} / {total} 项", + "notIncluded": "此备份未包含该项。", + "selectAtLeastOne": "请至少选择一项要恢复的内容。" } }, "apiKey": { diff --git a/docs/en/use/webui.md b/docs/en/use/webui.md index 837e43b903..c902cf11a2 100644 --- a/docs/en/use/webui.md +++ b/docs/en/use/webui.md @@ -99,6 +99,14 @@ Global settings are under `Settings` at the bottom of the sidebar: System configuration changes save automatically. Check for a successful save message and restart AstrBot if the page indicates that a restart is required. +Under `Settings → Maintenance`, open the backup dialog to export or restore data. Selection is grouped into **Main Data**, **Plugins & Extensions**, and **Temporary Files**. Use the checkbox to select a whole group, or click its heading to expand individual items and their descriptions. Partially selected groups show an indeterminate checkbox and a selection count. Export selects Main Data and Plugins & Extensions by default. Temporary Files are optional: these include message downloads and files produced by media processing and tools, and usually do not need to be backed up. + +After choosing or uploading a backup to restore, all available items are selected by default. Missing items are disabled; groups containing incomplete items expand automatically and explain why those items cannot be restored. Review the selected scope and the list of data that will be cleared and replaced before confirming. Restore attachments together with the main database when their records are needed. + +The Attachments option includes ordinary attachment files and legacy WebChat images. Conversation and attachment table records are included in Main Database. Legacy WebChat images from older backups retain their filenames and location when restored. Upload fragments are excluded from backups and are not restored from older backups, even when Temporary Files is selected. + +Review any warnings and the restored-data summary when the task finishes, including after a partial failure. Only one backup or restore task can run at a time. Import verification may temporarily delay dashboard responses. Selective backups require a version of AstrBot that supports selective restoration. + ## Plugins Select `Extensions` in the sidebar. The top tabs are `Plugins`, `Skills`, `MCP Servers`, and `Handlers`. Within `Plugins`, switch between `Installed` and `AstrBot Plugin Market` to view local and market plugins. diff --git a/docs/zh/use/webui.md b/docs/zh/use/webui.md index b5d38e1242..6c1aa8d273 100644 --- a/docs/zh/use/webui.md +++ b/docs/zh/use/webui.md @@ -99,6 +99,14 @@ ChatUI 支持以下常用能力： 系统配置修改后自动保存，请确认保存成功提示；如果页面提示需要重启，再按提示重启 AstrBot。 +在 `设置 → 维护` 中打开备份对话框，可导出或恢复数据。选择范围分为 **主要数据**、**插件与扩展** 和 **临时文件** 三组。勾选框用于整组选取，点击组标题可展开细项及说明；部分选中时，组勾选框显示半选状态，并显示已选数量。导出默认勾选主要数据和插件与扩展，临时文件按需选择。临时文件包含消息下载、媒体处理和工具运行产生的文件，通常无需备份。 + +选择或上传要恢复的备份后，默认勾选所有可用项。备份未包含的项不可选；存在条目缺失的组会自动展开，并说明无法恢复的原因。确认前请检查所选范围和即将被清空、替换的数据清单；需要附件记录时，应同时恢复附件与主数据库。 + +附件文件包含普通附件和旧版 WebChat 图片，对话及附件表记录包含在主数据库中。恢复旧版备份中的 WebChat 图片时，会保留原文件名和位置。上传分片不参与备份，也不会从旧版备份中恢复，勾选“临时文件”也不包含这些分片。 + +任务结束后请检查警告和已恢复的数据摘要，包括部分失败的情况。同一时间只允许运行一个备份或恢复任务。导入校验期间，仪表盘响应可能暂时延迟。选择性备份需要使用支持选择性恢复的 AstrBot 版本。 + ## 插件 点击左栏 `插件`，顶部可切换 `插件`、`技能`、`MCP` 和 `管理行为`。在 `插件` 标签内，通过 `已安装` 与 `插件市场` 切换查看本地插件和市场插件。 diff --git a/openspec/openapi-v1.yaml b/openspec/openapi-v1.yaml index 1893f1bfa1..1003442387 100644 --- a/openspec/openapi-v1.yaml +++ b/openspec/openapi-v1.yaml @@ -6434,14 +6434,11 @@ components: BackupExportRequest: type: object properties: - include: - type: array - items: - type: string - exclude: + components: type: array items: type: string + description: Component ids to export. Omitted exports everything. additionalProperties: false BackupUploadRequest: @@ -6532,6 +6529,11 @@ components: confirmed: type: boolean default: true + components: + type: array + items: + type: string + description: Component ids to restore. Omitted restores all available. additionalProperties: false UpdateRequest: diff --git a/pyproject.toml b/pyproject.toml index 14406011ac..91d0df63c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,6 +72,7 @@ dependencies = [ "pyotp>=2.9.0", "reportlab>=5.0.0", "xlrd>=2.0.2", + "ijson>=3.4,<4", ] [dependency-groups] diff --git a/requirements.txt b/requirements.txt index 9bcaa93a61..9fe0ff81fa 100644 --- a/requirements.txt +++ b/requirements.txt @@ -61,3 +61,4 @@ python-docx>=1.2.0 pyotp>=2.9.0 reportlab>=5.0.0 xlrd>=2.0.2 +ijson>=3.4,<4 diff --git a/tests/test_backup.py b/tests/test_backup.py index 069b61d6d6..d6a16ac802 100644 --- a/tests/test_backup.py +++ b/tests/test_backup.py @@ -1,11 +1,14 @@ """备份功能单元测试""" +import hashlib import json import os import re +import struct import zipfile +import zlib from datetime import datetime -from pathlib import Path +from pathlib import Path, PureWindowsPath from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -221,6 +224,10 @@ async def test_export_all_creates_zip( result = MagicMock() result.scalars.return_value.all.return_value = [] session.execute = AsyncMock(return_value=result) + stream = MagicMock() + stream.mappings.return_value.__aiter__.return_value = [] + stream.close = AsyncMock() + session.stream = AsyncMock(return_value=stream) mock_main_db.get_db.return_value = AsyncMock( __aenter__=AsyncMock(return_value=session), @@ -274,7 +281,7 @@ def test_validate_version_major_diff_rejected(self): # 使用一个明显不同的主版本 manifest = {"astrbot_version": "0.0.1"} - with pytest.raises(ValueError, match="主版本不兼容"): + with pytest.raises(ValueError, match="Incompatible major version"): importer._validate_version(manifest) def test_validate_version_minor_diff_allowed(self): @@ -294,7 +301,7 @@ def test_validate_version_missing(self): importer = AstrBotImporter(main_db=MagicMock()) manifest = {} - with pytest.raises(ValueError, match="缺少版本信息"): + with pytest.raises(ValueError, match="missing version information"): importer._validate_version(manifest) def test_convert_datetime_fields(self): @@ -480,7 +487,7 @@ def test_merge_platform_stats_rows_warns_on_invalid_count(self): warning_mock.call_count == PLATFORM_STATS_INVALID_COUNT_WARN_LIMIT + 1 ) assert any( - "告警已达到上限" in str(call.args[0]) + "warning limit reached" in str(call.args[0]) for call in warning_mock.call_args_list ) @@ -616,7 +623,7 @@ async def test_import_file_not_exists(self, mock_main_db, tmp_path): result = await importer.import_all(str(tmp_path / "nonexistent.zip")) assert result.success is False - assert any("不存在" in err for err in result.errors) + assert any("does not exist" in err for err in result.errors) @pytest.mark.asyncio async def test_import_invalid_zip(self, mock_main_db, tmp_path): @@ -629,7 +636,7 @@ async def test_import_invalid_zip(self, mock_main_db, tmp_path): result = await importer.import_all(str(invalid_zip)) assert result.success is False - assert any("无效" in err or "ZIP" in err for err in result.errors) + assert any("Invalid" in err or "ZIP" in err for err in result.errors) @pytest.mark.asyncio async def test_import_missing_manifest(self, mock_main_db, tmp_path): @@ -663,13 +670,17 @@ async def test_import_major_version_mismatch(self, mock_main_db, tmp_path): result = await importer.import_all(str(zip_path)) assert result.success is False - assert any("主版本不兼容" in err for err in result.errors) + assert any("Incompatible major version" in err for err in result.errors) @pytest.mark.asyncio async def test_import_replace_fails_when_clear_main_db_fails( self, mock_main_db, tmp_path ): - """测试 replace 模式下主库清空失败会直接终止导入""" + """测试 replace 模式下主库清空失败会直接终止导入 + + 清表已并入 _import_main_database 的导入事务（原子性），因此 + DatabaseClearError 现在从 _import_main_database 抛出。 + """ zip_path = tmp_path / "valid_backup.zip" manifest = { "version": "1.1", @@ -682,17 +693,19 @@ async def test_import_replace_fails_when_clear_main_db_fails( zf.writestr("databases/main_db.json", json.dumps(main_data)) importer = AstrBotImporter(main_db=mock_main_db) - importer._clear_main_db = AsyncMock( - side_effect=DatabaseClearError("清空表 platform_stats 失败: db locked") + importer._import_main_database = AsyncMock( + side_effect=DatabaseClearError( + "Failed to clear table platform_stats: db locked" + ) ) - importer._import_main_database = AsyncMock(return_value={}) result = await importer.import_all(str(zip_path), mode="replace") assert result.success is False - assert any("清空主数据库失败" in err for err in result.errors) - assert any("清空表 platform_stats 失败" in err for err in result.errors) - importer._import_main_database.assert_not_awaited() + assert any("Failed to clear main database" in err for err in result.errors) + assert any( + "Failed to clear table platform_stats" in err for err in result.errors + ) class TestSecureFilename: @@ -871,7 +884,7 @@ def test_pre_check_file_not_exists(self, mock_main_db): result = importer.pre_check("/nonexistent/file.zip") assert result.valid is False - assert "不存在" in result.error + assert "does not exist" in result.error def test_pre_check_invalid_zip(self, mock_main_db, tmp_path): """测试预检查无效的 ZIP 文件""" @@ -882,7 +895,7 @@ def test_pre_check_invalid_zip(self, mock_main_db, tmp_path): result = importer.pre_check(str(invalid_zip)) assert result.valid is False - assert "ZIP" in result.error or "无效" in result.error + assert "ZIP" in result.error or "Invalid" in result.error def test_pre_check_missing_manifest(self, mock_main_db, tmp_path): """测试预检查缺少 manifest 的 ZIP 文件""" @@ -897,20 +910,24 @@ def test_pre_check_missing_manifest(self, mock_main_db, tmp_path): assert "manifest" in result.error.lower() def test_pre_check_version_match(self, mock_main_db, tmp_path): - """测试预检查版本匹配""" + """测试预检查版本匹配 + + 摘要字段按 ZIP 实际条目推导（has_knowledge_bases / has_config + 不再是 manifest 自报字段），因此 zip 内需要放入对应条目。 + """ zip_path = tmp_path / "backup.zip" manifest = { "version": "1.1", "astrbot_version": VERSION, "created_at": "2024-01-01T12:00:00", "tables": {"platform_stats": 1}, - "has_knowledge_bases": True, - "has_config": True, "directories": ["plugins"], } with zipfile.ZipFile(zip_path, "w") as zf: zf.writestr("manifest.json", json.dumps(manifest)) + zf.writestr("databases/kb_metadata.json", json.dumps({})) + zf.writestr("config/cmd_config.json", json.dumps({})) importer = AstrBotImporter(main_db=mock_main_db) result = importer.pre_check(str(zip_path)) @@ -921,6 +938,29 @@ def test_pre_check_version_match(self, mock_main_db, tmp_path): assert result.backup_version == VERSION # confirm_message 现在由前端生成，后端不再生成 assert result.backup_summary["has_knowledge_bases"] is True + assert result.backup_summary["has_config"] is True + + def test_pre_check_summary_derived_from_real_entries(self, mock_main_db, tmp_path): + """摘要只信实际条目：manifest 自报布尔字段不被采信（幽灵字段回归）""" + zip_path = tmp_path / "backup.zip" + manifest = { + "version": "1.1", + "astrbot_version": VERSION, + "tables": {}, + "has_knowledge_bases": True, # 自报字段：不应影响推导结果 + "has_config": True, + } + + with zipfile.ZipFile(zip_path, "w") as zf: + zf.writestr("manifest.json", json.dumps(manifest)) + + importer = AstrBotImporter(main_db=mock_main_db) + result = importer.pre_check(str(zip_path)) + + assert result.valid is True + # zip 中没有 KB/配置条目，即使 manifest 自报也为 False + assert result.backup_summary["has_knowledge_bases"] is False + assert result.backup_summary["has_config"] is False def test_pre_check_minor_version_diff(self, mock_main_db, tmp_path): """测试预检查小版本差异""" @@ -1069,6 +1109,10 @@ async def test_export_import_roundtrip(self, tmp_path): result = MagicMock() result.scalars.return_value.all.return_value = [] session.execute = AsyncMock(return_value=result) + stream = MagicMock() + stream.mappings.return_value.__aiter__.return_value = [] + stream.close = AsyncMock() + session.stream = AsyncMock(return_value=stream) mock_db.get_db.return_value = AsyncMock( __aenter__=AsyncMock(return_value=session), @@ -1458,3 +1502,997 @@ async def test_short_write_chunk_is_rejected(self, backup_service): {"upload_id": upload_id}, owner="alice" ) assert result["size"] == 100 + + +def _make_working_mock_db(): + """Mock DB whose get_db()/session.begin() support the async CM protocol.""" + session = AsyncMock() + session.begin = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=session), + __aexit__=AsyncMock(return_value=None), + ) + ) + db = MagicMock() + db.get_db = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=session), + __aexit__=AsyncMock(return_value=None), + ) + ) + return db, session + + +def _sha256(data: bytes) -> str: + return f"sha256:{hashlib.sha256(data).hexdigest()}" + + +def _component_checksum(entries: dict[str, str]) -> str: + """Replicate the exporter's per-component digest for hand-built zips.""" + lines = sorted(f"{p}:{h}" for p, h in entries.items()) + return "sha256:" + hashlib.sha256("\n".join(lines).encode("utf-8")).hexdigest() + + +class TestSelectiveExport: + """选择性导出测试""" + + @pytest.mark.parametrize( + ("source", "relative_name", "archive_prefix"), + [ + ("webchat", "imgs/legacy.png", "directories/webchat"), + ("plugins", "example/main.py", "directories/plugins"), + ("kb_media", "media/image.png", "files/kb_media/kb1"), + ], + ) + def test_windows_archive_paths_match_checksums( + self, tmp_path, monkeypatch, source, relative_name, archive_prefix + ): + """Keep ZIP names and checksum keys portable for Windows source paths.""" + root = tmp_path / source + file_path = root / relative_name + file_path.parent.mkdir(parents=True) + file_path.write_bytes(b"backup-content") + exporter = AstrBotExporter(main_db=MagicMock()) + monkeypatch.setattr( + "astrbot.core.backup.exporter.get_backup_directories", + lambda: {source: root}, + ) + archive = tmp_path / "backup.zip" + relative_to = Path.relative_to + with zipfile.ZipFile(archive, "w") as zf, monkeypatch.context() as context: + # Simulate Windows separators while retaining real local file I/O. + context.setattr( + Path, + "relative_to", + lambda path, *other: PureWindowsPath(*relative_to(path, *other).parts), + ) + if source == "kb_media": + helper = MagicMock(kb_dir=root, kb_medias_dir=root / "media") + exporter._export_kb_media_files(zf, helper, "kb1") + else: + exporter._export_directories(zf, [source]) + + entry = f"{archive_prefix}/{relative_name}" + with zipfile.ZipFile(archive) as zf: + assert zf.namelist() == [entry] + assert exporter._checksums == {entry: _sha256(zf.read(entry))} + + @pytest.mark.asyncio + async def test_export_selective_components(self, temp_backup_dir, temp_data_dir): + """只导出勾选组件，manifest 记录实际写入的组件与聚合 hash""" + db, _ = _make_working_mock_db() + exporter = AstrBotExporter( + main_db=db, + kb_manager=None, + config_path=str(temp_data_dir / "cmd_config.json"), + ) + + zip_path = await exporter.export_all( + output_dir=str(temp_backup_dir), components=["cmd_config"] + ) + + with zipfile.ZipFile(zip_path, "r") as zf: + namelist = zf.namelist() + manifest = json.loads(zf.read("manifest.json")) + + assert "databases/main_db.json" not in namelist + assert "config/cmd_config.json" in namelist + assert manifest["components"] == ["cmd_config"] + assert set(manifest["component_checksums"]) == {"cmd_config"} + assert "config/cmd_config.json" in manifest["checksums"] + assert manifest["version"] == "1.2" + assert exporter.exported_components == ["cmd_config"] + + @pytest.mark.asyncio + async def test_export_invalid_components_rejected( + self, temp_backup_dir, temp_data_dir + ): + """空列表或全无效组件 id 被拒绝""" + exporter = AstrBotExporter( + main_db=MagicMock(), + kb_manager=None, + config_path=str(temp_data_dir / "cmd_config.json"), + ) + with pytest.raises(ValueError): + await exporter.export_all(output_dir=str(temp_backup_dir), components=[]) + with pytest.raises(ValueError): + await exporter.export_all( + output_dir=str(temp_backup_dir), components=["nope"] + ) + + @pytest.mark.asyncio + async def test_export_mid_write_failure_cleans_up_zip( + self, temp_backup_dir, temp_data_dir + ): + """回归：条目写入中途失败 -> 整个导出失败，半成品 ZIP 被清理""" + src = temp_data_dir / "att.bin" + src.write_bytes(b"x" * (2 << 20)) + + db, _ = _make_working_mock_db() + exporter = AstrBotExporter( + main_db=db, + kb_manager=None, + config_path=str(temp_data_dir / "cmd_config.json"), + ) + exporter._export_attachment_records = AsyncMock( + return_value=[{"path": str(src), "attachment_id": "x"}] + ) + + real_open = zipfile.ZipFile.open + + class _Boom: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def write(self, data): + raise OSError("disk on fire") + + def fake_open(zf, name, mode="r", *args, **kwargs): + if mode == "w": + return _Boom() + return real_open(zf, name, mode, *args, **kwargs) + + with patch.object(zipfile.ZipFile, "open", fake_open): + with pytest.raises(RuntimeError, match="mid-write"): + await exporter.export_all( + output_dir=str(temp_backup_dir), components=["attachments"] + ) + + # 半成品 ZIP 已被清理，不留无法通过完整性校验的产物 + assert list(temp_backup_dir.glob("*.zip")) == [] + + +class TestSelectiveImport: + """选择性导入与两阶段预检测试""" + + async def _full_backup(self, tmp_path, temp_data_dir): + """用真实 exporter 造一份 database + cmd_config 备份""" + db, _ = _make_working_mock_db() + exporter = AstrBotExporter( + main_db=db, + kb_manager=None, + config_path=str(temp_data_dir / "cmd_config.json"), + ) + exporter._export_main_database = AsyncMock( + return_value={"platform_stats": [], "conversations": [], "attachments": []} + ) + return await exporter.export_all( + output_dir=str(tmp_path / "bk"), components=["database", "cmd_config"] + ) + + @pytest.mark.asyncio + async def test_import_selective_restore_only_selected( + self, tmp_path, temp_data_dir + ): + """只恢复勾选组件：配置被替换，主库完全不被触碰""" + zip_path = await self._full_backup(tmp_path, temp_data_dir) + + target = tmp_path / "restore_cfg.json" + target.write_text(json.dumps({"old": "config"})) + db, session = _make_working_mock_db() + importer = AstrBotImporter(main_db=db, kb_manager=None, config_path=str(target)) + + result = await importer.import_all(zip_path, components=["cmd_config"]) + + assert result.success, result.errors + assert json.loads(target.read_text()) == {"test": "config"} + assert result.imported_files.get("config") == 1 + session.execute.assert_not_awaited() # 主库未被触碰 + + @pytest.mark.asyncio + async def test_import_empty_components_rejected(self, tmp_path, temp_data_dir): + """components=[] 显式拒绝（与 None 的全量语义区分）""" + zip_path = await self._full_backup(tmp_path, temp_data_dir) + db, _ = _make_working_mock_db() + importer = AstrBotImporter(main_db=db, kb_manager=None) + + result = await importer.import_all(zip_path, components=[]) + assert result.success is False + assert result.errors + + @pytest.mark.asyncio + async def test_import_all_invalid_components_rejected( + self, tmp_path, temp_data_dir + ): + """请求的组件全部无效时报错且不执行任何修改""" + zip_path = await self._full_backup(tmp_path, temp_data_dir) + db, _ = _make_working_mock_db() + importer = AstrBotImporter(main_db=db, kb_manager=None) + + result = await importer.import_all(zip_path, components=["nope1", "nope2"]) + assert result.success is False + assert any( + "None of the requested components can be restored" in e + for e in result.errors + ) + + @pytest.mark.asyncio + async def test_import_corrupt_config_aborts_zero_modification( + self, tmp_path, temp_data_dir + ): + """配置条目 hash 不匹配（硬失败）-> 中止，已有文件零改动""" + zip_path = await self._full_backup(tmp_path, temp_data_dir) + bad = tmp_path / "bad.zip" + with zipfile.ZipFile(zip_path) as zin, zipfile.ZipFile(bad, "w") as zout: + for item in zin.namelist(): + data = zin.read(item) + if item == "config/cmd_config.json": + data = b'{"tampered": true}' + zout.writestr(item, data) + + victim = tmp_path / "victim.json" + victim.write_text(json.dumps({"precious": "data"})) + db, _ = _make_working_mock_db() + importer = AstrBotImporter(main_db=db, kb_manager=None, config_path=str(victim)) + + result = await importer.import_all(str(bad), components=["cmd_config"]) + assert result.success is False + assert any("checksum" in e for e in result.errors) + assert json.loads(victim.read_text()) == {"precious": "data"} + + async def _legacy_zip(self, tmp_path, main_data: dict) -> str: + """造一份无 checksum 的旧格式（v1.1）备份""" + zip_path = tmp_path / "legacy.zip" + manifest = { + "version": "1.1", + "astrbot_version": VERSION, + "tables": {"main_db": list(main_data.keys())}, + } + with zipfile.ZipFile(zip_path, "w") as zf: + zf.writestr("manifest.json", json.dumps(manifest)) + zf.writestr("databases/main_db.json", json.dumps(main_data)) + return str(zip_path) + + @pytest.mark.asyncio + async def test_import_invalid_datetime_aborts_before_modification(self, tmp_path): + """非法日期：严格归一化 + model_validate 在清库前拦截（零修改）""" + zip_path = await self._legacy_zip( + tmp_path, + { + "conversations": [ + { + "conversation_id": "c1", + "platform_id": "p", + "user_id": "u", + "created_at": "not-a-date", + } + ] + }, + ) + db, session = _make_working_mock_db() + importer = AstrBotImporter(main_db=db, kb_manager=None) + + result = await importer.import_all(zip_path, components=["database"]) + assert result.success is False + assert any("Record validation failed" in e for e in result.errors) + session.execute.assert_not_awaited() # 清库未发生 + + @pytest.mark.asyncio + async def test_import_missing_required_field_aborts(self, tmp_path): + """缺必需字段：普通构造能过、model_validate 拒绝（零修改）""" + zip_path = await self._legacy_zip(tmp_path, {"conversations": [{}]}) + db, session = _make_working_mock_db() + importer = AstrBotImporter(main_db=db, kb_manager=None) + + result = await importer.import_all(zip_path, components=["database"]) + assert result.success is False + assert any("Record validation failed" in e for e in result.errors) + session.execute.assert_not_awaited() + + @pytest.mark.asyncio + async def test_import_broken_component_three_states(self, tmp_path, temp_data_dir): + """回归：UI 排除 broken 后返回 warning；默认恢复与显式选中 broken 中止""" + zip_path = await self._full_backup(tmp_path, temp_data_dir) + # 删掉 main_db.json 条目，制造"已声明但损坏"的 database + broken_zip = tmp_path / "broken.zip" + with zipfile.ZipFile(zip_path) as zin, zipfile.ZipFile(broken_zip, "w") as zout: + for item in zin.namelist(): + if item == "databases/main_db.json": + continue + zout.writestr(item, zin.read(item)) + + db, _ = _make_working_mock_db() + importer = AstrBotImporter( + main_db=db, + kb_manager=None, + config_path=str(tmp_path / "cfg.json"), + ) + + # 默认恢复：遇 broken 中止 + result = await importer.import_all(str(broken_zip)) + assert result.success is False + assert any("missing entries" in e for e in result.errors) + + # 显式排除 broken：恢复可用组件，result 带 warning 注明（不静默） + result = await importer.import_all(str(broken_zip), components=["cmd_config"]) + assert result.success, result.errors + assert any("excluded from this restore" in w for w in result.warnings) + + # 显式选中 broken 硬失败组件：修改前中止，不允许降格跳过 + result = await importer.import_all( + str(broken_zip), components=["database", "cmd_config"] + ) + assert result.success is False + assert any("missing entries" in e for e in result.errors) + + @pytest.mark.asyncio + async def test_import_legacy_roundtrip_from_real_export( + self, tmp_path, temp_data_dir + ): + """真实 exporter 产物降级为旧格式后仍可恢复（缺失 checksum 走降级警告）""" + zip_path = await self._full_backup(tmp_path, temp_data_dir) + + # 把 v1.2 manifest 降级成 v1.1：去掉新字段和 checksums + legacy = tmp_path / "legacy_full.zip" + with zipfile.ZipFile(zip_path) as zin, zipfile.ZipFile(legacy, "w") as zout: + for item in zin.namelist(): + data = zin.read(item) + if item == "manifest.json": + manifest = json.loads(data) + manifest["version"] = "1.1" + for key in ("components", "component_checksums", "checksums"): + manifest.pop(key, None) + data = json.dumps(manifest).encode() + zout.writestr(item, data) + + target = tmp_path / "restore.json" + db, _ = _make_working_mock_db() + importer = AstrBotImporter(main_db=db, kb_manager=None, config_path=str(target)) + + result = await importer.import_all(str(legacy)) + assert result.success, result.errors + assert json.loads(target.read_text()) == {"test": "config"} + assert result.imported_files.get("config") == 1 + # 旧格式降级有聚合警告，不静默 + assert any("no checksum" in w for w in result.warnings) + + @pytest.mark.asyncio + async def test_import_main_db_single_transaction(self): + """清表与插入在同一 session.begin() 事务内（原子性结构验证）""" + db, session = _make_working_mock_db() + importer = AstrBotImporter(main_db=db, kb_manager=None) + + await importer._import_main_database({"platform_stats": []}, clear=True) + + assert session.begin.call_count == 1 + # 13 张表各一次 delete，空数据无插入 + assert session.execute.await_count == len(MAIN_DB_MODELS) + session.add.assert_not_called() + + @pytest.mark.asyncio + async def test_import_main_db_clear_failure_raises_clear_error(self): + """清表失败抛 DatabaseClearError（事务回滚，旧数据保留）""" + db, session = _make_working_mock_db() + session.execute = AsyncMock(side_effect=Exception("db locked")) + importer = AstrBotImporter(main_db=db, kb_manager=None) + + with pytest.raises(DatabaseClearError): + await importer._import_main_database({"platform_stats": []}, clear=True) + + @pytest.mark.asyncio + async def test_import_corrupt_attachment_preserves_existing( + self, tmp_path, temp_data_dir + ): + """坏附件（软失败）：跳过且结果带 warning，恢复前的原文件不受损""" + attachments_dir = temp_data_dir / "attachments" + victim = attachments_dir / "abc.jpg" + victim.write_bytes(b"original-bytes") + + zip_path = tmp_path / "att.zip" + manifest = { + "version": "1.2", + "astrbot_version": VERSION, + "components": ["attachments"], + "checksums": {"files/attachments/abc.jpg": "sha256:wrong"}, + "component_checksums": {"attachments": "sha256:whatever"}, + "directories": [], + } + with zipfile.ZipFile(zip_path, "w") as zf: + zf.writestr("manifest.json", json.dumps(manifest)) + zf.writestr("files/attachments/abc.jpg", b"tampered-bytes") + + db, _ = _make_working_mock_db() + importer = AstrBotImporter( + main_db=db, + kb_manager=None, + config_path=str(temp_data_dir / "cmd_config.json"), + ) + + result = await importer.import_all(str(zip_path), components=["attachments"]) + assert result.success, result.errors + assert any("verification failed" in w for w in result.warnings) + # 原文件未被覆盖、未被删除 + assert victim.read_bytes() == b"original-bytes" + + @pytest.mark.asyncio + async def test_v12_missing_checksum_soft_component_skipped(self, tmp_path): + """v1.2 软失败组件条目缺 checksum（清单级错误）-> 整个组件跳过并记 error""" + zip_path = tmp_path / "v12_bad.zip" + manifest = { + "version": "1.2", + "astrbot_version": VERSION, + "components": ["attachments"], + "checksums": {}, # 条目缺 checksum：v1.2 不允许降级 + "component_checksums": {}, + "directories": [], + } + with zipfile.ZipFile(zip_path, "w") as zf: + zf.writestr("manifest.json", json.dumps(manifest)) + zf.writestr("files/attachments/a.jpg", b"data") + + db, _ = _make_working_mock_db() + importer = AstrBotImporter( + main_db=db, + kb_manager=None, + config_path=str(tmp_path / "cfg.json"), + ) + + result = await importer.import_all(str(zip_path), components=["attachments"]) + assert result.success is False + assert any("checksum" in e for e in result.errors) + + +class TestImportTaskResultRetention: + """服务层：失败任务保留完整结果""" + + @pytest.fixture + def backup_service(self, tmp_path): + service = BackupService(db=MagicMock(), core_lifecycle=MagicMock()) + service.backup_dir = str(tmp_path / "backups") + service.data_dir = str(tmp_path / "data") + return service + + @pytest.mark.asyncio + async def test_failed_import_task_keeps_full_result(self, backup_service, tmp_path): + """导入失败时 result（含 warnings/errors）完整保留，可通过进度接口查询""" + zip_path = tmp_path / "broken.zip" + manifest = { + "version": "1.2", + "astrbot_version": VERSION, + "components": ["database"], # 声明了 database 但条目缺失 + "checksums": {}, + "component_checksums": {}, + "directories": [], + } + with zipfile.ZipFile(zip_path, "w") as zf: + zf.writestr("manifest.json", json.dumps(manifest)) + + backup_service._init_task("t1", "import") + await backup_service.background_import_task("t1", str(zip_path)) + + progress = backup_service.get_progress("t1") + assert progress["status"] == "failed" + # 完整 result 保留：errors 说明中止原因 + assert progress["result"] is not None + assert any("missing entries" in e for e in progress["result"]["errors"]) + assert progress["error"] + + +class TestPreVerifyEdgeCases: + """预检边界情况回归测试（评审修正 2/3/4/5）""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "bad_entry", ["files/attachments/a.bin", "config/cmd_config.json"] + ) + async def test_deflate_error_classified_by_component(self, tmp_path, bad_entry): + """Preserve existing files and apply the component policy to broken DEFLATE.""" + entries = { + "files/attachments/a.bin": b"attachment data" * 20, + "config/cmd_config.json": b'{"restored": true}', + } + checksums = {name: _sha256(content) for name, content in entries.items()} + manifest = { + "version": "1.2", + "astrbot_version": VERSION, + "components": ["attachments", "cmd_config"], + "checksums": checksums, + "component_checksums": { + "attachments": _component_checksum( + {"files/attachments/a.bin": checksums["files/attachments/a.bin"]} + ), + "cmd_config": _component_checksum( + {"config/cmd_config.json": checksums["config/cmd_config.json"]} + ), + }, + } + zip_path = tmp_path / "deflate.zip" + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: + zf.writestr("manifest.json", json.dumps(manifest)) + for name, content in entries.items(): + zf.writestr(name, content) + header_offset = zf.getinfo(bad_entry).header_offset + + archive = bytearray(zip_path.read_bytes()) + name_size, extra_size = struct.unpack_from(" 32 * 1024 * 1024 + tracemalloc.start() + try: + data = resources.read_backup_json(archive, "databases/main_db.json") + assert isinstance(data, resources.BackupTableStream) + assert sum(1 for _ in data.get("conversations")) == count + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert peak < 8 * 1024 * 1024 + + +@pytest.mark.parametrize( + "name,limit", + [ + ("databases/main_db.json", "MAX_DATABASE_JSON_BYTES"), + ("databases/kb_metadata.json", "MAX_DATABASE_JSON_BYTES"), + ("databases/kb_test/documents.json", "MAX_KB_DOCUMENT_JSON_BYTES"), + ], +) +def test_streaming_entries_use_their_own_limits(tmp_path, monkeypatch, name, limit): + path = tmp_path / "typed.zip" + content = json.dumps({"documents": [{"text": "x" * 100}]}) + with zipfile.ZipFile(path, "w") as archive: + archive.writestr(name, content) + monkeypatch.setattr(resources, "MAX_JSON_BYTES", 16) + monkeypatch.setattr(resources, limit, 1024) + with resources.open_backup(path) as archive: + data = resources.read_backup_json(archive, name) + assert list(data.get("documents")) == [{"text": "x" * 100}] + monkeypatch.setattr(resources, limit, 32) + with pytest.raises(ValueError, match="size limit"): + resources.read_backup_json(archive, name) + + +def test_record_scanner_handles_escapes_and_utf8_across_chunks(): + data = {"rows": [{"text": '汉字😀\\"[]{}\\\\end', "nested": {"x": [1, 2]}}]} + source = io.BytesIO(json.dumps(data, ensure_ascii=False).encode()) + reader = resources._RecordLimitedReader(source) + assert list(ijson.items(reader, "rows.item", buf_size=3)) == data["rows"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend_name", sorted({"python", ijson.backend})) +@pytest.mark.parametrize("case", ["string", "record", "depth", "scalar", "truncated"]) +async def test_bad_stream_fails_before_any_database_write( + tmp_path, monkeypatch, case, backend_name +): + monkeypatch.setattr( + resources.ijson, + "basic_parse_coro", + ijson.get_backend(backend_name).basic_parse_coro, + ) + path = tmp_path / "bad.zip" + if case == "string": + payload = b'{"conversations":[{"text":"' + b"x" * 2048 + monkeypatch.setattr(resources, "MAX_JSON_RECORD_BYTES", 256) + elif case == "record": + payload = json.dumps( + {"conversations": [{str(i): "x" * 30 for i in range(30)}]} + ).encode() + monkeypatch.setattr(resources, "MAX_JSON_RECORD_BYTES", 256) + elif case == "depth": + payload = b'{"conversations":[{"text":' + b"[" * 80 + b"0" + b"]" * 80 + b"}]}" + elif case == "scalar": + payload = b'{"conversations":' + b"1" * 2048 + monkeypatch.setattr(resources, "MAX_JSON_RECORD_BYTES", 256) + else: + good = {"conversation_id": "one", "platform_id": "p", "user_id": "u"} + payload = b'{"conversations":[' + json.dumps(good).encode() + b"," + with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as archive: + archive.writestr("manifest.json", json.dumps({"astrbot_version": VERSION})) + archive.writestr("databases/main_db.json", payload) + db = MagicMock() + result = await AstrBotImporter(db).import_all(str(path), components=["database"]) + assert not result.success + db.get_db.assert_not_called() + if case != "truncated": + assert any("limit" in error for error in result.errors) + + +def test_row_batches_are_limited_by_bytes_as_well_as_count(monkeypatch): + monkeypatch.setattr(resources, "MAX_JSON_RECORD_BYTES", 128) + rows = [{"text": "x" * 70} for _ in range(3)] + batches = list(resources.backup_row_batches(iter(rows))) + assert [len(batch) for batch in batches] == [1, 1, 1] + assert [row for batch in batches for row in batch] == rows + + +def test_streamed_statistics_merge_matches_legacy_merge(): + rows = [ + { + "id": 1, + "timestamp": "2025-01-01T00:00:00Z", + "platform_id": "p", + "platform_type": "web", + "count": 2, + }, + { + "id": 2, + "timestamp": "2025-01-02T00:00:00Z", + "platform_id": "p", + "platform_type": "web", + "count": "3", + }, + { + "id": 3, + "timestamp": "2025-01-01T00:00:00+00:00", + "platform_id": "p", + "platform_type": "web", + "count": 4, + }, + { + "id": 4, + "timestamp": "", + "platform_id": "p", + "platform_type": "web", + "count": 1, + }, + { + "id": 5, + "timestamp": "", + "platform_id": "p", + "platform_type": "web", + "count": 1, + }, + ] + importer = AstrBotImporter(MagicMock()) + assert list( + importer._preprocess_main_table_rows("platform_stats", iter(rows)) + ) == importer._merge_platform_stats_rows(rows) + + +@pytest.mark.asyncio +async def test_streamed_export_restore_roundtrip(tmp_path, monkeypatch): + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'roundtrip.sqlite'}") + sessions = async_sessionmaker(engine, expire_on_commit=False) + db = SimpleNamespace(get_db=sessions) + models = {"conversations": ConversationV2} + monkeypatch.setattr("astrbot.core.backup.exporter.MAIN_DB_MODELS", models) + monkeypatch.setattr("astrbot.core.backup.importer.MAIN_DB_MODELS", models) + try: + async with engine.begin() as conn: + await conn.run_sync(ConversationV2.__table__.create) + async with sessions.begin() as session: + session.add_all( + ConversationV2( + conversation_id=str(i), + platform_id="p", + user_id="u", + content=[{"role": "user", "content": "汉字😀"}], + ) + for i in range(1001) + ) + exporter = AstrBotExporter(db) + path = await exporter.export_all(str(tmp_path), components=["database"]) + async with sessions.begin() as session: + await session.execute(ConversationV2.__table__.delete()) + result = await AstrBotImporter(db).import_all(path, components=["database"]) + assert result.success, result.errors + assert result.imported_tables["conversations"] == 1001 + async with sessions() as session: + records = (await session.execute(select(ConversationV2))).scalars().all() + assert len(records) == 1001 + assert records[0].content == [{"role": "user", "content": "汉字😀"}] + finally: + await engine.dispose() + + +@pytest.mark.parametrize("limit", ["MAX_DIRECTORY_BYTES", "MAX_ENTRIES"]) +def test_directory_limit_checked_before_zip_allocation(tmp_path, monkeypatch, limit): + path = tmp_path / "backup.zip" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("manifest.json", "{}") + monkeypatch.setattr(resources, limit, 0) + parse = MagicMock(side_effect=AssertionError("ZIP metadata must not be allocated")) + monkeypatch.setattr(zipfile.ZipFile, "_RealGetContents", parse) + with pytest.raises(ValueError, match="directory"): + with resources.open_backup(path): + pass + parse.assert_not_called() + + +def test_json_read_is_bounded_even_with_incorrect_size(monkeypatch): + monkeypatch.setattr(resources, "MAX_JSON_BYTES", 64) + archive = MagicMock() + archive.getinfo.return_value.file_size = 0 + source = io.BytesIO(b" " * 65) + archive.open.return_value = source + with pytest.raises(ValueError, match="size limit"): + resources.read_backup_json(archive, "config/cmd_config.json") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit", ["manifest", "entry", "expanded"]) +async def test_resource_limits_fail_before_restore(tmp_path, monkeypatch, limit): + path = tmp_path / "backup.zip" + config = tmp_path / "config.json" + config.write_text('{"old": true}') + manifest = {"version": "1.1", "astrbot_version": VERSION} + with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as archive: + archive.writestr("manifest.json", json.dumps(manifest)) + archive.writestr("databases/main_db.json", '{"conversations": []}') + archive.writestr("config/cmd_config.json", '{"new": true}') + if limit == "manifest": + monkeypatch.setattr(resources, "MAX_MANIFEST_BYTES", 8) + elif limit == "entry": + monkeypatch.setattr(resources, "MAX_DATABASE_JSON_BYTES", 8) + else: + monkeypatch.setattr(resources, "MAX_EXTRACTED_BYTES", 8) + db = MagicMock() + importer = AstrBotImporter(db, config_path=str(config)) + result = await importer.import_all(str(path)) + assert not result.success + assert any("limit" in error for error in result.errors) + db.get_db.assert_not_called() + assert config.read_text() == '{"old": true}' + + +def test_manifest_limit_applies_to_listing_and_precheck(tmp_path, monkeypatch): + path = tmp_path / "backup.zip" + with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as archive: + archive.writestr("manifest.json", json.dumps({"padding": "x" * 1000})) + monkeypatch.setattr(resources, "MAX_MANIFEST_BYTES", 100) + service = BackupService(MagicMock(), MagicMock()) + assert service.get_backup_manifest(str(path)) is None + check = AstrBotImporter(MagicMock()).pre_check(str(path)) + assert not check.can_import + assert "size limit" in check.error + + +@pytest.mark.asyncio +async def test_attachment_hints_do_not_read_oversized_unselected_database( + tmp_path, monkeypatch +): + path = tmp_path / "backup.zip" + with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as archive: + archive.writestr("manifest.json", json.dumps({"astrbot_version": VERSION})) + archive.writestr("databases/main_db.json", json.dumps({"padding": "x" * 1000})) + archive.writestr("files/attachments/test.txt", "attachment") + monkeypatch.setattr(resources, "MAX_DATABASE_JSON_BYTES", 100) + importer = AstrBotImporter(MagicMock()) + importer._import_attachments = AsyncMock(return_value=1) + result = await importer.import_all(str(path), components=["attachments"]) + assert result.success + assert any("path hints skipped" in warning for warning in result.warnings) + assert importer._import_attachments.call_args.args[1] == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("first", ["import", "export"]) +@pytest.mark.parametrize("status", ["pending", "processing"]) +async def test_busy_task_rejects_both_operations(tmp_path, first, status): + service = BackupService(MagicMock(), MagicMock()) + service.backup_dir = str(tmp_path) + (tmp_path / "backup.zip").touch() + service._init_task("active", first, status) + with pytest.raises(BackupServiceError, match="正在运行"): + service.export_backup() + with pytest.raises(BackupServiceError, match="正在运行"): + service.import_backup({"filename": "backup.zip", "confirmed": True}) + assert list(service.backup_tasks) == ["active"] + service._set_task_result("active", "failed") + service._init_task("next", "export") + assert service.backup_tasks["next"]["status"] == "pending" + + +@pytest.mark.asyncio +async def test_export_worker_yields_and_cancellation_waits_for_writer(tmp_path): + exporter = AstrBotExporter(MagicMock()) + entered = threading.Event() + release = threading.Event() + finished = threading.Event() + path = tmp_path / "output.zip" + + def write(archive): + entered.set() + assert release.wait(5) + archive.writestr("entry", "complete") + finished.set() + + with zipfile.ZipFile(path, "w") as archive: + task = asyncio.create_task(exporter._run_io(write, archive)) + try: + for _ in range(100): + if entered.is_set(): + break + await asyncio.sleep(0.01) + assert entered.is_set() + task.cancel() + await asyncio.sleep(0.02) + assert not task.done() + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert finished.is_set() + with zipfile.ZipFile(path) as archive: + assert archive.read("entry") == b"complete" + + +@pytest.mark.asyncio +async def test_export_json_limit_removes_partial_archive(tmp_path, monkeypatch): + exporter = AstrBotExporter(MagicMock()) + exporter._export_main_database = AsyncMock( + return_value={"conversations": [{"data": "x" * 100}]} + ) + monkeypatch.setattr(resources, "MAX_DATABASE_JSON_BYTES", 50) + with pytest.raises(ValueError, match="size limit"): + await exporter.export_all(str(tmp_path), components=["database"]) + assert list(tmp_path.glob("*.zip")) == [] + + +@pytest.mark.asyncio +async def test_export_archive_limit_removes_unrestorable_backup(tmp_path, monkeypatch): + exporter = AstrBotExporter(MagicMock()) + exporter._export_main_database = AsyncMock(return_value={"conversations": []}) + monkeypatch.setattr(resources, "MAX_ENTRIES", 1) + with pytest.raises(ValueError, match="resource limit"): + await exporter.export_all(str(tmp_path), components=["database"]) + assert list(tmp_path.glob("*.zip")) == [] + + +@pytest.mark.asyncio +async def test_batched_flush_preserves_whole_restore_transaction(tmp_path, monkeypatch): + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'db.sqlite'}") + sessions = async_sessionmaker(engine, expire_on_commit=False) + try: + async with engine.begin() as conn: + await conn.run_sync(ConversationV2.__table__.create) + async with sessions.begin() as session: + session.add( + ConversationV2(conversation_id="old", platform_id="p", user_id="u") + ) + monkeypatch.setattr( + "astrbot.core.backup.importer.MAIN_DB_MODELS", + {"conversations": ConversationV2}, + ) + session = sessions() + counts = [] + event.listen( + session.sync_session, + "before_flush", + lambda current, *_: counts.append(len(current.new)), + ) + importer = AstrBotImporter(SimpleNamespace(get_db=lambda: session)) + rows = [ + {"conversation_id": str(i), "platform_id": "p", "user_id": "u"} + for i in range(501) + ] + rows.append(rows[0].copy()) + with pytest.raises(IntegrityError): + await importer._import_main_database({"conversations": rows}, clear=True) + assert counts == [500, 2] + async with sessions() as session: + ids = ( + (await session.execute(select(ConversationV2.conversation_id))) + .scalars() + .all() + ) + assert ids == ["old"] + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_database_export_streams_rows_and_converts_off_loop( + tmp_path, monkeypatch +): + def reject_orm_json(raw): + raise AssertionError("JSON must not be decoded by the ORM on the event loop") + + engine = create_async_engine( + f"sqlite+aiosqlite:///{tmp_path / 'export.sqlite'}", + json_deserializer=reject_orm_json, + ) + sessions = async_sessionmaker(engine, expire_on_commit=False) + exporter = AstrBotExporter(MagicMock()) + convert = exporter._model_to_dict + thread_ids = set() + decode_threads = set() + loads = json.loads + + def track_decode(raw, *args, **kwargs): + if (isinstance(raw, str) and "export-thread" in raw) or ( + isinstance(raw, bytes) and b"export-thread" in raw + ): + decode_threads.add(threading.get_ident()) + return loads(raw, *args, **kwargs) + + monkeypatch.setattr(json, "loads", track_decode) + monkeypatch.setattr("astrbot.core.backup.exporter.MAX_JSON_RECORD_BYTES", 4096) + + def track_conversion(record): + thread_ids.add(threading.get_ident()) + return convert(record) + + exporter._model_to_dict = track_conversion + try: + async with engine.begin() as conn: + await conn.run_sync(ConversationV2.__table__.create) + async with sessions.begin() as session: + session.add_all( + ConversationV2( + conversation_id=str(i), + platform_id="p", + user_id="u", + content=[{"marker": "export-thread", "text": "x" * 500}], + ) + for i in range(1001) + ) + batches = [ + batch + async for batch in exporter._export_records( + sessions, ConversationV2, exporter._model_to_dict + ) + ] + rows = [row for batch in batches for row in batch] + assert max(map(len, batches)) < 500 + assert {row["conversation_id"] for row in rows} == {str(i) for i in range(1001)} + assert thread_ids + assert threading.get_ident() not in thread_ids + assert decode_threads and threading.get_ident() not in decode_threads + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_kb_document_restore_uses_bounded_batches(tmp_path, monkeypatch): + storage = MagicMock() + storage.initialize = AsyncMock() + storage.insert_documents_batch = AsyncMock() + storage.close = AsyncMock() + monkeypatch.setattr( + "astrbot.core.db.vec_db.faiss_impl.document_storage.DocumentStorage", + MagicMock(return_value=storage), + ) + importer = AstrBotImporter(MagicMock()) + importer.kb_root_dir = str(tmp_path) + documents = [ + {"doc_id": str(i), "text": "text", "metadata": "{}"} for i in range(1001) + ] + await importer._import_kb_documents("kb", {"documents": documents}) + assert [ + len(call.kwargs["doc_ids"]) + for call in storage.insert_documents_batch.await_args_list + ] == [500, 500, 1] + storage.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_streamed_kb_documents_preserve_text_and_metadata(tmp_path): + from astrbot.core.db.vec_db.faiss_impl.document_storage import DocumentStorage + + source = DocumentStorage(str(tmp_path / "source.sqlite")) + restored = None + try: + await source.initialize() + await source.insert_documents_batch( + doc_ids=["one", "two"], + texts=["汉字😀\x00tail", "second"], + metadatas=[{"source": "文档"}, {"n": 2}], + ) + exporter = AstrBotExporter(MagicMock()) + helper = SimpleNamespace(vec_db=SimpleNamespace(document_storage=source)) + data = await exporter._export_kb_documents(helper) + path = tmp_path / "kb.zip" + name = "databases/kb_test/documents.json" + with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as archive: + await exporter._write_table_dump(archive, name, data) + importer = AstrBotImporter(MagicMock()) + importer.kb_root_dir = str(tmp_path / "restored") + (tmp_path / "restored" / "test").mkdir(parents=True) + with resources.open_backup(path) as archive: + document_stream = resources.read_backup_json(archive, name) + await importer._import_kb_documents("test", document_stream) + restored = DocumentStorage(str(tmp_path / "restored" / "test" / "doc.db")) + await restored.initialize() + docs = await restored.get_documents({}, limit=None) + assert {doc["doc_id"]: doc["text"] for doc in docs} == { + "one": "汉字😀\x00tail", + "two": "second", + } + assert {doc["doc_id"]: json.loads(doc["metadata"]) for doc in docs} == { + "one": {"source": "文档"}, + "two": {"n": 2}, + } + finally: + await source.close() + if restored is not None: + await restored.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("backend_name", sorted({"python", ijson.backend})) +async def test_compact_record_and_large_integer_export_restore( + tmp_path, monkeypatch, backend_name +): + monkeypatch.setattr( + resources.ijson, + "basic_parse_coro", + ijson.get_backend(backend_name).basic_parse_coro, + ) + monkeypatch.setattr(resources, "MAX_JSON_RECORD_BYTES", 2048) + monkeypatch.setattr("astrbot.core.backup.exporter.MAX_JSON_RECORD_BYTES", 2048) + row = { + "conversation_id": "x", + "platform_id": "p", + "user_id": "u", + "content": [2**80, 1.25] + [0] * 800, + } + assert len(json.dumps(row)) > 2048 + assert len(json.dumps(row, separators=(",", ":"))) < 2048 + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'compact.sqlite'}") + sessions = async_sessionmaker(engine, expire_on_commit=False) + monkeypatch.setattr( + "astrbot.core.backup.importer.MAIN_DB_MODELS", {"conversations": ConversationV2} + ) + monkeypatch.setattr( + "astrbot.core.backup.exporter.MAIN_DB_MODELS", {"conversations": ConversationV2} + ) + try: + async with engine.begin() as conn: + await conn.run_sync(ConversationV2.__table__.create) + async with sessions.begin() as session: + session.add(ConversationV2(**row)) + path = await AstrBotExporter(SimpleNamespace(get_db=sessions)).export_all( + str(tmp_path), components=["database"] + ) + result = await AstrBotImporter(SimpleNamespace(get_db=sessions)).import_all( + path, components=["database"] + ) + assert result.success, result.errors + async with sessions() as session: + restored = (await session.execute(select(ConversationV2))).scalar_one() + assert restored.content == row["content"] + assert type(restored.content[0]) is int + assert type(restored.content[1]) is float + finally: + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", ["depth", "nonfinite"]) +async def test_export_rejects_unreadable_json_before_success(tmp_path, case): + value = float("nan") + if case == "depth": + value = 0 + for _ in range(70): + value = [value] + exporter = AstrBotExporter(MagicMock()) + exporter._export_main_database = AsyncMock( + return_value={ + "conversations": [ + { + "conversation_id": "x", + "platform_id": "p", + "user_id": "u", + "content": [value], + } + ] + } + ) + with pytest.raises(ValueError): + await exporter.export_all(str(tmp_path), components=["database"]) + assert not list(tmp_path.glob("*.zip")) + assert exporter.exported_components == [] + + +@pytest.mark.asyncio +async def test_kb_clear_failure_rolls_back_before_removing_files(tmp_path, monkeypatch): + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'kb.sqlite'}") + sessions = async_sessionmaker(engine, expire_on_commit=False) + kb_dir = tmp_path / "old-kb" + kb_dir.mkdir() + original = kb_dir / "original.txt" + original.write_text("keep") + helper = SimpleNamespace(kb_dir=kb_dir, terminate=AsyncMock()) + manager = SimpleNamespace( + kb_db=SimpleNamespace(get_db=sessions), + kb_insts={"old": helper}, + load_kbs=AsyncMock(), + ) + monkeypatch.setattr( + "astrbot.core.backup.importer.KB_METADATA_MODELS", + {"knowledge_bases": KnowledgeBase, "missing_table": ConversationV2}, + ) + try: + async with engine.begin() as conn: + await conn.run_sync(KnowledgeBase.__table__.create) + async with sessions.begin() as session: + session.add(KnowledgeBase(kb_id="old", kb_name="original")) + result = ImportResult() + with zipfile.ZipFile(tmp_path / "kb.zip", "w") as archive: + with pytest.raises(DatabaseClearError): + await AstrBotImporter(MagicMock(), manager)._import_knowledge_bases( + archive, + {"knowledge_bases": [{"kb_id": "new", "kb_name": "new"}]}, + result, + clear=True, + ) + async with sessions() as session: + assert ( + await session.execute(select(KnowledgeBase.kb_id)) + ).scalars().all() == ["old"] + assert original.read_text() == "keep" + assert result.imported_tables == {} + helper.terminate.assert_not_awaited() + manager.load_kbs.assert_not_awaited() + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_database_attachment_restore_parses_main_dump_only_twice( + tmp_path, monkeypatch +): + engine = create_async_engine( + f"sqlite+aiosqlite:///{tmp_path / 'attachments.sqlite'}" + ) + sessions = async_sessionmaker(engine, expire_on_commit=False) + models = {"conversations": ConversationV2, "attachments": Attachment} + monkeypatch.setattr("astrbot.core.backup.importer.MAIN_DB_MODELS", models) + destination = tmp_path / "attachments" / "nested" / "original.txt" + row = { + "attachment_id": "att", + "path": str(destination), + "type": "file", + "mime_type": "text/plain", + } + path = tmp_path / "restore.zip" + with zipfile.ZipFile(path, "w") as archive: + archive.writestr("manifest.json", json.dumps({"astrbot_version": VERSION})) + archive.writestr( + "databases/main_db.json", + json.dumps( + { + "conversations": [ + {"conversation_id": "x", "platform_id": "p", "user_id": "u"} + ], + "attachments": [row], + } + ), + ) + archive.writestr("files/attachments/att.txt", b"restored") + passes = [] + records = resources.BackupTableStream._records + + def count_passes(stream): + passes.append(stream.name) + yield from records(stream) + + monkeypatch.setattr(resources.BackupTableStream, "_records", count_passes) + try: + async with engine.begin() as conn: + for model in models.values(): + await conn.run_sync(model.__table__.create) + result = await AstrBotImporter( + SimpleNamespace(get_db=sessions), + config_path=str(tmp_path / "cmd_config.json"), + ).import_all(str(path), components=["database", "attachments"]) + assert result.success, result.errors + assert passes == ["databases/main_db.json"] * 2 + assert destination.read_bytes() == b"restored" + assert result.imported_files["attachments"] == 1 + finally: + await engine.dispose() + + +def test_manifest_kb_inventory_uses_only_written_entries(): + manager = SimpleNamespace(kb_insts={"not-exported": MagicMock()}) + exporter = AstrBotExporter(MagicMock(), manager) + exporter._checksums = { + "databases/kb_written/documents.json": "sha256:docs", + "files/kb_media/written/images/saved.png": "sha256:media", + } + manifest = exporter._generate_manifest({}, {}) + assert manifest["tables"]["kb_documents"] == {"written": "documents"} + assert manifest["files"]["kb_media"] == {"written": ["saved.png"]} + + +@pytest.mark.parametrize("backend_name", sorted({"python", ijson.backend})) +@pytest.mark.parametrize( + "value", + [ + None, + True, + 2**80, + -1.25e-8, + '汉字😀\x00\\"', + [[], {}, 1, {"x": [False, None, "汉字"]}], + {"nested": {"a": 1, "b": [1, 2]}, "empty": {}}, + ], +) +def test_raw_json_check_matches_compact_encoding(monkeypatch, backend_name, value): + monkeypatch.setattr( + resources.ijson, "basic_parse", ijson.get_backend(backend_name).basic_parse + ) + raw = json.dumps(value, ensure_ascii=True, indent=2).encode() + expected = len( + json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode() + ) + assert resources.check_backup_json_field(raw, expected) == expected + with pytest.raises(ValueError, match="size limit"): + resources.check_backup_json_field(raw, expected - 1) + + +@pytest.mark.asyncio +async def test_export_restore_preserves_nul_in_text_and_autostring( + tmp_path, monkeypatch +): + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'nul.sqlite'}") + sessions = async_sessionmaker(engine, expire_on_commit=False) + for module in ("exporter", "importer"): + monkeypatch.setattr( + f"astrbot.core.backup.{module}.MAIN_DB_MODELS", {"personas": Persona} + ) + prompt = "before\x00汉字😀after" + persona_id = "id\x00suffix" + try: + async with engine.begin() as conn: + await conn.run_sync(Persona.__table__.create) + async with sessions.begin() as session: + session.add(Persona(persona_id=persona_id, system_prompt=prompt)) + db = SimpleNamespace(get_db=sessions) + path = await AstrBotExporter(db).export_all( + str(tmp_path), components=["database"] + ) + result = await AstrBotImporter(db).import_all(path, components=["database"]) + assert result.success, result.errors + async with sessions() as session: + restored = (await session.execute(select(Persona))).scalar_one() + assert restored.system_prompt == prompt + assert restored.persona_id == persona_id + finally: + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", ["size", "combined", "depth"]) +async def test_export_checks_json_before_materializing_containers( + tmp_path, monkeypatch, case +): + limit = 256 * 1024 + monkeypatch.setattr("astrbot.core.backup.exporter.MAX_JSON_RECORD_BYTES", limit) + monkeypatch.setattr( + "astrbot.core.backup.exporter.MAIN_DB_MODELS", {"personas": Persona} + ) + if case == "depth": + first, second = "[" * 70 + "0" + "]" * 70, "null" + else: + count = 200000 if case == "size" else 50000 + first = "[" + ", ".join(["{}"] * count) + "]" + second = first if case == "combined" else "null" + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'limited.sqlite'}") + sessions = async_sessionmaker(engine, expire_on_commit=False) + decoded = [] + loads = json.loads + + def track_decode(raw, *args, **kwargs): + if isinstance(raw, bytes) and raw.startswith(b"["): + decoded.append(len(raw)) + return loads(raw, *args, **kwargs) + + try: + async with engine.begin() as conn: + await conn.run_sync(Persona.__table__.create) + async with sessions.begin() as session: + session.add(Persona(persona_id="x", system_prompt="prompt")) + async with sessions.begin() as session: + await session.execute( + text("UPDATE personas SET begin_dialogs=:first, tools=:second"), + {"first": first, "second": second}, + ) + monkeypatch.setattr(json, "loads", track_decode) + tracemalloc.start() + try: + with pytest.raises(ValueError, match="limit"): + await AstrBotExporter(SimpleNamespace(get_db=sessions)).export_all( + str(tmp_path), components=["database"] + ) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert decoded == [] + assert peak < 8 * 1024 * 1024 + assert not list(tmp_path.glob("*.zip")) + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_raw_autostring_fetch_is_bounded(tmp_path, monkeypatch): + monkeypatch.setattr("astrbot.core.backup.exporter.MAX_JSON_RECORD_BYTES", 256) + engine = create_async_engine( + f"sqlite+aiosqlite:///{tmp_path / 'autostring.sqlite'}" + ) + sessions = async_sessionmaker(engine, expire_on_commit=False) + # Inspect values delivered by the query, independently of the exporter's + # later size check. AutoString must not bypass the SQL read bound. + fetched = [] + try: + async with engine.begin() as conn: + await conn.run_sync(ConversationV2.__table__.create) + async with sessions.begin() as session: + session.add( + ConversationV2( + conversation_id="x", platform_id="p" * 10000, user_id="u" + ) + ) + session = sessions() + stream = session.stream + + async def capture(statement, *args, **kwargs): + async with sessions() as probe: + result = await probe.stream(statement) + try: + fetched.extend(await result.mappings().all()) + finally: + await result.close() + return await stream(statement, *args, **kwargs) + + monkeypatch.setattr(session, "stream", capture) + exporter = AstrBotExporter(MagicMock()) + with pytest.raises(ValueError, match="size limit"): + async for _ in exporter._export_records( + lambda: session, ConversationV2, exporter._model_to_dict + ): + pass + assert len(fetched[0]["platform_id"]) == 257 + assert isinstance(fetched[0]["platform_id"], bytes) + finally: + await engine.dispose()