117 lines
3.5 KiB
Python
117 lines
3.5 KiB
Python
# storage/init_storage.py
|
||
"""初始化 RustFS 对象存储"""
|
||
import asyncio
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# 添加项目根目录和 src 目录到 Python 路径
|
||
project_root = Path(__file__).parent.parent.parent
|
||
src_root = Path(__file__).parent.parent
|
||
sys.path.insert(0, str(project_root))
|
||
sys.path.insert(0, str(src_root))
|
||
|
||
from moldinsight.storage.rustfs_storage import rustfs_manager
|
||
from shared.config.settings import settings
|
||
from shared.utils.logger import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
async def init_rustfs_storage():
|
||
"""初始化 RustFS 对象存储"""
|
||
try:
|
||
# 连接到 RustFS (S3v4 API)
|
||
await rustfs_manager.connect(
|
||
endpoint=settings.RUSTFS_ENDPOINT,
|
||
access_key=settings.RUSTFS_ACCESS_KEY,
|
||
secret_key=settings.RUSTFS_SECRET_KEY,
|
||
timeout=settings.RUSTFS_TIMEOUT
|
||
)
|
||
|
||
logger.info("RustFS 对象存储初始化完成")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"RustFS 对象存储初始化失败: {e}")
|
||
return False
|
||
|
||
|
||
async def rustfs_startup_hook():
|
||
"""app_factory startup_hooks 注入点（D3 收敛）：moldinsight/unified 入口把
|
||
本钩子传入 create_app(startup_hooks=[rustfs_startup_hook])，RustFS 连接接线
|
||
归属 moldinsight 层，平台工厂不再持有模块专属依赖（原 connect_rustfs 已移除）。
|
||
"""
|
||
ok = await init_rustfs_storage()
|
||
print(f"[{'OK' if ok else 'WARN'}] RustFS")
|
||
|
||
|
||
async def test_storage():
|
||
"""测试 RustFS 对象存储功能"""
|
||
try:
|
||
import json
|
||
|
||
# 测试上传 JSON
|
||
test_data = {"test": True, "timestamp": "2024-01-01", "storage": "rustfs"}
|
||
result = await rustfs_manager.upload_json_data(
|
||
file_type='stp_files',
|
||
json_data=test_data,
|
||
file_hash='test-hash'
|
||
)
|
||
|
||
logger.info(f"RustFS 测试上传成功: {result['object_key']}")
|
||
|
||
# 测试下载
|
||
downloaded_bytes = await rustfs_manager.download_file(
|
||
file_type='stp_files',
|
||
object_key=result['object_key']
|
||
)
|
||
downloaded_data = json.loads(downloaded_bytes.decode('utf-8'))
|
||
logger.info(f"RustFS 测试下载成功: {downloaded_data}")
|
||
|
||
# 测试预签名 URL
|
||
url = await rustfs_manager.generate_presigned_url(
|
||
file_type='stp_files',
|
||
object_key=result['object_key'],
|
||
expires=3600
|
||
)
|
||
logger.info(f"RustFS 预签名URL: {url}")
|
||
|
||
# 清理测试文件
|
||
await rustfs_manager.delete_file(
|
||
file_type='stp_files',
|
||
object_key=result['object_key']
|
||
)
|
||
logger.info("RustFS 测试文件已清理")
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"RustFS 存储测试失败: {e}")
|
||
return False
|
||
|
||
|
||
if __name__ == "__main__":
|
||
async def main():
|
||
try:
|
||
print("=== 初始化 RustFS 对象存储 ===")
|
||
# 初始化存储
|
||
init_result = await init_rustfs_storage()
|
||
if init_result:
|
||
print("[OK] RustFS 连接成功")
|
||
|
||
print("\n=== 测试 RustFS 功能 ===")
|
||
# 运行测试
|
||
test_result = await test_storage()
|
||
if test_result:
|
||
print("[OK] RustFS 测试全部通过")
|
||
else:
|
||
print("[FAIL] RustFS 测试失败")
|
||
|
||
finally:
|
||
# 关闭连接
|
||
await rustfs_manager.close()
|
||
print("\n=== 连接已关闭 ===")
|
||
|
||
# 运行主函数
|
||
asyncio.run(main())
|