github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/client/interceptor.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"google.golang.org/grpc"
     8  	"google.golang.org/grpc/metadata"
     9  
    10  	"github.com/milvus-io/milvus-sdk-go/v2/internal/utils/crypto"
    11  )
    12  
    13  const (
    14  	authorizationHeader = `authorization`
    15  
    16  	identifierHeader = `identifier`
    17  
    18  	databaseHeader = `dbname`
    19  )
    20  
    21  // authenticationInterceptor appends credential into context metadata
    22  func authenticationInterceptor(ctx context.Context, username, password string) context.Context {
    23  	if username != "" || password != "" {
    24  		value := crypto.Base64Encode(fmt.Sprintf("%s:%s", username, password))
    25  		return metadata.AppendToOutgoingContext(ctx, authorizationHeader, value)
    26  	}
    27  	return ctx
    28  }
    29  
    30  func apiKeyInterceptor(ctx context.Context, apiKey string) context.Context {
    31  	if apiKey != "" {
    32  		value := crypto.Base64Encode(apiKey)
    33  		return metadata.AppendToOutgoingContext(ctx, authorizationHeader, value)
    34  	}
    35  	return ctx
    36  }
    37  
    38  func identifierInterceptor(ctx context.Context, identifierGetter func() string) context.Context {
    39  	identifier := identifierGetter()
    40  	if identifier != "" {
    41  		ctx = metadata.AppendToOutgoingContext(ctx, identifierHeader, identifier)
    42  	}
    43  	return ctx
    44  }
    45  
    46  // databaseNameInterceptor appends the dbName into metadata.
    47  func databaseNameInterceptor(ctx context.Context, dbNameGetter func() string) context.Context {
    48  	dbname := dbNameGetter()
    49  	if dbname != "" {
    50  		ctx = metadata.AppendToOutgoingContext(ctx, databaseHeader, dbname)
    51  	}
    52  	return ctx
    53  }
    54  
    55  // createMetaDataUnaryInterceptor creates a unary interceptor for metadata information.
    56  func createMetaDataUnaryInterceptor(cfg *Config) grpc.UnaryClientInterceptor {
    57  	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    58  		ctx = authenticationInterceptor(ctx, cfg.Username, cfg.Password)
    59  		ctx = apiKeyInterceptor(ctx, cfg.APIKey)
    60  		ctx = identifierInterceptor(ctx, func() string {
    61  			return cfg.Identifier
    62  		})
    63  		ctx = databaseNameInterceptor(ctx, func() string {
    64  			return cfg.DBName
    65  		})
    66  		return invoker(ctx, method, req, reply, cc, opts...)
    67  	}
    68  }