Add progress to file download

Gopkg.toml

@@ -26,3 +26,7 @@
 
 [[constraint]]
   name = "github.com/hashicorp/go-version"
+
+[[constraint]]
+  name = "github.com/dustin/go-humanize"
+  version = "1.0.0"

README.md

@@ -86,6 +86,7 @@ The supported options are:
   saved in bash history.
 - `--github-api-version` (**Optional**): Used when fetching an artifact from a GitHub Enterprise instance.
   Defaults to `v3`. This is ignored when fetching from GitHub.com.
+- `--progress` (**Optional**): Used when fetching a big file and want to see progress on the fetch.
 
 The supported arguments are:
 

checksum_test.go

@@ -37,7 +37,7 @@ func TestVerifyReleaseAsset(t *testing.T) {
 		t.Fatalf("Failed to parse sample release asset GitHub URL into Fetch GitHubRepo struct: %s", err)
 	}
 
-	assetPaths, fetchErr := downloadReleaseAssets(SAMPLE_RELEASE_ASSET_NAME, tmpDir, githubRepo, SAMPLE_RELEASE_ASSET_VERSION)
+	assetPaths, fetchErr := downloadReleaseAssets(SAMPLE_RELEASE_ASSET_NAME, tmpDir, githubRepo, SAMPLE_RELEASE_ASSET_VERSION, false)
 	if fetchErr != nil {
 		t.Fatalf("Failed to download release asset: %s", fetchErr)
 	}
@@ -72,7 +72,7 @@ func TestVerifyChecksumOfReleaseAsset(t *testing.T) {
 		t.Fatalf("Failed to parse sample release asset GitHub URL into Fetch GitHubRepo struct: %s", err)
 	}
 
-	assetPaths, fetchErr := downloadReleaseAssets(SAMPLE_RELEASE_ASSET_REGEX, tmpDir, githubRepo, SAMPLE_RELEASE_ASSET_VERSION)
+	assetPaths, fetchErr := downloadReleaseAssets(SAMPLE_RELEASE_ASSET_REGEX, tmpDir, githubRepo, SAMPLE_RELEASE_ASSET_VERSION, false)
 	if fetchErr != nil {
 		t.Fatalf("Failed to download release asset: %s", fetchErr)
 	}

github.go

