metadata: add iterator method
Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
This commit is contained in:
parent
dcca28944e
commit
8a2c4c511e
49
metadata/context.go
Normal file
49
metadata/context.go
Normal file
@ -0,0 +1,49 @@
|
||||
// Package metadata is a way of defining message headers
|
||||
package metadata
|
||||
|
||||
import (
|
||||
"context"
|
||||
)
|
||||
|
||||
// FromContext returns metadata from the given context
|
||||
func FromContext(ctx context.Context) (Metadata, bool) {
|
||||
if ctx == nil {
|
||||
return nil, false
|
||||
}
|
||||
md, ok := ctx.Value(metadataKey{}).(Metadata)
|
||||
if !ok {
|
||||
return nil, ok
|
||||
}
|
||||
nmd := Copy(md)
|
||||
return nmd, ok
|
||||
}
|
||||
|
||||
// NewContext creates a new context with the given metadata
|
||||
func NewContext(ctx context.Context, md Metadata) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
return context.WithValue(ctx, metadataKey{}, Copy(md))
|
||||
}
|
||||
|
||||
// MergeContext merges metadata to existing metadata, overwriting if specified
|
||||
func MergeContext(ctx context.Context, pmd Metadata, overwrite bool) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
md, ok := FromContext(ctx)
|
||||
if !ok {
|
||||
return context.WithValue(ctx, metadataKey{}, Copy(pmd))
|
||||
}
|
||||
nmd := Copy(md)
|
||||
for key, val := range pmd {
|
||||
if _, ok := nmd[key]; ok && !overwrite {
|
||||
// skip
|
||||
} else if val != "" {
|
||||
nmd.Set(key, val)
|
||||
} else {
|
||||
nmd.Del(key)
|
||||
}
|
||||
}
|
||||
return context.WithValue(ctx, metadataKey{}, nmd)
|
||||
}
|
@ -4,6 +4,7 @@ package metadata
|
||||
import (
|
||||
"context"
|
||||
"net/textproto"
|
||||
"sort"
|
||||
)
|
||||
|
||||
type metadataKey struct{}
|
||||
@ -18,6 +19,35 @@ var (
|
||||
DefaultMetadataSize = 6
|
||||
)
|
||||
|
||||
type Iterator struct {
|
||||
cur int
|
||||
cnt int
|
||||
keys []string
|
||||
md Metadata
|
||||
}
|
||||
|
||||
func (iter *Iterator) Next(k, v *string) bool {
|
||||
if iter.cur+1 > iter.cnt {
|
||||
return false
|
||||
}
|
||||
|
||||
*k = iter.keys[iter.cur]
|
||||
*v = iter.md[*k]
|
||||
iter.cur++
|
||||
return true
|
||||
}
|
||||
|
||||
// Iterate returns run user func with map key, val sorted by key
|
||||
func (md Metadata) Iterator() *Iterator {
|
||||
iter := &Iterator{md: md, cnt: len(md)}
|
||||
iter.keys = make([]string, 0, iter.cnt)
|
||||
for k := range md {
|
||||
iter.keys = append(iter.keys, k)
|
||||
}
|
||||
sort.Strings(iter.keys)
|
||||
return iter
|
||||
}
|
||||
|
||||
// Get returns value from metadata by key
|
||||
func (md Metadata) Get(key string) (string, bool) {
|
||||
// fast path
|
||||
@ -83,19 +113,6 @@ func Get(ctx context.Context, key string) (string, bool) {
|
||||
return md.Get(key)
|
||||
}
|
||||
|
||||
// FromContext returns metadata from the given context
|
||||
func FromContext(ctx context.Context) (Metadata, bool) {
|
||||
if ctx == nil {
|
||||
return nil, false
|
||||
}
|
||||
md, ok := ctx.Value(metadataKey{}).(Metadata)
|
||||
if !ok {
|
||||
return nil, ok
|
||||
}
|
||||
nmd := Copy(md)
|
||||
return nmd, ok
|
||||
}
|
||||
|
||||
// New return new sized metadata
|
||||
func New(size int) Metadata {
|
||||
if size == 0 {
|
||||
@ -103,33 +120,3 @@ func New(size int) Metadata {
|
||||
}
|
||||
return make(Metadata, size)
|
||||
}
|
||||
|
||||
// NewContext creates a new context with the given metadata
|
||||
func NewContext(ctx context.Context, md Metadata) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
return context.WithValue(ctx, metadataKey{}, Copy(md))
|
||||
}
|
||||
|
||||
// MergeContext merges metadata to existing metadata, overwriting if specified
|
||||
func MergeContext(ctx context.Context, pmd Metadata, overwrite bool) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
md, ok := FromContext(ctx)
|
||||
if !ok {
|
||||
return context.WithValue(ctx, metadataKey{}, Copy(pmd))
|
||||
}
|
||||
nmd := Copy(md)
|
||||
for key, val := range pmd {
|
||||
if _, ok := nmd[key]; ok && !overwrite {
|
||||
// skip
|
||||
} else if val != "" {
|
||||
nmd.Set(key, val)
|
||||
} else {
|
||||
nmd.Del(key)
|
||||
}
|
||||
}
|
||||
return context.WithValue(ctx, metadataKey{}, nmd)
|
||||
}
|
||||
|
@ -2,10 +2,26 @@ package metadata
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestIterator(t *testing.T) {
|
||||
md := Metadata{
|
||||
"1Last": "last",
|
||||
"2First": "first",
|
||||
"3Second": "second",
|
||||
}
|
||||
|
||||
iter := md.Iterator()
|
||||
var k, v string
|
||||
|
||||
for iter.Next(&k, &v) {
|
||||
fmt.Printf("k: %s, v: %s\n", k, v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMedataCanonicalKey(t *testing.T) {
|
||||
ctx := Set(context.TODO(), "x-request-id", "12345")
|
||||
v, ok := Get(ctx, "x-request-id")
|
||||
|
Loading…
Reference in New Issue
Block a user