From 902429b03068064c654458efb67709b73b8e60b2 Mon Sep 17 00:00:00 2001
From: jchappelow <140431406+jchappelow@users.noreply.github.com>
Date: Fri, 6 Oct 2023 10:23:47 -0500
Subject: [PATCH] kwil-cli,pkg/client: add --nonce override flag to database
 cmds (#337)

This adds the ability to set a nonce when making a transaction.

Rather than giving every method a nonce input arg, I've added a
WithNonceFunc option to pkg/client.Client. This is perhaps a
little strange but the cascade from a new argument was rather
large, and it seemed a little out of place as the other arguments
pertain to the action rather than the transaction containing it.
Can be switched if we'd prefer.
---
 cmd/kwil-cli/cmds/common/prompt/prompt.go    | 11 -------
 cmd/kwil-cli/cmds/database/batch.go          |  6 ++--
 cmd/kwil-cli/cmds/database/cmd.go            | 31 +++++++++++++-----
 cmd/kwil-cli/cmds/database/deploy.go         |  4 +--
 cmd/kwil-cli/cmds/database/drop.go           |  4 +--
 cmd/kwil-cli/cmds/database/execute.go        |  8 +++--
 internal/controller/grpc/txsvc/v1/account.go |  2 +-
 pkg/balances/account.go                      |  2 +-
 pkg/balances/sql.go                          |  5 +--
 pkg/client/client.go                         | 20 ++++++------
 pkg/client/opts.go                           | 12 +++++++
 pkg/client/tx.go                             | 33 +++++++++++++-------
 test/driver/client_driver.go                 |  2 +-
 13 files changed, 85 insertions(+), 55 deletions(-)

diff --git a/cmd/kwil-cli/cmds/common/prompt/prompt.go b/cmd/kwil-cli/cmds/common/prompt/prompt.go
index c185200fe..d3aa84b91 100644
--- a/cmd/kwil-cli/cmds/common/prompt/prompt.go
+++ b/cmd/kwil-cli/cmds/common/prompt/prompt.go
@@ -5,7 +5,6 @@ import (
 	"strings"
 
 	"github.com/manifoldco/promptui"
-	"github.com/spf13/cobra"
 )
 
 type Prompter struct {
@@ -50,13 +49,3 @@ func (p Prompter) Run() (string, error) {
 
 	return prompt.Run()
 }
-
-func ConfirmPrompt() bool {
-	prompt := promptui.Select{
-		Label: "Are you sure?",
-		Items: []string{"Apply", "Abort"},
-	}
-	_, result, err := prompt.Run()
-	cobra.CheckErr(err)
-	return result == "Apply"
-}
diff --git a/cmd/kwil-cli/cmds/database/batch.go b/cmd/kwil-cli/cmds/database/batch.go
index 822701efc..b31d2dd0b 100644
--- a/cmd/kwil-cli/cmds/database/batch.go
+++ b/cmd/kwil-cli/cmds/database/batch.go
@@ -36,7 +36,7 @@ The execution is treated as a single transaction, and will either succeed or fai
 		RunE: func(cmd *cobra.Command, args []string) error {
 			var resp []byte
 
-			err := common.DialClient(cmd.Context(), 0, func(ctx context.Context, client *client.Client, conf *config.KwilCliConfig) error {
+			err := common.DialClient(cmd.Context(), 0, func(ctx context.Context, cl *client.Client, conf *config.KwilCliConfig) error {
 				dbid, err := getSelectedDbid(cmd, conf)
 				if err != nil {
 					return err
@@ -61,7 +61,7 @@ The execution is treated as a single transaction, and will either succeed or fai
 					return fmt.Errorf("error building inputs: %w", err)
 				}
 
-				actionStructure, err := getAction(ctx, client, dbid, action)
+				actionStructure, err := getAction(ctx, cl, dbid, action)
 				if err != nil {
 					return fmt.Errorf("error getting action: %w", err)
 				}
@@ -71,7 +71,7 @@ The execution is treated as a single transaction, and will either succeed or fai
 					return fmt.Errorf("error creating action inputs: %w", err)
 				}
 
-				resp, err = client.ExecuteAction(ctx, dbid, strings.ToLower(action), tuples...)
+				resp, err = cl.ExecuteAction(ctx, dbid, strings.ToLower(action), tuples, client.WithNonce(nonceOverride))
 				if err != nil {
 					return fmt.Errorf("error executing action: %w", err)
 				}
diff --git a/cmd/kwil-cli/cmds/database/cmd.go b/cmd/kwil-cli/cmds/database/cmd.go
index 9b795cffb..eabc8db40 100644
--- a/cmd/kwil-cli/cmds/database/cmd.go
+++ b/cmd/kwil-cli/cmds/database/cmd.go
@@ -5,25 +5,40 @@ import (
 )
 
 var (
-	rootCmd = &cobra.Command{
+	dbCmd = &cobra.Command{
 		Use:     "database",
 		Aliases: []string{"db"},
 		Short:   "manage databases",
 		Long:    "Database is a command that contains subcommands for interacting with databases",
 	}
+
+	nonceOverride int64
 )
 
 func NewCmdDatabase() *cobra.Command {
-	rootCmd.AddCommand(
+	// readOnlyCmds do not create a transaction.
+	readOnlyCmds := []*cobra.Command{
+		listCmd(),
+		readSchemaCmd(),
+		queryCmd(),
+		callCmd(), // no tx, but may required key for signature, for now
+	}
+	dbCmd.AddCommand(readOnlyCmds...)
+
+	// writeCmds create a transactions, requiring a private key for signing/
+	writeCmds := []*cobra.Command{
 		deployCmd(),
 		dropCmd(),
-		readSchemaCmd(),
 		executeCmd(),
-		listCmd(),
 		batchCmd(),
-		queryCmd(),
-		callCmd(),
-	)
+	}
+	dbCmd.AddCommand(writeCmds...)
+
+	// The write commands may also specify a nonce to use instead of asking the
+	// node for the latest confirmed nonce.
+	for _, cmd := range writeCmds {
+		cmd.Flags().Int64VarP(&nonceOverride, "nonce", "N", -1, "nonce override (-1 means request from server)")
+	}
 
-	return rootCmd
+	return dbCmd
 }
diff --git a/cmd/kwil-cli/cmds/database/deploy.go b/cmd/kwil-cli/cmds/database/deploy.go
index 4b654a076..aeab1e684 100644
--- a/cmd/kwil-cli/cmds/database/deploy.go
+++ b/cmd/kwil-cli/cmds/database/deploy.go
@@ -29,7 +29,7 @@ func deployCmd() *cobra.Command {
 		RunE: func(cmd *cobra.Command, args []string) error {
 			var txHash []byte
 
-			err := common.DialClient(cmd.Context(), 0, func(ctx context.Context, client *client.Client, conf *config.KwilCliConfig) error {
+			err := common.DialClient(cmd.Context(), 0, func(ctx context.Context, cl *client.Client, conf *config.KwilCliConfig) error {
 				// read in the file
 				file, err := os.Open(filePath)
 				if err != nil {
@@ -49,7 +49,7 @@ func deployCmd() *cobra.Command {
 					return fmt.Errorf("failed to unmarshal file: %w", err)
 				}
 
-				txHash, err = client.DeployDatabase(ctx, db)
+				txHash, err = cl.DeployDatabase(ctx, db, client.WithNonce(nonceOverride))
 				return err
 			})
 
diff --git a/cmd/kwil-cli/cmds/database/drop.go b/cmd/kwil-cli/cmds/database/drop.go
index 6f7a03117..f6b613279 100644
--- a/cmd/kwil-cli/cmds/database/drop.go
+++ b/cmd/kwil-cli/cmds/database/drop.go
@@ -20,9 +20,9 @@ func dropCmd() *cobra.Command {
 		RunE: func(cmd *cobra.Command, args []string) error {
 			var resp []byte
 
-			err := common.DialClient(cmd.Context(), 0, func(ctx context.Context, client *client.Client, conf *config.KwilCliConfig) error {
+			err := common.DialClient(cmd.Context(), 0, func(ctx context.Context, cl *client.Client, conf *config.KwilCliConfig) error {
 				var _err error
-				resp, _err = client.DropDatabase(ctx, args[0])
+				resp, _err = cl.DropDatabase(ctx, args[0], client.WithNonce(nonceOverride))
 				if _err != nil {
 					return fmt.Errorf("error dropping database: %w", _err)
 				}
diff --git a/cmd/kwil-cli/cmds/database/execute.go b/cmd/kwil-cli/cmds/database/execute.go
index 58fec19fb..cc5e5d532 100644
--- a/cmd/kwil-cli/cmds/database/execute.go
+++ b/cmd/kwil-cli/cmds/database/execute.go
@@ -43,7 +43,7 @@ OR
 		RunE: func(cmd *cobra.Command, args []string) error {
 			var resp []byte
 
-			err := common.DialClient(cmd.Context(), 0, func(ctx context.Context, client *client.Client, conf *config.KwilCliConfig) error {
+			err := common.DialClient(cmd.Context(), 0, func(ctx context.Context, cl *client.Client, conf *config.KwilCliConfig) error {
 				dbId, err := getSelectedDbid(cmd, conf)
 				if err != nil {
 					return fmt.Errorf("target database not properly specified: %w", err)
@@ -51,7 +51,7 @@ OR
 
 				lowerName := strings.ToLower(actionName)
 
-				actionStructure, err := getAction(ctx, client, dbId, lowerName)
+				actionStructure, err := getAction(ctx, cl, dbId, lowerName)
 				if err != nil {
 					return fmt.Errorf("error getting action: %w", err)
 				}
@@ -61,7 +61,9 @@ OR
 					return fmt.Errorf("error getting inputs: %w", err)
 				}
 
-				resp, err = client.ExecuteAction(ctx, dbId, lowerName, inputs...)
+				// Could actually just directly pass nonce to the client method,
+				// but those methods don't need tx details in the inputs.
+				resp, err = cl.ExecuteAction(ctx, dbId, lowerName, inputs, client.WithNonce(nonceOverride))
 				if err != nil {
 					return fmt.Errorf("error executing database: %w", err)
 				}
diff --git a/internal/controller/grpc/txsvc/v1/account.go b/internal/controller/grpc/txsvc/v1/account.go
index eebf557c1..f97d05970 100644
--- a/internal/controller/grpc/txsvc/v1/account.go
+++ b/internal/controller/grpc/txsvc/v1/account.go
@@ -14,7 +14,7 @@ func (s *Service) GetAccount(ctx context.Context, req *txpb.GetAccountRequest) (
 
 	return &txpb.GetAccountResponse{
 		Account: &txpb.Account{
-			PublicKey: acc.PublicKey,
+			PublicKey: acc.PublicKey, // nil for non-existent account
 			Nonce:     acc.Nonce,
 			Balance:   acc.Balance.String(),
 		},
diff --git a/pkg/balances/account.go b/pkg/balances/account.go
index cb79cbcd2..b21f8e8c7 100644
--- a/pkg/balances/account.go
+++ b/pkg/balances/account.go
@@ -8,7 +8,7 @@ import (
 // emptyAccount returns an empty account with a balance of 0 and a nonce of 0.
 func emptyAccount() *Account {
 	return &Account{
-		PublicKey: nil,
+		PublicKey: nil, // do not change unless callers are updated
 		Balance:   big.NewInt(0),
 		Nonce:     0,
 	}
diff --git a/pkg/balances/sql.go b/pkg/balances/sql.go
index 4e4ecaade..00f42b64f 100644
--- a/pkg/balances/sql.go
+++ b/pkg/balances/sql.go
@@ -71,8 +71,9 @@ func (a *AccountStore) createAccount(ctx context.Context, pubKey []byte) error {
 	})
 }
 
-// getAccountReadOnly gets an account using a read-only connection.
-// it will not show uncommitted changes.
+// getAccountReadOnly gets an account using a read-only connection. it will not
+// show uncommitted changes. If the account does not exist, no error is
+// returned, but an account with a nil pubkey is returned.
 func (a *AccountStore) getAccountReadOnly(ctx context.Context, pubKey []byte) (*Account, error) {
 	results, err := a.db.Query(ctx, sqlGetAccount, map[string]interface{}{
 		"$public_key": pubKey,
diff --git a/pkg/client/client.go b/pkg/client/client.go
index 7fe9443e9..4cdd061a6 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -90,8 +90,8 @@ func (c *Client) GetSchema(ctx context.Context, dbid string) (*transactions.Sche
 }
 
 // DeployDatabase deploys a schema
-func (c *Client) DeployDatabase(ctx context.Context, payload *transactions.Schema) (transactions.TxHash, error) {
-	tx, err := c.newTx(ctx, payload)
+func (c *Client) DeployDatabase(ctx context.Context, payload *transactions.Schema, opts ...TxOpt) (transactions.TxHash, error) {
+	tx, err := c.newTx(ctx, payload, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -104,12 +104,12 @@ func (c *Client) DeployDatabase(ctx context.Context, payload *transactions.Schem
 }
 
 // DropDatabase drops a database
-func (c *Client) DropDatabase(ctx context.Context, name string) (transactions.TxHash, error) {
+func (c *Client) DropDatabase(ctx context.Context, name string, opts ...TxOpt) (transactions.TxHash, error) {
 	identifier := &transactions.DropSchema{
 		DBID: engineUtils.GenerateDBID(name, c.Signer.PublicKey()),
 	}
 
-	tx, err := c.newTx(ctx, identifier)
+	tx, err := c.newTx(ctx, identifier, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -129,7 +129,7 @@ func (c *Client) DropDatabase(ctx context.Context, name string) (transactions.Tx
 // ExecuteAction executes an action.
 // It returns the receipt, as well as outputs which is the decoded body of the receipt.
 // It can take any number of inputs, and if multiple tuples of inputs are passed, it will execute them transactionally.
-func (c *Client) ExecuteAction(ctx context.Context, dbid string, action string, tuples ...[]any) (transactions.TxHash, error) {
+func (c *Client) ExecuteAction(ctx context.Context, dbid string, action string, tuples [][]any, opts ...TxOpt) (transactions.TxHash, error) {
 	stringTuples, err := convertTuples(tuples)
 	if err != nil {
 		return nil, err
@@ -141,7 +141,7 @@ func (c *Client) ExecuteAction(ctx context.Context, dbid string, action string,
 		Arguments: stringTuples,
 	}
 
-	tx, err := c.newTx(ctx, executionBody)
+	tx, err := c.newTx(ctx, executionBody, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -270,7 +270,7 @@ func (c *Client) CurrentValidators(ctx context.Context) ([]*validators.Validator
 	return c.transportClient.CurrentValidators(ctx)
 }
 
-func (c *Client) ApproveValidator(ctx context.Context, joiner []byte) ([]byte, error) {
+func (c *Client) ApproveValidator(ctx context.Context, joiner []byte, opts ...TxOpt) ([]byte, error) {
 	_, err := crypto.Ed25519PublicKeyFromBytes(joiner)
 	if err != nil {
 		return nil, fmt.Errorf("invalid candidate validator public key: %w", err)
@@ -278,7 +278,7 @@ func (c *Client) ApproveValidator(ctx context.Context, joiner []byte) ([]byte, e
 	payload := &transactions.ValidatorApprove{
 		Candidate: joiner,
 	}
-	tx, err := c.newTx(ctx, payload)
+	tx, err := c.newTx(ctx, payload, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -299,7 +299,7 @@ func (c *Client) ValidatorLeave(ctx context.Context) ([]byte, error) {
 	return c.validatorUpdate(ctx, 0)
 }
 
-func (c *Client) validatorUpdate(ctx context.Context, power int64) ([]byte, error) {
+func (c *Client) validatorUpdate(ctx context.Context, power int64, opts ...TxOpt) ([]byte, error) {
 	pubKey := c.Signer.PublicKey()
 
 	var payload transactions.Payload
@@ -314,7 +314,7 @@ func (c *Client) validatorUpdate(ctx context.Context, power int64) ([]byte, erro
 		}
 	}
 
-	tx, err := c.newTx(ctx, payload)
+	tx, err := c.newTx(ctx, payload, opts...)
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/client/opts.go b/pkg/client/opts.go
index 494e19c76..3c7501867 100644
--- a/pkg/client/opts.go
+++ b/pkg/client/opts.go
@@ -51,3 +51,15 @@ func Authenticated(shouldSign bool) CallOpt {
 		o.forceAuthenticated = &copied
 	}
 }
+
+type txOptions struct {
+	nonce int64
+}
+
+type TxOpt func(*txOptions)
+
+func WithNonce(nonce int64) TxOpt {
+	return func(o *txOptions) {
+		o.nonce = nonce
+	}
+}
diff --git a/pkg/client/tx.go b/pkg/client/tx.go
index 2d7dd0843..76e3314ff 100644
--- a/pkg/client/tx.go
+++ b/pkg/client/tx.go
@@ -3,26 +3,37 @@ package client
 import (
 	"context"
 	"fmt"
-	"math/big"
 
-	"github.com/kwilteam/kwil-db/pkg/balances"
 	"github.com/kwilteam/kwil-db/pkg/transactions"
 )
 
 // newTx creates a new Transaction signed by the Client's Signer
-func (c *Client) newTx(ctx context.Context, data transactions.Payload) (*transactions.Transaction, error) {
-	// get nonce from address
-	acc, err := c.transportClient.GetAccount(ctx, c.Signer.PublicKey())
-	if err != nil {
-		acc = &balances.Account{
-			PublicKey: c.Signer.PublicKey(),
-			Nonce:     0,
-			Balance:   big.NewInt(0),
+func (c *Client) newTx(ctx context.Context, data transactions.Payload, opts ...TxOpt) (*transactions.Transaction, error) {
+	txOpts := &txOptions{}
+	for _, opt := range opts {
+		opt(txOpts)
+	}
+
+	var nonce uint64
+	if txOpts.nonce > 0 {
+		nonce = uint64(txOpts.nonce)
+	} else {
+		// Get the latest nonce for the account, if it exists.
+		acc, err := c.transportClient.GetAccount(ctx, c.Signer.PublicKey())
+		if err != nil {
+			return nil, err
+		}
+		// NOTE: an error type would be more robust signalling of a non-existent
+		// account, but presently a nil pubkey is set by pkg/balances.
+		if len(acc.PublicKey) > 0 {
+			nonce = uint64(acc.Nonce + 1)
+		} else {
+			nonce = 1
 		}
 	}
 
 	// build transaction
-	tx, err := transactions.CreateTransaction(data, uint64(acc.Nonce+1))
+	tx, err := transactions.CreateTransaction(data, nonce)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create transaction: %w", err)
 	}
diff --git a/test/driver/client_driver.go b/test/driver/client_driver.go
index 9055e0019..2e2445cc6 100644
--- a/test/driver/client_driver.go
+++ b/test/driver/client_driver.go
@@ -113,7 +113,7 @@ func (d *KwildClientDriver) DatabaseExists(ctx context.Context, dbid string) err
 }
 
 func (d *KwildClientDriver) ExecuteAction(ctx context.Context, dbid string, actionName string, actionInputs ...[]any) ([]byte, error) {
-	rec, err := d.clt.ExecuteAction(ctx, dbid, actionName, actionInputs...)
+	rec, err := d.clt.ExecuteAction(ctx, dbid, actionName, actionInputs)
 	if err != nil {
 		return nil, fmt.Errorf("error executing query: %w", err)
 	}