26F26g

"""
Milvus 연결 진단 스크립트
실행: python scripts/diagnose_milvus.py
"""
import socket
import sys
import time
import requests
from config import settings


def check_tcp(host: str, port: int, timeout: int = 5) -> bool:
    """TCP 레벨 연결 확인"""
    try:
        sock = socket.create_connection((host, port), timeout=timeout)
        sock.close()
        print(f"  ✓ TCP {host}:{port} 연결 성공")
        return True
    except socket.timeout:
        print(f"  ✗ TCP {host}:{port} 타임아웃 ({timeout}s)")
    except ConnectionRefusedError:
        print(f"  ✗ TCP {host}:{port} 연결 거부 (포트 닫힘)")
    except socket.gaierror as e:
        print(f"  ✗ DNS 해석 실패: {host} → {e}")
    return False


def check_milvus_http(host: str) -> None:
    """Milvus REST API (9091) health check"""
    url = f"http://{host}:9091/healthz"
    try:
        r = requests.get(url, timeout=5)
        print(f"  ✓ HTTP healthz: {r.status_code} {r.text[:100]}")
    except Exception as e:
        print(f"  ✗ HTTP healthz 실패: {e}")
        print(f"    (Milvus 2.3+ 는 9091 포트로 health check 가능)")


def check_pymilvus_version() -> None:
    """pymilvus 버전 확인"""
    try:
        import pymilvus
        print(f"  ✓ pymilvus 버전: {pymilvus.__version__}")
        # 버전별 권장 Milvus 서버 버전
        ver = tuple(int(x) for x in pymilvus.__version__.split(".")[:2])
        if ver >= (2, 4):
            print(f"    → Milvus 서버 2.4.x 필요")
        elif ver >= (2, 3):
            print(f"    → Milvus 서버 2.3.x 필요")
        elif ver >= (2, 2):
            print(f"    → Milvus 서버 2.2.x 필요")
    except ImportError:
        print("  ✗ pymilvus 미설치")


def try_connect_grpc(host: str, port: int) -> bool:
    """gRPC 방식 연결 시도"""
    from pymilvus import connections
    try:
        connections.connect(
            alias="grpc_test",
            host=host,
            port=port,
            timeout=10,
        )
        print(f"  ✓ gRPC 연결 성공 ({host}:{port})")
        connections.disconnect("grpc_test")
        return True
    except Exception as e:
        print(f"  ✗ gRPC 연결 실패: {e}")
        return False


def try_connect_uri(host: str, grpc_port: int) -> bool:
    """URI 방식 연결 시도 (pymilvus 2.3+)"""
    from pymilvus import MilvusClient
    uri = f"http://{host}:{grpc_port}"
    try:
        client = MilvusClient(uri=uri)
        print(f"  ✓ MilvusClient URI 연결 성공: {uri}")
        return True
    except Exception as e:
        print(f"  ✗ MilvusClient URI 실패 ({uri}): {e}")

    # gRPC URI
    uri = f"grpc://{host}:{grpc_port}"
    try:
        client = MilvusClient(uri=uri)
        print(f"  ✓ MilvusClient gRPC URI 성공: {uri}")
        return True
    except Exception as e:
        print(f"  ✗ MilvusClient gRPC URI 실패 ({uri}): {e}")

    return False


def try_connect_with_auth(host: str, port: int, user: str, password: str) -> bool:
    """인증 포함 연결 시도"""
    from pymilvus import connections
    if not user:
        print("  - 인증 정보 없음, 스킵")
        return False
    try:
        connections.connect(
            alias="auth_test",
            host=host,
            port=port,
            user=user,
            password=password,
            timeout=10,
        )
        print(f"  ✓ 인증 포함 연결 성공")
        connections.disconnect("auth_test")
        return True
    except Exception as e:
        print(f"  ✗ 인증 포함 연결 실패: {e}")
        return False


if __name__ == "__main__":
    host = settings.MILVUS_HOST
    port = settings.MILVUS_PORT

    print("=" * 55)
    print(f"Milvus 연결 진단")
    print(f"  설정 HOST: {host}")
    print(f"  설정 PORT: {port}")
    print("=" * 55)

    print("\n[1] pymilvus 버전 확인")
    check_pymilvus_version()

    print(f"\n[2] TCP 연결 확인")
    tcp_ok = check_tcp(host, port)
    if not tcp_ok:
        print("\n  → 가능한 원인:")
        print("    a) MILVUS_HOST 값이 잘못됨 (kubectl get svc 로 확인)")
        print("    b) MILVUS_PORT가 잘못됨 (기본 19530, Helm에 따라 다를 수 있음)")
        print("    c) Milvus Pod가 아직 기동 중 (Running 상태인지 확인)")
        print("    d) 네트워크 정책으로 포트 막힘 (NetworkPolicy 확인)")
        print("\n  → 아래 명령으로 실제 서비스 확인:")
        print("    kubectl get svc -A | grep milvus")
        print("    kubectl get pods -A | grep milvus")

    print(f"\n[3] Milvus HTTP healthz 확인 (포트 9091)")
    check_milvus_http(host)

    print(f"\n[4] gRPC 연결 시도")
    grpc_ok = try_connect_grpc(host, port)

    if not grpc_ok:
        print(f"\n[5] URI 방식 연결 시도 (pymilvus 2.3+)")
        try_connect_uri(host, port)

        print(f"\n[6] 인증 포함 연결 시도")
        try_connect_with_auth(
            host, port,
            settings.MILVUS_USER,
            settings.MILVUS_PASSWORD,
        )

    print("\n" + "=" * 55)
    if grpc_ok:
        print("✅ 연결 성공 — init_milvus.py 실행 가능")
    else:
        print("❌ 연결 실패 — 위 진단 결과를 확인하세요")
    print("=" * 55)