From b177afc3263f5f33f2182f58c6156b4d77115a0b Mon Sep 17 00:00:00 2001
From: tnextday <fw2k4@163.com>
Date: Tue, 15 Dec 2015 00:14:02 +0800
Subject: [PATCH] `weed download` command use stream download the large file.

---
 go/operation/chunked_file.go |   4 --
 go/util/http_util.go         |   5 +-
 go/weed/download.go          | 100 ++++++++++++++++++++++-------------
 3 files changed, 66 insertions(+), 43 deletions(-)

diff --git a/go/operation/chunked_file.go b/go/operation/chunked_file.go
index 5d7f99176..70564cbd2 100644
--- a/go/operation/chunked_file.go
+++ b/go/operation/chunked_file.go
@@ -83,10 +83,6 @@ func (cm *ChunkManifest) DeleteChunks(master string) error {
 	return nil
 }
 
-//func (cm *ChunkManifest) StoredHelper() error {
-//	return nil
-//}
-
 func readChunkNeedle(fileUrl string, w io.Writer, offset int64) (written int64, e error) {
 	req, err := http.NewRequest("GET", fileUrl, nil)
 	if err != nil {
diff --git a/go/util/http_util.go b/go/util/http_util.go
index d56aaa39a..f80ab0c24 100644
--- a/go/util/http_util.go
+++ b/go/util/http_util.go
@@ -136,12 +136,11 @@ func GetUrlStream(url string, values url.Values, readFn func(io.Reader) error) e
 	return readFn(r.Body)
 }
 
-func DownloadUrl(fileUrl string) (filename string, content []byte, e error) {
+func DownloadUrl(fileUrl string) (filename string, rc io.ReadCloser, e error) {
 	response, err := client.Get(fileUrl)
 	if err != nil {
 		return "", nil, err
 	}
-	defer response.Body.Close()
 	contentDisposition := response.Header["Content-Disposition"]
 	if len(contentDisposition) > 0 {
 		if strings.HasPrefix(contentDisposition[0], "filename=") {
@@ -149,7 +148,7 @@ func DownloadUrl(fileUrl string) (filename string, content []byte, e error) {
 			filename = strings.Trim(filename, "\"")
 		}
 	}
-	content, e = ioutil.ReadAll(response.Body)
+	rc = response.Body
 	return
 }
 
diff --git a/go/weed/download.go b/go/weed/download.go
index 3c55b3a34..392b65edb 100644
--- a/go/weed/download.go
+++ b/go/weed/download.go
@@ -3,9 +3,11 @@ package main
 import (
 	"fmt"
 	"io"
-	"io/ioutil"
 	"os"
 	"path"
+
+	"io/ioutil"
+
 	"strings"
 
 	"github.com/chrislusf/seaweedfs/go/operation"
@@ -43,50 +45,76 @@ var cmdDownload = &Command{
 
 func runDownload(cmd *Command, args []string) bool {
 	for _, fid := range args {
-		filename, content, e := fetchFileId(*d.server, fid)
-		if e != nil {
-			fmt.Println("Fetch Error:", e)
-			continue
-		}
-		if filename == "" {
-			filename = fid
-		}
-		if strings.HasSuffix(filename, "-list") {
-			filename = filename[0 : len(filename)-len("-list")]
-			fids := strings.Split(string(content), "\n")
-			f, err := os.OpenFile(path.Join(*d.dir, filename), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, os.ModePerm)
-			if err != nil {
-				fmt.Println("File Creation Error:", e)
-				continue
-			}
-			defer f.Close()
-			for _, partId := range fids {
-				var n int
-				_, part, err := fetchFileId(*d.server, partId)
-				if err == nil {
-					n, err = f.Write(part)
-				}
-				if err == nil && n < len(part) {
-					err = io.ErrShortWrite
-				}
-				if err != nil {
-					fmt.Println("File Write Error:", err)
-					break
-				}
-			}
-		} else {
-			ioutil.WriteFile(path.Join(*d.dir, filename), content, os.ModePerm)
+		if e := downloadToFile(*d.server, fid, *d.dir); e != nil {
+			fmt.Println("Download Error:", e)
 		}
 	}
 	return true
 }
 
-func fetchFileId(server string, fileId string) (filename string, content []byte, e error) {
+func downloadToFile(server, fileId, saveDir string) error {
+	fileUrl, lookupError := operation.LookupFileId(server, fileId)
+	if lookupError != nil {
+		return lookupError
+	}
+	filename, rc, err := util.DownloadUrl(fileUrl)
+	if err != nil {
+		return err
+	}
+	defer rc.Close()
+	if filename == "" {
+		filename = fileId
+	}
+	isFileList := false
+	if strings.HasSuffix(filename, "-list") {
+		// old command compatible
+		isFileList = true
+		filename = filename[0 : len(filename)-len("-list")]
+	}
+	f, err := os.OpenFile(path.Join(saveDir, filename), os.O_WRONLY|os.O_CREATE|os.O_TRUNC, os.ModePerm)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	if isFileList {
+		content, err := ioutil.ReadAll(rc)
+		if err != nil {
+			return err
+		}
+		fids := strings.Split(string(content), "\n")
+		for _, partId := range fids {
+			var n int
+			_, part, err := fetchContent(*d.server, partId)
+			if err == nil {
+				n, err = f.Write(part)
+			}
+			if err == nil && n < len(part) {
+				err = io.ErrShortWrite
+			}
+			if err != nil {
+				return err
+			}
+		}
+	} else {
+		if _, err = io.Copy(f, rc); err != nil {
+			return err
+		}
+
+	}
+	return nil
+}
+
+func fetchContent(server string, fileId string) (filename string, content []byte, e error) {
 	fileUrl, lookupError := operation.LookupFileId(server, fileId)
 	if lookupError != nil {
 		return "", nil, lookupError
 	}
-	filename, content, e = util.DownloadUrl(fileUrl)
+	var rc io.ReadCloser
+	if filename, rc, e = util.DownloadUrl(fileUrl); e != nil {
+		return "", nil, e
+	}
+	content, e = ioutil.ReadAll(rc)
+	rc.Close()
 	return
 }