oss download with progress bar and range

download.go

package common

import (
	"io"
	"os"
	"strconv"

	pb "gopkg.in/cheggaaa/pb.v1"

	"github.com/sirupsen/logrus"

	"github.com/aliyun/aliyun-oss-go-sdk/oss"
)

const defaultOssEndpoint = "a"
const defaultOssAccessKeyID = "b"
const defaultOSSAccessKeySecret = "c"

// DefaultPartSize ...
//const DefaultPartSize = 1024 * 1024
const DefaultPartSize = 1

// DefaultOssClient ...
var DefaultOssClient *oss.Client

// OssDownload ...
func OssDownload(bucketName, objectKey, filePath string) (err error) {
	logrus.WithFields(logrus.Fields{
		"objectKey":  objectKey,
		"bucketName": bucketName,
		"filePath":   filePath,
	}).Debug("download oss file")
	bucket, err := DefaultOssClient.Bucket(bucketName)
	if err != nil {
		return
	}

	total, err := getObjectSize(bucket, objectKey)
	if err != nil {
		return
	}

	out, transmitted, err := openDownloadFile(filePath, total)
	if err != nil {
		return
	}
	defer out.Close()

	logrus.WithField("transmitted", transmitted).Debug("file stat")

	if transmitted == total {
		logrus.Warn("download already completed")
		return
	}

	body, err := bucket.GetObject(objectKey, oss.Range(transmitted, total-1))
	if err != nil {
		return
	}
	defer body.Close()

	bar := pb.New64(total).SetUnits(pb.U_BYTES).Set64(transmitted)
	bar.Start()

	writer := io.MultiWriter(out, bar)

	io.Copy(writer, body)
	bar.Finish()
	return
}

func openDownloadFile(filePath string, total int64) (file *os.File, transmitted int64, err error) {
	fi, err := os.Stat(filePath)

	if os.IsNotExist(err) {
		file, err = os.Create(filePath)
		return
	}

	if fi.Size() > total {
		file, err = os.Create(filePath)
		return
	}

	file, err = os.OpenFile(filePath, os.O_APPEND|os.O_WRONLY, 0666)
	if err != nil {
		return
	}

	fi, err = file.Stat()
	if err != nil {
		return
	}
	transmitted = fi.Size()
	return
}

func getObjectSize(bucket *oss.Bucket, objectKey string) (size int64, err error) {
	meta, err := bucket.GetObjectMeta(objectKey)
	if err != nil {
		return
	}
	size, err = strconv.ParseInt(meta.Get("Content-Length"), 10, 0)
	return
}

func getEnvWithDefault(name, defaults string) string {
	result := os.Getenv(name)
	if result != "" {
		return result
	}
	logrus.WithFields(logrus.Fields{
		"name":  name,
		"value": defaults,
	}).Warn("using default for env")
	return defaults
}