feat: field mask util.
This commit is contained in:
210
fieldmaskutil/fieldmaskutil.go
Normal file
210
fieldmaskutil/fieldmaskutil.go
Normal file
@@ -0,0 +1,210 @@
|
||||
package fieldmaskutil
|
||||
|
||||
import (
|
||||
"google.golang.org/protobuf/proto"
|
||||
"google.golang.org/protobuf/reflect/protoreflect"
|
||||
)
|
||||
|
||||
// Filter keeps the msg fields that are listed in the paths and clears all the rest.
|
||||
//
|
||||
// This is a handy wrapper for NestedMask.Filter method.
|
||||
// If the same paths are used to process multiple proto messages use NestedMask.Filter method directly.
|
||||
func Filter(msg proto.Message, paths []string) {
|
||||
NestedMaskFromPaths(paths).Filter(msg)
|
||||
}
|
||||
|
||||
// Prune clears all the fields listed in paths from the given msg.
|
||||
//
|
||||
// This is a handy wrapper for NestedMask.Prune method.
|
||||
// If the same paths are used to process multiple proto messages use NestedMask.Filter method directly.
|
||||
func Prune(msg proto.Message, paths []string) {
|
||||
NestedMaskFromPaths(paths).Prune(msg)
|
||||
}
|
||||
|
||||
// Overwrite overwrites all the fields listed in paths in the dest msg using values from src msg.
|
||||
//
|
||||
// This is a handy wrapper for NestedMask.Overwrite method.
|
||||
// If the same paths are used to process multiple proto messages use NestedMask.Overwrite method directly.
|
||||
func Overwrite(src, dest proto.Message, paths []string) {
|
||||
NestedMaskFromPaths(paths).Overwrite(src, dest)
|
||||
}
|
||||
|
||||
// NestedMask represents a field mask as a recursive map.
|
||||
type NestedMask map[string]NestedMask
|
||||
|
||||
// NestedMaskFromPaths creates an instance of NestedMask for the given paths.
|
||||
func NestedMaskFromPaths(paths []string) NestedMask {
|
||||
mask := make(NestedMask)
|
||||
for _, path := range paths {
|
||||
curr := mask
|
||||
var letters []rune
|
||||
for _, letter := range path {
|
||||
if letter == '.' {
|
||||
if len(letters) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
key := string(letters)
|
||||
c, ok := curr[key]
|
||||
if !ok {
|
||||
c = make(NestedMask)
|
||||
curr[key] = c
|
||||
}
|
||||
curr = c
|
||||
letters = nil
|
||||
continue
|
||||
}
|
||||
letters = append(letters, letter)
|
||||
}
|
||||
if len(letters) != 0 {
|
||||
key := string(letters)
|
||||
if _, ok := curr[key]; !ok {
|
||||
curr[key] = make(NestedMask)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return mask
|
||||
}
|
||||
|
||||
// Filter keeps the msg fields that are listed in the paths and clears all the rest.
|
||||
//
|
||||
// If the mask is empty then all the fields are kept.
|
||||
// Paths are assumed to be valid and normalized otherwise the function may panic.
|
||||
// See google.golang.org/protobuf/types/known/fieldmaskpb for details.
|
||||
func (mask NestedMask) Filter(msg proto.Message) {
|
||||
if len(mask) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
rft := msg.ProtoReflect()
|
||||
rft.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool {
|
||||
m, ok := mask[string(fd.Name())]
|
||||
if ok {
|
||||
if len(m) == 0 {
|
||||
return true
|
||||
}
|
||||
|
||||
if fd.IsMap() {
|
||||
xmap := rft.Get(fd).Map()
|
||||
xmap.Range(func(mk protoreflect.MapKey, mv protoreflect.Value) bool {
|
||||
if mi, ok := m[mk.String()]; ok {
|
||||
if i, ok := mv.Interface().(protoreflect.Message); ok && len(mi) > 0 {
|
||||
mi.Filter(i.Interface())
|
||||
}
|
||||
} else {
|
||||
xmap.Clear(mk)
|
||||
}
|
||||
|
||||
return true
|
||||
})
|
||||
} else if fd.IsList() {
|
||||
list := rft.Get(fd).List()
|
||||
for i := 0; i < list.Len(); i++ {
|
||||
m.Filter(list.Get(i).Message().Interface())
|
||||
}
|
||||
} else if fd.Kind() == protoreflect.MessageKind {
|
||||
m.Filter(rft.Get(fd).Message().Interface())
|
||||
}
|
||||
} else {
|
||||
rft.Clear(fd)
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// Prune clears all the fields listed in paths from the given msg.
|
||||
//
|
||||
// All other fields are kept untouched. If the mask is empty no fields are cleared.
|
||||
// This operation is the opposite of NestedMask.Filter.
|
||||
// Paths are assumed to be valid and normalized otherwise the function may panic.
|
||||
// See google.golang.org/protobuf/types/known/fieldmaskpb for details.
|
||||
func (mask NestedMask) Prune(msg proto.Message) {
|
||||
if len(mask) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
rft := msg.ProtoReflect()
|
||||
rft.Range(func(fd protoreflect.FieldDescriptor, _ protoreflect.Value) bool {
|
||||
m, ok := mask[string(fd.Name())]
|
||||
if ok {
|
||||
if len(m) == 0 {
|
||||
rft.Clear(fd)
|
||||
return true
|
||||
}
|
||||
|
||||
if fd.IsMap() {
|
||||
xmap := rft.Get(fd).Map()
|
||||
xmap.Range(func(mk protoreflect.MapKey, mv protoreflect.Value) bool {
|
||||
if mi, ok := m[mk.String()]; ok {
|
||||
if i, ok := mv.Interface().(protoreflect.Message); ok && len(mi) > 0 {
|
||||
mi.Prune(i.Interface())
|
||||
} else {
|
||||
xmap.Clear(mk)
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
})
|
||||
} else if fd.IsList() {
|
||||
list := rft.Get(fd).List()
|
||||
for i := 0; i < list.Len(); i++ {
|
||||
m.Prune(list.Get(i).Message().Interface())
|
||||
}
|
||||
} else if fd.Kind() == protoreflect.MessageKind {
|
||||
m.Prune(rft.Get(fd).Message().Interface())
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// Overwrite overwrites all the fields listed in paths in the dest msg using values from src msg.
|
||||
//
|
||||
// All other fields are kept untouched. If the mask is empty, no fields are overwritten.
|
||||
// Supports scalars, messages, repeated fields, and maps.
|
||||
// If the parent of the field is nil message, the parent is initiated before overwriting the field
|
||||
// If the field in src is empty value, the field in dest is cleared.
|
||||
// Paths are assumed to be valid and normalized otherwise the function may panic.
|
||||
func (mask NestedMask) Overwrite(src, dest proto.Message) {
|
||||
mask.overwrite(src.ProtoReflect(), dest.ProtoReflect())
|
||||
}
|
||||
|
||||
func (mask NestedMask) overwrite(src, dest protoreflect.Message) {
|
||||
for k, v := range mask {
|
||||
srcFD := src.Descriptor().Fields().ByName(protoreflect.Name(k))
|
||||
destFD := dest.Descriptor().Fields().ByName(protoreflect.Name(k))
|
||||
if srcFD == nil || destFD == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// Leaf mask -> copy value from src to dest
|
||||
if len(v) == 0 {
|
||||
if srcFD.Kind() == destFD.Kind() { // TODO: Full type equality check
|
||||
val := src.Get(srcFD)
|
||||
if isValid(srcFD, val) {
|
||||
dest.Set(destFD, val)
|
||||
} else {
|
||||
dest.Clear(destFD)
|
||||
}
|
||||
}
|
||||
} else if srcFD.Kind() == protoreflect.MessageKind {
|
||||
// If dest field is nil
|
||||
if !dest.Get(destFD).Message().IsValid() {
|
||||
dest.Set(destFD, protoreflect.ValueOf(dest.Get(destFD).Message().New()))
|
||||
}
|
||||
v.overwrite(src.Get(srcFD).Message(), dest.Get(destFD).Message())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func isValid(fd protoreflect.FieldDescriptor, val protoreflect.Value) bool {
|
||||
if fd.IsMap() {
|
||||
return val.Map().IsValid()
|
||||
} else if fd.IsList() {
|
||||
return val.List().IsValid()
|
||||
} else if fd.Message() != nil {
|
||||
return val.Message().IsValid()
|
||||
}
|
||||
return true
|
||||
}
|
||||
59
fieldmaskutil/fieldmaskutil_test.go
Normal file
59
fieldmaskutil/fieldmaskutil_test.go
Normal file
@@ -0,0 +1,59 @@
|
||||
package fieldmaskutil
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func Test_NestedMaskFromPaths(t *testing.T) {
|
||||
type args struct {
|
||||
paths []string
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
args args
|
||||
want NestedMask
|
||||
}{
|
||||
{
|
||||
name: "no nested fields",
|
||||
args: args{paths: []string{"a", "b", "c"}},
|
||||
want: NestedMask{"a": NestedMask{}, "b": NestedMask{}, "c": NestedMask{}},
|
||||
},
|
||||
{
|
||||
name: "with nested fields",
|
||||
args: args{paths: []string{"aaa.bb.c", "dd.e", "f"}},
|
||||
want: NestedMask{
|
||||
"aaa": NestedMask{"bb": NestedMask{"c": NestedMask{}}},
|
||||
"dd": NestedMask{"e": NestedMask{}},
|
||||
"f": NestedMask{}},
|
||||
},
|
||||
{
|
||||
name: "single field",
|
||||
args: args{paths: []string{"a"}},
|
||||
want: NestedMask{"a": NestedMask{}},
|
||||
},
|
||||
{
|
||||
name: "empty fields",
|
||||
args: args{paths: []string{}},
|
||||
want: NestedMask{},
|
||||
},
|
||||
{
|
||||
name: "invalid input",
|
||||
args: args{paths: []string{".", "..", "..."}},
|
||||
want: NestedMask{},
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := NestedMaskFromPaths(tt.args.paths); !reflect.DeepEqual(got, tt.want) {
|
||||
t.Errorf("NestedMaskFromPaths() = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkNestedMaskFromPaths(b *testing.B) {
|
||||
for i := 0; i < b.N; i++ {
|
||||
NestedMaskFromPaths([]string{"aaa.bbb.c.d.e.f", "aa.b.cc.ddddddd", "e", "f", "g.h.i.j.k"})
|
||||
}
|
||||
}
|
||||
1
go.mod
1
go.mod
@@ -10,6 +10,7 @@ require (
|
||||
github.com/stretchr/testify v1.8.4
|
||||
golang.org/x/crypto v0.14.0
|
||||
golang.org/x/exp v0.0.0-20231006140011-7918f672742d
|
||||
google.golang.org/protobuf v1.31.0
|
||||
)
|
||||
|
||||
require (
|
||||
|
||||
7
go.sum
7
go.sum
@@ -3,6 +3,9 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/gobwas/glob v0.2.3 h1:A4xDbljILXROh+kObIiy5kIaPYD8e96x1tgBhUI5J+Y=
|
||||
github.com/gobwas/glob v0.2.3/go.mod h1:d3Ez4x06l9bZtSvzIay5+Yzi0fmZzPgnTbPcKjJAkT8=
|
||||
github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
|
||||
github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
|
||||
github.com/google/go-cmp v0.5.8 h1:e6P7q2lk1O+qJJb4BtCQXlK8vWEO8V1ZeuEdJNOqZyg=
|
||||
github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4=
|
||||
github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/gosimple/slug v1.13.1 h1:bQ+kpX9Qa6tHRaK+fZR0A0M2Kd7Pa5eHPPsb1JpHD+Q=
|
||||
@@ -30,6 +33,10 @@ golang.org/x/crypto v0.14.0 h1:wBqGXzWJW6m1XrIKlAH0Hs1JJ7+9KBwnIO8v66Q9cHc=
|
||||
golang.org/x/crypto v0.14.0/go.mod h1:MVFd36DqK4CsrnJYDkBA3VC4m2GkXAM0PvzMCn4JQf4=
|
||||
golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM8rJBtfilJ2qTU199MI=
|
||||
golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
|
||||
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
|
||||
google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
|
||||
google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
|
||||
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
|
||||
|
||||
Reference in New Issue
Block a user