103 lines
3.6 KiB
Python
103 lines
3.6 KiB
Python
"""就绪探针：MySQL / Redis / Milvus 必须是真实探测结果。
|
||
|
||
B3 修复：原先 ``checks["milvus"] = True`` 是硬编码，健康检查永远报告 Milvus 正常，
|
||
属于误导性探针。现在改为真实探测（连接 + ``get_server_version``），并满足：
|
||
|
||
- 探测有超时（默认 2s），超时返回 ``False`` + ``timeout`` 状态；
|
||
- 外部依赖不可用时不抛异常，只返回 ``False`` + 明确状态，整体降级为 ``degraded``；
|
||
- pymilvus 是可选依赖（且无类型存根），缺失时返回 ``False`` + ``client_missing``。
|
||
"""
|
||
|
||
import asyncio
|
||
from typing import Any
|
||
|
||
from sqlalchemy import text
|
||
|
||
from app.core.config import get_settings
|
||
from app.infrastructure.db import SessionFactory
|
||
|
||
MILVUS_PROBE_TIMEOUT_SECONDS = 2.0
|
||
|
||
|
||
def _create_milvus_client(uri: str, token: str) -> Any:
|
||
"""延迟导入并建连：pymilvus 为可选依赖，缺失时由调用方转为明确状态。"""
|
||
from pymilvus import AsyncMilvusClient # type: ignore[import-untyped]
|
||
|
||
return AsyncMilvusClient(uri=uri, token=token or None)
|
||
|
||
|
||
class HealthService:
|
||
def __init__(self, *, milvus_timeout_seconds: float = MILVUS_PROBE_TIMEOUT_SECONDS) -> None:
|
||
self._milvus_timeout_seconds = milvus_timeout_seconds
|
||
|
||
async def ready(self) -> dict[str, Any]:
|
||
checks: dict[str, bool] = {}
|
||
probes: dict[str, str] = {}
|
||
try:
|
||
async with SessionFactory() as session:
|
||
await session.execute(text("SELECT 1"))
|
||
checks["mysql"] = True
|
||
except Exception:
|
||
checks["mysql"] = False
|
||
probes["mysql"] = "unavailable"
|
||
checks["redis"] = await self._redis_ready()
|
||
milvus_ready, milvus_state = await self._milvus_ready()
|
||
checks["milvus"] = milvus_ready
|
||
probes["milvus"] = milvus_state
|
||
return {
|
||
"status": "ready" if all(checks.values()) else "degraded",
|
||
"checks": checks,
|
||
"probes": probes,
|
||
}
|
||
|
||
async def _redis_ready(self) -> bool:
|
||
client = None
|
||
try:
|
||
from redis.asyncio import Redis
|
||
|
||
settings = get_settings()
|
||
client = Redis.from_url(
|
||
settings.redis_url,
|
||
socket_connect_timeout=settings.redis_connect_timeout_seconds,
|
||
socket_timeout=settings.redis_connect_timeout_seconds,
|
||
decode_responses=True,
|
||
)
|
||
return bool(await client.ping())
|
||
except Exception:
|
||
return False
|
||
finally:
|
||
if client is not None:
|
||
try:
|
||
await client.aclose()
|
||
except Exception:
|
||
pass
|
||
|
||
async def _milvus_ready(self) -> tuple[bool, str]:
|
||
"""真实探测 Milvus：状态取值 ok / unavailable / client_missing / timeout。"""
|
||
settings = get_settings()
|
||
client: Any = None
|
||
try:
|
||
async with asyncio.timeout(self._milvus_timeout_seconds):
|
||
client = await asyncio.to_thread(
|
||
_create_milvus_client, settings.resolved_milvus_uri, settings.milvus_token
|
||
)
|
||
await client.get_server_version()
|
||
return True, "ok"
|
||
except TimeoutError:
|
||
return False, "timeout"
|
||
except ImportError:
|
||
return False, "client_missing"
|
||
except Exception:
|
||
return False, "unavailable"
|
||
finally:
|
||
await self._close_milvus_client(client)
|
||
|
||
@staticmethod
|
||
async def _close_milvus_client(client: Any) -> None:
|
||
if client is None:
|
||
return
|
||
try:
|
||
await client.close()
|
||
except Exception:
|
||
pass
|