59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
from typing import Optional
|
|
|
|
from pymilvus import MilvusClient as PyMilvusClient
|
|
|
|
from app.conf.app_config import MilvusConfig, app_config
|
|
|
|
|
|
class MilvusClientWrapper:
|
|
"""Milvus 客户端单例包装器"""
|
|
_instance = None
|
|
_initialized = False
|
|
|
|
def __new__(cls, milvus_config: MilvusConfig = None):
|
|
if cls._instance is None:
|
|
cls._instance = super(MilvusClientWrapper, cls).__new__(cls)
|
|
return cls._instance
|
|
|
|
def __init__(self, milvus_config: MilvusConfig = None):
|
|
if self._initialized:
|
|
return
|
|
|
|
if milvus_config is None:
|
|
milvus_config = app_config.milvus
|
|
|
|
self.milvus_config = milvus_config
|
|
self.client: Optional[PyMilvusClient] = None
|
|
self._initialized = True
|
|
|
|
def init(self):
|
|
if self.client is not None:
|
|
return
|
|
|
|
uri = f"http://{self.milvus_config.host}:{self.milvus_config.port}"
|
|
self.client = PyMilvusClient(
|
|
uri=uri,
|
|
user=self.milvus_config.user,
|
|
password=self.milvus_config.password
|
|
)
|
|
|
|
def close(self):
|
|
if self.client is not None:
|
|
self.client.close()
|
|
self.client = None
|
|
|
|
def has_collection(self, collection_name: str) -> bool:
|
|
self._ensure_connection()
|
|
return self.client.has_collection(collection_name)
|
|
|
|
def drop_collection(self, collection_name: str):
|
|
self._ensure_connection()
|
|
if self.client.has_collection(collection_name):
|
|
self.client.drop_collection(collection_name)
|
|
|
|
def _ensure_connection(self):
|
|
if self.client is None:
|
|
self.init()
|
|
|
|
|
|
milvus_client = MilvusClientWrapper() |