From c3feff71d2774dffd135416854e93b38a96c8a2a Mon Sep 17 00:00:00 2001 From: Marin Dzhigarov Date: Mon, 23 Dec 2024 10:15:02 +0200 Subject: [PATCH] Adds ability to change the print format to yaml for kube_capture --- k8s/object_writer.go | 12 ++++++++--- k8s/result_writer.go | 12 ++++++++++- starlark/kube_capture.go | 15 +++++++++++--- starlark/kube_capture_test.go | 39 ++++++++++++++++++++++++++++++++++- 4 files changed, 70 insertions(+), 8 deletions(-) diff --git a/k8s/object_writer.go b/k8s/object_writer.go index 5419f943..8dd38d52 100644 --- a/k8s/object_writer.go +++ b/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) } diff --git a/k8s/result_writer.go b/k8s/result_writer.go index 08e74492..9e081f11 100644 --- a/k8s/result_writer.go +++ b/k8s/result_writer.go @@ -6,6 +6,7 @@ import ( "os" "path/filepath" + "k8s.io/cli-runtime/pkg/printers" "k8s.io/client-go/rest" ) @@ -13,9 +14,10 @@ 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) { @@ -23,8 +25,15 @@ func NewResultWriter(workdir, what string, restApi rest.Interface) (*ResultWrite } 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 { diff --git a/starlark/kube_capture.go b/starlark/kube_capture.go index e50494e7..aeaf7577 100644 --- a/starlark/kube_capture.go +++ b/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") } diff --git a/starlark/kube_capture_test.go b/starlark/kube_capture_test.go index 92e65c9b..8d83ecd3 100644 --- a/starlark/kube_capture_test.go +++ b/starlark/kube_capture_test.go @@ -93,6 +93,37 @@ 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 { @@ -100,6 +131,7 @@ func TestKubeCapture(t *testing.T) { []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()) + } }, }, {