Files
geMoldInsight/scripts/migrations/migrate_rustfs_structure.py
T
2026-04-21 14:20:20 +08:00

378 lines
14 KiB
Python

#!/usr/bin/env python3
"""
RustFS 存储结构迁移脚本
将旧的多桶结构迁移到新的单桶结构:
旧结构: moldinsight-geometry, moldinsight-stp-files, moldinsight-mold-cavities, moldinsight-html-files, moldinsight-user-files
新结构: moldinsight-storage/
├── moldinsight/stp-files/{uuid}.stp
├── moldinsight/geometry/{file_hash}.json
├── moldinsight/mold-cavities/{file_hash}.json
├── moldinsight/html/{file_hash}.json
└── moldinsight/user-files/{uuid}.{ext}
注意:这是rustFS,而不是minio,只是用了minio的通用S3接口
"""
import asyncio
import sys
import os
from pathlib import Path
from typing import Dict, List, Optional
import json
from datetime import datetime
# 添加项目根目录和 src 目录到 Python 路径
project_root = Path(__file__).parent
src_root = project_root / "src"
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(src_root))
from storage.rustfs_storage import RustFSManager
from config.settings import settings
from utils.logger import get_logger
logger = get_logger(__name__)
class RustFSMigration:
"""RustFS 存储结构迁移器"""
def __init__(self):
self.rustfs = RustFSManager()
self.is_connected = False
# 旧桶名称映射
self.old_buckets = {
'stp_files': 'moldinsight-stp-files',
'geometry_data': 'moldinsight-geometry',
'mold_cavities': 'moldinsight-mold-cavities',
'html_files': 'moldinsight-html-files',
'user_files': 'moldinsight-user-files'
}
# 新桶名称
self.new_bucket = 'moldinsight'
# 文件类型前缀映射
self.file_type_mapping = {
'stp_files': 'stp-files',
'geometry_data': 'geometry',
'mold_cavities': 'mold-cavities',
'html_files': 'html',
'user_files': 'user-files'
}
async def connect(self):
"""连接到 RustFS"""
try:
await self.rustfs.connect(
endpoint=settings.RUSTFS_ENDPOINT,
access_key=settings.RUSTFS_ACCESS_KEY,
secret_key=settings.RUSTFS_SECRET_KEY,
timeout=settings.RUSTFS_TIMEOUT
)
self.is_connected = True
logger.info("RustFS 连接成功")
return True
except Exception as e:
logger.error(f"RustFS 连接失败: {e}")
return False
async def close(self):
"""关闭连接"""
await self.rustfs.close()
self.is_connected = False
logger.info("RustFS 连接已关闭")
async def list_old_buckets(self) -> Dict[str, List[Dict]]:
"""列出所有旧桶及其文件"""
if not self.is_connected:
raise RuntimeError("RustFS 未连接")
buckets_info = {}
for file_type, bucket_name in self.old_buckets.items():
try:
# 检查桶是否存在
if not self.rustfs.client.bucket_exists(bucket_name):
logger.info(f"桶不存在: {bucket_name}")
buckets_info[bucket_name] = []
continue
# 列出桶中所有文件
objects = self.rustfs.client.list_objects(bucket_name, recursive=True)
files = []
for obj in objects:
files.append({
'object_key': obj.object_name,
'size': obj.size,
'last_modified': obj.last_modified,
'etag': obj.etag
})
buckets_info[bucket_name] = files
logger.info(f"桶 {bucket_name} 包含 {len(files)} 个文件")
except Exception as e:
logger.error(f"列出桶 {bucket_name} 失败: {e}")
buckets_info[bucket_name] = []
return buckets_info
async def ensure_new_bucket(self):
"""确保新桶存在"""
if not self.is_connected:
raise RuntimeError("RustFS 未连接")
try:
if not self.rustfs.client.bucket_exists(self.new_bucket):
self.rustfs.client.make_bucket(self.new_bucket)
logger.info(f"创建新桶: {self.new_bucket}")
else:
logger.info(f"新桶已存在: {self.new_bucket}")
except Exception as e:
logger.error(f"确保新桶存在失败: {e}")
raise
async def migrate_file(self, old_bucket: str, old_object_key: str, file_type: str) -> bool:
"""迁移单个文件到新结构"""
try:
# 下载旧文件
response = self.rustfs.client.get_object(old_bucket, old_object_key)
file_data = response.read()
response.close()
response.release_conn()
# 生成新对象键
if file_type in ['stp_files', 'user_files']:
# STP文件和用户文件:使用UUID格式
import uuid
unique_id = str(uuid.uuid4())
ext = Path(old_object_key).suffix or ('.stp' if file_type == 'stp_files' else '')
new_object_key = f"moldinsight/{self.file_type_mapping[file_type]}/{unique_id}{ext}"
else:
# JSON数据文件:使用文件哈希格式
import hashlib
file_hash = hashlib.sha256(file_data).hexdigest()
new_object_key = f"moldinsight/{self.file_type_mapping[file_type]}/{file_hash}.json"
# 上传到新桶
self.rustfs.client.put_object(
self.new_bucket,
new_object_key,
data=file_data,
length=len(file_data),
content_type='application/octet-stream'
)
logger.info(f"文件迁移成功: {old_bucket}/{old_object_key} -> {self.new_bucket}/{new_object_key}")
return True
except Exception as e:
logger.error(f"文件迁移失败 {old_bucket}/{old_object_key}: {e}")
return False
async def migrate_bucket(self, old_bucket: str, file_type: str) -> Dict[str, any]:
"""迁移整个桶"""
if not self.is_connected:
raise RuntimeError("RustFS 未连接")
migration_result = {
'total_files': 0,
'successful': 0,
'failed': 0,
'failed_files': []
}
try:
# 检查桶是否存在
if not self.rustfs.client.bucket_exists(old_bucket):
logger.info(f"桶不存在,跳过迁移: {old_bucket}")
return migration_result
# 列出桶中所有文件
objects = self.rustfs.client.list_objects(old_bucket, recursive=True)
files = list(objects)
migration_result['total_files'] = len(files)
logger.info(f"开始迁移桶 {old_bucket}, 包含 {len(files)} 个文件")
for obj in files:
success = await self.migrate_file(old_bucket, obj.object_name, file_type)
if success:
migration_result['successful'] += 1
else:
migration_result['failed'] += 1
migration_result['failed_files'].append(obj.object_name)
logger.info(f"桶 {old_bucket} 迁移完成: 成功 {migration_result['successful']}, 失败 {migration_result['failed']}")
except Exception as e:
logger.error(f"桶 {old_bucket} 迁移失败: {e}")
return migration_result
async def delete_old_buckets(self) -> Dict[str, bool]:
"""删除所有旧桶(可选操作)"""
if not self.is_connected:
raise RuntimeError("RustFS 未连接")
deletion_results = {}
for file_type, bucket_name in self.old_buckets.items():
try:
# 检查桶是否存在
if not self.rustfs.client.bucket_exists(bucket_name):
logger.info(f"桶不存在,跳过删除: {bucket_name}")
deletion_results[bucket_name] = True
continue
# 删除桶中所有文件
objects = self.rustfs.client.list_objects(bucket_name, recursive=True)
for obj in objects:
self.rustfs.client.remove_object(bucket_name, obj.object_name)
# 删除空桶
self.rustfs.client.remove_bucket(bucket_name)
deletion_results[bucket_name] = True
logger.info(f"桶删除成功: {bucket_name}")
except Exception as e:
deletion_results[bucket_name] = False
logger.error(f"桶删除失败 {bucket_name}: {e}")
return deletion_results
async def run_migration(self, delete_old_buckets: bool = False) -> Dict[str, any]:
"""运行完整迁移流程"""
migration_summary = {
'start_time': datetime.now().isoformat(),
'connection_status': False,
'old_buckets_info': {},
'migration_results': {},
'deletion_results': {},
'end_time': None,
'status': 'failed'
}
try:
# 1. 连接
logger.info("=== RustFS 存储结构迁移开始 ===")
connection_result = await self.connect()
if not connection_result:
raise RuntimeError("无法连接到 RustFS")
migration_summary['connection_status'] = True
# 2. 列出旧桶信息
logger.info("1. 检查旧桶结构...")
old_buckets_info = await self.list_old_buckets()
migration_summary['old_buckets_info'] = old_buckets_info
# 3. 确保新桶存在
logger.info("2. 确保新桶存在...")
await self.ensure_new_bucket()
# 4. 执行迁移
logger.info("3. 开始迁移文件...")
migration_results = {}
for file_type, bucket_name in self.old_buckets.items():
if old_buckets_info.get(bucket_name):
logger.info(f"迁移桶: {bucket_name}")
result = await self.migrate_bucket(bucket_name, file_type)
migration_results[bucket_name] = result
else:
logger.info(f"跳过空桶: {bucket_name}")
migration_summary['migration_results'] = migration_results
# 5. 可选:删除旧桶
if delete_old_buckets:
logger.info("4. 删除旧桶...")
deletion_results = await self.delete_old_buckets()
migration_summary['deletion_results'] = deletion_results
else:
logger.info("4. 保留旧桶(跳过删除)")
# 6. 完成
migration_summary['status'] = 'completed'
migration_summary['end_time'] = datetime.now().isoformat()
logger.info("=== RustFS 存储结构迁移完成 ===")
return migration_summary
except Exception as e:
migration_summary['error'] = str(e)
migration_summary['end_time'] = datetime.now().isoformat()
logger.error(f"迁移失败: {e}")
return migration_summary
finally:
await self.close()
async def main():
"""主函数"""
migration = RustFSMigration()
print("=== RustFS 存储结构迁移工具 ===")
print("注意:这是rustFS,而不是minio,只是用了minio的通用S3接口")
print()
# 询问是否删除旧桶
delete_old = input("是否在迁移完成后删除旧桶?(y/N): ").strip().lower() == 'y'
print("\n开始迁移...")
# 运行迁移
result = await migration.run_migration(delete_old_buckets=delete_old)
# 输出结果摘要
print("\n=== 迁移结果摘要 ===")
print(f"状态: {result['status']}")
print(f"开始时间: {result['start_time']}")
print(f"结束时间: {result['end_time']}")
if 'error' in result:
print(f"错误: {result['error']}")
# 旧桶信息
print("\n--- 旧桶信息 ---")
for bucket_name, files in result['old_buckets_info'].items():
print(f"{bucket_name}: {len(files)} 个文件")
# 迁移结果
print("\n--- 迁移结果 ---")
total_files = 0
total_success = 0
total_failed = 0
for bucket_name, migration_result in result['migration_results'].items():
print(f"{bucket_name}:")
print(f" 总文件数: {migration_result['total_files']}")
print(f" 成功: {migration_result['successful']}")
print(f" 失败: {migration_result['failed']}")
total_files += migration_result['total_files']
total_success += migration_result['successful']
total_failed += migration_result['failed']
print(f"\n总计: {total_files} 个文件, 成功 {total_success}, 失败 {total_failed}")
# 删除结果(如果执行了删除)
if result['deletion_results']:
print("\n--- 旧桶删除结果 ---")
for bucket_name, success in result['deletion_results'].items():
status = "成功" if success else "失败"
print(f"{bucket_name}: {status}")
print("\n=== 迁移完成 ===")
if __name__ == "__main__":
asyncio.run(main())