client: use set instead of append for grpc headers

Signed-off-by: huanghaoyuanhhy <[email protected]>

Author: huanghaoyuanhhy

internal/client/milvus/grpc.go

@@ -317,19 +317,31 @@ func NewGrpc(cfg *paramtable.MilvusConfig) (*GrpcClient, error) {
 	return cli, nil
 }
 
-func (g *GrpcClient) newCtx(ctx context.Context) context.Context {
+func (g *GrpcClient) newAuthMD(ctx context.Context) metadata.MD {
+	md := metadata.MD{}
+	if outgoingMD, ok := metadata.FromOutgoingContext(ctx); ok {
+		md = outgoingMD.Copy()
+	}
+
 	if g.auth != "" {
-		ctx = metadata.AppendToOutgoingContext(ctx, authorizationHeader, g.auth)
+		md.Set(authorizationHeader, g.auth)
 	}
 	if g.identifier != "" {
-		ctx = metadata.AppendToOutgoingContext(ctx, identifierHeader, g.identifier)
+		md.Set(identifierHeader, g.identifier)
 	}
-	return ctx
+
+	return md
+}
+
+func (g *GrpcClient) newCtx(ctx context.Context) context.Context {
+	return metadata.NewOutgoingContext(ctx, g.newAuthMD(ctx))
 }
 
 func (g *GrpcClient) newCtxWithDB(ctx context.Context, db string) context.Context {
-	ctx = g.newCtx(ctx)
-	return metadata.AppendToOutgoingContext(ctx, databaseHeader, db)
+	md := g.newAuthMD(ctx)
+	md.Set(databaseHeader, db)
+
+	return metadata.NewOutgoingContext(ctx, md)
 }
 
 func (g *GrpcClient) connect(ctx context.Context) error {

internal/client/milvus/grpc_test.go

@@ -76,20 +76,49 @@ func TestIsRateLimitError(t *testing.T) {
 }
 
 func TestGrpcClient_newCtx(t *testing.T) {
-	cli := &GrpcClient{auth: "auth", identifier: "identifier"}
-	ctx := cli.newCtx(context.Background())
-	md, ok := metadata.FromOutgoingContext(ctx)
-	assert.True(t, ok)
-	assert.Equal(t, "auth", md.Get(authorizationHeader)[0])
-	assert.Equal(t, "identifier", md.Get(identifierHeader)[0])
+	t.Run("Normal", func(t *testing.T) {
+		cli := &GrpcClient{auth: "auth", identifier: "identifier"}
+		ctx := cli.newCtx(context.Background())
+		md, ok := metadata.FromOutgoingContext(ctx)
+		assert.True(t, ok)
+		assert.Equal(t, "auth", md.Get(authorizationHeader)[0])
+		assert.Len(t, md.Get(authorizationHeader), 1)
+		assert.Equal(t, "identifier", md.Get(identifierHeader)[0])
+		assert.Len(t, md.Get(identifierHeader), 1)
+	})
+
+	t.Run("SetMultipleTimes", func(t *testing.T) {
+		cli := &GrpcClient{auth: "auth", identifier: "identifier"}
+		ctx := cli.newCtx(context.Background())
+		ctx = cli.newCtx(ctx)
+		md, ok := metadata.FromOutgoingContext(ctx)
+		assert.True(t, ok)
+		assert.Equal(t, "auth", md.Get(authorizationHeader)[0])
+		assert.Len(t, md.Get(authorizationHeader), 1)
+		assert.Equal(t, "identifier", md.Get(identifierHeader)[0])
+		assert.Len(t, md.Get(identifierHeader), 1)
+	})
+
 }
 
 func TestGrpcClient_newCtxWithDB(t *testing.T) {
-	cli := &GrpcClient{}
-	ctx := cli.newCtxWithDB(context.Background(), "db")
-	md, ok := metadata.FromOutgoingContext(ctx)
-	assert.True(t, ok)
-	assert.Equal(t, "db", md.Get(databaseHeader)[0])
+	t.Run("Normal", func(t *testing.T) {
+		cli := &GrpcClient{}
+		ctx := cli.newCtxWithDB(context.Background(), "db")
+		md, ok := metadata.FromOutgoingContext(ctx)
+		assert.True(t, ok)
+		assert.Equal(t, "db", md.Get(databaseHeader)[0])
+	})
+
+	t.Run("SetMultipleTimes", func(t *testing.T) {
+		cli := &GrpcClient{}
+		ctx := cli.newCtxWithDB(context.Background(), "db")
+		ctx = cli.newCtxWithDB(ctx, "db2")
+		md, ok := metadata.FromOutgoingContext(ctx)
+		assert.True(t, ok)
+		assert.Equal(t, "db2", md.Get(databaseHeader)[0])
+		assert.Len(t, md.Get(databaseHeader), 1)
+	})
 }
 
 func TestGrpcClient_HasFeature(t *testing.T) {