47 lines
970 B
Go
47 lines
970 B
Go
package models
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
|
|
"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)
|
|
m.Client, err = client.NewGrpcClient(
|
|
context.Background(),
|
|
fmt.Sprintf("%s:%d", host, port),
|
|
)
|
|
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
|
|
}
|