Mario Finelli Blog

golang: upload to amazon s3 with progress bar

3 July 2020

For a side-project that I'm working on I wanted to implement uploading to AWS S3 with a progress bar. There's an offical example provided by the SDK but it is not a progress bar it just dumps the progress to the terminal ina very noisy way. It also has a problem that the progress starts at 50%. Fortunately, someone has fixed this issue (though it has not been merged to master yet). So I combined the above two code samples along with the mpb library to make it all work.

package main

import (
        "fmt"
        "github.com/aws/aws-sdk-go/aws"
        "github.com/aws/aws-sdk-go/aws/session"
        "github.com/aws/aws-sdk-go/service/s3/s3manager"
        "github.com/vbauerster/mpb/v5"
        "github.com/vbauerster/mpb/v5/decor"
        "os"
        "sync"
)

type CustomReader struct {
        fp      *os.File
        size    int64
        read    int64
        bar     *mpb.Bar
        signMap map[int64]struct{}
        mux     sync.Mutex
}

func (r *CustomReader) Read(p []byte) (int, error) {
        return r.fp.Read(p)
}

func (r *CustomReader) ReadAt(p []byte, off int64) (int, error) {
        n, err := r.fp.ReadAt(p, off)
        if err != nil {
                return n, err
        }

        r.bar.SetTotal(r.size, false)

        r.mux.Lock()
        // Ignore the first signature call
        if _, ok := r.signMap[off]; ok {
                r.read += int64(n)
                r.bar.SetCurrent(r.read)
        } else {
                r.signMap[off] = struct{}{}
        }
        r.mux.Unlock()

        return n, err
}

func (r *CustomReader) Seek(offset int64, whence int) (int64, error) {
        return r.fp.Seek(offset, whence)
}

func main() {
        sess, err := session.NewSession(&aws.Config{
                Region: aws.String("us-east-1"),
        })

        if err != nil {
                fmt.Println(err)
                os.Exit(1)
        }

        file, err := os.Open(os.Args[1:][0])

        if err != nil {
                fmt.Println(err)
                os.Exit(1)
        }

        fileInfo, err := file.Stat()

        if err != nil {
                fmt.Println(err)
                os.Exit(1)
        }

        p := mpb.New()

        reader := &CustomReader{
                fp:      file,
                size:    fileInfo.Size(),
                signMap: map[int64]struct{}{},
                bar: p.AddBar(fileInfo.Size(),
                        mpb.PrependDecorators(
                                decor.Name("uploading..."),
                                decor.Percentage(decor.WCSyncSpace),
                        ),
                ),
        }

        uploader := s3manager.NewUploader(sess, func(u *s3manager.Uploader) {
        })

        result, err := uploader.Upload(&s3manager.UploadInput{
                Bucket: aws.String("your-bucket-name"),
                Key:    aws.String("test-object"),
                Body:   reader,
        })

        if err != nil {
                fmt.Println(err)
                os.Exit(1)
        }

        fmt.Println(result)
}

The above program when compiled (go build -o main main.go) will upload the argument that you pass to the bucket your-bucket-name as the object test-object with a progress bar so that you can see how far along it is.