diff --git a/Gopkg.lock b/Gopkg.lock index 0c29a1b..11cdd71 100644 --- a/Gopkg.lock +++ b/Gopkg.lock @@ -9,6 +9,14 @@ revision = "8991bc29aa16c548c550c7ff78260e27b9ab7c73" version = "v1.1.1" +[[projects]] + digest = "1:f1a75a8e00244e5ea77ff274baa9559eb877437b240ee7b278f3fc560d9f08bf" + name = "github.com/dustin/go-humanize" + packages = ["."] + pruneopts = "" + revision = "9f541cc9db5d55bce703bd99987c9d5cb8eea45e" + version = "v1.0.0" + [[projects]] digest = "1:7398a6dcdc9283fdbd3ee279c7b80fa81e7ad9d6f9e71bc8e6394810870cc54d" name = "github.com/hashicorp/go-version" @@ -44,6 +52,7 @@ analyzer-name = "dep" analyzer-version = 1 input-imports = [ + "github.com/dustin/go-humanize", "github.com/hashicorp/go-version", "github.com/stretchr/testify/assert", "gopkg.in/urfave/cli.v1", diff --git a/Gopkg.toml b/Gopkg.toml index fc8bd92..79fd464 100644 --- a/Gopkg.toml +++ b/Gopkg.toml @@ -26,3 +26,7 @@ [[constraint]] name = "github.com/hashicorp/go-version" + +[[constraint]] + name = "github.com/dustin/go-humanize" + version = "1.0.0" diff --git a/README.md b/README.md index cf07986..9854087 100644 --- a/README.md +++ b/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: diff --git a/checksum_test.go b/checksum_test.go index 48fd0ed..82627ad 100644 --- a/checksum_test.go +++ b/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) } diff --git a/github.go b/github.go index 104e1fa..982f39f 100644 --- a/github.go +++ b/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) } diff --git a/github_test.go b/github_test.go index 61e208a..9b5f36f 100644 --- a/github_test.go +++ b/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()) } diff --git a/main.go b/main.go index de7955b..38420ca 100644 --- a/main.go +++ b/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 { diff --git a/main_test.go b/main_test.go index 5b35205..3c07e8e 100644 --- a/main_test.go +++ b/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 {