46 lines
1.0 KiB
Go
46 lines
1.0 KiB
Go
package models
|
|
|
|
import (
|
|
"context"
|
|
"log"
|
|
"strconv"
|
|
|
|
"github.com/milvus-io/milvus-sdk-go/v2/client"
|
|
)
|
|
|
|
type MilvusConnection struct {
|
|
Client client.Client
|
|
}
|
|
|
|
func (m *MilvusConnection) GetClient() client.Client {
|
|
return m.Client
|
|
}
|
|
|
|
func (m *MilvusConnection) Init() (err error) {
|
|
log.Println("Milvus connection init")
|
|
host := Viper.Get("milvus.host").(string)
|
|
port := Viper.Get("milvus.port").(int)
|
|
address_of_milvus := host + ":" + strconv.Itoa(port)
|
|
log.Println("Milvus address:", address_of_milvus)
|
|
m.Client, err = client.NewGrpcClient(context.Background(), address_of_milvus)
|
|
if err != nil {
|
|
log.Println("Milvus connection failed:", err)
|
|
return
|
|
}
|
|
log.Println("Milvus connection success")
|
|
return
|
|
}
|
|
|
|
func (m *MilvusConnection) GetCollection(collection_name string) (collection *client.Client, err error) {
|
|
if m.Client == nil {
|
|
m.Init()
|
|
}
|
|
err = m.Client.LoadCollection(context.Background(), collection_name, false)
|
|
if err != nil {
|
|
log.Println("Milvus load collection failed:", err)
|
|
return
|
|
}
|
|
collection = &m.Client
|
|
return
|
|
}
|