databroker: add support for field masks on Put (#3210)

* databroker: add support for field masks on Put

* return errors

* clean up go.mod
This commit is contained in:
Caleb Doxsey 2022-03-29 16:36:40 -06:00 committed by GitHub
parent 8fc5dbf4c5
commit 2dc778035d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 381 additions and 134 deletions

View file

@ -11,12 +11,14 @@ import (
"github.com/google/btree"
"github.com/rs/zerolog"
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/types/known/fieldmaskpb"
"google.golang.org/protobuf/types/known/timestamppb"
"github.com/pomerium/pomerium/internal/log"
"github.com/pomerium/pomerium/internal/signal"
"github.com/pomerium/pomerium/pkg/cryptutil"
"github.com/pomerium/pomerium/pkg/grpc/databroker"
"github.com/pomerium/pomerium/pkg/protoutil"
"github.com/pomerium/pomerium/pkg/storage"
)
@ -204,7 +206,11 @@ func (backend *Backend) Lease(_ context.Context, leaseName, leaseID string, ttl
}
// Put puts a record into the in-memory store.
func (backend *Backend) Put(ctx context.Context, record *databroker.Record) (serverVersion uint64, err error) {
func (backend *Backend) Put(
ctx context.Context,
record *databroker.Record,
mask *fieldmaskpb.FieldMask,
) (serverVersion uint64, err error) {
if record == nil {
return backend.serverVersion, fmt.Errorf("records cannot be nil")
}
@ -219,14 +225,24 @@ func (backend *Backend) Put(ctx context.Context, record *databroker.Record) (ser
defer backend.mu.Unlock()
defer backend.onChange.Broadcast(ctx)
backend.recordChange(record)
c, ok := backend.lookup[record.GetType()]
if !ok {
c = NewRecordCollection()
backend.lookup[record.GetType()] = c
}
if mask != nil {
oldRecord := c.Get(record.GetId())
if oldRecord != nil {
record.Data, err = protoutil.MergeAnyWithFieldMask(oldRecord.Data, record.Data, mask)
if err != nil {
return serverVersion, err
}
}
}
backend.recordChange(record)
if record.GetDeletedAt() != nil {
c.Delete(record.GetId())
} else {