Implement the GetBalances RPC

go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/dynamodb v1.59.0
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.102.2
 	github.com/code-payments/code-vm-indexer v1.2.0
-	github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54
+	github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260914142118-62b4738e3521
 	github.com/emirpasic/gods v1.12.0
 	github.com/envoyproxy/protoc-gen-validate v1.3.3
 	github.com/golang/protobuf v1.5.4

go.sum

@@ -80,8 +80,8 @@ github.com/cockroachdb/apd v1.1.0 h1:3LFP3629v+1aKXU5Q37mxmRxX/pIu1nijXydLShEq5I
 github.com/cockroachdb/apd v1.1.0/go.mod h1:8Sl8LxpKi29FqWXR16WEFZRNSz3SoPzUzeMeY4+DwBQ=
 github.com/code-payments/code-vm-indexer v1.2.0 h1:rSHpBMiT9BKgmKcXg/VIoi/h0t7jNxGx07Qz59m+6Q0=
 github.com/code-payments/code-vm-indexer v1.2.0/go.mod h1:vn91YN2qNqb+gGJeZe2+l+TNxVmEEiRHXXnIn2Y40h8=
-github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54 h1:idVNVzNMZWunI9RlDid6GNto1+z8HVApQhHqmfRIAPc=
-github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260911150858-07d584b02a54/go.mod h1:tw6BooY5a8l6CtSZnKOruyKII0W04n89pcM4BizrgG8=
+github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260914142118-62b4738e3521 h1:b6uE5r53BiWXhjJ9UquZzNYsrvLl6nERUdXIuuqYGwQ=
+github.com/code-payments/ocp-protobuf-api v1.16.1-0.20260914142118-62b4738e3521/go.mod h1:tw6BooY5a8l6CtSZnKOruyKII0W04n89pcM4BizrgG8=
 github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6 h1:NmTXa/uVnDyp0TY5MKi197+3HWcnYWfnHGyaFthlnGw=
 github.com/containerd/continuity v0.0.0-20190827140505-75bee3e2ccb6/go.mod h1:GL3xCUCBDV3CZiTSEKksMWbLE66hEyuu9qyDOOqM47Y=
 github.com/coreos/bbolt v1.3.2/go.mod h1:iRUV2dpdMOn7Bo10OQBFzIJO9kkE559Wcmn+qkEiiKk=

ocp/balance/calculator.go

