code.vegaprotocol.io/vega@v0.79.0/datanode/service/accounts.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package service
    17  
    18  import (
    19  	"context"
    20  
    21  	"code.vegaprotocol.io/vega/datanode/entities"
    22  	"code.vegaprotocol.io/vega/datanode/utils"
    23  	"code.vegaprotocol.io/vega/logging"
    24  	"code.vegaprotocol.io/vega/protos/vega"
    25  )
    26  
    27  type AccountStore interface {
    28  	// Use get by raw ID to avoid using the AccountID type because mockgen does not support it
    29  	GetByRawID(ctx context.Context, id string) (entities.Account, error)
    30  	GetAll(ctx context.Context) ([]entities.Account, error)
    31  	GetByTxHash(ctx context.Context, txHash entities.TxHash) ([]entities.Account, error)
    32  	Obtain(ctx context.Context, a *entities.Account) error
    33  	Query(ctx context.Context, filter entities.AccountFilter) ([]entities.Account, error)
    34  	QueryBalances(ctx context.Context, filter entities.AccountFilter, pagination entities.CursorPagination) ([]entities.AccountBalance, entities.PageInfo, error)
    35  	GetBalancesByTxHash(ctx context.Context, txHash entities.TxHash) ([]entities.AccountBalance, error)
    36  }
    37  
    38  type BalanceStore interface {
    39  	Flush(ctx context.Context) ([]entities.AccountBalance, error)
    40  	Add(b entities.AccountBalance) error
    41  	Query(ctx context.Context, filter entities.AccountFilter, dateRange entities.DateRange, pagination entities.CursorPagination) (*[]entities.AggregatedBalance, entities.PageInfo, error)
    42  }
    43  
    44  type Account struct {
    45  	aStore    AccountStore
    46  	bStore    BalanceStore
    47  	bObserver utils.Observer[entities.AccountBalance]
    48  }
    49  
    50  func NewAccount(aStore AccountStore, bStore BalanceStore, log *logging.Logger) *Account {
    51  	return &Account{
    52  		aStore:    aStore,
    53  		bStore:    bStore,
    54  		bObserver: utils.NewObserver[entities.AccountBalance]("account_balance", log, 0, 0),
    55  	}
    56  }
    57  
    58  func (a *Account) GetByID(ctx context.Context, id entities.AccountID) (entities.Account, error) {
    59  	return a.aStore.GetByRawID(ctx, id.String())
    60  }
    61  
    62  func (a *Account) GetAll(ctx context.Context) ([]entities.Account, error) {
    63  	return a.aStore.GetAll(ctx)
    64  }
    65  
    66  func (a *Account) Obtain(ctx context.Context, acc *entities.Account) error {
    67  	return a.aStore.Obtain(ctx, acc)
    68  }
    69  
    70  func (a *Account) Query(ctx context.Context, filter entities.AccountFilter) ([]entities.Account, error) {
    71  	return a.aStore.Query(ctx, filter)
    72  }
    73  
    74  func (a *Account) QueryBalances(ctx context.Context, filter entities.AccountFilter, pagination entities.CursorPagination) ([]entities.AccountBalance, entities.PageInfo, error) {
    75  	return a.aStore.QueryBalances(ctx, filter, pagination)
    76  }
    77  
    78  func (a *Account) GetByTxHash(ctx context.Context, txHash entities.TxHash) ([]entities.Account, error) {
    79  	return a.aStore.GetByTxHash(ctx, txHash)
    80  }
    81  
    82  func (a *Account) GetBalancesByTxHash(ctx context.Context, txHash entities.TxHash) ([]entities.AccountBalance, error) {
    83  	return a.aStore.GetBalancesByTxHash(ctx, txHash)
    84  }
    85  
    86  func (a *Account) AddAccountBalance(b entities.AccountBalance) error {
    87  	return a.bStore.Add(b)
    88  }
    89  
    90  func (a *Account) Flush(ctx context.Context) error {
    91  	flushed, err := a.bStore.Flush(ctx)
    92  	if err != nil {
    93  		return err
    94  	}
    95  	a.bObserver.Notify(flushed)
    96  	return nil
    97  }
    98  
    99  func (a *Account) Unsubscribe(ctx context.Context, ref uint64) error {
   100  	return a.bObserver.Unsubscribe(ctx, ref)
   101  }
   102  
   103  func (a *Account) QueryAggregatedBalances(ctx context.Context, filter entities.AccountFilter, dateRange entities.DateRange, pagination entities.CursorPagination) (*[]entities.AggregatedBalance, entities.PageInfo, error) {
   104  	return a.bStore.Query(ctx, filter, dateRange, pagination)
   105  }
   106  
   107  func (a *Account) ObserveAccountBalances(ctx context.Context, retries int, marketID string,
   108  	asset string, ty vega.AccountType, partyIDs map[string]string,
   109  ) (accountCh <-chan []entities.AccountBalance, ref uint64) {
   110  	ch, ref := a.bObserver.Observe(ctx,
   111  		retries,
   112  		func(ab entities.AccountBalance) bool {
   113  			_, partyOK := partyIDs[ab.PartyID.String()]
   114  
   115  			return (len(marketID) == 0 || marketID == ab.MarketID.String()) &&
   116  				(partyOK) &&
   117  				(len(asset) == 0 || asset == ab.AssetID.String()) &&
   118  				(ty == vega.AccountType_ACCOUNT_TYPE_UNSPECIFIED || ty == ab.Type)
   119  		})
   120  	return ch, ref
   121  }