From d9252e934972db957be8763de7a6adad3776546f Mon Sep 17 00:00:00 2001
From: Hieu Vu <72878483+hieuvubk@users.noreply.github.com>
Date: Fri, 14 Jun 2024 15:25:24 +0700
Subject: [PATCH] print output depend on FlagOutput

---
 server/v2/cometbft/commands.go | 54 ++++++++++++++++++++++------------
 1 file changed, 35 insertions(+), 19 deletions(-)

diff --git a/server/v2/cometbft/commands.go b/server/v2/cometbft/commands.go
index fe3d513feddc..b635ca505c36 100644
--- a/server/v2/cometbft/commands.go
+++ b/server/v2/cometbft/commands.go
@@ -72,10 +72,7 @@ func (s *CometBFTServer[T]) StatusCommand() *cobra.Command {
 				return err
 			}
 
-			cmd.Println(string(output))
-
-			// TODO: figure out yaml and json output
-			return nil
+			return printOutput(cmd, output)
 		},
 	}
 
@@ -215,15 +212,12 @@ for. Each module documents its respective events under 'xx_events.md'.
 				return err
 			}
 
-			// return clientCtx.PrintProto(blocks) // TODO: previously we had this, but I think it can be replaced with a simple json marshal.
-			// We are missing YAML output tho.
 			bz, err := protojson.Marshal(blocks)
 			if err != nil {
 				return err
 			}
 
-			_, err = cmd.OutOrStdout().Write(bz)
-			return err
+			return printOutput(cmd, bz)
 		},
 	}
 
@@ -282,15 +276,12 @@ $ %s query block --%s=%s <hash>
 					return fmt.Errorf("no block found with height %s", args[0])
 				}
 
-				// return clientCtx.PrintProto(output)
-
 				bz, err := json.Marshal(output)
 				if err != nil {
 					return err
 				}
 
-				_, err = cmd.OutOrStdout().Write(bz)
-				return err
+				return printOutput(cmd, bz)
 
 			case auth.TypeHash:
 
@@ -308,14 +299,12 @@ $ %s query block --%s=%s <hash>
 					return fmt.Errorf("no block found with hash %s", args[0])
 				}
 
-				// return clientCtx.PrintProto(output)
 				bz, err := json.Marshal(output)
 				if err != nil {
 					return err
 				}
 
-				_, err = cmd.OutOrStdout().Write(bz)
-				return err
+				return printOutput(cmd, bz)
 
 			default:
 				return fmt.Errorf("unknown --%s value %s", auth.FlagType, typ)
@@ -371,10 +360,7 @@ func (s *CometBFTServer[T]) QueryBlockResultsCmd() *cobra.Command {
 				return err
 			}
 
-			cmd.Println(string(blockResStr))
-
-			// TODO: figure out yaml and json output
-			return nil
+			return printOutput(cmd, blockResStr)
 		},
 	}
 
@@ -425,3 +411,33 @@ func (s *CometBFTServer[T]) BootstrapStateCmd() *cobra.Command {
 
 	return cmd
 }
+
+func printOutput(cmd *cobra.Command, out []byte) error {
+	// Get flags output
+	outFlag, err := cmd.Flags().GetString(flags.FlagOutput) 
+	if err != nil {
+		return err
+	}
+
+	if outFlag == "text" {
+		out, err = yaml.JSONToYAML(out)
+		if err != nil {
+			return err
+		}
+	}
+
+	writer := cmd.OutOrStdout()
+	_, err = writer.Write(out)
+	if err != nil {
+		return err
+	}
+
+	if outFlag != "text" {
+		// append new-line for formats besides YAML
+		_, err = writer.Write([]byte("\n"))
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}