nmvuong92 · September 26, 2024 12:09 · Sep 18, 2024 · Sep 8, 2023 · Sep 7, 2023
diff --git a/golang s3 download with progressbar.go b/golang s3 download with progressbar.go
@@ -1,6 +1,7 @@
 package awss3
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"io"
@@ -9,91 +10,112 @@ import (
 	"path/filepath"
 	"time"
 
-  "github.com/schollz/progressbar/v3"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/aws/aws-sdk-go/service/s3/s3manager"
+	"github.com/schollz/progressbar/v3"
 )
 
 type S3Bucket struct {
 	BucketName string
 	PathPrefix string
 	Region     string
+	session    *session.Session
 }
-func (s *S3Bucket) Download(s3filename string, destPath string) (err error) {
+
+func NewS3Bucket(bucketName, pathPrefix, region string) (*S3Bucket, error) {
 	sess, err := session.NewSession(&aws.Config{
-		Region: aws.String(s.Region),
+		Region: aws.String(region),
 	})
-	if utils.CheckFileExist(destPath) {
-		_ = os.Remove(destPath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create session: %w", err)
+	}
+
+	return &S3Bucket{
+		BucketName: bucketName,
+		PathPrefix: pathPrefix,
+		Region:     region,
+		session:    sess,
+	}, nil
+}
+
+func (s *S3Bucket) Download(ctx context.Context, s3filename, destPath string) error {
+	if err := os.RemoveAll(destPath); err != nil {
+		return fmt.Errorf("failed to remove existing file: %w", err)
 	}
+
 	destFile, err := os.Create(destPath)
 	if err != nil {
-		fmt.Println(err)
+		return fmt.Errorf("failed to create destination file: %w", err)
 	}
 	defer destFile.Close()
 
-	s3object, err := s.GetObject(s.PathPrefix + "/" + s3filename)
+	s3object, err := s.GetObject(ctx, filepath.Join(s.PathPrefix, s3filename))
+	if err != nil {
+		return fmt.Errorf("failed to get S3 object: %w", err)
+	}
+
 	bar := progressbar.DefaultBytes(
 		*s3object.ContentLength,
 		"[info] downloading s3: "+s3filename,
 	)
-	progressbarTicker := time.NewTicker(200 * time.Millisecond)
-	quitProgressBarChannel := make(chan interface{})
-	go func() {
-		for {
-			select {
-			case <-progressbarTicker.C:
-				localObject, err2 := os.Stat(destPath)
-				if err2 != nil {
-					return
-				}
-				_ = bar.Set64(localObject.Size())
-			case <-quitProgressBarChannel:
-				progressbarTicker.Stop()
-				return
-			}
-		}
-	}()
 
-	downloader := s3manager.NewDownloader(sess, func(d *s3manager.Downloader) {
+	downloader := s3manager.NewDownloader(s.session, func(d *s3manager.Downloader) {
 		d.PartSize = 64 * 1024 * 1024
 		d.Concurrency = 6
 	})
 
-	_, err = downloader.Download(destFile,
+	_, err = downloader.DownloadWithContext(ctx, destFile,
 		&s3.GetObjectInput{
 			Bucket: aws.String(s.BucketName),
-			Key:    aws.String(s.PathPrefix + "/" + s3filename),
+			Key:    aws.String(filepath.Join(s.PathPrefix, s3filename)),
 		})
-
-	defer func() {
-		time.Sleep(200 * time.Millisecond)
-		quitProgressBarChannel <- true
-	}()
-
 	if err != nil {
-		fmt.Println(err)
+		return fmt.Errorf("failed to download file: %w", err)
 	}
 
-	return
+	go s.updateProgressBar(ctx, bar, destPath)
+
+	return nil
 }
 
-func (s *S3Bucket) GetObject(objectKey string) (result *s3.GetObjectOutput, err error) {
-	client := &http.Client{Transport: &http.Transport{
-		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
-	}}
-	sess, err := session.NewSession(&aws.Config{
-		Region: aws.String(s.Region),
-	})
-	svc := s3.New(sess, &aws.Config{
+func (s *S3Bucket) updateProgressBar(ctx context.Context, bar *progressbar.ProgressBar, filePath string) {
+	ticker := time.NewTicker(200 * time.Millisecond)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			info, err := os.Stat(filePath)
+			if err != nil {
+				continue
+			}
+			_ = bar.Set64(info.Size())
+		}
+	}
+}
+
+func (s *S3Bucket) GetObject(ctx context.Context, objectKey string) (*s3.GetObjectOutput, error) {
+	client := &http.Client{
+		Transport: &http.Transport{
+			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+		},
+	}
+
+	svc := s3.New(s.session, &aws.Config{
 		HTTPClient: client,
 	})
-	result, err = svc.GetObject(&s3.GetObjectInput{
+
+	result, err := svc.GetObjectWithContext(ctx, &s3.GetObjectInput{
 		Bucket: aws.String(s.BucketName),
 		Key:    aws.String(objectKey),
 	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to get object: %w", err)
+	}
 
-	return
-}
+	return result, nil
+}
diff --git a/golang s3 download with progressbar.go b/golang s3 download with progressbar.go
@@ -39,7 +39,7 @@ func (s *S3Bucket) Download(s3filename string, destPath string) (err error) {
 		*s3object.ContentLength,
 		"[info] downloading s3: "+s3filename,
 	)
-	progressbarTicker := time.NewTicker(1 * time.Second)
+	progressbarTicker := time.NewTicker(200 * time.Millisecond)
 	quitProgressBarChannel := make(chan interface{})
 	go func() {
 		for {
@@ -69,6 +69,7 @@ func (s *S3Bucket) Download(s3filename string, destPath string) (err error) {
 		})
 
 	defer func() {
+		time.Sleep(200 * time.Millisecond)
 		quitProgressBarChannel <- true
 	}()
 

diff --git a/golang s3 download with progressbar.go b/golang s3 download with progressbar.go
@@ -0,0 +1,98 @@
+package awss3
+
+import (
+	"crypto/tls"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"time"
+
+  "github.com/schollz/progressbar/v3"
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/s3"
+	"github.com/aws/aws-sdk-go/service/s3/s3manager"
+)
+
+type S3Bucket struct {
+	BucketName string
+	PathPrefix string
+	Region     string
+}
+func (s *S3Bucket) Download(s3filename string, destPath string) (err error) {
+	sess, err := session.NewSession(&aws.Config{
+		Region: aws.String(s.Region),
+	})
+	if utils.CheckFileExist(destPath) {
+		_ = os.Remove(destPath)
+	}
+	destFile, err := os.Create(destPath)
+	if err != nil {
+		fmt.Println(err)
+	}
+	defer destFile.Close()
+
+	s3object, err := s.GetObject(s.PathPrefix + "/" + s3filename)
+	bar := progressbar.DefaultBytes(
+		*s3object.ContentLength,
+		"[info] downloading s3: "+s3filename,
+	)
+	progressbarTicker := time.NewTicker(1 * time.Second)
+	quitProgressBarChannel := make(chan interface{})
+	go func() {
+		for {
+			select {
+			case <-progressbarTicker.C:
+				localObject, err2 := os.Stat(destPath)
+				if err2 != nil {
+					return
+				}
+				_ = bar.Set64(localObject.Size())
+			case <-quitProgressBarChannel:
+				progressbarTicker.Stop()
+				return
+			}
+		}
+	}()
+
+	downloader := s3manager.NewDownloader(sess, func(d *s3manager.Downloader) {
+		d.PartSize = 64 * 1024 * 1024
+		d.Concurrency = 6
+	})
+
+	_, err = downloader.Download(destFile,
+		&s3.GetObjectInput{
+			Bucket: aws.String(s.BucketName),
+			Key:    aws.String(s.PathPrefix + "/" + s3filename),
+		})
+
+	defer func() {
+		quitProgressBarChannel <- true
+	}()
+
+	if err != nil {
+		fmt.Println(err)
+	}
+
+	return
+}
+
+func (s *S3Bucket) GetObject(objectKey string) (result *s3.GetObjectOutput, err error) {
+	client := &http.Client{Transport: &http.Transport{
+		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
+	}}
+	sess, err := session.NewSession(&aws.Config{
+		Region: aws.String(s.Region),
+	})
+	svc := s3.New(sess, &aws.Config{
+		HTTPClient: client,
+	})
+	result, err = svc.GetObject(&s3.GetObjectInput{
+		Bucket: aws.String(s.BucketName),
+		Key:    aws.String(objectKey),
+	})
+
+	return
+}