Adds ability to change the print format to yaml for kube_capture

k8s/object_writer.go

@@ -12,6 +12,7 @@ import (
 
 type ObjectWriter struct {
 	writeDir string
+	printer  printers.ResourcePrinter
 }
 
 func (w ObjectWriter) Write(result SearchResult) (string, error) {
@@ -35,7 +36,13 @@ func (w ObjectWriter) Write(result SearchResult) (string, error) {
 	}
 
 	now := time.Now().Format("200601021504.0000")
-	path := filepath.Join(w.writeDir, fmt.Sprintf("%s-%s.json", result.ResourceName, now))
+	var extension string
+	if _, ok := w.printer.(*printers.JSONPrinter); ok {
+		extension = "json"
+	} else {
+		extension = "yaml"
+	}
+	path := filepath.Join(w.writeDir, fmt.Sprintf("%s-%s.%s", result.ResourceName, now, extension))
 
 	file, err := os.Create(path)
 	if err != nil {
@@ -45,8 +52,7 @@ func (w ObjectWriter) Write(result SearchResult) (string, error) {
 
 	logrus.Debugf("objectWriter: saving %s search results to: %s", result.ResourceName, path)
 
-	printer := new(printers.JSONPrinter)
-	if err := printer.PrintObj(result.List, file); err != nil {
+	if err := w.printer.PrintObj(result.List, file); err != nil {
 		if wErr := writeError(err, file); wErr != nil {
 			return "", fmt.Errorf("objectWriter: failed to write previous err [%s] to file: %s", err, wErr)
 		}

k8s/result_writer.go

@@ -6,25 +6,34 @@ import (
 	"os"
 	"path/filepath"
 
+	"k8s.io/cli-runtime/pkg/printers"
 	"k8s.io/client-go/rest"
 )
 
 type ResultWriter struct {
 	workdir   string
 	writeLogs bool
 	restApi   rest.Interface
+	printer   printers.ResourcePrinter
 }
 
-func NewResultWriter(workdir, what string, restApi rest.Interface) (*ResultWriter, error) {
+func NewResultWriter(workdir, what, outputFormat string, restApi rest.Interface) (*ResultWriter, error) {
 	var err error
 	workdir = filepath.Join(workdir, BaseDirname)
 	if err := os.MkdirAll(workdir, 0744); err != nil && !os.IsExist(err) {
 		return nil, err
 	}
 
 	writeLogs := what == "logs" || what == "all"
+	var printer printers.ResourcePrinter
+	if outputFormat == "json" {
+		printer = &printers.JSONPrinter{}
+	} else if outputFormat == "yaml" {
+		printer = &printers.YAMLPrinter{}
+	}
 	return &ResultWriter{
 		workdir:   workdir,
+		printer:   printer,
 		writeLogs: writeLogs,
 		restApi:   restApi,
 	}, err
@@ -44,6 +53,7 @@ func (w *ResultWriter) Write(ctx context.Context, searchResults []SearchResult)
 	for _, result := range searchResults {
 		objWriter := ObjectWriter{
 			writeDir: w.workdir,
+			printer:  w.printer,
 		}
 		writeDir, err := objWriter.Write(result)
 		if err != nil {

starlark/kube_capture.go

@@ -22,10 +22,13 @@ func KubeCaptureFn(thread *starlark.Thread, _ *starlark.Builtin, args starlark.T
 	var groups, categories, kinds, namespaces, versions, names, labels, containers *starlark.List
 	var kubeConfig *starlarkstruct.Struct
 	var what string
+	var outputFormat string
+	logrus.Info(kwargs)
 
 	if err := starlark.UnpackArgs(
 		identifiers.kubeCapture, args, kwargs,
 		"what", &what,
+		"output_format?", &outputFormat,
 		"groups?", &groups,
 		"categories?", &categories,
 		"kinds?", &kinds,
@@ -61,7 +64,7 @@ func KubeCaptureFn(thread *starlark.Thread, _ *starlark.Builtin, args starlark.T
 	data := thread.Local(identifiers.crashdCfg)
 	cfg, _ := data.(*starlarkstruct.Struct)
 	workDirVal, _ := cfg.Attr("workdir")
-	resultDir, err := write(ctx, trimQuotes(workDirVal.String()), what, client, k8s.SearchParams{
+	resultDir, err := write(ctx, trimQuotes(workDirVal.String()), what, outputFormat, client, k8s.SearchParams{
 		Groups:     toSlice(groups),
 		Categories: toSlice(categories),
 		Kinds:      toSlice(kinds),
@@ -85,7 +88,7 @@ func KubeCaptureFn(thread *starlark.Thread, _ *starlark.Builtin, args starlark.T
 		}), nil
 }
 
-func write(ctx context.Context, workdir, what string, client *k8s.Client, params k8s.SearchParams) (string, error) {
+func write(ctx context.Context, workdir, what, outputFormat string, client *k8s.Client, params k8s.SearchParams) (string, error) {
 
 	logrus.Debugf("kube_capture(what=%s)", what)
 	switch what {
@@ -97,13 +100,19 @@ func write(ctx context.Context, workdir, what string, client *k8s.Client, params
 	default:
 		return "", errors.Errorf("don't know how to get: %s", what)
 	}
+	if outputFormat == "" {
+		outputFormat = "json"
+	}
+	if outputFormat != "json" && outputFormat != "yaml" {
+		return "", errors.Errorf("unsupported output format: %s", outputFormat)
+	}
 
 	searchResults, err := client.Search(ctx, params)
 	if err != nil {
 		return "", err
 	}
 
-	resultWriter, err := k8s.NewResultWriter(workdir, what, client.CoreRest)
+	resultWriter, err := k8s.NewResultWriter(workdir, what, outputFormat, client.CoreRest)
 	if err != nil {
 		return "", errors.Wrap(err, "failed to initialize writer")
 	}

starlark/kube_capture_test.go

@@ -93,13 +93,45 @@ func TestKubeCapture(t *testing.T) {
 				}
 			},
 		},
+		{
+			name: "test with invalid output format",
+			kwargs: func(t *testing.T) []starlark.Tuple {
+				return []starlark.Tuple{
+					[]starlark.Value{starlark.String("what"), starlark.String("objects")},
+					[]starlark.Value{starlark.String("groups"), starlark.NewList([]starlark.Value{starlark.String("core")})},
+					[]starlark.Value{starlark.String("kinds"), starlark.NewList([]starlark.Value{starlark.String("services")})},
+					[]starlark.Value{starlark.String("namespaces"), starlark.NewList([]starlark.Value{starlark.String("default"), starlark.String("kube-system")})},
+					[]starlark.Value{starlark.String("output_format"), starlark.String("xml")},
+				}
+			},
+			eval: func(t *testing.T, kwargs []starlark.Tuple) {
+				val, err := KubeCaptureFn(newTestThreadLocal(t), nil, nil, kwargs)
+				if err != nil {
+					t.Fatalf("failed to execute: %s", err)
+				}
+				resultStruct, ok := val.(*starlarkstruct.Struct)
+				if !ok {
+					t.Fatalf("expecting type *starlarkstruct.Struct, got %T", val)
+				}
+
+				errVal, err := resultStruct.Attr("error")
+				if err != nil {
+					t.Error(err)
+				}
+				resultErr := errVal.(starlark.String).GoString()
+				if resultErr != "unsupported output format: xml" {
+					t.Fatalf("Expected error \"unsupported output format: xml\" but got: \"%s\"", resultErr)
+				}
+			},
+		},
 		{
 			name: "test for non-namespaced objects",
 			kwargs: func(t *testing.T) []starlark.Tuple {
 				return []starlark.Tuple{
 					[]starlark.Value{starlark.String("what"), starlark.String("objects")},
 					[]starlark.Value{starlark.String("groups"), starlark.NewList([]starlark.Value{starlark.String("core")})},
 					[]starlark.Value{starlark.String("kinds"), starlark.NewList([]starlark.Value{starlark.String("nodes")})},
+					[]starlark.Value{starlark.String("output_format"), starlark.String("yaml")},
 				}
 			},
 			eval: func(t *testing.T, kwargs []starlark.Tuple) {
@@ -150,7 +182,9 @@ func TestKubeCapture(t *testing.T) {
 				if !strings.Contains(files[0].Name(), "nodes-") {
 					t.Errorf("expecting to find a node output file, but fond: %s", files[0].Name())
 				}
-
+				if !strings.HasSuffix(files[0].Name(), ".yaml") {
+					t.Errorf("expecting to find a yaml file, but found: %s", files[0].Name())
+				}
 			},
 		},
 		{
@@ -461,6 +495,9 @@ kube_data = kube_capture(what="objects", groups=["core"], kinds=["nodes"])`, wor
 				if !strings.Contains(files[0].Name(), "nodes-") {
 					t.Errorf("expecting to find a node output file, but fond: %s", files[0].Name())
 				}
+				if !strings.HasSuffix(files[0].Name(), ".json") {
+					t.Errorf("expecting to find a json file, but found: %s", files[0].Name())
+				}
 			},
 		},
 		{