From 03896e83330d2bbb8d0f104c61db9d2bbd2d3dad Mon Sep 17 00:00:00 2001
From: Dzung Do <neitdung@gmail.com>
Date: Tue, 7 May 2024 17:36:14 +0700
Subject: [PATCH] Add balances cmd to query multi chain balances

---
 cmd/flags.go |  9 +++++++
 cmd/query.go | 73 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 82 insertions(+)

diff --git a/cmd/flags.go b/cmd/flags.go
index 90b6eeb87..7fbb99bd7 100644
--- a/cmd/flags.go
+++ b/cmd/flags.go
@@ -50,6 +50,7 @@ const (
 	flagInitialBlockHistory            = "block-history"
 	flagFlushInterval                  = "flush-interval"
 	flagMemo                           = "memo"
+	flagKeyName                        = "key-name"
 	flagFilterRule                     = "filter-rule"
 	flagFilterChannels                 = "filter-channels"
 	flagSrcChainID                     = "src-chain-id"
@@ -477,6 +478,14 @@ func memoFlag(v *viper.Viper, cmd *cobra.Command) *cobra.Command {
 	return cmd
 }
 
+func keyNameFlag(v *viper.Viper, cmd *cobra.Command) *cobra.Command {
+	cmd.Flags().String(flagKeyName, "", "a key from the keychain associated with a particular chain")
+	if err := v.BindPFlag(flagMemo, cmd.Flags().Lookup(flagKeyName)); err != nil {
+		panic(err)
+	}
+	return cmd
+}
+
 func OverwriteConfigFlag(v *viper.Viper, cmd *cobra.Command) *cobra.Command {
 	cmd.Flags().BoolP(flagOverwriteConfig, "o", false,
 		"overwrite already configured paths - will clear channel filter(s)")
diff --git a/cmd/query.go b/cmd/query.go
index bc8c6376b..4d9a8ffec 100644
--- a/cmd/query.go
+++ b/cmd/query.go
@@ -34,6 +34,7 @@ func queryCmd(a *appState) *cobra.Command {
 		queryUnrelayedAcknowledgements(a),
 		lineBreakCommand(),
 		queryBalanceCmd(a),
+		queryBalancesCmd(a),
 		queryHeaderCmd(a),
 		queryNodeStateCmd(a),
 		queryTxs(a),
@@ -325,6 +326,78 @@ $ %s query balance ibc-0 testkey`,
 	return cmd
 }
 
+func queryBalancesCmd(a *appState) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "balances [chain-name...]",
+		Short: "query the relayer's account balances on given networks by chain-ID",
+		Args:  withUsage(cobra.MinimumNArgs(1)),
+		Example: strings.TrimSpace(fmt.Sprintf(`
+$ %s query balances ibc-0 ibc-1
+$ %s query balances ibc-0 ibc-1 --key-name=test`,
+			appName, appName,
+		)),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			keyName, _ := cmd.Flags().GetString(flagKeyName)
+
+			data := map[string]string{}
+			for _, arg := range args {
+				chain, ok := a.config.Chains[arg]
+				if !ok {
+					return errChainNotFound(args[0])
+				}
+
+				chainKey := keyName
+				if chainKey == "" {
+					chainKey = chain.ChainProvider.Key()
+				}
+
+				showDenoms, err := cmd.Flags().GetBool(flagIBCDenoms)
+				if err != nil {
+					return err
+				}
+
+				if !chain.ChainProvider.KeyExists(chainKey) {
+					return errKeyDoesntExist(chainKey)
+				}
+				addr, err := chain.ChainProvider.ShowAddress(chainKey)
+				if err != nil {
+					return err
+				}
+
+				coins, err := relayer.QueryBalance(cmd.Context(), chain, addr, showDenoms)
+				if err != nil {
+					return err
+				}
+
+				data[addr] = coins.String()
+			}
+			// Convert the map to a JSON string
+			jsonOutput, err := json.Marshal(data)
+			if err != nil {
+				return err
+			}
+
+			output, _ := cmd.Flags().GetString(flagOutput)
+			switch output {
+			case formatJson:
+				fmt.Fprint(cmd.OutOrStdout(), string(jsonOutput))
+			case formatLegacy:
+				fallthrough
+			default:
+				for addr, balance := range data {
+					fmt.Fprintf(cmd.OutOrStdout(), "address {%s} balance {%s} \n", addr, balance)
+				}
+			}
+			return nil
+		},
+	}
+
+	cmd = addOutputFlag(a.viper, cmd)
+	cmd = ibcDenomFlags(a.viper, cmd)
+	cmd = keyNameFlag(a.viper, cmd)
+	return cmd
+}
+
 func queryHeaderCmd(a *appState) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "header chain_name [height]",