metadata: add iterator method

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
This commit is contained in:
Василий Толстов 2021-01-21 18:35:31 +03:00
parent dcca28944e
commit 8a2c4c511e
3 changed files with 95 additions and 43 deletions

49
metadata/context.go Normal file
View 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)
}

View File

@ -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)
}

View File

@ -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")