milvus init

models/milvus.go

@@ -0,0 +1,45 @@
+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
+}