Repository: milvus-io/milvus-sdk-go
Pull request #473 (2023-05-29 08:00:28)

[2.2.x] Support API token config
/kind improvement


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/473#" title="Author self-approved">congqixia</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/v2.2.x/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":[]} -->
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/473?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) Report
> Merging [#473](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/473?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (63e8299) into [v2.2.x](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/commit/19d5b0b41a07192e598aabb0fcfe01ee7407106b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (19d5b0b) will **decrease** coverage by `0.20%`.
> The diff coverage is `30.00%`.

[![Impacted file tree graph](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/473/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://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/473?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)

```diff
@@            Coverage Diff             @@
##           v2.2.x     #473      +/-   ##
==========================================
- Coverage   81.28%   81.09%   -0.20%     
==========================================
  Files          33       33              
  Lines        2683     2692       +9     
==========================================
+ Hits         2181     2183       +2     
- Misses        403      409       +6     
- Partials       99      100       +1     
```


| [Impacted Files](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/473?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) | Coverage Δ | |
|---|---|---|
| [client/config.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/473?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NvbmZpZy5nbw==) | `60.56% <0.00%> (-3.62%)` | :arrow_down: |
| [client/interceptor.go](https://app.codecov.io/gh/milvus-io/milvus-sdk-go/pull/473?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2ludGVyY2VwdG9yLmdv) | `47.82% <50.00%> (-2.18%)` | :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/473#" title="Author self-approved">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/473#pullrequestreview-1449088501" title="LGTM">yah01</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/v2.2.x/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/config.go b/client/config.go
index 83cdb5e..1f1bd67 100644
--- a/client/config.go
+++ b/client/config.go
@@ -49,6 +49,7 @@ type Config struct {
 	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.
 
@@ -162,6 +163,13 @@ func (c *syncConfig) getDialOption() []grpc.DialOption {
 		)
 	}
 
+	// Construct api token
+	if c.cfg.APIKey != "" {
+		options = append(options, grpc.WithChainUnaryInterceptor(
+			createAPIKeyUnaryInteceptor(c.cfg.APIKey),
+		))
+	}
+
 	// Construct DBName field
 	dbNameGetter := func() string {
 		c.Lock()
diff --git a/client/interceptor.go b/client/interceptor.go
index 795009c..c4009c2 100644
--- a/client/interceptor.go
+++ b/client/interceptor.go
@@ -9,13 +9,22 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
+const (
+	authorizationHeader = `authorization`
+)
+
 // authenticationInterceptor appends credential into context metadata
 func authenticationInterceptor(ctx context.Context, username, password string) context.Context {
 	value := crypto.Base64Encode(fmt.Sprintf("%s:%s", username, password))
-	return metadata.AppendToOutgoingContext(ctx, "authorization", value)
+	return metadata.AppendToOutgoingContext(ctx, authorizationHeader, value)
 }
 
-// CreateAuthenticationUnaryInterceptor creates a unary interceptor for authentication
+func apiKeyInterceptor(ctx context.Context, apiKey string) context.Context {
+	value := crypto.Base64Encode(fmt.Sprintf("Bearer: %s", apiKey))
+	return metadata.AppendToOutgoingContext(ctx, authorizationHeader, value)
+}
+
+// CreateAuthenticationUnaryInterceptor creates a unary interceptor for basic authentication.
 func createAuthenticationUnaryInterceptor(username, password string) grpc.UnaryClientInterceptor {
 	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 		ctx = authenticationInterceptor(ctx, username, password)
@@ -23,6 +32,14 @@ func createAuthenticationUnaryInterceptor(username, password string) grpc.UnaryC
 	}
 }
 
+// createAPIKeyUnaryInteceptor creates a unary inteceptor for api key authentication.
+func createAPIKeyUnaryInteceptor(apiKey string) grpc.UnaryClientInterceptor {
+	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+		ctx = apiKeyInterceptor(ctx, apiKey)
+		return invoker(ctx, method, req, reply, cc, opts...)
+	}
+}
+
 // createAuthenticationStreamInterceptor creates a stream interceptor for authentication
 func createAuthenticationStreamInterceptor(username, password string) grpc.StreamClientInterceptor {
 	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {


Test patch:
diff --git a/client/config_test.go b/client/config_test.go
index 77b9527..0f33ed6 100644
--- a/client/config_test.go
+++ b/client/config_test.go
@@ -59,6 +59,7 @@ func TestClientConfig(t *testing.T) {
 		t,
 		&Config{
 			Address: "https://xxxx-xxxxxxxxxxxxxx-xxxxxx.aws-us-west-2.vectordb-sit.zillizcloud.com:19530/database_name",
+			APIKey:  "test-token",
 		},
 		"xxxx-xxxxxxxxxxxxxx-xxxxxx.aws-us-west-2.vectordb-sit.zillizcloud.com:19530", "database_name", true, false,
 	)
diff --git a/client/interceptor_test.go b/client/interceptor_test.go
index 0847719..cdd9379 100644
--- a/client/interceptor_test.go
+++ b/client/interceptor_test.go
@@ -32,3 +32,13 @@ func TestDBNameInterceptor(t *testing.T) {
 	assert.Equal(t, 1, len(md["dbname"]))
 	assert.Equal(t, testDBName, md["dbname"][0])
 }
+
+func TestAPIInterceptor(t *testing.T) {
+	ctx := context.Background()
+	ctx = apiKeyInterceptor(ctx, "test-token")
+	md, ok := metadata.FromOutgoingContext(ctx)
+	value := crypto.Base64Encode(fmt.Sprintf("Bearer: %s", "test-token"))
+	assert.True(t, ok)
+	assert.Equal(t, 1, len(md[authorizationHeader]))
+	assert.Equal(t, value, md[authorizationHeader][0])
+}