# 通过pymilvus sdk来操作数据库

from pymilvus import (
    connections,
    db,
    connections,
    Collection,
    utility,
)


def test_connection(
    db_name: str = "blog",
    collection_name: str = "articles",
    fied_name: str = "tile_vector",
) -> bool:
    try:
        # 建立与milvus的连接
        connections.connect(
            "default", host="49.233.56.55", port="19530", timeout=10, db_name=db_name
        )
        # 使用或者切换到db_name
        db.using_database(db_name)
        if not utility.has_collection(collection_name):
            raise ValueError(f"集合{collection_name}不存在，无法创建索引")
        # 获取集合对象
        collection = Collection(collection_name)
        # 判断是否已经有索引
        has_index = bool(collection.indexes)
        if has_index:
            print("所以已经存在，跳过创建")
        else:
            index_params = {
                "index_type": "IVF_FLAT",
                "metric_type": "L2",
                "params": {"nlist": 128},
            }
            # 创建索引
            collection.create_index(field_name=fied_name, index_params=index_params)
            print("索引创建成功")
        # 加载结合到内存
        collection.load()
        print("集合已经加载到内存，可供查询")

    except Exception as exc:
        print("索引管理失败:", exc)
        raise
    finally:
        connections.disconnect("default")


if __name__ == "__main__":
    if not test_connection():
        raise SystemExit(1)