118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
"""文件变更追踪 — 基于 SHA256 + mtime 判断文件是否需要重新入库."""
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
from threading import Lock
|
||
from typing import TypedDict
|
||
|
||
logger = logging.getLogger("md-vector-db")
|
||
|
||
|
||
class FileRecord(TypedDict):
|
||
"""追踪记录."""
|
||
|
||
hash: str
|
||
mtime: float
|
||
size: int
|
||
ingested_at: str
|
||
|
||
|
||
class FileTracker:
|
||
"""文件入库追踪器 — JSON 文件持久化.
|
||
|
||
通过比对文件内容的 SHA256 哈希判断文件是否变更。
|
||
使用 threading.Lock 保护 JSON 文件的并发读写。
|
||
"""
|
||
|
||
def __init__(self, db_path: str = "./ingest_tracker.json"):
|
||
self._db_path = Path(db_path)
|
||
self._lock = Lock()
|
||
self._records: dict[str, FileRecord] = {}
|
||
self._load()
|
||
|
||
# -- 公开 API --
|
||
|
||
def is_stale(self, file_path: str) -> bool:
|
||
"""检查文件是否需要重新入库.
|
||
|
||
Returns:
|
||
True: 文件不存在 / 无记录 / 内容已变更
|
||
False: 文件未变更且记录存在
|
||
"""
|
||
path = Path(file_path)
|
||
if not path.exists():
|
||
return True
|
||
record = self.get_record(file_path)
|
||
if record is None:
|
||
return True
|
||
current_hash = self.compute_hash(file_path)
|
||
return current_hash != record["hash"]
|
||
|
||
def get_record(self, file_path: str) -> FileRecord | None:
|
||
"""获取文件的追踪记录（无记录返回 None）."""
|
||
return self._records.get(self._abs_key(file_path))
|
||
|
||
def mark_ingested(self, file_path: str) -> None:
|
||
"""标记文件已入库（创建或更新追踪记录）."""
|
||
path = Path(file_path)
|
||
if not path.exists():
|
||
logger.warning("标记已入库时文件不存在: %s", file_path)
|
||
return
|
||
stat = path.stat()
|
||
record: FileRecord = {
|
||
"hash": self.compute_hash(file_path),
|
||
"mtime": stat.st_mtime,
|
||
"size": stat.st_size,
|
||
"ingested_at": self._now_iso(),
|
||
}
|
||
with self._lock:
|
||
self._records[self._abs_key(file_path)] = record
|
||
self._save()
|
||
|
||
def remove_record(self, file_path: str) -> None:
|
||
"""移除文件的追踪记录."""
|
||
key = self._abs_key(file_path)
|
||
with self._lock:
|
||
if key in self._records:
|
||
del self._records[key]
|
||
self._save()
|
||
|
||
@staticmethod
|
||
def compute_hash(file_path: str) -> str:
|
||
"""计算文件 SHA256 哈希（分块读取，适合大文件）."""
|
||
sha = hashlib.sha256()
|
||
with open(file_path, "rb") as f:
|
||
while chunk := f.read(8192):
|
||
sha.update(chunk)
|
||
return sha.hexdigest()
|
||
|
||
# -- 内部 --
|
||
|
||
def _abs_key(self, file_path: str) -> str:
|
||
"""生成标准化 key — 使用绝对路径."""
|
||
return str(Path(file_path).resolve())
|
||
|
||
@staticmethod
|
||
def _now_iso() -> str:
|
||
return datetime.now(timezone.utc).isoformat()
|
||
|
||
def _load(self) -> None:
|
||
if self._db_path.exists():
|
||
try:
|
||
with open(self._db_path, "r", encoding="utf-8") as f:
|
||
self._records = json.load(f)
|
||
except (json.JSONDecodeError, OSError) as e:
|
||
logger.warning("tracker 文件损坏，重置为空: %s", e)
|
||
self._records = {}
|
||
|
||
def _save(self) -> None:
|
||
try:
|
||
with open(self._db_path, "w", encoding="utf-8") as f:
|
||
json.dump(self._records, f, ensure_ascii=False, indent=2)
|
||
except OSError as e:
|
||
logger.error("无法写入 tracker 文件: %s", e)
|