Repository: milvus-io/milvus-sdk-go
Pull request #304 (2022-08-17 13:54:49)

To compatiable with auth interceptor for Grpc opts
Signed-off-by: yun.zhang <yun.zhang@zilliz.com>
/kind improvement

sre-ci-robot: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/304#" title="Author self-approved">jaime0815</a>*
To complete the [pull request process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process), please assign **sunby** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @sunby` in a comment when ready.

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["sunby"]} -->
codecov[bot]: # [Codecov](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/304?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) Report
> Merging [#304](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/304?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (8106202) into [master](https://codecov.io/gh/milvus-io/milvus-sdk-go/commit/ef83b3742a9da1349d323327bee367519cb86056?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (ef83b37) will **decrease** coverage by `0.42%`.
> The diff coverage is `54.16%`.

[![Impacted file tree graph](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/304/graphs/tree.svg?width=650&height=150&src=pr&token=WJK43PV390&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/304?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)

```diff
@@            Coverage Diff             @@
##           master     #304      +/-   ##
==========================================
- Coverage   90.32%   89.90%   -0.43%     
==========================================
  Files          24       24              
  Lines        2668     2665       -3     
==========================================
- Hits         2410     2396      -14     
- Misses        177      187      +10     
- Partials       81       82       +1     
```


| [Impacted Files](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/304?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) | Coverage Δ | |
|---|---|---|
| [client/client.go](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/304/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NsaWVudC5nbw==) | `59.25% <52.17%> (-40.75%)` | :arrow_down: |
| [client/client\_grpc\_collection.go](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/304/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NsaWVudF9ncnBjX2NvbGxlY3Rpb24uZ28=) | `85.47% <100.00%> (-1.41%)` | :arrow_down: |

sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/304#pullrequestreview-1076642812" title="LGTM">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/304#" title="Author self-approved">jaime0815</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)~~ [congqixia]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/client/client.go b/client/client.go
index 7b69710..f48ea87 100644
--- a/client/client.go
+++ b/client/client.go
@@ -14,8 +14,17 @@ package client
 
 import (
 	"context"
+	"crypto/tls"
+	"math"
 	"time"
 
+	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
+	"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"
+
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 
 	"google.golang.org/grpc"
@@ -140,21 +149,92 @@ type SearchResult struct {
 	Err         error           // search error if any
 }
 
-// alias type for context field key
-type grpcKey int
-
-const (
-	dialOption grpcKey = 1
-)
+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,
+	}),
+}
 
 // NewGrpcClient create client with grpc addr
 // the `Connect` API will be called for you
 // dialOptions contains the dial option(s) that control the grpc dialing process
 func NewGrpcClient(ctx context.Context, addr string, dialOptions ...grpc.DialOption) (Client, error) {
 	c := &grpcClient{}
-	err := c.connect(ctx, addr, dialOptions...)
+	if len(dialOptions) == 0 {
+		return NewDefaultGrpcClient(ctx, addr)
+	}
+
+	if err := c.connect(ctx, addr, dialOptions...); err != nil {
+		return nil, err
+	}
+
+	return c, nil
+}
+
+func NewDefaultGrpcClient(ctx context.Context, addr string) (Client, error) {
+	c := &grpcClient{}
+	defaultOpts := append(DefaultGrpcOpts,
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithUnaryInterceptor(
+			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)),
+		),
+	)
+	err := c.connect(ctx, addr, defaultOpts...)
 	if err != nil {
 		return nil, err
 	}
 	return c, nil
 }
+
+// NewDefaultGrpcClientWithTLSAuth enable transport security
+func NewDefaultGrpcClientWithTLSAuth(ctx context.Context, addr, username, password string) (Client, error) {
+	defaultOpts := getDefaultAuthOpts(username, password, true)
+	return NewGrpcClient(ctx, addr, defaultOpts...)
+}
+
+// NewDefaultGrpcClientWithAuth  disable transport security
+func NewDefaultGrpcClientWithAuth(ctx context.Context, addr, username, password string) (Client, error) {
+	defaultOpts := getDefaultAuthOpts(username, password, false)
+	return NewGrpcClient(ctx, addr, defaultOpts...)
+}
+
+func getDefaultAuthOpts(username, password string, enableTLS bool) []grpc.DialOption {
+	var credential credentials.TransportCredentials
+	if enableTLS {
+		credential = credentials.NewTLS(&tls.Config{})
+	} else {
+		credential = insecure.NewCredentials()
+	}
+
+	defaultOpts := append(DefaultGrpcOpts,
+		grpc.WithTransportCredentials(credential),
+		grpc.WithChainUnaryInterceptor(
+			CreateAuthenticationUnaryInterceptor(username, password),
+			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)),
+		),
+		grpc.WithStreamInterceptor(CreateAuthenticationStreamInterceptor(username, password)),
+	)
+	return defaultOpts
+}
diff --git a/client/client_grpc_collection.go b/client/client_grpc_collection.go
index 80d5b9b..558effe 100644
--- a/client/client_grpc_collection.go
+++ b/client/client_grpc_collection.go
@@ -15,17 +15,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"math"
-
 	"time"
 
-	"google.golang.org/grpc/backoff"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/credentials/insecure"
-	"google.golang.org/grpc/keepalive"
-
 	"github.com/golang/protobuf/proto"
-	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
 	"google.golang.org/grpc"
 
@@ -41,50 +33,16 @@ type grpcClient struct {
 
 // connect connect to service
 func (c *grpcClient) connect(ctx context.Context, addr string, opts ...grpc.DialOption) error {
-	if err := c.establishConnection(addr, opts...); err != nil {
-		return err
-	}
-
-	c.service = server.NewMilvusServiceClient(c.conn)
-	return nil
-}
-
-func (c *grpcClient) establishConnection(addr string, opts ...grpc.DialOption) error {
 	if addr == "" {
 		return fmt.Errorf("address is empty")
 	}
-
-	defaultOpts := []grpc.DialOption{
-		grpc.WithTransportCredentials(insecure.NewCredentials()),
-		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,
-		}),
-		grpc.WithUnaryInterceptor(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))),
-	}
-
-	opts = append(opts, defaultOpts...)
 	conn, err := grpc.Dial(addr, opts...)
 	if err != nil {
 		return err
 	}
 
 	c.conn = conn
+	c.service = server.NewMilvusServiceClient(c.conn)
 	return nil
 }
 
diff --git a/examples/auth/auth.go b/examples/auth/auth.go
index aea71b3..9f9b822 100644
--- a/examples/auth/auth.go
+++ b/examples/auth/auth.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/milvus-io/milvus-sdk-go/v2/client"
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
-	"google.golang.org/grpc"
 )
 
 const (
@@ -22,31 +21,8 @@ const (
 
 func main() {
 	ctx := context.Background()
-	// create user credential if needed
-	c, err := client.NewGrpcClient(ctx, addr)
-	if err != nil {
-		log.Fatal("failed to connect:", err)
-	}
-	users, err := c.ListCredUsers(ctx)
-	registered := false
-	for _, user := range users {
-		if user == username {
-			registered = true
-			break
-		}
-	}
-	if !registered {
-		if err = c.CreateCredential(ctx, username, password); err != nil {
-			log.Fatal("failed to create credential:", err)
-		}
-	}
-	c.Close()
-
-	// client with auth interceptor
-	unaryAuthInterceptor := client.CreateAuthenticationUnaryInterceptor(username, password)
-	streamAuthInterceptor := client.CreateAuthenticationStreamInterceptor(username, password)
-	c, err = client.NewGrpcClient(ctx, addr, grpc.WithInsecure(),
-		grpc.WithUnaryInterceptor(unaryAuthInterceptor), grpc.WithStreamInterceptor(streamAuthInterceptor))
+	// create grpc client that tls is enabled
+	c, err := client.NewDefaultGrpcClientWithTLSAuth(ctx, addr, username, password)
 	if err != nil {
 		log.Fatal("failed to connect:", err)
 	}


Test patch:
diff --git a/client/client_test.go b/client/client_test.go
index bbbda51..4b2a876 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -203,11 +203,11 @@ func TestGrpcClientRetryPolicy(t *testing.T) {
 	}()
 	defer s.Stop()
 
-	grpcClient := &grpcClient{}
-	grpcClient.establishConnection(address)
-	defer grpcClient.Close()
+	client, err := NewGrpcClient(context.TODO(), address)
+	assert.Nil(t, err)
+	defer client.Close()
 
-	greeterClient := helloworld.NewGreeterClient(grpcClient.conn)
+	greeterClient := helloworld.NewGreeterClient(client.(*grpcClient).conn)
 	ctx := context.Background()
 	name := fmt.Sprintf("hello world %d", time.Now().Second())
 	res, err := greeterClient.SayHello(ctx, &helloworld.HelloRequest{Name: name})
@@ -215,6 +215,16 @@ func TestGrpcClientRetryPolicy(t *testing.T) {
 	assert.Equal(t, res.Message, strings.ToUpper(name))
 }
 
+func TestClient_getDefaultAuthOpts(t *testing.T) {
+	username := "u"
+	password := "pwd"
+	defaultOpts := getDefaultAuthOpts(username, password, true)
+	assert.True(t, len(defaultOpts) == 6)
+
+	defaultOpts = getDefaultAuthOpts(username, password, false)
+	assert.True(t, len(defaultOpts) == 6)
+}
+
 func successStatus() (*common.Status, error) {
 	return &common.Status{ErrorCode: common.ErrorCode_Success}, nil
 }