@@ -185,21 +185,56 @@ func BatchCalculateFromCache(ctx context.Context, data ocp_data.Provider, tokenA
 //
 // Note: Use this method when calculating balances for accounts that are managed by
 // Code (ie. Timelock account) and operate within the L2 system.
-func BatchCalculateFromCacheByOwner(ctx context.Context, data ocp_data.Provider, owner *common.Account) (map[string]*Balance, error) {
-	tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateFromCacheByOwner")
-	tracer.AddAttribute("owner", owner.PublicKey().ToBase58())
+func BatchCalculateFromCacheByOwner(ctx context.Context, data ocp_data.Provider, owner *common.Account, mints ...*common.Account) (map[string]*Balance, error) {
+	balancesByOwner, err := BatchCalculateFromCacheByOwners(ctx, data, []*common.Account{owner}, mints)
+	if err != nil {
+		return nil, err
+	}
+
+	res, ok := balancesByOwner[owner.PublicKey().ToBase58()]
+	if !ok {
+		return make(map[string]*Balance), nil
+	}
+	return res, nil
+}
+
+// BatchCalculateFromCacheByOwners is BatchCalculateFromCacheByOwner for a set
+// of owners in a single ledger read. The result is keyed by owner, then by
+// token account. Owners the ledger holds nothing for are omitted. An empty
+// mints slice includes every mint, otherwise the read is limited to the
+// provided mints.
+func BatchCalculateFromCacheByOwners(ctx context.Context, data ocp_data.Provider, owners, mints []*common.Account) (map[string]map[string]*Balance, error) {
+	tracer := metrics.TraceMethodCall(ctx, metricsPackageName, "BatchCalculateFromCacheByOwners")
+	tracer.AddAttribute("owner_count", len(owners))
+	tracer.AddAttribute("mint_count", len(mints))
 	defer tracer.End()
 
-	balanceRecords, err := data.GetAllBalancesByOwner(ctx, owner.PublicKey().ToBase58())
-	if err != nil && err != balance.ErrRecordNotFound {
+	ownerAddresses := make([]string, len(owners))
+	for i, owner := range owners {
+		ownerAddresses[i] = owner.PublicKey().ToBase58()
+	}
+
+	mintAddresses := make([]string, len(mints))
+	for i, mint := range mints {
+		mintAddresses[i] = mint.PublicKey().ToBase58()
+	}
+
+	balanceRecordsByOwner, err := data.GetAllBalancesByOwnerBatch(ctx, ownerAddresses, mintAddresses)
+	if err != nil {
 		tracer.OnError(err)
 		return nil, err
 	}
 
-	res := make(map[string]*Balance, len(balanceRecords))
-	for _, balanceRecord := range balanceRecords {
-		if cached, ok := balanceFromRecord(balanceRecord); ok {
-			res[balanceRecord.TokenAccount] = cached
+	res := make(map[string]map[string]*Balance, len(balanceRecordsByOwner))
+	for owner, balanceRecords := range balanceRecordsByOwner {
+		balancesByTokenAccount := make(map[string]*Balance, len(balanceRecords))
+		for _, balanceRecord := range balanceRecords {
+			if cached, ok := balanceFromRecord(balanceRecord); ok {
+				balancesByTokenAccount[balanceRecord.TokenAccount] = cached
+			}
+		}
+		if len(balancesByTokenAccount) > 0 {
+			res[owner] = balancesByTokenAccount
 		}
 	}
 	return res, nil

ocp/balance/calculator_test.go

@@ -120,6 +120,62 @@ func TestDefaultCalculationMethods_ByOwner(t *testing.T) {
 	assert.Empty(t, balanceByAccount)
 }
 
+func TestDefaultCalculationMethods_ByOwners(t *testing.T) {
+	env := setupBalanceTestEnv(t)
+
+	coreVmConfig := testutil.NewRandomVmConfig(t, true)
+	otherVmConfig := testutil.NewRandomVmConfig(t, false)
+
+	owner1 := testutil.NewRandomAccount(t)
+	owner2 := testutil.NewRandomAccount(t)
+	unknownOwner := testutil.NewRandomAccount(t)
+
+	owner1CoreMint := newBalanceTestAccountForOwner(t, env, owner1, coreVmConfig)
+	owner1OtherMint := newBalanceTestAccountForOwner(t, env, owner1, otherVmConfig)
+	owner2CoreMint := newBalanceTestAccountForOwner(t, env, owner2, coreVmConfig)
+	owner2Unlocked := newBalanceTestAccountForOwner(t, env, owner2, otherVmConfig)
+
+	saveBalanceTestRecord(t, env, owner1CoreMint, &balance.Record{Quarks: 42, IsOpen: true, IsLocked: true})
+	saveBalanceTestRecord(t, env, owner1OtherMint, &balance.Record{Quarks: 33, IsOpen: true, IsLocked: true})
+	saveBalanceTestRecord(t, env, owner2CoreMint, &balance.Record{Quarks: 7, IsOpen: true, IsLocked: true})
+	saveBalanceTestRecord(t, env, owner2Unlocked, &balance.Record{Quarks: 99, IsOpen: true})
+
+	// Every owner's managed accounts are reported in a single read. Unlocked
+	// accounts and owners the ledger holds nothing for are omitted.
+	balanceByOwnerAndAccount, err := BatchCalculateFromCacheByOwners(env.ctx, env.data, []*common.Account{owner1, owner2, unknownOwner}, nil)
+	require.NoError(t, err)
+	assert.Equal(t, map[string]map[string]*Balance{
+		owner1.PublicKey().ToBase58(): {
+			owner1CoreMint.tokenAccount.PublicKey().ToBase58():  {MintAccount: owner1CoreMint.mint(), Quarks: 42},
+			owner1OtherMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: owner1OtherMint.mint(), Quarks: 33},
+		},
+		owner2.PublicKey().ToBase58(): {
+			owner2CoreMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: owner2CoreMint.mint(), Quarks: 7},
+		},
+	}, balanceByOwnerAndAccount)
+
+	// A mint filter limits the read to the provided mints
+	balanceByOwnerAndAccount, err = BatchCalculateFromCacheByOwners(env.ctx, env.data, []*common.Account{owner1, owner2}, []*common.Account{otherVmConfig.Mint})
+	require.NoError(t, err)
+	assert.Equal(t, map[string]map[string]*Balance{
+		owner1.PublicKey().ToBase58(): {
+			owner1OtherMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: owner1OtherMint.mint(), Quarks: 33},
+		},
+	}, balanceByOwnerAndAccount)
+
+	// The single owner variant applies the same filter
+	balanceByAccount, err := BatchCalculateFromCacheByOwner(env.ctx, env.data, owner1, otherVmConfig.Mint)
+	require.NoError(t, err)
+	assert.Equal(t, map[string]*Balance{
+		owner1OtherMint.tokenAccount.PublicKey().ToBase58(): {MintAccount: owner1OtherMint.mint(), Quarks: 33},
+	}, balanceByAccount)
+
+	// A filter naming a mint nobody holds is empty rather than an error
+	balanceByOwnerAndAccount, err = BatchCalculateFromCacheByOwners(env.ctx, env.data, []*common.Account{owner1, owner2}, []*common.Account{testutil.NewRandomAccount(t)})
+	require.NoError(t, err)
+	assert.Empty(t, balanceByOwnerAndAccount)
+}
+
 func TestDefaultCalculation_ExternalAccount(t *testing.T) {
 	env := setupBalanceTestEnv(t)
 	externalAccount := testutil.NewRandomAccount(t)

ocp/data/balance/memory/store.go

@@ -90,6 +90,40 @@ func (s *store) GetAllByOwner(_ context.Context, owner string) ([]*balance.Recor
 	})
 }
 
+// GetAllByOwnerBatch implements balance.Store.GetAllByOwnerBatch
+func (s *store) GetAllByOwnerBatch(_ context.Context, owners, mints []string) (map[string][]*balance.Record, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	ownerSet := make(map[string]struct{}, len(owners))
+	for _, owner := range owners {
+		ownerSet[owner] = struct{}{}
+	}
+
+	var mintSet map[string]struct{}
+	if len(mints) > 0 {
+		mintSet = make(map[string]struct{}, len(mints))
+		for _, mint := range mints {
+			mintSet[mint] = struct{}{}
+		}
+	}
+
+	res := make(map[string][]*balance.Record)
+	for _, item := range s.balanceRecords {
+		if _, ok := ownerSet[item.OwnerAccount]; !ok {
+			continue
+		}
+		if mintSet != nil {
+			if _, ok := mintSet[item.MintAccount]; !ok {
+				continue
+			}
+		}
+		cloned := item.Clone()
+		res[item.OwnerAccount] = append(res[item.OwnerAccount], &cloned)
+	}
+	return res, nil
+}
+
 // GetAllByOwnerAndMint implements balance.Store.GetAllByOwnerAndMint
 func (s *store) GetAllByOwnerAndMint(_ context.Context, owner, mint string) ([]*balance.Record, error) {
 	s.mu.Lock()

ocp/data/balance/postgres/model.go

@@ -136,6 +136,30 @@ func dbGetBatch(ctx context.Context, db *sqlx.DB, tokenAccounts ...string) ([]*m
 	return res, nil
 }
 
+func dbGetAllByOwnerBatch(ctx context.Context, db *sqlx.DB, owners, mints []string) ([]*model, error) {
+	res := []*model{}
+	if len(owners) == 0 {
+		return res, nil
+	}
+
+	query := `SELECT ` + allColumns + ` FROM ` + tableName + `
+		WHERE owner_account = ANY($1)`
+	args := []any{owners}
+	if len(mints) > 0 {
+		query += ` AND mint_account = ANY($2)`
+		args = append(args, mints)
+	}
+	query += ` ORDER BY id ASC`
+
+	err := pgutil.ExecuteInTx(ctx, db, sql.LevelDefault, func(tx *sqlx.Tx) error {
+		return tx.SelectContext(ctx, &res, query, args...)
+	})
+	if err != nil && !pgutil.IsNoRows(err) {
+		return nil, err
+	}
+	return res, nil
+}
+
 func dbGetAllByOwner(ctx context.Context, db *sqlx.DB, owner string, mint *string) ([]*model, error) {
 	res := []*model{}
 

ocp/data/balance/postgres/store.go

@@ -68,6 +68,20 @@ func (s *store) GetAllByOwner(ctx context.Context, owner string) ([]*balance.Rec
 	return fromModels(models), nil
 }
 
+// GetAllByOwnerBatch implements balance.Store.GetAllByOwnerBatch
+func (s *store) GetAllByOwnerBatch(ctx context.Context, owners, mints []string) (map[string][]*balance.Record, error) {
+	models, err := dbGetAllByOwnerBatch(ctx, s.db, owners, mints)
+	if err != nil {
+		return nil, err
+	}
+
+	res := make(map[string][]*balance.Record)
+	for _, model := range models {
+		res[model.OwnerAccount] = append(res[model.OwnerAccount], fromModel(model))
+	}
+	return res, nil
+}
+
 // GetAllByOwnerAndMint implements balance.Store.GetAllByOwnerAndMint
 func (s *store) GetAllByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) {
 	models, err := dbGetAllByOwner(ctx, s.db, owner, &mint)

ocp/data/balance/store.go

@@ -51,6 +51,12 @@ type Store interface {
 	// ErrRecordNotFound is returned if no records exist.
 	GetAllByOwner(ctx context.Context, owner string) ([]*Record, error)
 
+	// GetAllByOwnerBatch gets all balance records for a set of owners, keyed
+	// by owner. An empty mints slice includes every mint, otherwise only
+	// records for the provided mints are returned. Owners without records
+	// are omitted from the result.
+	GetAllByOwnerBatch(ctx context.Context, owners, mints []string) (map[string][]*Record, error)
+
 	// GetAllByOwnerAndMint gets all balance records for an owner and mint.
 	//
 	// ErrRecordNotFound is returned if no records exist.

ocp/data/balance/tests/tests.go

@@ -45,6 +45,10 @@ func testRecordHappyPath(t *testing.T, s balance.Store) {
 		require.NoError(t, err)
 		assert.Empty(t, batch)
 
+		byOwnerBatch, err := s.GetAllByOwnerBatch(ctx, []string{"owner_1", "owner_2"}, nil)
+		require.NoError(t, err)
+		assert.Empty(t, byOwnerBatch)
+
 		start := time.Now()
 
 		expected := &balance.Record{
@@ -95,6 +99,25 @@ func testRecordHappyPath(t *testing.T, s balance.Store) {
 		assert.Equal(t, "token_account_1", byOwner[0].TokenAccount)
 		assert.Equal(t, "token_account_2", byOwner[1].TokenAccount)
 
+		byOwnerBatch, err = s.GetAllByOwnerBatch(ctx, []string{"owner_1", "owner_2", "owner_3"}, nil)
+		require.NoError(t, err)
+		require.Len(t, byOwnerBatch, 2)
+		require.Len(t, byOwnerBatch["owner_1"], 2)
+		assertEquivalentRecords(t, &cloned, byOwnerBatch["owner_1"][0])
+		assert.Equal(t, "token_account_2", byOwnerBatch["owner_1"][1].TokenAccount)
+		require.Len(t, byOwnerBatch["owner_2"], 1)
+		assert.Equal(t, "token_account_3", byOwnerBatch["owner_2"][0].TokenAccount)
+
+		byOwnerBatch, err = s.GetAllByOwnerBatch(ctx, []string{"owner_1", "owner_2"}, []string{"mint_2"})
+		require.NoError(t, err)
+		require.Len(t, byOwnerBatch, 1)
+		require.Len(t, byOwnerBatch["owner_1"], 1)
+		assert.Equal(t, "token_account_2", byOwnerBatch["owner_1"][0].TokenAccount)
+
+		byOwnerBatch, err = s.GetAllByOwnerBatch(ctx, []string{"owner_1", "owner_2"}, []string{"mint_3"})
+		require.NoError(t, err)
+		assert.Empty(t, byOwnerBatch)
+
 		byOwnerAndMint, err := s.GetAllByOwnerAndMint(ctx, "owner_1", "mint_2")
 		require.NoError(t, err)
 		require.Len(t, byOwnerAndMint, 1)

ocp/data/internal.go

@@ -125,6 +125,7 @@ type DatabaseData interface {
 	GetBalance(ctx context.Context, tokenAccount string) (*balance.Record, error)
 	GetBalanceBatch(ctx context.Context, tokenAccounts ...string) (map[string]*balance.Record, error)
 	GetAllBalancesByOwner(ctx context.Context, owner string) ([]*balance.Record, error)
+	GetAllBalancesByOwnerBatch(ctx context.Context, owners, mints []string) (map[string][]*balance.Record, error)
 	GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error)
 	GetAllLockedBalancesByMint(ctx context.Context, mint string, minQuarks uint64, cursor query.Cursor, limit uint64, direction query.Ordering) ([]*balance.Record, error)
 	CountLockedBalancesByMint(ctx context.Context, mint string, minQuarks uint64) (uint64, error)
@@ -477,6 +478,9 @@ func (dp *DatabaseProvider) GetBalanceBatch(ctx context.Context, tokenAccounts .
 func (dp *DatabaseProvider) GetAllBalancesByOwner(ctx context.Context, owner string) ([]*balance.Record, error) {
 	return dp.balance.GetAllByOwner(ctx, owner)
 }
+func (dp *DatabaseProvider) GetAllBalancesByOwnerBatch(ctx context.Context, owners, mints []string) (map[string][]*balance.Record, error) {
+	return dp.balance.GetAllByOwnerBatch(ctx, owners, mints)
+}
 func (dp *DatabaseProvider) GetAllBalancesByOwnerAndMint(ctx context.Context, owner, mint string) ([]*balance.Record, error) {
 	return dp.balance.GetAllByOwnerAndMint(ctx, owner, mint)
 }