[Persistence] Consolidate common behaviour between Pool and Account into a shared interface (Issue #102)

persistence/account.go

@@ -5,208 +5,71 @@ import (
 	"math/big"
 
 	"github.com/pokt-network/pocket/persistence/types"
-	"github.com/pokt-network/pocket/shared/converters"
 	coreTypes "github.com/pokt-network/pocket/shared/core/types"
-
-	"github.com/jackc/pgx/v5"
-)
-
-// TODO(https://github.com/pokt-network/pocket/issues/102): Generalize Pool and Account operations.
-
-const (
-	defaultAccountAmountStr string = "0"
 )
 
 // --- Account Functions ---
 
 func (p PostgresContext) GetAccountAmount(address []byte, height int64) (amount string, err error) {
-	return p.getAccountAmountStr(hex.EncodeToString(address), height)
-}
-
-func (p PostgresContext) getAccountAmountStr(address string, height int64) (amount string, err error) {
-	ctx, tx, err := p.getCtxAndTx()
-	if err != nil {
-		return
-	}
-
-	amount = defaultAccountAmountStr
-	if err = tx.QueryRow(ctx, types.GetAccountAmountQuery(address, height)).Scan(&amount); err != pgx.ErrNoRows {
-		return
-	}
-
-	return amount, nil
+	return p.getAccountAmount(types.Account, hex.EncodeToString(address), height)
 }
 
 func (p PostgresContext) AddAccountAmount(address []byte, amount string) error {
-	return p.operationAccountAmount(address, amount, func(orig *big.Int, delta *big.Int) error {
+	return p.operationAccountAmount(types.Account, hex.EncodeToString(address), amount, func(orig, delta *big.Int) error {
 		orig.Add(orig, delta)
 		return nil
 	})
 }
 
 func (p PostgresContext) SubtractAccountAmount(address []byte, amount string) error {
-	return p.operationAccountAmount(address, amount, func(orig *big.Int, delta *big.Int) error {
+	return p.operationAccountAmount(types.Account, hex.EncodeToString(address), amount, func(orig, delta *big.Int) error {
 		orig.Sub(orig, delta)
 		return nil
 	})
 }
 
-// DISCUSS(team): If we are okay with `GetAccountAmount` return 0 as a default, this function can leverage
-// `operationAccountAmount` with `*orig = *delta` and make everything much simpler.
 func (p PostgresContext) SetAccountAmount(address []byte, amount string) error {
-	ctx, tx, err := p.getCtxAndTx()
-	if err != nil {
-		return err
-	}
-	height, err := p.GetHeight()
-	if err != nil {
-		return err
-	}
-	// DISCUSS(team): Do we want to panic if `amount < 0` here?
-	if _, err = tx.Exec(ctx, types.InsertAccountAmountQuery(hex.EncodeToString(address), amount, height)); err != nil {
-		return err
-	}
-	return nil
+	return p.operationAccountAmount(types.Account, hex.EncodeToString(address), amount, func(orig, amount *big.Int) error {
+		orig.Set(amount)
+		return nil
+	})
 }
 
 func (p PostgresContext) GetAccountsUpdated(height int64) (accounts []*coreTypes.Account, err error) {
-	return p.getPoolOrAccUpdatedInternal(types.GetAccountsUpdatedAtHeightQuery(height))
-}
-
-func (p *PostgresContext) operationAccountAmount(address []byte, deltaAmount string, op func(*big.Int, *big.Int) error) error {
-	return p.operationPoolOrAccAmount(hex.EncodeToString(address), deltaAmount, op, p.getAccountAmountStr, types.InsertAccountAmountQuery)
+	return p.getAccountsUpdated(types.Account, height)
 }
 
 // --- Pool Functions ---
 
-// TODO(andrew): remove address param
-func (p PostgresContext) InsertPool(name string, address []byte, amount string) error {
-	ctx, tx, err := p.getCtxAndTx()
-	if err != nil {
-		return err
-	}
-	height, err := p.GetHeight()
-	if err != nil {
-		return err
-	}
-	if _, err = tx.Exec(ctx, types.InsertPoolAmountQuery(name, amount, height)); err != nil {
-		return err
-	}
-	return nil
+func (p PostgresContext) InsertPool(name string, amount string) error {
+	return p.insertAccount(types.Pool, name, amount)
 }
 
 func (p PostgresContext) GetPoolAmount(name string, height int64) (amount string, err error) {
-	ctx, tx, err := p.getCtxAndTx()
-	if err != nil {
-		return
-	}
-
-	amount = defaultAccountAmountStr
-	if err = tx.QueryRow(ctx, types.GetPoolAmountQuery(name, height)).Scan(&amount); err != pgx.ErrNoRows {
-		return
-	}
-
-	return amount, nil
+	return p.getAccountAmount(types.Pool, name, height)
 }
 
 func (p PostgresContext) AddPoolAmount(name string, amount string) error {
-	return p.operationPoolAmount(name, amount, func(s *big.Int, s1 *big.Int) error {
-		s.Add(s, s1)
+	return p.operationAccountAmount(types.Pool, name, amount, func(orig, delta *big.Int) error {
+		orig.Add(orig, delta)
 		return nil
 	})
 }
 
 func (p PostgresContext) SubtractPoolAmount(name string, amount string) error {
-	return p.operationPoolAmount(name, amount, func(s *big.Int, s1 *big.Int) error {
-		s.Sub(s, s1)
+	return p.operationAccountAmount(types.Pool, name, amount, func(orig, delta *big.Int) error {
+		orig.Sub(orig, delta)
 		return nil
 	})
 }
 
-// DISCUSS(team): If we are okay with `GetPoolAmount` return 0 as a default, this function can leverage
-//
-//	`operationPoolAmount` with `*orig = *delta` and make everything much simpler.
-//
-// DISCUSS(team): Do we have a use-case for this function?
 func (p PostgresContext) SetPoolAmount(name string, amount string) error {
-	ctx, tx, err := p.getCtxAndTx()
-	if err != nil {
-		return err
-	}
-	height, err := p.GetHeight()
-	if err != nil {
-		return err
-	}
-	if _, err = tx.Exec(ctx, types.InsertPoolAmountQuery(name, amount, height)); err != nil {
-		return err
-	}
-	return nil
-}
-
-func (p *PostgresContext) operationPoolAmount(name string, amount string, op func(*big.Int, *big.Int) error) error {
-	return p.operationPoolOrAccAmount(name, amount, op, p.GetPoolAmount, types.InsertPoolAmountQuery)
+	return p.operationAccountAmount(types.Pool, name, amount, func(orig, amount *big.Int) error {
+		orig.Set(amount)
+		return nil
+	})
 }
 
 func (p PostgresContext) GetPoolsUpdated(height int64) ([]*coreTypes.Account, error) {
-	return p.getPoolOrAccUpdatedInternal(types.GetPoolsUpdatedAtHeightQuery(height))
-}
-
-// Joint Pool & Account Helpers
-
-// Shared logic between `getPoolsUpdated` & `getAccountsUpdated` to keep explicit external interfaces
-func (p *PostgresContext) getPoolOrAccUpdatedInternal(query string) (accounts []*coreTypes.Account, err error) {
-	ctx, tx, err := p.getCtxAndTx()
-	if err != nil {
-		return
-	}
-
-	rows, err := tx.Query(ctx, query)
-	if err != nil {
-		return nil, err
-	}
-	defer rows.Close()
-
-	for rows.Next() {
-		account := new(coreTypes.Account)
-		if err = rows.Scan(&account.Address, &account.Amount); err != nil {
-			return nil, err
-		}
-		accounts = append(accounts, account)
-	}
-
-	return
-}
-
-func (p *PostgresContext) operationPoolOrAccAmount(
-	name, amount string,
-	op func(*big.Int, *big.Int) error,
-	getAmount func(string, int64) (string, error),
-	insert func(name, amount string, height int64) string) error {
-	ctx, tx, err := p.getCtxAndTx()
-	if err != nil {
-		return err
-	}
-	height, err := p.GetHeight()
-	if err != nil {
-		return err
-	}
-	originalAmount, err := getAmount(name, height)
-	if err != nil {
-		return err
-	}
-	originalAmountBig, err := converters.StringToBigInt(originalAmount)
-	if err != nil {
-		return err
-	}
-	amountBig, err := converters.StringToBigInt(amount)
-	if err != nil {
-		return err
-	}
-	if err := op(originalAmountBig, amountBig); err != nil {
-		return err
-	}
-	if _, err = tx.Exec(ctx, insert(name, converters.BigIntToString(originalAmountBig), height)); err != nil {
-		return err
-	}
-	return nil
+	return p.getAccountsUpdated(types.Pool, height)
 }

persistence/account_shared_sql.go

@@ -0,0 +1,101 @@
+package persistence
+
+import (
+	"github.com/jackc/pgx/v5"
+	"github.com/pokt-network/pocket/persistence/types"
+	"github.com/pokt-network/pocket/shared/converters"
+	coreTypes "github.com/pokt-network/pocket/shared/core/types"
+	"math/big"
+)
+
+const (
+	defaultAccountAmountStr = "0"
+)
+
+func (p *PostgresContext) getAccountAmount(accountSchema types.ProtocolAccountSchema, identifier string, height int64) (amount string, err error) {
+	ctx, tx, err := p.getCtxAndTx()
+	if err != nil {
+		return
+	}
+	amount = defaultAccountAmountStr
+	if err = tx.QueryRow(ctx, accountSchema.GetAccountAmountQuery(identifier, height)).Scan(&amount); err != pgx.ErrNoRows {
+		return
+	}
+
+	return amount, nil
+}
+
+func (p *PostgresContext) operationAccountAmount(
+	accountSchema types.ProtocolAccountSchema,
+	identifier, amount string,
+	op func(*big.Int, *big.Int) error,
+) error {
+	ctx, tx, err := p.getCtxAndTx()
+	if err != nil {
+		return err
+	}
+	height, err := p.GetHeight()
+	if err != nil {
+		return err
+	}
+	originalAmount, err := p.getAccountAmount(accountSchema, identifier, height)
+	if err != nil {
+		return err
+	}
+	originalAmountBig, err := converters.StringToBigInt(originalAmount)
+	if err != nil {
+		return err
+	}
+	amountBig, err := converters.StringToBigInt(amount)
+	if err != nil {
+		return err
+	}
+	if err := op(originalAmountBig, amountBig); err != nil {
+		return err
+	}
+	if _, err = tx.Exec(ctx, accountSchema.InsertAccountQuery(identifier, converters.BigIntToString(originalAmountBig), height)); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (p *PostgresContext) getAccountsUpdated(accountType types.ProtocolAccountSchema, height int64) (accounts []*coreTypes.Account, err error) {
+	query := accountType.GetAccountsUpdatedAtHeightQuery(height)
+
+	ctx, tx, err := p.getCtxAndTx()
+	if err != nil {
+		return
+	}
+
+	rows, err := tx.Query(ctx, query)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		acc := new(coreTypes.Account)
+		if err = rows.Scan(&acc.Address, &acc.Amount); err != nil {
+			return nil, err
+		}
+		accounts = append(accounts, acc)
+	}
+
+	return
+}
+
+func (p *PostgresContext) insertAccount(accountType types.ProtocolAccountSchema, identifier, amount string) error {
+	ctx, tx, err := p.getCtxAndTx()
+	if err != nil {
+		return err
+	}
+	height, err := p.GetHeight()
+	if err != nil {
+		return err
+	}
+	// DISCUSS(team): Do we want to panic if `amount < 0` here?
+	if _, err = tx.Exec(ctx, accountType.InsertAccountQuery(identifier, amount, height)); err != nil {
+		return err
+	}
+	return nil
+}

persistence/shared_sql.go → persistence/actor_shared_sql.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/hex"
 	"fmt"
-
 	"github.com/jackc/pgx/v5"
 	"github.com/pokt-network/pocket/persistence/types"
 	coreTypes "github.com/pokt-network/pocket/shared/core/types"

persistence/db.go

@@ -167,10 +167,10 @@ func initializeProtocolActorTables(ctx context.Context, db *pgx.Conn, actor type
 }
 
 func initializeAccountTables(ctx context.Context, db *pgx.Conn) error {
-	if _, err := db.Exec(ctx, fmt.Sprintf(`%s %s %s %s`, CreateTable, IfNotExists, types.AccountTableName, types.AccountTableSchema)); err != nil {
+	if _, err := db.Exec(ctx, fmt.Sprintf(`%s %s %s %s`, CreateTable, IfNotExists, types.AccountTableName, types.Account.GetTableSchema())); err != nil {
 		return err
 	}
-	if _, err := db.Exec(ctx, fmt.Sprintf(`%s %s %s %s`, CreateTable, IfNotExists, types.PoolTableName, types.PoolTableSchema)); err != nil {
+	if _, err := db.Exec(ctx, fmt.Sprintf(`%s %s %s %s`, CreateTable, IfNotExists, types.PoolTableName, types.Pool.GetTableSchema())); err != nil {
 		return err
 	}
 	return nil

persistence/debug.go

@@ -13,8 +13,8 @@ import (
 
 // A list of functions to clear data from the DB not associated with protocol actors
 var nonActorClearFunctions = []func() string{
-	types.ClearAllAccounts,
-	types.ClearAllPools,
+	types.Account.ClearAllAccounts,
+	types.Pool.ClearAllAccounts,
 	types.ClearAllGovParamsQuery,
 	types.ClearAllGovFlagsQuery,
 	types.ClearAllBlocksQuery,

persistence/genesis.go

@@ -46,8 +46,7 @@ func (m *persistenceModule) populateGenesisState(state *genesis.GenesisState) {
 		}
 	}
 	for _, pool := range state.GetPools() {
-		poolNameBytes := []byte(pool.GetAddress())
-		err = rwContext.InsertPool(pool.GetAddress(), poolNameBytes, pool.GetAmount())
+		err = rwContext.InsertPool(pool.GetAddress(), pool.GetAmount()) // pool.GetAddress() returns the pool's semantic name
 		if err != nil {
 			log.Fatalf("an error occurred inserting an pool in the genesis state: %s", err.Error())
 		}
@@ -142,7 +141,7 @@ func (p PostgresContext) GetAllAccounts(height int64) (accs []*coreTypes.Account
 	if err != nil {
 		return nil, err
 	}
-	rows, err := tx.Query(ctx, types.SelectAccounts(height, types.AccountTableName))
+	rows, err := tx.Query(ctx, types.Account.GetAllQuery(height))
 	if err != nil {
 		return nil, err
 	}
@@ -166,7 +165,7 @@ func (p PostgresContext) GetAllPools(height int64) (accs []*coreTypes.Account, e
 	if err != nil {
 		return nil, err
 	}
-	rows, err := tx.Query(ctx, types.SelectPools(height, types.PoolTableName))
+	rows, err := tx.Query(ctx, types.Pool.GetAllQuery(height))
 	if err != nil {
 		return nil, err
 	}