Mutual_Fund/config/database/milvus.py

"""Milvus 会话：全局单例 AsyncMilvusClient。

timeout 入参是数据操作默认超时（用 MILVUS_TIMEOUT）；
构造为急切连接、由注册表 init_db 做指数退避重试 + 严格模式兜底。
"""
from pymilvus import AsyncMilvusClient

from config.settings import settings

_client: AsyncMilvusClient | None = None


def client() -> AsyncMilvusClient:
    global _client
    if _client is None:
        _client = AsyncMilvusClient(
            uri=settings.milvus.uri,
            token=settings.milvus.token,
            user=settings.milvus.user,
            password=settings.milvus.password,
            db_name=settings.milvus.db_name,
            timeout=settings.milvus.timeout,
        )
    return _client


async def ensure_database(milvus_client: AsyncMilvusClient | None = None) -> None:
    """Create the configured project database only when it does not exist."""
    target = milvus_client or client()
    database_name = settings.milvus.db_name
    if not database_name:
        raise RuntimeError("MILVUS_DB must be configured")
    if database_name not in await target.list_databases():
        await target.create_database(database_name)


async def init_db() -> None:
    client()  # 急切建连；失败由注册表重试（config/database/__init__.py）


async def dispose() -> None:
    global _client
    if _client is not None:
        await _client.close()
        _client = None


async def check_health() -> None:
    await client().get_server_version()