Improve retrieve command with extraction option

Adds a flag to sonobuoy extract which allows the user
to choose to download and extract the data in a single
command. This is ideal especially when worrying about if
someone on windows 1) can untar files easily 2) can be
sure that untar will end up being aware that the tarball
uses slash pathing and the client is on Windows.

Fixes #1277

Signed-off-by: John Schnake <[email protected]>

cmd/sonobuoy/app/args.go

@@ -390,3 +390,11 @@ func AddDryRunFlag(flag *bool, flags *pflag.FlagSet) {
 func AddNodeSelectorsFlag(p *NodeSelectors, flags *pflag.FlagSet) {
 	flags.Var(p, "aggregator-node-selector", "Node selectors to add to the aggregator. Values can be given multiple times and are in the form key:value")
 }
+
+// AddExtractFlag adds a boolean flag to extract results instead of just downloading the tarball.
+func AddExtractFlag(flag *bool, flags *pflag.FlagSet) {
+	flags.BoolVarP(
+		flag, "extract", "x", false,
+		"If true, extracts the results instead of just downloading the results",
+	)
+}

cmd/sonobuoy/app/retrieve.go

@@ -18,6 +18,7 @@ package app
 
 import (
 	"fmt"
+	"io"
 	"os"
 
 	"github.com/pkg/errors"
@@ -33,64 +34,90 @@ var (
 )
 
 type retrieveFlags struct {
-	namespace string
-	kubecfg   Kubeconfig
+	namespace      string
+	kubecfg        Kubeconfig
+	extract        bool
+	outputLocation string
 }
 
-var rcvFlags retrieveFlags
-
 func NewCmdRetrieve() *cobra.Command {
+	rcvFlags := retrieveFlags{}
 	cmd := &cobra.Command{
 		Use:   "retrieve [path]",
 		Short: "Retrieves the results of a sonobuoy run to a specified path",
-		Run:   retrieveResults,
+		Run:   retrieveResultsCmd(&rcvFlags),
 		Args:  cobra.MaximumNArgs(1),
 	}
 
 	AddKubeconfigFlag(&rcvFlags.kubecfg, cmd.Flags())
 	AddNamespaceFlag(&rcvFlags.namespace, cmd.Flags())
-
+	AddExtractFlag(&rcvFlags.extract, cmd.Flags())
 	return cmd
 }
 
-func retrieveResults(cmd *cobra.Command, args []string) {
-	outDir := defaultOutDir
-	if len(args) > 0 {
-		outDir = args[0]
-	}
+func retrieveResultsCmd(opts *retrieveFlags) func(cmd *cobra.Command, args []string) {
+	return func(cmd *cobra.Command, args []string) {
+		opts.outputLocation = defaultOutDir
+		if len(args) > 0 {
+			opts.outputLocation = args[0]
+		}
 
-	sbc, err := getSonobuoyClientFromKubecfg(rcvFlags.kubecfg)
-	if err != nil {
-		errlog.LogError(errors.Wrap(err, "could not create sonobuoy client"))
-		os.Exit(1)
-	}
+		sbc, err := getSonobuoyClientFromKubecfg(opts.kubecfg)
+		if err != nil {
+			errlog.LogError(errors.Wrap(err, "could not create sonobuoy client"))
+			os.Exit(1)
+		}
 
-	// Get a reader that contains the tar output of the results directory.
-	reader, ec, err := sbc.RetrieveResults(&client.RetrieveConfig{Namespace: rcvFlags.namespace})
-	if err != nil {
-		errlog.LogError(err)
-		os.Exit(1)
+		// Get a reader that contains the tar output of the results directory.
+		reader, ec, err := sbc.RetrieveResults(&client.RetrieveConfig{Namespace: opts.namespace})
+		if err != nil {
+			errlog.LogError(err)
+			os.Exit(1)
+		}
+
+		err = retrieveResults(*opts, reader, ec)
+		if _, ok := err.(exec.CodeExitError); ok {
+			fmt.Fprintln(os.Stderr, "Results not ready yet. Check `sonobuoy status` for status.")
+			os.Exit(1)
+		} else if err != nil {
+			fmt.Fprintf(os.Stderr, "error retrieving results: %v\n", err)
+			os.Exit(2)
+		}
 	}
+}
 
+func retrieveResults(opts retrieveFlags, r io.Reader, ec <-chan error) error {
 	eg := &errgroup.Group{}
 	eg.Go(func() error { return <-ec })
 	eg.Go(func() error {
-		filesCreated, err := client.UntarAll(reader, outDir, "")
+		// This untars the request itself, which is tar'd as just part of the API request, not the sonobuoy logic.
+		filesCreated, err := client.UntarAll(r, opts.outputLocation, "")
 		if err != nil {
 			return err
 		}
-		for _, name := range filesCreated {
-			fmt.Println(name)
+		if !opts.extract {
+			// Only print the filename if not extracting. Allows capturing the filename for scripting.
+			for _, name := range filesCreated {
+				fmt.Println(name)
+			}
+			return nil
+		} else {
+			for _, filename := range filesCreated {
+				err := client.UntarFile(filename, opts.outputLocation, true)
+				if err != nil {
+					// Just log errors if it is just not cleaning up the file.
+					re, ok := err.(*client.DeletionError)
+					if ok {
+						errlog.LogError(re)
+					} else {
+						return err
+					}
+				}
+			}
+
+			return err
 		}
-		return nil
 	})
 
-	err = eg.Wait()
-	if _, ok := err.(exec.CodeExitError); ok {
-		fmt.Fprintln(os.Stderr, "Results not ready yet. Check `sonobuoy status` for status.")
-		os.Exit(1)
-	} else if err != nil {
-		fmt.Fprintf(os.Stderr, "error retrieving results: %v\n", err)
-		os.Exit(2)
-	}
+	return eg.Wait()
 }

pkg/client/retrieve.go

@@ -43,6 +43,15 @@ var (
 	}
 )
 
+type DeletionError struct {
+	filename string
+	err      error
+}
+
+func (d *DeletionError) Error() string {
+	return errors.Wrapf(d.err, "failed to delete file %q", d.filename).Error()
+}
+
 // RetrieveResults copies results from a sonobuoy run into a Reader in tar format.
 // It also returns a channel of errors, where any errors encountered when writing results
 // will be sent, and an error in the case where the config validation fails.
@@ -210,6 +219,27 @@ func UntarAll(reader io.Reader, destFile, prefix string) (filenames []string, re
 	return filenames, nil
 }
 
+// UntarFile untars the file, filename, into the given destination directory with the given prefix. If delete
+// is true, it deletes the original file after extraction.
+func UntarFile(filename string, destination string, delete bool) error {
+	f, err := os.Open(filename)
+	if err != nil {
+		return errors.Wrapf(err, "failed to open file %q for extraction after downloading it", filename)
+	}
+	defer f.Close()
+
+	_, err = UntarAll(f, destination, "")
+	if err != nil {
+		return errors.Wrapf(err, "failed to extract file %q after downloading it", filename)
+	}
+	if err := os.Remove(filename); err != nil {
+		// Returning a specific type of error so that consumers can ignore if desired. The file did
+		// get untar'd successfully and so this error has less significance.
+		return &DeletionError{filename: filename, err: err}
+	}
+	return nil
+}
+
 // dirExists checks if a path exists and is a directory.
 func dirExists(path string) (bool, error) {
 	fi, err := os.Stat(path)

test/integration/sonobuoy_integration_test.go

@@ -148,6 +148,48 @@ func TestSimpleRun(t *testing.T) {
 	mustRunSonobuoyCommandWithContext(ctx, t, args)
 }
 
+func TestRetrieveAndExtract(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	ns, cleanup := getNamespace(t)
+	defer cleanup()
+
+	args := fmt.Sprintf("run --image-pull-policy IfNotPresent --wait -p testImage/yaml/job-junit-passing-singlefile.yaml -n %v", ns)
+	mustRunSonobuoyCommandWithContext(ctx, t, args)
+
+	// Create tmpdir and extract contents into it
+	tmpdir, err := ioutil.TempDir("", "TestRetrieveAndExtract")
+	if err != nil {
+		t.Fatal("Failed to create tmp dir")
+	}
+	defer os.RemoveAll(tmpdir)
+	args = fmt.Sprintf("retrieve %v -n %v --extract", tmpdir, ns)
+	mustRunSonobuoyCommandWithContext(ctx, t, args)
+
+	// Check that the files are there. Lots of ways to test this but I'm simply going to check that we have
+	// a "lot" of files.
+	files := []string{}
+	err = filepath.Walk(tmpdir,
+		func(path string, info os.FileInfo, err error) error {
+			if err != nil {
+				return err
+			}
+			files = append(files, path)
+			return nil
+		})
+	if err != nil {
+		t.Fatalf("Failed to walk path to check: %v", err)
+	}
+
+	// Verbose logging here in case we want to just see if certain files were found. Can remove
+	// this and just log on error if it is too much.
+	t.Logf("Extracted files:\n%v", strings.Join(files, "\n\t-"))
+	if len(files) < 20 {
+		t.Errorf("Expected many files to be extracted into %v, but only got %v", tmpdir, len(files))
+	}
+}
+
 // TestQuick runs a real "--mode quick" check against the cluster to ensure that it passes.
 func TestQuick(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)