Improve tsh play JSON & YAML output (#18768)

Before this commit, events in the output of `tsh play --format=yaml` were
separated by a newline. This commit ensures that these events are now
separated by `---`, allowing the command output to play well with `yq`.

For consistency, the output of `tsh play --format=yaml <session-id>`
now mirrors what we get e.g. with `tctl nodes ls --format=yaml`.
The same applies to `tsh play --format json`.

lib/utils/jsontools.go

@@ -18,6 +18,7 @@ package utils
 
 import (
 	"bytes"
+	"encoding/json"
 	"io"
 	"reflect"
 	"unicode"
@@ -100,6 +101,14 @@ func FastMarshalIndent(v interface{}, prefix, indent string) ([]byte, error) {
 	return data, nil
 }
 
+// WriteJSON marshals multiple documents as a JSON list with indentation.
+func WriteJSON(w io.Writer, values interface{}) error {
+	encoder := json.NewEncoder(w)
+	encoder.SetIndent("", "    ")
+	err := encoder.Encode(values)
+	return trace.Wrap(err)
+}
+
 const yamlDocDelimiter = "---"
 
 // WriteYAML detects whether value is a list

tool/tctl/common/collection.go

@@ -17,7 +17,6 @@ limitations under the License.
 package common
 
 import (
-	"encoding/json"
 	"fmt"
 	"io"
 	"sort"
@@ -157,17 +156,12 @@ func (s *serverCollection) writeText(w io.Writer) error {
 	return trace.Wrap(err)
 }
 
-func (s *serverCollection) writeYaml(w io.Writer) error {
+func (s *serverCollection) writeYAML(w io.Writer) error {
 	return utils.WriteYAML(w, s.servers)
 }
 
 func (s *serverCollection) writeJSON(w io.Writer) error {
-	data, err := json.MarshalIndent(s.resources(), "", "    ")
-	if err != nil {
-		return trace.Wrap(err)
-	}
-	_, err = w.Write(data)
-	return trace.Wrap(err)
+	return utils.WriteJSON(w, s.servers)
 }
 
 type userCollection struct {
@@ -437,12 +431,7 @@ func formatLastHeartbeat(t time.Time) string {
 }
 
 func writeJSON(c ResourceCollection, w io.Writer) error {
-	data, err := json.MarshalIndent(c.resources(), "", "    ")
-	if err != nil {
-		return trace.Wrap(err)
-	}
-	_, err = w.Write(data)
-	return trace.Wrap(err)
+	return utils.WriteJSON(w, c.resources())
 }
 
 func writeYAML(c ResourceCollection, w io.Writer) error {
@@ -506,20 +495,11 @@ func (a *appServerCollection) writeText(w io.Writer) error {
 }
 
 func (a *appServerCollection) writeJSON(w io.Writer) error {
-	data, err := json.MarshalIndent(a.toMarshal(), "", "    ")
-	if err != nil {
-		return trace.Wrap(err)
-	}
-	_, err = w.Write(data)
-	return trace.Wrap(err)
-}
-
-func (a *appServerCollection) toMarshal() interface{} {
-	return a.servers
+	return utils.WriteJSON(w, a.servers)
 }
 
 func (a *appServerCollection) writeYAML(w io.Writer) error {
-	return utils.WriteYAML(w, a.toMarshal())
+	return utils.WriteYAML(w, a.servers)
 }
 
 type appCollection struct {
@@ -676,20 +656,11 @@ func (c *databaseServerCollection) writeText(w io.Writer) error {
 }
 
 func (c *databaseServerCollection) writeJSON(w io.Writer) error {
-	data, err := json.MarshalIndent(c.toMarshal(), "", "    ")
-	if err != nil {
-		return trace.Wrap(err)
-	}
-	_, err = w.Write(data)
-	return trace.Wrap(err)
-}
-
-func (c *databaseServerCollection) toMarshal() interface{} {
-	return c.servers
+	return utils.WriteJSON(w, c.servers)
 }
 
 func (c *databaseServerCollection) writeYAML(w io.Writer) error {
-	return utils.WriteYAML(w, c.toMarshal())
+	return utils.WriteYAML(w, c.servers)
 }
 
 type databaseCollection struct {
@@ -837,12 +808,7 @@ func (c *windowsDesktopAndServiceCollection) writeYAML(w io.Writer) error {
 }
 
 func (c *windowsDesktopAndServiceCollection) writeJSON(w io.Writer) error {
-	data, err := json.MarshalIndent(c.desktops, "", "    ")
-	if err != nil {
-		return trace.Wrap(err)
-	}
-	_, err = w.Write(data)
-	return trace.Wrap(err)
+	return utils.WriteJSON(w, c.desktops)
 }
 
 type tokenCollection struct {
@@ -911,12 +877,7 @@ func (c *kubeServerCollection) writeYAML(w io.Writer) error {
 }
 
 func (c *kubeServerCollection) writeJSON(w io.Writer) error {
-	data, err := json.MarshalIndent(c.servers, "", "    ")
-	if err != nil {
-		return trace.Wrap(err)
-	}
-	_, err = w.Write(data)
-	return trace.Wrap(err)
+	return utils.WriteJSON(w, c.servers)
 }
 
 type kubeClusterCollection struct {

tool/tctl/common/node_command.go

@@ -208,7 +208,7 @@ func (c *NodeCommand) Invite(ctx context.Context, client auth.ClientI) error {
 	return nil
 }
 
-// ListActive retreives the list of nodes who recently sent heartbeats to
+// ListActive retrieves the list of nodes who recently sent heartbeats to
 // to a cluster and prints it to stdout
 func (c *NodeCommand) ListActive(ctx context.Context, clt auth.ClientI) error {
 	labels, err := libclient.ParseLabelSpec(c.labels)
@@ -244,7 +244,7 @@ func (c *NodeCommand) ListActive(ctx context.Context, clt auth.ClientI) error {
 			return trace.Wrap(err)
 		}
 	case teleport.YAML:
-		if err := coll.writeYaml(os.Stdout); err != nil {
+		if err := coll.writeYAML(os.Stdout); err != nil {
 			return trace.Wrap(err)
 		}
 	case teleport.JSON:

tool/tsh/tsh.go

@@ -1271,18 +1271,21 @@ func exportSession(cf *CLIConf) error {
 		// when playing from a file, id is not included, this
 		// makes the outputs otherwise identical
 		delete(event, "id")
-		var e []byte
-		var err error
-		if format == teleport.JSON {
-			e, err = utils.FastMarshal(event)
-		} else {
-			e, err = yaml.Marshal(event)
-		}
-		if err != nil {
+	}
+
+	switch format {
+	case teleport.JSON:
+		if err := utils.WriteJSON(os.Stdout, events); err != nil {
 			return trace.Wrap(err)
 		}
-		fmt.Println(string(e))
+	case teleport.YAML:
+		if err := utils.WriteYAML(os.Stdout, events); err != nil {
+			return trace.Wrap(err)
+		}
+	default:
+		return trace.Errorf("Invalid format %s, only pty, json and yaml are supported", format)
 	}
+
 	return nil
 }