Merge pull request #442 from replicatedhq/divolgin/closer

Allow memory writers

Author: divolgin

pkg/collect/copy.go

@@ -138,6 +138,8 @@ func copyFilesFromPod(ctx context.Context, dstPath string, clientConfig *restcli
 		if err != nil {
 			return nil, nil, errors.Wrap(err, "failed to craete dest file")
 		}
+		defer result.CloseWriter(dstPath, filepath.Base(containerPath)+".tar", w)
+
 		stdoutWriter = w
 	}
 

pkg/collect/copy_from_host.go

@@ -315,6 +315,8 @@ func copyFilesFromHost(ctx context.Context, dstPath string, clientConfig *restcl
 		if err != nil {
 			return nil, nil, errors.Wrap(err, "failed to craete dest file")
 		}
+		defer result.CloseWriter(dstPath, "archive.tar", w)
+
 		stdoutWriter = w
 	}
 

pkg/collect/logs.go

@@ -125,7 +125,7 @@ func savePodLogs(ctx context.Context, bundlePath string, client *kubernetes.Clie
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get log writer")
 	}
-	defer logWriter.Close()
+	defer result.CloseWriter(bundlePath, fileKey+".log", logWriter)
 
 	_, err = io.Copy(logWriter, podLogs)
 	if err != nil {
@@ -145,7 +145,7 @@ func savePodLogs(ctx context.Context, bundlePath string, client *kubernetes.Clie
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to get previous log writer")
 	}
-	defer logWriter.Close()
+	defer result.CloseWriter(bundlePath, fileKey+"-previous.log", logWriter)
 
 	_, err = io.Copy(prevLogWriter, podLogs)
 	if err != nil {

pkg/collect/result.go

@@ -89,6 +89,10 @@ func (r CollectorResult) GetReader(bundlePath string, relativePath string) (io.R
 		return bytes.NewReader(r[relativePath]), nil
 	}
 
+	if bundlePath == "" {
+		return nil, errors.New("cannot create reader, bundle path is empty")
+	}
+
 	filename := filepath.Join(bundlePath, relativePath)
 	f, err := os.Open(filename)
 	if err != nil {
@@ -98,9 +102,10 @@ func (r CollectorResult) GetReader(bundlePath string, relativePath string) (io.R
 	return f, nil
 }
 
-func (r CollectorResult) GetWriter(bundlePath string, relativePath string) (io.WriteCloser, error) {
-	if r[relativePath] != nil {
-		return nil, errors.New("writing to memory is not implemented")
+func (r CollectorResult) GetWriter(bundlePath string, relativePath string) (io.Writer, error) {
+	if bundlePath == "" {
+		var b bytes.Buffer
+		return &b, nil
 	}
 
 	fileDir, _ := filepath.Split(relativePath)
@@ -120,6 +125,19 @@ func (r CollectorResult) GetWriter(bundlePath string, relativePath string) (io.W
 	return f, nil
 }
 
+func (r CollectorResult) CloseWriter(bundlePath string, relativePath string, writer interface{}) error {
+	if c, ok := writer.(io.Closer); ok {
+		return errors.Wrap(c.Close(), "failed to close writer")
+	}
+
+	if b, ok := writer.(*bytes.Buffer); ok {
+		r[relativePath] = b.Bytes()
+		return nil
+	}
+
+	return errors.Errorf("cannot close writer of type %T", writer)
+}
+
 func TarSupportBundleDir(bundlePath string, input CollectorResult, outputFilename string) error {
 	fileWriter, err := os.Create(outputFilename)
 	if err != nil {