"""
向量存储迁移和备份工具
提供向量存储系统间的数据迁移、备份恢复和数据验证功能。
支持内存存储、FAISS、Milvus等不同后端之间的无缝迁移。
"""
import hashlib
import json
import logging
import os
import shutil
from datetime import datetime
from typing import Any, Dict, Iterator, List, Optional, Protocol, runtime_checkable
from haystack import Document as HaystackDocument
from haystack.document_stores.in_memory import InMemoryDocumentStore
from document_stores.implementations.base_vector_store import BaseVectorStore, VectorStoreError
from document_stores.vector_store_factory import create_vector_store
@runtime_checkable
class _HaystackDocumentSource(Protocol):
def get_haystack_store(self) -> InMemoryDocumentStore: ...
@runtime_checkable
class _MappedDocumentSource(Protocol):
documents: Dict[str, HaystackDocument]
[docs]
class VectorStoreMigrator:
"""
向量存储迁移工具类。
提供完整的迁移、备份、恢复和验证功能。
"""
[docs]
def __init__(self) -> None:
"""初始化迁移工具"""
self.logger = logging.getLogger("VectorStoreMigrator")
[docs]
def migrate(
self,
source_config: Dict[str, Any],
target_config: Dict[str, Any],
batch_size: int = 1000,
verify_migration: bool = True,
backup_before_migration: bool = True,
) -> bool:
"""
执行向量存储迁移。
Args:
source_config: 源存储配置
target_config: 目标存储配置
batch_size: 批量迁移大小
verify_migration: 是否验证迁移结果
backup_before_migration: 迁移前是否备份
Returns:
迁移成功返回True，失败返回False
"""
try:
self.logger.info("开始向量存储迁移...")
# 创建源和目标存储
source_store = create_vector_store(source_config, auto_fallback=False)
target_store = create_vector_store(target_config, auto_fallback=False)
# 迁移前备份
backup_path = None
if backup_before_migration:
backup_path = self.backup_storage(
source_store, f"backup_before_migration_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
)
if not backup_path:
self.logger.error("迁移前备份失败")
return False
# 执行迁移
success = self._execute_migration(source_store, target_store, batch_size)
if not success:
self.logger.error("迁移执行失败")
return False
# 验证迁移结果
if verify_migration:
if not self.verify_migration(source_store, target_store):
self.logger.error("迁移验证失败")
return False
self.logger.info("向量存储迁移完成")
# 清理资源
source_store.close()
target_store.close()
return True
except Exception as e:
self.logger.error(f"迁移过程发生错误: {str(e)}")
return False
def _execute_migration(self, source_store: BaseVectorStore, target_store: BaseVectorStore, batch_size: int) -> bool:
"""
执行实际的迁移操作。
Args:
source_store: 源存储
target_store: 目标存储
batch_size: 批处理大小
Returns:
迁移成功返回True
"""
total_docs = source_store.get_document_count()
if total_docs == 0:
self.logger.info("源存储为空，无需迁移")
return True
self.logger.info(f"开始迁移 {total_docs} 个文档")
migrated_count = 0
try:
# 批量迁移文档
for batch_docs in self._get_documents_in_batches(source_store, batch_size):
if not batch_docs:
continue
# 添加到目标存储
if not target_store.add_documents(batch_docs):
self.logger.error(f"批量添加文档失败，已迁移 {migrated_count} 个文档")
return False
migrated_count += len(batch_docs)
self.logger.info(f"已迁移 {migrated_count}/{total_docs} 个文档")
self.logger.info(f"迁移完成，总共迁移 {migrated_count} 个文档")
return migrated_count == total_docs
except Exception as e:
self.logger.error(f"迁移执行失败: {str(e)}")
return False
def _get_documents_in_batches(self, store: BaseVectorStore, batch_size: int) -> Iterator[List[HaystackDocument]]:
"""
批量获取文档的生成器。
Args:
store: 向量存储
batch_size: 批处理大小
Yields:
文档批次列表
"""
if isinstance(store, _HaystackDocumentSource):
all_docs = store.get_haystack_store().filter_documents({})
elif isinstance(store, _MappedDocumentSource) and isinstance(store.documents, dict):
all_docs = list(store.documents.values())
else:
raise VectorStoreError("Storage backend does not expose documents for migration")
for i in range(0, len(all_docs), batch_size):
yield all_docs[i : i + batch_size]
[docs]
def backup_storage(self, store: BaseVectorStore, backup_name: str, backup_dir: str = "./backups") -> Optional[str]:
"""
备份向量存储。
Args:
store: 要备份的存储
backup_name: 备份名称
backup_dir: 备份目录
Returns:
备份路径，失败返回None
"""
try:
# 创建备份目录
os.makedirs(backup_dir, exist_ok=True)
backup_path = os.path.join(backup_dir, backup_name)
os.makedirs(backup_path, exist_ok=True)
self.logger.info(f"开始备份到 {backup_path}")
# 保存存储索引
if not store.save_index(backup_path):
self.logger.error("保存存储索引失败")
return None
# 保存配置和元数据
backup_metadata = {
"backup_time": datetime.now().isoformat(),
"storage_info": store.get_storage_info(),
"document_count": store.get_document_count(),
"backup_version": "1.0",
}
metadata_path = os.path.join(backup_path, "backup_metadata.json")
with open(metadata_path, "w", encoding="utf-8") as f:
json.dump(backup_metadata, f, indent=2, ensure_ascii=False)
# 创建备份校验和
checksum = self._calculate_backup_checksum(backup_path)
checksum_path = os.path.join(backup_path, "backup.checksum")
with open(checksum_path, "w") as f:
f.write(checksum)
self.logger.info(f"备份完成: {backup_path}")
return backup_path
except Exception as e:
self.logger.error(f"备份失败: {str(e)}")
return None
[docs]
def restore_storage(self, backup_path: str, target_config: Dict[str, Any], verify_backup: bool = True) -> bool:
"""
从备份恢复向量存储。
Args:
backup_path: 备份路径
target_config: 目标存储配置
verify_backup: 是否验证备份完整性
Returns:
恢复成功返回True
"""
try:
self.logger.info(f"开始从 {backup_path} 恢复存储")
# 验证备份完整性
if verify_backup and not self.verify_backup(backup_path):
self.logger.error("备份验证失败")
return False
# 加载备份元数据
metadata_path = os.path.join(backup_path, "backup_metadata.json")
if not os.path.exists(metadata_path):
self.logger.error("备份元数据文件不存在")
return False
with open(metadata_path, "r", encoding="utf-8") as f:
backup_metadata = json.load(f)
# 创建目标存储
target_store = create_vector_store(target_config, auto_fallback=False)
# 加载索引
if not target_store.load_index(backup_path):
self.logger.error("加载索引失败")
return False
# 验证恢复结果
restored_count = target_store.get_document_count()
expected_count = backup_metadata.get("document_count", 0)
if restored_count != expected_count:
self.logger.warning(f"恢复的文档数量 ({restored_count}) 与备份记录 ({expected_count}) 不匹配")
self.logger.info(f"存储恢复完成，恢复了 {restored_count} 个文档")
target_store.close()
return True
except Exception as e:
self.logger.error(f"恢复失败: {str(e)}")
return False
@staticmethod
def _documents_match(source: HaystackDocument, target: Optional[HaystackDocument]) -> bool:
return (
target is not None
and target.content == source.content
and target.meta == source.meta
and target.embedding == source.embedding
)
[docs]
def verify_migration(
self, source_store: BaseVectorStore, target_store: BaseVectorStore, sample_ratio: float = 0.1
) -> bool:
"""
验证迁移结果。
Args:
source_store: 源存储
target_store: 目标存储
sample_ratio: 抽样验证比例
Returns:
验证通过返回True
"""
try:
self.logger.info("开始验证迁移结果...")
# 验证文档数量
source_count = source_store.get_document_count()
target_count = target_store.get_document_count()
if source_count != target_count:
self.logger.error(f"文档数量不匹配: 源存储 {source_count}, 目标存储 {target_count}")
return False
if source_count == 0:
self.logger.info("存储为空，验证通过")
return True
# 抽样验证文档内容和向量
sample_size = max(1, int(source_count * sample_ratio))
self.logger.info(f"抽样验证 {sample_size} 个文档")
verified = 0
for batch in self._get_documents_in_batches(source_store, max(1, sample_size)):
for source_doc in batch:
target_doc = target_store.get_document_by_id(source_doc.id)
if not self._documents_match(source_doc, target_doc):
return False
verified += 1
if verified >= sample_size:
break
if verified >= sample_size:
break
success_rate = verified / sample_size
self.logger.info(f"迁移验证通过，成功率: {success_rate:.2%}")
return verified > 0
except Exception as e:
self.logger.error(f"迁移验证失败: {str(e)}")
return False
[docs]
def verify_backup(self, backup_path: str) -> bool:
"""
验证备份完整性。
Args:
backup_path: 备份路径
Returns:
验证通过返回True
"""
try:
# 检查必要文件
required_files = ["backup_metadata.json", "backup.checksum"]
for file_name in required_files:
file_path = os.path.join(backup_path, file_name)
if not os.path.exists(file_path):
self.logger.error(f"备份文件缺失: {file_name}")
return False
# 验证校验和
checksum_path = os.path.join(backup_path, "backup.checksum")
with open(checksum_path, "r") as f:
stored_checksum = f.read().strip()
calculated_checksum = self._calculate_backup_checksum(backup_path)
if stored_checksum != calculated_checksum:
self.logger.error("备份校验和不匹配")
return False
self.logger.info("备份验证通过")
return True
except Exception as e:
self.logger.error(f"备份验证失败: {str(e)}")
return False
def _calculate_backup_checksum(self, backup_path: str) -> str:
"""
计算备份的校验和。
Args:
backup_path: 备份路径
Returns:
校验和字符串
"""
hasher = hashlib.sha256()
# 遍历所有文件计算校验和
for root, dirs, files in os.walk(backup_path):
for file in sorted(files):
if file == "backup.checksum": # 跳过校验和文件本身
continue
file_path = os.path.join(root, file)
try:
with open(file_path, "rb") as f:
while chunk := f.read(8192):
hasher.update(chunk)
except Exception as e:
self.logger.warning(f"无法读取文件 {file_path}: {str(e)}")
return hasher.hexdigest()
[docs]
def list_backups(self, backup_dir: str = "./backups") -> List[Dict[str, Any]]:
"""
列出所有备份。
Args:
backup_dir: 备份目录
Returns:
备份信息列表
"""
backups: List[Dict[str, Any]] = []
if not os.path.exists(backup_dir):
return backups
try:
for backup_name in os.listdir(backup_dir):
backup_path = os.path.join(backup_dir, backup_name)
if not os.path.isdir(backup_path):
continue
metadata_path = os.path.join(backup_path, "backup_metadata.json")
if not os.path.exists(metadata_path):
continue
try:
with open(metadata_path, "r", encoding="utf-8") as f:
metadata = json.load(f)
backup_info = {
"name": backup_name,
"path": backup_path,
"backup_time": metadata.get("backup_time"),
"document_count": metadata.get("document_count", 0),
"storage_type": metadata.get("storage_info", {}).get("storage_type"),
"size_mb": self._get_directory_size(backup_path) / 1024 / 1024,
}
backups.append(backup_info)
except Exception as e:
self.logger.warning(f"无法读取备份元数据 {backup_name}: {str(e)}")
except Exception as e:
self.logger.error(f"列出备份失败: {str(e)}")
# 按备份时间排序
backups.sort(key=lambda x: x.get("backup_time", ""), reverse=True)
return backups
def _get_directory_size(self, path: str) -> int:
"""
获取目录大小（字节）。
Args:
path: 目录路径
Returns:
目录大小
"""
total_size = 0
try:
for dirpath, dirnames, filenames in os.walk(path):
for filename in filenames:
file_path = os.path.join(dirpath, filename)
if os.path.exists(file_path):
total_size += os.path.getsize(file_path)
except Exception as e:
self.logger.warning(f"计算目录大小失败: {str(e)}")
return total_size
[docs]
def cleanup_old_backups(self, backup_dir: str = "./backups", keep_count: int = 10) -> int:
"""
清理旧备份。
Args:
backup_dir: 备份目录
keep_count: 保留的备份数量
Returns:
删除的备份数量
"""
try:
backups = self.list_backups(backup_dir)
if len(backups) <= keep_count:
self.logger.info(f"备份数量 ({len(backups)}) 未超过保留限制 ({keep_count})")
return 0
# 删除最旧的备份
backups_to_delete = backups[keep_count:]
deleted_count = 0
for backup in backups_to_delete:
try:
shutil.rmtree(backup["path"])
self.logger.info(f"已删除旧备份: {backup['name']}")
deleted_count += 1
except Exception as e:
self.logger.error(f"删除备份失败 {backup['name']}: {str(e)}")
self.logger.info(f"清理完成，删除了 {deleted_count} 个旧备份")
return deleted_count
except Exception as e:
self.logger.error(f"清理备份失败: {str(e)}")
return 0
# 便捷函数
[docs]
def migrate_vector_storage(
source_config: Dict[str, Any], target_config: Dict[str, Any], batch_size: int = 1000
) -> bool:
"""
向量存储迁移的便捷函数。
Args:
source_config: 源存储配置
target_config: 目标存储配置
batch_size: 批处理大小
Returns:
迁移成功返回True
"""
migrator = VectorStoreMigrator()
return migrator.migrate(source_config, target_config, batch_size)
[docs]
def backup_vector_storage(store: BaseVectorStore, backup_name: str) -> Optional[str]:
"""
向量存储备份的便捷函数。
Args:
store: 要备份的存储
backup_name: 备份名称
Returns:
备份路径，失败返回None
"""
migrator = VectorStoreMigrator()
return migrator.backup_storage(store, backup_name)