@@ -9,6 +9,9 @@ import (
 	"net/url"
 	"os"
 	"regexp"
+	"strings"
+
+	"github.com/dustin/go-humanize"
 )
 
 type GitHubRepo struct {
@@ -156,14 +159,13 @@ func ParseUrlIntoGitHubRepo(url string, token string, instance GitHubInstance) (
 }
 
 // Download the release asset with the given id and return its body
-func DownloadReleaseAsset(repo GitHubRepo, assetId int, destPath string) *FetchError {
+func DownloadReleaseAsset(repo GitHubRepo, assetId int, destPath string, withProgress bool) *FetchError {
 	url := createGitHubRepoUrlForPath(repo, fmt.Sprintf("releases/assets/%d", assetId))
 	resp, err := callGitHubApi(repo, url, map[string]string{"Accept": "application/octet-stream"})
 	if err != nil {
 		return err
 	}
-
-	return writeResonseToDisk(resp, destPath)
+	return writeResonseToDisk(resp, destPath, withProgress)
 }
 
 // Get information about the GitHub release with the given tag
@@ -235,8 +237,39 @@ func callGitHubApi(repo GitHubRepo, path string, customHeaders map[string]string
 	return resp, nil
 }
 
+type writeCounter struct {
+	written uint64
+	suffix string // contains " / SIZE MB" if size is known, otherwise empty
+}
+
+func newWriteCounter(total int64) *writeCounter {
+	if total > 0 {
+		return &writeCounter{
+			suffix: fmt.Sprintf(" / %s", humanize.Bytes(uint64(total))),
+		}
+	}
+	return &writeCounter{}
+}
+
+func (wc *writeCounter) Write(p []byte) (int, error) {
+	n := len(p)
+	wc.written += uint64(n)
+	wc.PrintProgress()
+	return n, nil
+}
+
+func (wc writeCounter) PrintProgress() {
+	// Clear the line by using a character return to go back to the start and remove
+	// the remaining characters by filling it with spaces
+	fmt.Printf("\r%s", strings.Repeat(" ", 35))
+
+	// Return again and print current status of download
+	// We use the humanize package to print the bytes in a meaningful way (e.g. 10 MB)
+	fmt.Printf("\rDownloading... %s%s", humanize.Bytes(wc.written), wc.suffix)
+}
+
 // Write the body of the given HTTP response to disk at the given path
-func writeResonseToDisk(resp *http.Response, destPath string) *FetchError {
+func writeResonseToDisk(resp *http.Response, destPath string, withProgress bool) *FetchError {
 	out, err := os.Create(destPath)
 	if err != nil {
 		return wrapError(err)
@@ -245,6 +278,12 @@ func writeResonseToDisk(resp *http.Response, destPath string) *FetchError {
 	defer out.Close()
 	defer resp.Body.Close()
 
-	_, err = io.Copy(out, resp.Body)
+	var readCloser io.Reader
+	if withProgress{
+		readCloser = io.TeeReader(resp.Body, newWriteCounter(resp.ContentLength))
+	} else {
+		readCloser = resp.Body
+	}
+	_, err = io.Copy(out, readCloser)
 	return wrapError(err)
 }

github_test.go

@@ -274,9 +274,11 @@ func TestDownloadGitHubPulicReleaseAsset(t *testing.T) {
 		repoToken string
 		tag       string
 		assetId   int
+		progress  bool
 	}{
-		{"https://github.com/gruntwork-io/fetch-test-private", token, "v0.0.2", 1872521},
-		{"https://github.com/gruntwork-io/fetch-test-public", "", "v0.0.2", 1872641},
+		{"https://github.com/gruntwork-io/fetch-test-private", token, "v0.0.2", 1872521, false},
+		{"https://github.com/gruntwork-io/fetch-test-public", "", "v0.0.2", 1872641, false},
+		{"https://github.com/gruntwork-io/fetch-test-public", "", "v0.0.2", 1872641, true},
 	}
 
 	for _, tc := range cases {
@@ -290,7 +292,7 @@ func TestDownloadGitHubPulicReleaseAsset(t *testing.T) {
 			t.Fatalf("Failed to create temp file due to error: %s", tmpErr.Error())
 		}
 
-		if err := DownloadReleaseAsset(repo, tc.assetId, tmpFile.Name()); err != nil {
+		if err := DownloadReleaseAsset(repo, tc.assetId, tmpFile.Name(), tc.progress); err != nil {
 			t.Fatalf("Failed to download asset %d to %s from GitHub URL %s due to error: %s", tc.assetId, tmpFile.Name(), tc.repoUrl, err.Error())
 		}
 

main.go

@@ -28,6 +28,7 @@ type FetchOptions struct {
 	ReleaseAssetChecksumAlgo string
 	LocalDownloadPath        string
 	GithubApiVersion         string
+	WithProgress             bool
 }
 
 type AssetDownloadResult struct {
@@ -45,6 +46,7 @@ const OPTION_RELEASE_ASSET = "release-asset"
 const OPTION_RELEASE_ASSET_CHECKSUM = "release-asset-checksum"
 const OPTION_RELEASE_ASSET_CHECKSUM_ALGO = "release-asset-checksum-algo"
 const OPTION_GITHUB_API_VERSION = "github-api-version"
+const OPTION_WITH_PROGRESS = "progress"
 
 const ENV_VAR_GITHUB_TOKEN = "GITHUB_OAUTH_TOKEN"
 
@@ -98,6 +100,10 @@ func main() {
 			Value: "v3",
 			Usage: "The api version of the GitHub instance. If left blank, v3 will be used.\n\tThis will only be used if the repo url is not a github.com url.",
 		},
+		cli.BoolFlag{
+			Name:  OPTION_WITH_PROGRESS,
+			Usage: "Display progress on file downloads, especially useful for large files",
+		},
 	}
 
 	app.Action = runFetchWrapper
@@ -174,7 +180,7 @@ func runFetch(c *cli.Context) error {
 	}
 
 	// Download the requested release assets
-	assetPaths, err := downloadReleaseAssets(options.ReleaseAsset, options.LocalDownloadPath, repo, desiredTag)
+	assetPaths, err := downloadReleaseAssets(options.ReleaseAsset, options.LocalDownloadPath, repo, desiredTag, options.WithProgress)
 	if err != nil {
 		return err
 	}
@@ -217,6 +223,7 @@ func parseOptions(c *cli.Context) (FetchOptions, error) {
         OPTION_RELEASE_ASSET_CHECKSUM,
         OPTION_RELEASE_ASSET_CHECKSUM_ALGO,
         OPTION_GITHUB_API_VERSION,
+        OPTION_WITH_PROGRESS,
     }
 
     for _, option := range optionsList {
@@ -252,6 +259,7 @@ func parseOptions(c *cli.Context) (FetchOptions, error) {
 		ReleaseAssetChecksumAlgo: c.String(OPTION_RELEASE_ASSET_CHECKSUM_ALGO),
 		LocalDownloadPath:        localDownloadPath,
 		GithubApiVersion:         c.String(OPTION_GITHUB_API_VERSION),
+		WithProgress:             c.IsSet(OPTION_WITH_PROGRESS),
 	}, nil
 }
 
@@ -331,7 +339,7 @@ func downloadSourcePaths(sourcePaths []string, destPath string, githubRepo GitHu
 // were downloaded. For those that succeeded, the path they were downloaded to will be passed back
 // along with the error.
 // Returns the paths where the release assets were downloaded.
-func downloadReleaseAssets(assetRegex string, destPath string, githubRepo GitHubRepo, tag string) ([]string, error) {
+func downloadReleaseAssets(assetRegex string, destPath string, githubRepo GitHubRepo, tag string, withProgress bool) ([]string, error) {
 	var err error
 	var assetPaths []string
 
@@ -363,7 +371,7 @@ func downloadReleaseAssets(assetRegex string, destPath string, githubRepo GitHub
 
 			assetPath := path.Join(destPath, asset.Name)
 			fmt.Printf("Downloading release asset %s to %s\n", asset.Name, assetPath)
-			if downloadErr := DownloadReleaseAsset(githubRepo, asset.Id, assetPath); downloadErr == nil {
+			if downloadErr := DownloadReleaseAsset(githubRepo, asset.Id, assetPath, withProgress); downloadErr == nil {
 				fmt.Printf("Downloaded %s\n", assetPath)
 				results <- AssetDownloadResult{assetPath, nil}
 			} else {

main_test.go

@@ -25,7 +25,7 @@ func TestDownloadReleaseAssets(t *testing.T) {
 		t.Fatalf("Failed to parse sample release asset GitHub URL into Fetch GitHubRepo struct: %s", err)
 	}
 
-	assetPaths, fetchErr := downloadReleaseAssets(SAMPLE_RELEASE_ASSET_REGEX, tmpDir, githubRepo, SAMPLE_RELEASE_ASSET_VERSION)
+	assetPaths, fetchErr := downloadReleaseAssets(SAMPLE_RELEASE_ASSET_REGEX, tmpDir, githubRepo, SAMPLE_RELEASE_ASSET_VERSION, false)
 	if fetchErr != nil {
 		t.Fatalf("Failed to download release asset: %s", fetchErr)
 	}
@@ -55,7 +55,7 @@ func TestInvalidReleaseAssetsRegex(t *testing.T) {
 		t.Fatalf("Failed to parse sample release asset GitHub URL into Fetch GitHubRepo struct: %s", err)
 	}
 
-	_, fetchErr := downloadReleaseAssets("*", tmpDir, githubRepo, SAMPLE_RELEASE_ASSET_VERSION)
+	_, fetchErr := downloadReleaseAssets("*", tmpDir, githubRepo, SAMPLE_RELEASE_ASSET_VERSION, false)
 	if fetchErr == nil {
 		t.Fatalf("Expected error for invalid regex")
 	}
@@ -97,6 +97,9 @@ func TestEmptyOptionValues(t *testing.T) {
 			Name:  OPTION_GITHUB_API_VERSION,
 			Value: "v3",
 		},
+		cli.StringFlag{
+			Name: OPTION_WITH_PROGRESS,
+		},
 	}
 
     app.Action = func(c *cli.Context) error {
@@ -119,6 +122,7 @@ func TestEmptyOptionValues(t *testing.T) {
         OPTION_RELEASE_ASSET_CHECKSUM,
         OPTION_RELEASE_ASSET_CHECKSUM_ALGO,
         OPTION_GITHUB_API_VERSION,
+        OPTION_WITH_PROGRESS,
     }
 
     for _, option := range optionsList {