feat: Add milvusclient package and migrate GoSDK (#32907)

Related to #31293

Signed-off-by: Congqi Xia <[email protected]>

.github/workflows/main.yaml

@@ -7,6 +7,7 @@ on:
     paths:
       - 'scripts/**'
       - 'internal/**'
+      - 'client/**'
       - 'pkg/**'
       - 'cmd/**'
       - 'build/**'
@@ -24,6 +25,7 @@ on:
       - 'scripts/**'
       - 'internal/**'
       - 'pkg/**'
+      - 'client/**'
       - 'cmd/**'
       - 'build/**'
       - 'tests/integration/**' # run integration test

client/Makefile

@@ -0,0 +1,33 @@
+# Licensed to the LF AI & Data foundation under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+GO		  ?= go
+PWD 	  := $(shell pwd)
+GOPATH	:= $(shell $(GO) env GOPATH)
+SHELL 	:= /bin/bash
+OBJPREFIX := "github.com/milvus-io/milvus/cmd/milvus/v2"
+
+# TODO pass golangci-lint path
+lint:
+	@echo "Running lint checks..." 
+
+unittest:
+	@echo "Running unittests..."
+	@(env bash $(PWD)/scripts/run_unittest.sh)
+
+generate-mockery:
+	@echo "Generating mockery Milvus service server"
+	@../bin/mockery --srcpkg=github.com/milvus-io/milvus-proto/go-api/v2/milvuspb --name=MilvusServiceServer --filename=mock_milvus_server_test.go --output=. --outpkg=client --with-expecter 

client/OWNERS

@@ -0,0 +1,6 @@
+reviewers:
+  - congqixia
+
+approvers:
+  - maintainers
+

client/client.go

@@ -0,0 +1,157 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package client
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"strconv"
+	"time"
+
+	"github.com/gogo/status"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
+	"github.com/milvus-io/milvus/client/v2/common"
+	"github.com/milvus-io/milvus/client/v2/entity"
+	"github.com/milvus-io/milvus/pkg/util/merr"
+)
+
+type Client struct {
+	conn    *grpc.ClientConn
+	service milvuspb.MilvusServiceClient
+	config  *ClientConfig
+
+	collCache *CollectionCache
+}
+
+func New(ctx context.Context, config *ClientConfig) (*Client, error) {
+	if err := config.parse(); err != nil {
+		return nil, err
+	}
+
+	c := &Client{
+		config: config,
+	}
+
+	// Parse remote address.
+	addr := c.config.getParsedAddress()
+
+	// Parse grpc options
+	options := c.config.getDialOption()
+
+	// Connect the grpc server.
+	if err := c.connect(ctx, addr, options...); err != nil {
+		return nil, err
+	}
+
+	c.collCache = NewCollectionCache(func(ctx context.Context, collName string) (*entity.Collection, error) {
+		return c.DescribeCollection(ctx, NewDescribeCollectionOption(collName))
+	})
+
+	return c, nil
+}
+
+func (c *Client) Close(ctx context.Context) error {
+	if c.conn == nil {
+		return nil
+	}
+	err := c.conn.Close()
+	if err != nil {
+		return err
+	}
+	c.conn = nil
+	c.service = nil
+	return nil
+}
+
+func (c *Client) connect(ctx context.Context, addr string, options ...grpc.DialOption) error {
+	if addr == "" {
+		return fmt.Errorf("address is empty")
+	}
+	conn, err := grpc.DialContext(ctx, addr, options...)
+	if err != nil {
+		return err
+	}
+
+	c.conn = conn
+	c.service = milvuspb.NewMilvusServiceClient(c.conn)
+
+	if !c.config.DisableConn {
+		err = c.connectInternal(ctx)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (c *Client) connectInternal(ctx context.Context) error {
+	hostName, err := os.Hostname()
+	if err != nil {
+		return err
+	}
+
+	req := &milvuspb.ConnectRequest{
+		ClientInfo: &commonpb.ClientInfo{
+			SdkType:    "Golang",
+			SdkVersion: common.SDKVersion,
+			LocalTime:  time.Now().String(),
+			User:       c.config.Username,
+			Host:       hostName,
+		},
+	}
+
+	resp, err := c.service.Connect(ctx, req)
+	if err != nil {
+		status, ok := status.FromError(err)
+		if ok {
+			if status.Code() == codes.Unimplemented {
+				// disable unsupported feature
+				c.config.addFlags(
+					disableDatabase |
+						disableJSON |
+						disableParitionKey |
+						disableDynamicSchema)
+			}
+			return nil
+		}
+		return err
+	}
+
+	if !merr.Ok(resp.GetStatus()) {
+		return merr.Error(resp.GetStatus())
+	}
+
+	c.config.setServerInfo(resp.GetServerInfo().GetBuildTags())
+	c.config.setIdentifier(strconv.FormatInt(resp.GetIdentifier(), 10))
+
+	return nil
+}
+
+func (c *Client) callService(fn func(milvusService milvuspb.MilvusServiceClient) error) error {
+	service := c.service
+	if service == nil {
+		return merr.WrapErrServiceNotReady("SDK", 0, "not connected")
+	}
+
+	return fn(c.service)
+}

client/client_config.go

@@ -0,0 +1,182 @@
+package client
+
+import (
+	"crypto/tls"
+	"fmt"
+	"math"
+	"net/url"
+	"regexp"
+	"strings"
+	"time"
+
+	"github.com/cockroachdb/errors"
+	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/backoff"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/keepalive"
+)
+
+const (
+	disableDatabase uint64 = 1 << iota
+	disableJSON
+	disableDynamicSchema
+	disableParitionKey
+)
+
+var regexValidScheme = regexp.MustCompile(`^https?:\/\/`)
+
+// DefaultGrpcOpts is GRPC options for milvus client.
+var DefaultGrpcOpts = []grpc.DialOption{
+	grpc.WithBlock(),
+	grpc.WithKeepaliveParams(keepalive.ClientParameters{
+		Time:                5 * time.Second,
+		Timeout:             10 * time.Second,
+		PermitWithoutStream: true,
+	}),
+	grpc.WithConnectParams(grpc.ConnectParams{
+		Backoff: backoff.Config{
+			BaseDelay:  100 * time.Millisecond,
+			Multiplier: 1.6,
+			Jitter:     0.2,
+			MaxDelay:   3 * time.Second,
+		},
+		MinConnectTimeout: 3 * time.Second,
+	}),
+}
+
+// ClientConfig for milvus client.
+type ClientConfig struct {
+	Address  string // Remote address, "localhost:19530".
+	Username string // Username for auth.
+	Password string // Password for auth.
+	DBName   string // DBName for this client.
+
+	EnableTLSAuth bool   // Enable TLS Auth for transport security.
+	APIKey        string // API key
+
+	DialOptions []grpc.DialOption // Dial options for GRPC.
+
+	// RetryRateLimit *RetryRateLimitOption // option for retry on rate limit inteceptor
+
+	DisableConn bool
+
+	identifier    string // Identifier for this connection
+	ServerVersion string // ServerVersion
+	parsedAddress *url.URL
+	flags         uint64 // internal flags
+}
+
+func (cfg *ClientConfig) parse() error {
+	// Prepend default fake tcp:// scheme for remote address.
+	address := cfg.Address
+	if !regexValidScheme.MatchString(address) {
+		address = fmt.Sprintf("tcp://%s", address)
+	}
+
+	remoteURL, err := url.Parse(address)
+	if err != nil {
+		return errors.Wrap(err, "milvus address parse fail")
+	}
+	// Remote Host should never be empty.
+	if remoteURL.Host == "" {
+		return errors.New("empty remote host of milvus address")
+	}
+	// Use DBName in remote url path.
+	if cfg.DBName == "" {
+		cfg.DBName = strings.TrimLeft(remoteURL.Path, "/")
+	}
+	// Always enable tls auth for https remote url.
+	if remoteURL.Scheme == "https" {
+		cfg.EnableTLSAuth = true
+	}
+	if remoteURL.Port() == "" && cfg.EnableTLSAuth {
+		remoteURL.Host += ":443"
+	}
+	cfg.parsedAddress = remoteURL
+	return nil
+}
+
+// Get parsed remote milvus address, should be called after parse was called.
+func (c *ClientConfig) getParsedAddress() string {
+	return c.parsedAddress.Host
+}
+
+// useDatabase change the inner db name.
+func (c *ClientConfig) useDatabase(dbName string) {
+	c.DBName = dbName
+}
+
+// useDatabase change the inner db name.
+func (c *ClientConfig) setIdentifier(identifier string) {
+	c.identifier = identifier
+}
+
+func (c *ClientConfig) setServerInfo(serverInfo string) {
+	c.ServerVersion = serverInfo
+}
+
+// Get parsed grpc dial options, should be called after parse was called.
+func (c *ClientConfig) getDialOption() []grpc.DialOption {
+	options := c.DialOptions
+	if c.DialOptions == nil {
+		// Add default connection options.
+		options = make([]grpc.DialOption, len(DefaultGrpcOpts))
+		copy(options, DefaultGrpcOpts)
+	}
+
+	// Construct dial option.
+	if c.EnableTLSAuth {
+		options = append(options, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{})))
+	} else {
+		options = append(options, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	}
+
+	options = append(options,
+		grpc.WithChainUnaryInterceptor(grpc_retry.UnaryClientInterceptor(
+			grpc_retry.WithMax(6),
+			grpc_retry.WithBackoff(func(attempt uint) time.Duration {
+				return 60 * time.Millisecond * time.Duration(math.Pow(3, float64(attempt)))
+			}),
+			grpc_retry.WithCodes(codes.Unavailable, codes.ResourceExhausted)),
+		// c.getRetryOnRateLimitInterceptor(),
+		))
+
+	// options = append(options, grpc.WithChainUnaryInterceptor(
+	// 	createMetaDataUnaryInterceptor(c),
+	// ))
+	return options
+}
+
+// func (c *ClientConfig) getRetryOnRateLimitInterceptor() grpc.UnaryClientInterceptor {
+// 	if c.RetryRateLimit == nil {
+// 		c.RetryRateLimit = c.defaultRetryRateLimitOption()
+// 	}
+
+// 	return RetryOnRateLimitInterceptor(c.RetryRateLimit.MaxRetry, c.RetryRateLimit.MaxBackoff, func(ctx context.Context, attempt uint) time.Duration {
+// 		return 10 * time.Millisecond * time.Duration(math.Pow(3, float64(attempt)))
+// 	})
+// }
+
+// func (c *ClientConfig) defaultRetryRateLimitOption() *RetryRateLimitOption {
+// 	return &RetryRateLimitOption{
+// 		MaxRetry:   75,
+// 		MaxBackoff: 3 * time.Second,
+// 	}
+// }
+
+// addFlags set internal flags
+func (c *ClientConfig) addFlags(flags uint64) {
+	c.flags |= flags
+}
+
+// hasFlags check flags is set
+func (c *ClientConfig) hasFlags(flags uint64) bool {
+	return (c.flags & flags) > 0
+}
+
+func (c *ClientConfig) resetFlags(flags uint64) {
+	c.flags &= ^flags
+}