github.com/lastbackend/toolkit@v0.0.0-20241020043710-cafa37b95aad/pkg/context/metadata/metadata.go (about)

     1  package metadata
     2  
     3  import (
     4  	"context"
     5  )
     6  
     7  const emptyString = ""
     8  
     9  type MD map[string]string
    10  type mdKey struct{}
    11  
    12  func (m MD) Get(key string) (string, bool) {
    13  	val, ok := m[key]
    14  	if ok {
    15  		return val, ok
    16  	}
    17  	val, ok = m[key]
    18  	return val, ok
    19  }
    20  
    21  func (m MD) Set(key, val string) {
    22  	m[key] = val
    23  }
    24  
    25  func (m MD) Del(key string) {
    26  	delete(m, key)
    27  }
    28  
    29  func NewContext(ctx context.Context, md MD) context.Context {
    30  	return context.WithValue(ctx, mdKey{}, md)
    31  }
    32  
    33  func DeleteContextValue(ctx context.Context, k string) context.Context {
    34  	return SetContextValue(ctx, k, emptyString)
    35  }
    36  
    37  func SetContextValue(ctx context.Context, k, v string) context.Context {
    38  	md, ok := LoadFromContext(ctx)
    39  	if !ok {
    40  		md = make(MD, 0)
    41  	}
    42  	if len(v) != 0 {
    43  		md[k] = v
    44  	} else {
    45  		delete(md, k)
    46  	}
    47  	return context.WithValue(ctx, mdKey{}, md)
    48  }
    49  
    50  func GetContextValue(ctx context.Context, key string) (string, bool) {
    51  	md, ok := LoadFromContext(ctx)
    52  	if !ok {
    53  		return emptyString, ok
    54  	}
    55  	val, ok := md[key]
    56  	return val, ok
    57  }
    58  
    59  func LoadFromContext(ctx context.Context) (MD, bool) {
    60  	md, ok := ctx.Value(mdKey{}).(MD)
    61  	if !ok {
    62  		return nil, ok
    63  	}
    64  	newMD := make(MD, len(md))
    65  	for k, v := range md {
    66  		newMD[k] = v
    67  	}
    68  	return newMD, ok
    69  }
    70  
    71  func MergeContext(ctx context.Context, metadata MD, overwrite bool) context.Context {
    72  	if ctx == nil {
    73  		ctx = context.Background()
    74  	}
    75  	md, ok := ctx.Value(mdKey{}).(MD)
    76  	if !ok {
    77  		return ctx
    78  	}
    79  	cp := make(MD, len(md))
    80  	for k, v := range md {
    81  		cp[k] = v
    82  	}
    83  	for k, v := range metadata {
    84  		if _, ok := cp[k]; ok && !overwrite {
    85  			continue
    86  		}
    87  		if len(v) != 0 {
    88  			cp[k] = v
    89  		} else {
    90  			delete(cp, k)
    91  		}
    92  	}
    93  	return context.WithValue(ctx, mdKey{}, cp)
    94  }