github.com/annwntech/go-micro/v2@v2.9.5/metadata/metadata.go (about)

     1  // Package metadata is a way of defining message headers
     2  package metadata
     3  
     4  import (
     5  	"context"
     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  func (md Metadata) Get(key string) (string, bool) {
    15  	// attempt to get as is
    16  	val, ok := md[key]
    17  	if ok {
    18  		return val, ok
    19  	}
    20  
    21  	// attempt to get lower case
    22  	val, ok = md[strings.Title(key)]
    23  	return val, ok
    24  }
    25  
    26  func (md Metadata) Set(key, val string) {
    27  	md[key] = val
    28  }
    29  
    30  func (md Metadata) Delete(key string) {
    31  	// delete key as-is
    32  	delete(md, key)
    33  	// delete also Title key
    34  	delete(md, strings.Title(key))
    35  }
    36  
    37  // Copy makes a copy of the metadata
    38  func Copy(md Metadata) Metadata {
    39  	cmd := make(Metadata, len(md))
    40  	for k, v := range md {
    41  		cmd[k] = v
    42  	}
    43  	return cmd
    44  }
    45  
    46  // Delete key from metadata
    47  func Delete(ctx context.Context, k string) context.Context {
    48  	return Set(ctx, k, "")
    49  }
    50  
    51  // Set add key with val to metadata
    52  func Set(ctx context.Context, k, v string) context.Context {
    53  	md, ok := FromContext(ctx)
    54  	if !ok {
    55  		md = make(Metadata)
    56  	}
    57  	if v == "" {
    58  		delete(md, k)
    59  	} else {
    60  		md[k] = v
    61  	}
    62  	return context.WithValue(ctx, metadataKey{}, md)
    63  }
    64  
    65  // Get returns a single value from metadata in the context
    66  func Get(ctx context.Context, key string) (string, bool) {
    67  	md, ok := FromContext(ctx)
    68  	if !ok {
    69  		return "", ok
    70  	}
    71  	// attempt to get as is
    72  	val, ok := md[key]
    73  	if ok {
    74  		return val, ok
    75  	}
    76  
    77  	// attempt to get lower case
    78  	val, ok = md[strings.Title(key)]
    79  
    80  	return val, ok
    81  }
    82  
    83  // MergeContext merges metadata to existing metadata, overwriting if specified
    84  func MergeContext(ctx context.Context, patchMd Metadata, overwrite bool) context.Context {
    85  	if ctx == nil {
    86  		ctx = context.Background()
    87  	}
    88  	md, _ := ctx.Value(metadataKey{}).(Metadata)
    89  	cmd := make(Metadata, len(md))
    90  	for k, v := range md {
    91  		cmd[k] = v
    92  	}
    93  	for k, v := range patchMd {
    94  		if _, ok := cmd[k]; ok && !overwrite {
    95  			// skip
    96  		} else if v != "" {
    97  			cmd[k] = v
    98  		} else {
    99  			delete(cmd, k)
   100  		}
   101  	}
   102  	return context.WithValue(ctx, metadataKey{}, cmd)
   103  }