github.com/volts-dev/volts@v0.0.0-20240120094013-5e9c65924106/internal/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  type metadataKey struct{}
    10  
    11  // Metadata is our way of representing request headers internally.
    12  // They're used at the RPC level and translate back and forth
    13  // from Transport headers.
    14  type Metadata map[string]string
    15  
    16  func (md Metadata) Get(key string) (string, bool) {
    17  	// attempt to get as is
    18  	val, ok := md[key]
    19  	if ok {
    20  		return val, ok
    21  	}
    22  
    23  	// attempt to get lower case
    24  	val, ok = md[strings.Title(key)]
    25  	return val, ok
    26  }
    27  
    28  func (md Metadata) Set(key, val string) {
    29  	md[key] = val
    30  }
    31  
    32  func (md Metadata) Delete(key string) {
    33  	// delete key as-is
    34  	delete(md, key)
    35  	// delete also Title key
    36  	delete(md, strings.Title(key))
    37  }
    38  
    39  // Copy makes a copy of the metadata
    40  func Copy(md Metadata) Metadata {
    41  	cmd := make(Metadata, len(md))
    42  	for k, v := range md {
    43  		cmd[k] = v
    44  	}
    45  	return cmd
    46  }
    47  
    48  // Delete key from metadata
    49  func Delete(ctx context.Context, k string) context.Context {
    50  	return Set(ctx, k, "")
    51  }
    52  
    53  // Set add key with val to metadata
    54  func Set(ctx context.Context, k, v string) context.Context {
    55  	md, ok := FromContext(ctx)
    56  	if !ok {
    57  		md = make(Metadata)
    58  	}
    59  	if v == "" {
    60  		delete(md, k)
    61  	} else {
    62  		md[k] = v
    63  	}
    64  	return context.WithValue(ctx, metadataKey{}, md)
    65  }
    66  
    67  // Get returns a single value from metadata in the context
    68  func Get(ctx context.Context, key string) (string, bool) {
    69  	md, ok := FromContext(ctx)
    70  	if !ok {
    71  		return "", ok
    72  	}
    73  	// attempt to get as is
    74  	val, ok := md[key]
    75  	if ok {
    76  		return val, ok
    77  	}
    78  
    79  	// attempt to get lower case
    80  	val, ok = md[strings.Title(key)]
    81  
    82  	return val, ok
    83  }
    84  
    85  // FromContext returns metadata from the given context
    86  func FromContext(ctx context.Context) (Metadata, bool) {
    87  	md, ok := ctx.Value(metadataKey{}).(Metadata)
    88  	if !ok {
    89  		return nil, ok
    90  	}
    91  
    92  	// capitalise all values
    93  	newMD := make(Metadata, len(md))
    94  	for k, v := range md {
    95  		newMD[strings.Title(k)] = v
    96  	}
    97  
    98  	return newMD, ok
    99  }
   100  
   101  // NewContext creates a new context with the given metadata
   102  func NewContext(ctx context.Context, md Metadata) context.Context {
   103  	return context.WithValue(ctx, metadataKey{}, md)
   104  }
   105  
   106  // MergeContext merges metadata to existing metadata, overwriting if specified
   107  func MergeContext(ctx context.Context, patchMd Metadata, overwrite bool) context.Context {
   108  	if ctx == nil {
   109  		ctx = context.Background()
   110  	}
   111  	md, _ := ctx.Value(metadataKey{}).(Metadata)
   112  	cmd := make(Metadata, len(md))
   113  	for k, v := range md {
   114  		cmd[k] = v
   115  	}
   116  	for k, v := range patchMd {
   117  		if _, ok := cmd[k]; ok && !overwrite {
   118  			// skip
   119  		} else if v != "" {
   120  			cmd[k] = v
   121  		} else {
   122  			delete(cmd, k)
   123  		}
   124  	}
   125  	return context.WithValue(ctx, metadataKey{}, cmd)
   126  }