databroker: refactor databroker to sync all changes (#1879)

* refactor backend, implement encrypted store

* refactor in-memory store

* wip

* wip

* wip

* add syncer test

* fix redis expiry

* fix linting issues

* fix test by skipping non-config records

* fix backoff import

* fix init issues

* fix query

* wait for initial sync before starting directory sync

* add type to SyncLatest

* add more log messages, fix deadlock in in-memory store, always return server version from SyncLatest

* update sync types and tests

* add redis tests

* skip macos in github actions

* add comments to proto

* split getBackend into separate methods

* handle errors in initVersion

* return different error for not found vs other errors in get

* use exponential backoff for redis transaction retry

* rename raw to result

* use context instead of close channel

* store type urls as constants in databroker

* use timestampb instead of ptypes

* fix group merging not waiting

* change locked names

* update GetAll to return latest record version

* add method to grpcutil to get the type url for a protobuf type

pkg/grpc/config/config.go

@@ -1,3 +1,4 @@
+// Package config contains protobuf definitions for config.
 package config
 
 // IsSet returns true if one of the route redirect options has been chosen.

pkg/grpc/databroker/databroker.go

@@ -3,10 +3,9 @@ package databroker
 
 import (
 	"context"
+	"fmt"
 	"io"
 	"strings"
-
-	"google.golang.org/protobuf/proto"
 )
 
 // GetUserID gets the databroker user id from a provider user id.
@@ -37,19 +36,17 @@ func ApplyOffsetAndLimit(all []*Record, offset, limit int) (records []*Record, t
 	return records, len(all)
 }
 
-// InitialSync performs a sync with no_wait set to true and then returns all the results.
-func InitialSync(ctx context.Context, client DataBrokerServiceClient, in *SyncRequest) (*SyncResponse, error) {
-	dup := new(SyncRequest)
-	proto.Merge(dup, in)
-	dup.NoWait = true
-
-	stream, err := client.Sync(ctx, dup)
+// InitialSync performs a sync latest and then returns all the results.
+func InitialSync(
+	ctx context.Context,
+	client DataBrokerServiceClient,
+	req *SyncLatestRequest,
+) (records []*Record, recordVersion, serverVersion uint64, err error) {
+	stream, err := client.SyncLatest(ctx, req)
 	if err != nil {
-		return nil, err
+		return nil, 0, 0, err
 	}
 
-	finalRes := &SyncResponse{}
-
 loop:
 	for {
 		res, err := stream.Recv()
@@ -57,12 +54,19 @@ loop:
 		case err == io.EOF:
 			break loop
 		case err != nil:
-			return nil, err
+			return nil, 0, 0, err
 		}
 
-		finalRes.ServerVersion = res.GetServerVersion()
-		finalRes.Records = append(finalRes.Records, res.GetRecords()...)
+		switch res := res.GetResponse().(type) {
+		case *SyncLatestResponse_Versions:
+			recordVersion = res.Versions.GetLatestRecordVersion()
+			serverVersion = res.Versions.GetServerVersion()
+		case *SyncLatestResponse_Record:
+			records = append(records, res.Record)
+		default:
+			panic(fmt.Sprintf("unexpected response: %T", res))
+		}
 	}
 
-	return finalRes, nil
+	return records, recordVersion, serverVersion, nil
 }

pkg/grpc/databroker/databroker.proto

@@ -4,46 +4,27 @@ package databroker;
 option go_package = "github.com/pomerium/pomerium/pkg/grpc/databroker";
 
 import "google/protobuf/any.proto";
-import "google/protobuf/empty.proto";
 import "google/protobuf/timestamp.proto";
 
-message ServerVersion {
-  string version = 1;
-}
-
 message Record {
-  string version = 1;
+  uint64 version = 1;
   string type = 2;
   string id = 3;
   google.protobuf.Any data = 4;
-  google.protobuf.Timestamp created_at = 5;
-  google.protobuf.Timestamp modified_at = 6;
-  google.protobuf.Timestamp deleted_at = 7;
+  google.protobuf.Timestamp modified_at = 5;
+  google.protobuf.Timestamp deleted_at = 6;
 }
-
-message DeleteRequest {
-  string type = 1;
-  string id = 2;
+message Versions {
+  // the server version indicates the version of the server storing the data
+  uint64 server_version = 1;
+  uint64 latest_record_version = 2;
 }
 
 message GetRequest {
   string type = 1;
   string id = 2;
 }
-message GetResponse {
-  Record record = 1;
-}
-
-message GetAllRequest {
-  string type = 1;
-  string page_token = 2;
-}
-message GetAllResponse {
-  repeated Record records = 1;
-  string server_version = 2;
-  string record_version = 3;
-  string next_page_token = 4;
-}
+message GetResponse { Record record = 1; }
 
 message QueryRequest {
   string type = 1;
@@ -56,39 +37,39 @@ message QueryResponse {
   int64 total_count = 2;
 }
 
-message SetRequest {
-  string type = 1;
-  string id = 2;
-  google.protobuf.Any data = 3;
-}
-message SetResponse {
-  Record record = 1;
-  string server_version = 2;
+message PutRequest { Record record = 1; }
+message PutResponse {
+  uint64 server_version = 1;
+  Record record = 2;
 }
 
 message SyncRequest {
-  string server_version = 1;
-  string record_version = 2;
-  string type = 3;
-  bool no_wait = 4;
+  uint64 server_version = 1;
+  uint64 record_version = 2;
 }
 message SyncResponse {
-  string server_version = 1;
-  repeated Record records = 2;
+  uint64 server_version = 1;
+  Record record = 2;
 }
 
-message GetTypesResponse {
-  repeated string types = 1;
+message SyncLatestRequest { string type = 1; }
+message SyncLatestResponse {
+  oneof response {
+    Record record = 1;
+    Versions versions = 2;
+  }
 }
 
+// The DataBrokerService stores key-value data.
 service DataBrokerService {
-  rpc Delete(DeleteRequest) returns (google.protobuf.Empty);
+  // Get gets a record.
   rpc Get(GetRequest) returns (GetResponse);
-  rpc GetAll(GetAllRequest) returns (GetAllResponse);
+  // Put saves a record.
+  rpc Put(PutRequest) returns (PutResponse);
+  // Query queries for records.
   rpc Query(QueryRequest) returns (QueryResponse);
-  rpc Set(SetRequest) returns (SetResponse);
+  // Sync streams changes to records after the specified version.
   rpc Sync(SyncRequest) returns (stream SyncResponse);
-
-  rpc GetTypes(google.protobuf.Empty) returns (GetTypesResponse);
-  rpc SyncTypes(google.protobuf.Empty) returns (stream GetTypesResponse);
+  // SyncLatest streams the latest version of every record.
+  rpc SyncLatest(SyncLatestRequest) returns (stream SyncLatestResponse);
 }

pkg/grpc/databroker/databroker_test.go

@@ -61,18 +61,26 @@ func TestInitialSync(t *testing.T) {
 
 	r1 := new(Record)
 	r2 := new(Record)
-	r3 := new(Record)
 
 	m := &mockServer{
-		sync: func(req *SyncRequest, stream DataBrokerService_SyncServer) error {
-			assert.Equal(t, true, req.GetNoWait())
-			stream.Send(&SyncResponse{
-				ServerVersion: "a",
-				Records:       []*Record{r1, r2},
+		syncLatest: func(req *SyncLatestRequest, stream DataBrokerService_SyncLatestServer) error {
+			stream.Send(&SyncLatestResponse{
+				Response: &SyncLatestResponse_Record{
+					Record: r1,
+				},
 			})
-			stream.Send(&SyncResponse{
-				ServerVersion: "b",
-				Records:       []*Record{r3},
+			stream.Send(&SyncLatestResponse{
+				Response: &SyncLatestResponse_Record{
+					Record: r2,
+				},
+			})
+			stream.Send(&SyncLatestResponse{
+				Response: &SyncLatestResponse_Versions{
+					Versions: &Versions{
+						LatestRecordVersion: 2,
+						ServerVersion:       1,
+					},
+				},
 			})
 			return nil
 		},
@@ -90,20 +98,19 @@ func TestInitialSync(t *testing.T) {
 
 	c := NewDataBrokerServiceClient(cc)
 
-	res, err := InitialSync(ctx, c, &SyncRequest{
-		Type: "TEST",
-	})
+	records, recordVersion, serverVersion, err := InitialSync(ctx, c, new(SyncLatestRequest))
 	assert.NoError(t, err)
-	assert.Equal(t, "b", res.GetServerVersion())
-	assert.Equal(t, []*Record{r1, r2, r3}, res.GetRecords())
+	assert.Equal(t, uint64(2), recordVersion)
+	assert.Equal(t, uint64(1), serverVersion)
+	assert.Equal(t, []*Record{r1, r2}, records)
 }
 
 type mockServer struct {
 	DataBrokerServiceServer
 
-	sync func(*SyncRequest, DataBrokerService_SyncServer) error
+	syncLatest func(empty *SyncLatestRequest, server DataBrokerService_SyncLatestServer) error
 }
 
-func (m *mockServer) Sync(req *SyncRequest, stream DataBrokerService_SyncServer) error {
-	return m.sync(req, stream)
+func (m *mockServer) SyncLatest(req *SyncLatestRequest, stream DataBrokerService_SyncLatestServer) error {
+	return m.syncLatest(req, stream)
 }

pkg/grpc/databroker/syncer.go

@@ -0,0 +1,171 @@
+package databroker
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	backoff "github.com/cenkalti/backoff/v4"
+	"github.com/rs/zerolog"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/pomerium/pomerium/internal/log"
+)
+
+type syncerConfig struct {
+	typeURL string
+}
+
+// A SyncerOption customizes the syncer configuration.
+type SyncerOption func(cfg *syncerConfig)
+
+func getSyncerConfig(options ...SyncerOption) *syncerConfig {
+	cfg := new(syncerConfig)
+	for _, option := range options {
+		option(cfg)
+	}
+	return cfg
+}
+
+// WithTypeURL restricts the sync'd results to the given type.
+func WithTypeURL(typeURL string) SyncerOption {
+	return func(cfg *syncerConfig) {
+		cfg.typeURL = typeURL
+	}
+}
+
+// A SyncerHandler receives sync events from the Syncer.
+type SyncerHandler interface {
+	GetDataBrokerServiceClient() DataBrokerServiceClient
+	ClearRecords(ctx context.Context)
+	UpdateRecords(ctx context.Context, records []*Record)
+}
+
+// A Syncer is a helper type for working with Sync and SyncLatest. It will make a call to
+// SyncLatest to retrieve the latest version of the data, then begin syncing with a call
+// to Sync. If the server version changes `ClearRecords` will be called and the process
+// will start over.
+type Syncer struct {
+	cfg     *syncerConfig
+	handler SyncerHandler
+	backoff *backoff.ExponentialBackOff
+
+	recordVersion uint64
+	serverVersion uint64
+
+	closeCtx       context.Context
+	closeCtxCancel func()
+}
+
+// NewSyncer creates a new Syncer.
+func NewSyncer(handler SyncerHandler, options ...SyncerOption) *Syncer {
+	closeCtx, closeCtxCancel := context.WithCancel(context.Background())
+
+	bo := backoff.NewExponentialBackOff()
+	bo.MaxElapsedTime = 0
+	return &Syncer{
+		cfg:     getSyncerConfig(options...),
+		handler: handler,
+		backoff: bo,
+
+		closeCtx:       closeCtx,
+		closeCtxCancel: closeCtxCancel,
+	}
+}
+
+// Close closes the Syncer.
+func (syncer *Syncer) Close() error {
+	syncer.closeCtxCancel()
+	return nil
+}
+
+// Run runs the Syncer.
+func (syncer *Syncer) Run(ctx context.Context) error {
+	ctx, cancel := context.WithCancel(ctx)
+	go func() {
+		<-syncer.closeCtx.Done()
+		cancel()
+	}()
+
+	for {
+		var err error
+		if syncer.serverVersion == 0 {
+			err = syncer.init(ctx)
+		} else {
+			err = syncer.sync(ctx)
+		}
+
+		if err != nil {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-time.After(syncer.backoff.NextBackOff()):
+			}
+		}
+	}
+}
+
+func (syncer *Syncer) init(ctx context.Context) error {
+	syncer.log().Info().Msg("syncing latest records")
+	records, recordVersion, serverVersion, err := InitialSync(ctx, syncer.handler.GetDataBrokerServiceClient(), &SyncLatestRequest{
+		Type: syncer.cfg.typeURL,
+	})
+	if err != nil {
+		syncer.log().Error().Err(err).Msg("error during initial sync")
+		return err
+	}
+	syncer.backoff.Reset()
+
+	// reset the records as we have to sync latest
+	syncer.handler.ClearRecords(ctx)
+
+	syncer.recordVersion = recordVersion
+	syncer.serverVersion = serverVersion
+	syncer.handler.UpdateRecords(ctx, records)
+
+	return nil
+}
+
+func (syncer *Syncer) sync(ctx context.Context) error {
+	stream, err := syncer.handler.GetDataBrokerServiceClient().Sync(ctx, &SyncRequest{
+		ServerVersion: syncer.serverVersion,
+		RecordVersion: syncer.recordVersion,
+	})
+	if err != nil {
+		syncer.log().Error().Err(err).Msg("error during sync")
+		return err
+	}
+
+	for {
+		res, err := stream.Recv()
+		if status.Code(err) == codes.Aborted {
+			syncer.log().Error().Err(err).Msg("aborted sync due to mismatched server version")
+			// server version changed, so re-init
+			syncer.serverVersion = 0
+			return nil
+		} else if err != nil {
+			return err
+		}
+
+		if syncer.recordVersion != res.GetRecord().GetVersion()-1 {
+			syncer.log().Error().Err(err).
+				Uint64("received", res.GetRecord().GetVersion()).
+				Msg("aborted sync due to missing record")
+			syncer.serverVersion = 0
+			return fmt.Errorf("missing record version")
+		}
+		syncer.recordVersion = res.GetRecord().GetVersion()
+		if syncer.cfg.typeURL == "" || syncer.cfg.typeURL == res.GetRecord().GetType() {
+			syncer.handler.UpdateRecords(ctx, []*Record{res.GetRecord()})
+		}
+	}
+}
+
+func (syncer *Syncer) log() *zerolog.Logger {
+	l := log.With().Str("service", "syncer").
+		Str("type", syncer.cfg.typeURL).
+		Uint64("server_version", syncer.serverVersion).
+		Uint64("record_version", syncer.recordVersion).Logger()
+	return &l
+}

pkg/grpc/databroker/syncer_test.go

@@ -0,0 +1,222 @@
+package databroker
+
+import (
+	"context"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+	"google.golang.org/grpc/test/bufconn"
+
+	"github.com/pomerium/pomerium/internal/testutil"
+)
+
+type testSyncerHandler struct {
+	getDataBrokerServiceClient func() DataBrokerServiceClient
+	clearRecords               func(ctx context.Context)
+	updateRecords              func(ctx context.Context, records []*Record)
+}
+
+func (t testSyncerHandler) GetDataBrokerServiceClient() DataBrokerServiceClient {
+	return t.getDataBrokerServiceClient()
+}
+
+func (t testSyncerHandler) ClearRecords(ctx context.Context) {
+	t.clearRecords(ctx)
+}
+
+func (t testSyncerHandler) UpdateRecords(ctx context.Context, records []*Record) {
+	t.updateRecords(ctx, records)
+}
+
+type testServer struct {
+	DataBrokerServiceServer
+	sync       func(request *SyncRequest, server DataBrokerService_SyncServer) error
+	syncLatest func(req *SyncLatestRequest, server DataBrokerService_SyncLatestServer) error
+}
+
+func (t testServer) Sync(request *SyncRequest, server DataBrokerService_SyncServer) error {
+	return t.sync(request, server)
+}
+
+func (t testServer) SyncLatest(req *SyncLatestRequest, server DataBrokerService_SyncLatestServer) error {
+	return t.syncLatest(req, server)
+}
+
+func TestSyncer(t *testing.T) {
+	ctx := context.Background()
+	ctx, clearTimeout := context.WithTimeout(ctx, time.Second*10)
+	defer clearTimeout()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	lis := bufconn.Listen(1)
+	r1 := &Record{Version: 1000, Id: "r1"}
+	r2 := &Record{Version: 1001, Id: "r2"}
+	r3 := &Record{Version: 1002, Id: "r3"}
+	r5 := &Record{Version: 1004, Id: "r5"}
+
+	syncCount := 0
+	syncLatestCount := 0
+
+	gs := grpc.NewServer()
+	RegisterDataBrokerServiceServer(gs, testServer{
+		sync: func(request *SyncRequest, server DataBrokerService_SyncServer) error {
+			syncCount++
+			switch syncCount {
+			case 1:
+				return status.Error(codes.Internal, "SOME INTERNAL ERROR")
+			case 2:
+				return status.Error(codes.Aborted, "ABORTED")
+			case 3:
+				_ = server.Send(&SyncResponse{
+					ServerVersion: 2001,
+					Record:        r3,
+				})
+				_ = server.Send(&SyncResponse{
+					ServerVersion: 2001,
+					Record:        r5,
+				})
+			case 4:
+				select {} // block forever
+			default:
+				t.Fatal("unexpected call to sync", request)
+			}
+			return nil
+		},
+		syncLatest: func(req *SyncLatestRequest, server DataBrokerService_SyncLatestServer) error {
+			syncLatestCount++
+			switch syncLatestCount {
+			case 1:
+				_ = server.Send(&SyncLatestResponse{
+					Response: &SyncLatestResponse_Record{
+						Record: r1,
+					},
+				})
+				_ = server.Send(&SyncLatestResponse{
+					Response: &SyncLatestResponse_Versions{
+						Versions: &Versions{
+							LatestRecordVersion: r1.Version,
+							ServerVersion:       2000,
+						},
+					},
+				})
+			case 2:
+				_ = server.Send(&SyncLatestResponse{
+					Response: &SyncLatestResponse_Record{
+						Record: r2,
+					},
+				})
+				_ = server.Send(&SyncLatestResponse{
+					Response: &SyncLatestResponse_Versions{
+						Versions: &Versions{
+							LatestRecordVersion: r2.Version,
+							ServerVersion:       2001,
+						},
+					},
+				})
+			case 3:
+				return status.Error(codes.Internal, "SOME INTERNAL ERROR")
+			case 4:
+				_ = server.Send(&SyncLatestResponse{
+					Response: &SyncLatestResponse_Record{
+						Record: r3,
+					},
+				})
+				_ = server.Send(&SyncLatestResponse{
+					Response: &SyncLatestResponse_Record{
+						Record: r5,
+					},
+				})
+				_ = server.Send(&SyncLatestResponse{
+					Response: &SyncLatestResponse_Versions{
+						Versions: &Versions{
+							LatestRecordVersion: r5.Version,
+							ServerVersion:       2001,
+						},
+					},
+				})
+			default:
+				t.Fatal("unexpected call to sync latest")
+			}
+			return nil
+		},
+	})
+	go func() { _ = gs.Serve(lis) }()
+
+	gc, err := grpc.DialContext(ctx, "bufnet",
+		grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
+			return lis.Dial()
+		}),
+		grpc.WithInsecure())
+	require.NoError(t, err)
+	defer func() { _ = gc.Close() }()
+
+	clearCh := make(chan struct{})
+	updateCh := make(chan []*Record)
+	syncer := NewSyncer(testSyncerHandler{
+		getDataBrokerServiceClient: func() DataBrokerServiceClient {
+			return NewDataBrokerServiceClient(gc)
+		},
+		clearRecords: func(ctx context.Context) {
+			clearCh <- struct{}{}
+		},
+		updateRecords: func(ctx context.Context, records []*Record) {
+			updateCh <- records
+		},
+	})
+	go func() { _ = syncer.Run(ctx) }()
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("1. expected call to clear records")
+	case <-clearCh:
+	}
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("2. expected call to update records")
+	case records := <-updateCh:
+		testutil.AssertProtoJSONEqual(t, `[{"id": "r1", "version": "1000"}]`, records)
+	}
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("3. expected call to clear records due to server version change")
+	case <-clearCh:
+	}
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("4. expected call to update records")
+	case records := <-updateCh:
+		testutil.AssertProtoJSONEqual(t, `[{"id": "r2", "version": "1001"}]`, records)
+	}
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("5. expected call to update records from sync")
+	case records := <-updateCh:
+		testutil.AssertProtoJSONEqual(t, `[{"id": "r3", "version": "1002"}]`, records)
+	}
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("6. expected call to clear records due to skipped version")
+	case <-clearCh:
+	}
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("7. expected call to update records")
+	case records := <-updateCh:
+		testutil.AssertProtoJSONEqual(t, `[{"id": "r3", "version": "1002"}, {"id": "r5", "version": "1004"}]`, records)
+	}
+
+	assert.NoError(t, syncer.Close())
+}

pkg/grpc/session/session.go

@@ -17,9 +17,13 @@ import (
 // Delete deletes a session from the databroker.
 func Delete(ctx context.Context, client databroker.DataBrokerServiceClient, sessionID string) error {
 	any, _ := ptypes.MarshalAny(new(Session))
-	_, err := client.Delete(ctx, &databroker.DeleteRequest{
-		Type: any.GetTypeUrl(),
-		Id:   sessionID,
+	_, err := client.Put(ctx, &databroker.PutRequest{
+		Record: &databroker.Record{
+			Type:      any.GetTypeUrl(),
+			Id:        sessionID,
+			Data:      any,
+			DeletedAt: timestamppb.Now(),
+		},
 	})
 	return err
 }
@@ -44,13 +48,15 @@ func Get(ctx context.Context, client databroker.DataBrokerServiceClient, session
 	return &s, nil
 }
 
-// Set sets a session in the databroker.
-func Set(ctx context.Context, client databroker.DataBrokerServiceClient, s *Session) (*databroker.SetResponse, error) {
+// Put sets a session in the databroker.
+func Put(ctx context.Context, client databroker.DataBrokerServiceClient, s *Session) (*databroker.PutResponse, error) {
 	any, _ := anypb.New(s)
-	res, err := client.Set(ctx, &databroker.SetRequest{
-		Type: any.GetTypeUrl(),
-		Id:   s.Id,
-		Data: any,
+	res, err := client.Put(ctx, &databroker.PutRequest{
+		Record: &databroker.Record{
+			Type: any.GetTypeUrl(),
+			Id:   s.Id,
+			Data: any,
+		},
 	})
 	return res, err
 }

pkg/grpc/user/user.go

@@ -32,13 +32,15 @@ func Get(ctx context.Context, client databroker.DataBrokerServiceClient, userID
 	return &u, nil
 }
 
-// Set sets a user in the databroker.
-func Set(ctx context.Context, client databroker.DataBrokerServiceClient, u *User) (*databroker.Record, error) {
+// Put sets a user in the databroker.
+func Put(ctx context.Context, client databroker.DataBrokerServiceClient, u *User) (*databroker.Record, error) {
 	any, _ := anypb.New(u)
-	res, err := client.Set(ctx, &databroker.SetRequest{
-		Type: any.GetTypeUrl(),
-		Id:   u.Id,
-		Data: any,
+	res, err := client.Put(ctx, &databroker.PutRequest{
+		Record: &databroker.Record{
+			Type: any.GetTypeUrl(),
+			Id:   u.Id,
+			Data: any,
+		},
 	})
 	if err != nil {
 		return nil, err
@@ -46,13 +48,15 @@ func Set(ctx context.Context, client databroker.DataBrokerServiceClient, u *User
 	return res.GetRecord(), nil
 }
 
-// SetServiceAccount sets a service account in the databroker.
-func SetServiceAccount(ctx context.Context, client databroker.DataBrokerServiceClient, sa *ServiceAccount) (*databroker.Record, error) {
+// PutServiceAccount sets a service account in the databroker.
+func PutServiceAccount(ctx context.Context, client databroker.DataBrokerServiceClient, sa *ServiceAccount) (*databroker.Record, error) {
 	any, _ := anypb.New(sa)
-	res, err := client.Set(ctx, &databroker.SetRequest{
-		Type: any.GetTypeUrl(),
-		Id:   sa.GetId(),
-		Data: any,
+	res, err := client.Put(ctx, &databroker.PutRequest{
+		Record: &databroker.Record{
+			Type: any.GetTypeUrl(),
+			Id:   sa.GetId(),
+			Data: any,
+		},
 	})
 	if err != nil {
 		return nil, err