Update go sdk v2

Signed-off-by: wayblink <[email protected]>
zilliztech · Nov 29, 2024 · c682188 · c682188
1 parent 98f0526
commit c682188
Show file tree

Hide file tree

Showing 5 changed files with 487 additions and 90 deletions.
diff --git a/core/backup_context.go b/core/backup_context.go
@@ -143,14 +143,12 @@ func CreateBackupContext(ctx context.Context, params paramtable.BackupParams) *B
 
 func (b *BackupContext) getMilvusClient() *MilvusClient {
 	if b.milvusClient == nil {
-		milvusClient, err := CreateMilvusClient(b.ctx, b.params)
+		milvusClient, err := newMilvusClient(b.ctx, b.params)
 		if err != nil {
 			log.Error("failed to initial milvus client", zap.Error(err))
 			panic(err)
 		}
-		b.milvusClient = &MilvusClient{
-			client: milvusClient,
-		}
+		b.milvusClient = milvusClient
 	}
 	return b.milvusClient
 }

diff --git a/core/backup_impl_create_backup.go b/core/backup_impl_create_backup.go
@@ -151,7 +151,7 @@ func (b *BackupContext) parseBackupCollections(request *backuppb.CreateBackupReq
 		}
 		for db, collections := range dbCollections {
 			if len(collections) == 0 {
-				collections, err := b.getMilvusClient().ListCollections(b.ctx, db)
+				collections, err := b.getMilvusClient().ListCollectionsV2(b.ctx, db)
 				if err != nil {
 					log.Error("fail in ListCollections", zap.Error(err))
 					return nil, err
@@ -176,7 +176,7 @@ func (b *BackupContext) parseBackupCollections(request *backuppb.CreateBackupReq
 			// compatible to milvus under v2.2.8 without database support
 			if strings.Contains(err.Error(), "feature not supported") {
 				// default database only
-				collections, err := b.getMilvusClient().ListCollections(b.ctx, "default")
+				collections, err := b.getMilvusClient().ListCollectionsV2(b.ctx, "default")
 				if err != nil {
 					log.Error("fail in ListCollections", zap.Error(err))
 					return nil, err
@@ -190,7 +190,7 @@ func (b *BackupContext) parseBackupCollections(request *backuppb.CreateBackupReq
 			}
 		} else {
 			for _, db := range dbs {
-				collections, err := b.getMilvusClient().ListCollections(b.ctx, db.Name)
+				collections, err := b.getMilvusClient().ListCollectionsV2(b.ctx, db.Name)
 				if err != nil {
 					log.Error("fail in ListCollections", zap.Error(err))
 					return nil, err

diff --git a/core/milvus_sdk_wrapper.go b/core/milvus_sdk_wrapper.go
@@ -5,17 +5,66 @@ import (
 	"sync"
 	"time"
 
+	"github.com/cockroachdb/errors"
+	"go.uber.org/zap"
+
 	"github.com/milvus-io/milvus-proto/go-api/v2/msgpb"
 	gomilvus "github.com/milvus-io/milvus-sdk-go/v2/client"
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 
+	entityV2 "github.com/milvus-io/milvus/client/v2/entity"
+	milvusClientV2 "github.com/milvus-io/milvus/client/v2/milvusclient"
+
+	"github.com/zilliztech/milvus-backup/core/paramtable"
+	"github.com/zilliztech/milvus-backup/internal/log"
 	"github.com/zilliztech/milvus-backup/internal/util/retry"
 )
 
 // MilvusClient wrap db into milvus API to make it thread safe
 type MilvusClient struct {
 	mu     sync.Mutex
 	client gomilvus.Client
+	// sdk v2
+	milvusClientV2 *milvusClientV2.Client
+}
+
+func newMilvusClient(ctx context.Context, params paramtable.BackupParams) (*MilvusClient, error) {
+	milvusEndpoint := params.MilvusCfg.Address + ":" + params.MilvusCfg.Port
+	log.Debug("Start Milvus client", zap.String("endpoint", milvusEndpoint))
+	var c gomilvus.Client
+	var err error
+	if params.MilvusCfg.AuthorizationEnabled && params.MilvusCfg.User != "" && params.MilvusCfg.Password != "" {
+		if params.MilvusCfg.TLSMode == 0 {
+			c, err = gomilvus.NewDefaultGrpcClientWithAuth(ctx, milvusEndpoint, params.MilvusCfg.User, params.MilvusCfg.Password)
+		} else if params.MilvusCfg.TLSMode == 1 || params.MilvusCfg.TLSMode == 2 {
+			c, err = gomilvus.NewDefaultGrpcClientWithTLSAuth(ctx, milvusEndpoint, params.MilvusCfg.User, params.MilvusCfg.Password)
+		} else {
+			log.Error("milvus.TLSMode is not illegal, support value 0, 1, 2")
+			return nil, errors.New("milvus.TLSMode is not illegal, support value 0, 1, 2")
+		}
+	} else {
+		c, err = gomilvus.NewGrpcClient(ctx, milvusEndpoint)
+	}
+	if err != nil {
+		log.Error("failed to connect to milvus", zap.Error(err))
+		return nil, err
+	}
+
+	milvusClientV2, err := milvusClientV2.New(ctx, &milvusClientV2.ClientConfig{
+		Address:       milvusEndpoint,
+		Username:      params.MilvusCfg.User,
+		Password:      params.MilvusCfg.Password,
+		EnableTLSAuth: params.MilvusCfg.AuthorizationEnabled,
+	})
+	if err != nil {
+		log.Error("failed to initial milvus client v2", zap.Error(err))
+		return nil, err
+	}
+
+	return &MilvusClient{
+		client:         c,
+		milvusClientV2: milvusClientV2,
+	}, nil
 }
 
 func (m *MilvusClient) Close() error {
@@ -104,6 +153,31 @@ func (m *MilvusClient) ListCollections(ctx context.Context, db string) ([]*entit
 	return m.client.ListCollections(ctx)
 }
 
+func (m *MilvusClient) ListCollectionsV2(ctx context.Context, db string) ([]*entityV2.Collection, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	err := m.milvusClientV2.UsingDatabase(ctx, milvusClientV2.NewUsingDatabaseOption(db))
+	if err != nil {
+		return nil, err
+	}
+
+	collections, err := m.milvusClientV2.ListCollections(ctx, milvusClientV2.NewListCollectionOption())
+	if err != nil {
+		return nil, err
+	}
+
+	collectionEntities := make([]*entityV2.Collection, 0)
+	for _, collection := range collections {
+		coll, err := m.milvusClientV2.DescribeCollection(ctx, milvusClientV2.NewDescribeCollectionOption(collection))
+		if err != nil {
+			return nil, err
+		}
+		collectionEntities = append(collectionEntities, coll)
+	}
+
+	return collectionEntities, nil
+}
+
 func (m *MilvusClient) HasCollection(ctx context.Context, db, collName string) (bool, error) {
 	m.mu.Lock()
 	defer m.mu.Unlock()

diff --git a/go.mod b/go.mod
@@ -1,6 +1,6 @@
 module github.com/zilliztech/milvus-backup
 
-go 1.18
+go 1.21
 
 require (
 	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.7.0
@@ -10,54 +10,67 @@ require (
 	github.com/blang/semver/v4 v4.0.0
 	github.com/cockroachdb/errors v1.9.1
 	github.com/gin-gonic/gin v1.8.1
-	github.com/golang/protobuf v1.5.2
-	github.com/google/btree v1.0.1
-	github.com/google/uuid v1.3.0
+	github.com/golang/protobuf v1.5.4
+	github.com/google/btree v1.1.2
+	github.com/google/uuid v1.6.0
 	github.com/json-iterator/go v1.1.12
 	github.com/lingdor/stackerror v0.0.0-20191119040541-976d8885ed76
 	//github.com/milvus-io/milvus-proto/go-api/v2 v2.4.3
-	github.com/milvus-io/milvus-sdk-go/v2 v2.4.2-0.20241009111120-8008f14f2c16
+	github.com/milvus-io/milvus-sdk-go/v2 v2.4.2-0.20241114020221-3d596ff9fe9f
+	github.com/milvus-io/milvus/client/v2 v2.0.0-20241120124432-f3a36f8a299b
 	github.com/minio/minio-go/v7 v7.0.61
 	github.com/pkg/errors v0.9.1
-	github.com/samber/lo v1.39.0
+	github.com/samber/lo v1.47.0
 	github.com/sony/sonyflake v1.1.0
 	github.com/spf13/cast v1.3.1
 	github.com/spf13/cobra v1.5.0
 	github.com/spf13/viper v1.8.1
-	github.com/stretchr/testify v1.8.4
+	github.com/stretchr/testify v1.9.0
 	github.com/swaggo/files v1.0.0
 	github.com/swaggo/gin-swagger v1.5.3
 	github.com/swaggo/swag v1.16.1
 	github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common v1.0.865
-	github.com/uber/jaeger-client-go v2.25.0+incompatible
+	github.com/uber/jaeger-client-go v2.30.0+incompatible
 	go.uber.org/atomic v1.10.0
-	go.uber.org/zap v1.17.0
-	golang.org/x/oauth2 v0.0.0-20210402161424-2e8d93401602
+	go.uber.org/zap v1.27.0
+	golang.org/x/oauth2 v0.20.0
 	golang.org/x/sync v0.8.0
-	golang.org/x/time v0.0.0-20201208040808-7e3f01d25324
-	google.golang.org/grpc v1.48.0
+	golang.org/x/time v0.5.0
+	google.golang.org/grpc v1.65.0
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
-replace github.com/milvus-io/milvus-sdk-go/v2 => github.com/wayblink/milvus-sdk-go/v2 v2.3.0-beta4.0.20241030091852-d6eb85c1a8ff
+require github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20241108105827-266fb751b620
 
-require github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240909041258-8f8ca67816cd
+replace github.com/milvus-io/milvus-sdk-go/v2 => github.com/wayblink/milvus-sdk-go/v2 v2.3.0-beta4.0.20241030091852-d6eb85c1a8ff
 
 require (
-	cloud.google.com/go v0.81.0 // indirect
+	cloud.google.com/go/compute/metadata v0.3.0 // indirect
 	github.com/Azure/azure-sdk-for-go/sdk/internal v1.3.0 // indirect
 	github.com/AzureAD/microsoft-authentication-library-for-go v1.1.1 // indirect
 	github.com/KyleBanks/depth v1.2.1 // indirect
 	github.com/alibabacloud-go/debug v0.0.0-20190504072949-9472017b5c68 // indirect
 	github.com/alibabacloud-go/tea v1.1.8 // indirect
+	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/cenkalti/backoff/v4 v4.2.1 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/cilium/ebpf v0.11.0 // indirect
 	github.com/cockroachdb/logtags v0.0.0-20211118104740-dabe8e521a4f // indirect
 	github.com/cockroachdb/redact v1.1.3 // indirect
+	github.com/containerd/cgroups/v3 v3.0.3 // indirect
+	github.com/coreos/go-semver v0.3.0 // indirect
+	github.com/coreos/go-systemd/v22 v22.3.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/docker/go-units v0.4.0 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
+	github.com/form3tech-oss/jwt-go v3.2.3+incompatible // indirect
 	github.com/fsnotify/fsnotify v1.4.9 // indirect
 	github.com/getsentry/sentry-go v0.12.0 // indirect
 	github.com/gin-contrib/sse v0.1.0 // indirect
+	github.com/go-logr/logr v1.4.2 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
+	github.com/go-ole/go-ole v1.2.6 // indirect
 	github.com/go-openapi/jsonpointer v0.20.0 // indirect
 	github.com/go-openapi/jsonreference v0.20.2 // indirect
 	github.com/go-openapi/spec v0.20.9 // indirect
@@ -66,51 +79,97 @@ require (
 	github.com/go-playground/universal-translator v0.18.0 // indirect
 	github.com/go-playground/validator/v10 v10.10.0 // indirect
 	github.com/goccy/go-json v0.9.7 // indirect
+	github.com/godbus/dbus/v5 v5.0.4 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang-jwt/jwt/v5 v5.0.0 // indirect
+	github.com/gorilla/websocket v1.4.2 // indirect
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 // indirect
+	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 // indirect
+	github.com/grpc-ecosystem/grpc-gateway v1.16.0 // indirect
+	github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect
+	github.com/jonboulle/clockwork v0.2.2 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
-	github.com/klauspost/compress v1.16.7 // indirect
+	github.com/klauspost/compress v1.17.7 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.5 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/kylelemons/godebug v1.1.0 // indirect
 	github.com/leodido/go-urn v1.2.1 // indirect
+	github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
 	github.com/magiconair/properties v1.8.5 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
-	github.com/mattn/go-isatty v0.0.14 // indirect
+	github.com/mattn/go-isatty v0.0.19 // indirect
+	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
+	github.com/milvus-io/milvus/pkg v0.0.2-0.20241111021426-5e90f348fcbb // indirect
 	github.com/minio/md5-simd v1.1.2 // indirect
 	github.com/minio/sha256-simd v1.0.1 // indirect
 	github.com/mitchellh/mapstructure v1.4.1 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/opencontainers/runtime-spec v1.0.2 // indirect
+	github.com/panjf2000/ants/v2 v2.7.2 // indirect
 	github.com/pelletier/go-toml v1.9.3 // indirect
 	github.com/pelletier/go-toml/v2 v2.0.1 // indirect
 	github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/rogpeppe/go-internal v1.9.0 // indirect
+	github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
+	github.com/prometheus/client_golang v1.14.0 // indirect
+	github.com/prometheus/client_model v0.3.0 // indirect
+	github.com/prometheus/common v0.42.0 // indirect
+	github.com/prometheus/procfs v0.9.0 // indirect
+	github.com/rogpeppe/go-internal v1.10.0 // indirect
 	github.com/rs/xid v1.5.0 // indirect
+	github.com/shirou/gopsutil/v3 v3.22.9 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
+	github.com/soheilhy/cmux v0.1.5 // indirect
+	github.com/spaolacci/murmur3 v1.1.0 // indirect
 	github.com/spf13/afero v1.6.0 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/subosito/gotenv v1.2.0 // indirect
-	github.com/tidwall/gjson v1.14.4 // indirect
+	github.com/tidwall/gjson v1.17.1 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.0 // indirect
+	github.com/tklauser/go-sysconf v0.3.10 // indirect
+	github.com/tklauser/numcpus v0.4.0 // indirect
+	github.com/tmc/grpc-websocket-proxy v0.0.0-20201229170055-e5319fda7802 // indirect
 	github.com/ugorji/go/codec v1.2.7 // indirect
-	go.uber.org/multierr v1.6.0 // indirect
-	golang.org/x/crypto v0.14.0 // indirect
-	golang.org/x/exp v0.0.0-20220303212507-bbda1eaf7a17 // indirect
-	golang.org/x/net v0.17.0 // indirect
-	golang.org/x/sys v0.13.0 // indirect
-	golang.org/x/text v0.13.0 // indirect
-	golang.org/x/tools v0.11.0 // indirect
-	google.golang.org/appengine v1.6.7 // indirect
-	google.golang.org/genproto v0.0.0-20220503193339-ba3ae3f07e29 // indirect
-	google.golang.org/protobuf v1.33.0 // indirect
+	github.com/x448/float16 v0.8.4 // indirect
+	github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2 // indirect
+	github.com/yusufpapurcu/wmi v1.2.2 // indirect
+	go.etcd.io/bbolt v1.3.6 // indirect
+	go.etcd.io/etcd/api/v3 v3.5.5 // indirect
+	go.etcd.io/etcd/client/pkg/v3 v3.5.5 // indirect
+	go.etcd.io/etcd/client/v2 v2.305.5 // indirect
+	go.etcd.io/etcd/client/v3 v3.5.5 // indirect
+	go.etcd.io/etcd/pkg/v3 v3.5.5 // indirect
+	go.etcd.io/etcd/raft/v3 v3.5.5 // indirect
+	go.etcd.io/etcd/server/v3 v3.5.5 // indirect
+	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 // indirect
+	go.opentelemetry.io/otel v1.28.0 // indirect
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.20.0 // indirect
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.20.0 // indirect
+	go.opentelemetry.io/otel/metric v1.28.0 // indirect
+	go.opentelemetry.io/otel/sdk v1.28.0 // indirect
+	go.opentelemetry.io/otel/trace v1.28.0 // indirect
+	go.opentelemetry.io/proto/otlp v1.0.0 // indirect
+	go.uber.org/automaxprocs v1.5.3 // indirect
+	go.uber.org/multierr v1.11.0 // indirect
+	golang.org/x/crypto v0.25.0 // indirect
+	golang.org/x/exp v0.0.0-20230224173230-c95f2b4c22f2 // indirect
+	golang.org/x/net v0.27.0 // indirect
+	golang.org/x/sys v0.22.0 // indirect
+	golang.org/x/text v0.16.0 // indirect
+	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
+	google.golang.org/genproto v0.0.0-20231106174013-bbf56f31fb17 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20240528184218-531527333157 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20240730163845-b1a4ccb954bf // indirect
+	google.golang.org/protobuf v1.34.2 // indirect
+	gopkg.in/inf.v0 v0.9.1 // indirect
 	gopkg.in/ini.v1 v1.67.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
+	k8s.io/apimachinery v0.28.6 // indirect
+	sigs.k8s.io/yaml v1.3.0 // indirect
 )