github.com/songzhibin97/gkit@v1.2.13/internal/metadata/metadata.go (about)

     1  package metadata
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strings"
     7  )
     8  
     9  // Metadata is our way of representing request headers internally.
    10  // They're used at the RPC level and translate back and forth
    11  // from Transport headers.
    12  type Metadata map[string]string
    13  
    14  // NewMetadata creates an MD from a given key-values map.
    15  func NewMetadata(mds ...map[string]string) Metadata {
    16  	md := Metadata{}
    17  	for _, m := range mds {
    18  		for k, v := range m {
    19  			md.Set(k, v)
    20  		}
    21  	}
    22  	return md
    23  }
    24  
    25  // GetValue returns the value associated with the passed key.
    26  func (m Metadata) GetValue(key string) string {
    27  	k := strings.ToLower(key)
    28  	return m[k]
    29  }
    30  
    31  // Set stores the key-value pair.
    32  func (m Metadata) Set(key string, value string) {
    33  	if key == "" || value == "" {
    34  		return
    35  	}
    36  	k := strings.ToLower(key)
    37  	m[k] = value
    38  }
    39  
    40  // Range iterate over element in metadata.
    41  func (m Metadata) Range(f func(k, v string) bool) {
    42  	for k, v := range m {
    43  		ret := f(k, v)
    44  		if !ret {
    45  			break
    46  		}
    47  	}
    48  }
    49  
    50  // Clone returns a deep copy of Metadata
    51  func (m Metadata) Clone() Metadata {
    52  	md := Metadata{}
    53  	for k, v := range m {
    54  		md[k] = v
    55  	}
    56  	return md
    57  }
    58  
    59  type serverMetadataKey struct{}
    60  
    61  // NewServerContext creates a new context with client md attached.
    62  func NewServerContext(ctx context.Context, md Metadata) context.Context {
    63  	return context.WithValue(ctx, serverMetadataKey{}, md)
    64  }
    65  
    66  // FromServerContext returns the server metadata in ctx if it exists.
    67  func FromServerContext(ctx context.Context) (Metadata, bool) {
    68  	md, ok := ctx.Value(serverMetadataKey{}).(Metadata)
    69  	return md, ok
    70  }
    71  
    72  type clientMetadataKey struct{}
    73  
    74  // NewClientContext creates a new context with client md attached.
    75  func NewClientContext(ctx context.Context, md Metadata) context.Context {
    76  	return context.WithValue(ctx, clientMetadataKey{}, md)
    77  }
    78  
    79  // FromClientContext returns the client metadata in ctx if it exists.
    80  func FromClientContext(ctx context.Context) (Metadata, bool) {
    81  	md, ok := ctx.Value(clientMetadataKey{}).(Metadata)
    82  	return md, ok
    83  }
    84  
    85  // AppendToClientContext returns a new context with the provided kv merged
    86  // with any existing metadata in the context.
    87  func AppendToClientContext(ctx context.Context, kv ...string) context.Context {
    88  	if len(kv)%2 == 1 {
    89  		panic(fmt.Sprintf("metadata: AppendToOutgoingContext got an odd number of input pairs for metadata: %d", len(kv)))
    90  	}
    91  	md, _ := FromClientContext(ctx)
    92  	md = md.Clone()
    93  	for i := 0; i < len(kv); i += 2 {
    94  		md.Set(kv[i], kv[i+1])
    95  	}
    96  	return NewClientContext(ctx, md)
    97  }
    98  
    99  // MergeToClientContext merge new metadata into ctx.
   100  func MergeToClientContext(ctx context.Context, cmd Metadata) context.Context {
   101  	md, _ := FromClientContext(ctx)
   102  	md = md.Clone()
   103  	for k, v := range cmd {
   104  		md[k] = v
   105  	}
   106  	return NewClientContext(ctx, md)
   107  }