From 9bdbfe20a5f15ef12b641a234fda6a598e6fce28 Mon Sep 17 00:00:00 2001
From: Viacheslav Poturaev <vearutop@gmail.com>
Date: Sat, 14 May 2022 13:54:15 +0200
Subject: [PATCH] Add a flag to save downloaded S3 resource to a file (#24)

---
 s3/cmd.go |  4 ++++
 s3/job.go | 22 ++++++++++++++++++++--
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/s3/cmd.go b/s3/cmd.go
index 75a6c63..08ed691 100644
--- a/s3/cmd.go
+++ b/s3/cmd.go
@@ -16,6 +16,8 @@ type Flags struct {
 	Bucket       string
 	Key          string
 	PathStyle    bool
+
+	Save string
 }
 
 // AddCommand registers curl command into CLI app.
@@ -37,6 +39,8 @@ func AddCommand(lf *loadgen.Flags) {
 	s3.Flag("path-style", "To use path-style addressing, i.e., `http://s3.amazonaws.com/BUCKET/KEY`.").
 		BoolVar(&f.PathStyle)
 
+	s3.Flag("save", "Path to local file to save the entry.").StringVar(&f.Save)
+
 	s3.Action(func(kp *kingpin.ParseContext) error {
 		return run(*lf, f)
 	})
diff --git a/s3/job.go b/s3/job.go
index ced3121..0eed176 100644
--- a/s3/job.go
+++ b/s3/job.go
@@ -2,6 +2,8 @@ package s3
 
 import (
 	"fmt"
+	"io"
+	"os"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -70,10 +72,26 @@ func (nopWriterAt) WriteAt(p []byte, off int64) (n int, err error) {
 	return len(p), nil
 }
 
-func (j *jobProducer) Job(_ int) (time.Duration, error) {
+func (j *jobProducer) Job(i int) (time.Duration, error) {
 	start := time.Now()
+	w := io.WriterAt(nopWriterAt{})
 
-	n, err := j.dl.Download(nopWriterAt{}, &s3.GetObjectInput{
+	if i == 0 && j.f.Save != "" {
+		f, err := os.Create(j.f.Save)
+		if err != nil {
+			return 0, fmt.Errorf("failed to create file to save S3 result: %w", err)
+		}
+
+		w = f
+
+		defer func() {
+			if err := f.Close(); err != nil {
+				println("failed to close file:", err)
+			}
+		}()
+	}
+
+	n, err := j.dl.Download(w, &s3.GetObjectInput{
 		Bucket: aws.String(j.f.Bucket),
 		Key:    aws.String(j.f.Key),
 	})