259 lines
9.2 KiB
Python
259 lines
9.2 KiB
Python
"""外部系统大响应的 task 私有缓存与结构化分段读取。"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import re
|
||
import time
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
from uuid import uuid4
|
||
|
||
from core.file_store import atomic_write_text
|
||
|
||
RESULT_CACHE_SUBDIR = ".zcbot_external_results"
|
||
RESULT_TTL_SECONDS = 24 * 60 * 60
|
||
MAX_STORED_RESULT_BYTES = 10 * 1024 * 1024
|
||
MAX_TASK_CACHE_BYTES = 50 * 1024 * 1024
|
||
MAX_USER_CACHE_BYTES = 200 * 1024 * 1024
|
||
_REF_RE = re.compile(r"^extres_[0-9a-f]{32}$")
|
||
_PREVIEW_MAX_DEPTH = 6
|
||
_PREVIEW_MAX_NODES = 120
|
||
_PREVIEW_MAX_DICT_KEYS = 30
|
||
_PREVIEW_MAX_LIST_ITEMS = 5
|
||
_PREVIEW_MAX_STRING_CHARS = 500
|
||
|
||
|
||
class ExternalResultError(RuntimeError):
|
||
pass
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class StoredExternalResult:
|
||
result_ref: str
|
||
original_bytes: int
|
||
expires_in_seconds: int
|
||
|
||
|
||
def _json_bytes(value: Any) -> bytes:
|
||
return json.dumps(value, ensure_ascii=False, default=str).encode("utf-8")
|
||
|
||
|
||
def _pointer_escape(value: str) -> str:
|
||
return value.replace("~", "~0").replace("/", "~1")
|
||
|
||
|
||
def build_result_preview(value: Any) -> tuple[Any, list[dict[str, Any]]]:
|
||
"""生成合法、定界的 JSON 预览，并列出可继续分段读取的位置。"""
|
||
reads: list[dict[str, Any]] = []
|
||
budget = [_PREVIEW_MAX_NODES]
|
||
|
||
def visit(item: Any, pointer: str, depth: int) -> Any:
|
||
if budget[0] <= 0:
|
||
reads.append({"json_pointer": pointer, "reason": "node_budget"})
|
||
return {"$preview_truncated": True}
|
||
budget[0] -= 1
|
||
if depth > _PREVIEW_MAX_DEPTH:
|
||
reads.append({"json_pointer": pointer, "reason": "depth"})
|
||
return {"$preview_truncated": True}
|
||
if isinstance(item, dict):
|
||
pairs = list(item.items())
|
||
shown = pairs[:_PREVIEW_MAX_DICT_KEYS]
|
||
result: Any = {
|
||
str(key): visit(
|
||
child,
|
||
pointer + "/" + _pointer_escape(str(key)),
|
||
depth + 1,
|
||
)
|
||
for key, child in shown
|
||
}
|
||
if len(pairs) > len(shown):
|
||
result["$preview_omitted_keys"] = len(pairs) - len(shown)
|
||
reads.append({
|
||
"json_pointer": pointer,
|
||
"reason": "keys",
|
||
"total_keys": len(pairs),
|
||
})
|
||
return result
|
||
if isinstance(item, list):
|
||
shown = item[:_PREVIEW_MAX_LIST_ITEMS]
|
||
result = [
|
||
visit(child, pointer + "/" + str(index), depth + 1)
|
||
for index, child in enumerate(shown)
|
||
]
|
||
if len(item) > len(shown):
|
||
reads.append({
|
||
"json_pointer": pointer,
|
||
"reason": "items",
|
||
"total_items": len(item),
|
||
"next_offset": len(shown),
|
||
})
|
||
return result
|
||
if isinstance(item, str) and len(item) > _PREVIEW_MAX_STRING_CHARS:
|
||
reads.append({
|
||
"json_pointer": pointer,
|
||
"reason": "string",
|
||
"total_chars": len(item),
|
||
})
|
||
return item[:_PREVIEW_MAX_STRING_CHARS] + "...[preview]"
|
||
return item
|
||
|
||
return visit(value, "", 0), reads
|
||
|
||
|
||
def resolve_json_pointer(value: Any, pointer: str) -> Any:
|
||
"""解析 RFC 6901 JSON Pointer；空串表示完整根值。"""
|
||
pointer = (pointer or "").strip()
|
||
if not pointer:
|
||
return value
|
||
if not pointer.startswith("/"):
|
||
raise ExternalResultError("json_pointer 必须为空或以 / 开头")
|
||
current = value
|
||
for raw_part in pointer[1:].split("/"):
|
||
part = raw_part.replace("~1", "/").replace("~0", "~")
|
||
if isinstance(current, dict):
|
||
if part not in current:
|
||
raise ExternalResultError(f"json_pointer 不存在: {pointer}")
|
||
current = current[part]
|
||
elif isinstance(current, list):
|
||
try:
|
||
if not part.isdigit():
|
||
raise ValueError(part)
|
||
index = int(part)
|
||
current = current[index]
|
||
except (ValueError, IndexError) as exc:
|
||
raise ExternalResultError(f"json_pointer 不存在: {pointer}") from exc
|
||
else:
|
||
raise ExternalResultError(f"json_pointer 不存在: {pointer}")
|
||
return current
|
||
|
||
|
||
class ExternalResultStore:
|
||
def __init__(self, user_root: Path, task_id: str):
|
||
self.cache_root = Path(user_root) / RESULT_CACHE_SUBDIR
|
||
self.root = self.cache_root / str(task_id)
|
||
|
||
def _path(self, result_ref: str) -> Path:
|
||
if not _REF_RE.fullmatch(result_ref or ""):
|
||
raise ExternalResultError("result_ref 无效")
|
||
return self.root / f"{result_ref}.json"
|
||
|
||
def _sweep(self) -> None:
|
||
if not self.cache_root.is_dir():
|
||
return
|
||
cutoff = time.time() - RESULT_TTL_SECONDS
|
||
for path in self.cache_root.glob("*/extres_*.json"):
|
||
try:
|
||
if path.stat().st_mtime < cutoff:
|
||
path.unlink()
|
||
try:
|
||
path.parent.rmdir()
|
||
except OSError:
|
||
pass
|
||
except OSError:
|
||
continue
|
||
|
||
def store(self, system_id: str, result: dict[str, Any]) -> StoredExternalResult:
|
||
encoded = _json_bytes(result)
|
||
if len(encoded) > MAX_STORED_RESULT_BYTES:
|
||
raise ExternalResultError(
|
||
f"外部系统响应超过安全保存上限({MAX_STORED_RESULT_BYTES} bytes)"
|
||
)
|
||
self.root.mkdir(parents=True, exist_ok=True)
|
||
self._sweep()
|
||
existing = sorted(
|
||
self.root.glob("extres_*.json"),
|
||
key=lambda path: path.stat().st_mtime,
|
||
)
|
||
total = sum(path.stat().st_size for path in existing)
|
||
while existing and total + len(encoded) > MAX_TASK_CACHE_BYTES:
|
||
oldest = existing.pop(0)
|
||
try:
|
||
size = oldest.stat().st_size
|
||
oldest.unlink()
|
||
total -= size
|
||
except OSError:
|
||
continue
|
||
if total + len(encoded) > MAX_TASK_CACHE_BYTES:
|
||
raise ExternalResultError("当前任务的外部结果缓存已达上限")
|
||
user_files = sorted(
|
||
self.cache_root.glob("*/extres_*.json"),
|
||
key=lambda path: path.stat().st_mtime,
|
||
)
|
||
user_total = sum(path.stat().st_size for path in user_files)
|
||
while user_files and user_total + len(encoded) > MAX_USER_CACHE_BYTES:
|
||
oldest = user_files.pop(0)
|
||
try:
|
||
size = oldest.stat().st_size
|
||
oldest.unlink()
|
||
user_total -= size
|
||
except OSError:
|
||
continue
|
||
if user_total + len(encoded) > MAX_USER_CACHE_BYTES:
|
||
raise ExternalResultError("当前用户的外部结果缓存已达上限")
|
||
result_ref = f"extres_{uuid4().hex}"
|
||
envelope = {
|
||
"version": 1,
|
||
"created_at": int(time.time()),
|
||
"system_id": str(system_id),
|
||
"result": result,
|
||
}
|
||
atomic_write_text(self._path(result_ref), json.dumps(
|
||
envelope, ensure_ascii=False, default=str
|
||
))
|
||
return StoredExternalResult(
|
||
result_ref=result_ref,
|
||
original_bytes=len(encoded),
|
||
expires_in_seconds=RESULT_TTL_SECONDS,
|
||
)
|
||
|
||
def load(self, result_ref: str) -> dict[str, Any]:
|
||
self._sweep()
|
||
path = self._path(result_ref)
|
||
try:
|
||
if time.time() - path.stat().st_mtime > RESULT_TTL_SECONDS:
|
||
path.unlink()
|
||
raise ExternalResultError("result_ref 已过期")
|
||
payload = json.loads(path.read_text(encoding="utf-8"))
|
||
except FileNotFoundError as exc:
|
||
raise ExternalResultError("result_ref 不存在或已过期") from exc
|
||
except (OSError, ValueError) as exc:
|
||
raise ExternalResultError("result_ref 无法读取") from exc
|
||
if not isinstance(payload, dict) or not isinstance(payload.get("result"), dict):
|
||
raise ExternalResultError("result_ref 内容无效")
|
||
return payload
|
||
|
||
|
||
def project_result(
|
||
value: Any,
|
||
*,
|
||
offset: int = 0,
|
||
limit: int = 50,
|
||
fields: list[str] | None = None,
|
||
) -> dict[str, Any]:
|
||
"""对缓存节点分页并可选投影字典字段。"""
|
||
offset = max(0, int(offset))
|
||
limit = max(1, min(int(limit), 200))
|
||
if fields is not None and not isinstance(fields, list):
|
||
raise ExternalResultError("fields 必须是字符串数组")
|
||
wanted = [str(field) for field in (fields or [])[:50]]
|
||
|
||
def select(item: Any) -> Any:
|
||
if not wanted or not isinstance(item, dict):
|
||
return item
|
||
return {field: item.get(field) for field in wanted if field in item}
|
||
|
||
if isinstance(value, list):
|
||
page = [select(item) for item in value[offset : offset + limit]]
|
||
return {
|
||
"type": "array",
|
||
"total_items": len(value),
|
||
"offset": offset,
|
||
"limit": limit,
|
||
"has_more": offset + len(page) < len(value),
|
||
"data": page,
|
||
}
|
||
if isinstance(value, dict):
|
||
return {"type": "object", "data": select(value)}
|
||
return {"type": type(value).__name__, "data": value}
|