<commit_before>package s3resource

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
	"github.com/rlmcpherson/s3gof3r"
)

//go:generate counterfeiter . S3Client

type S3Client interface {
	BucketFiles(bucketName string, prefixHint string) ([]string, error)

	UploadFile(bucketName string, remotePath string, localPath string) error
	DownloadFile(bucketName string, remotePath string, localPath string) error

	URL(bucketName string, remotePath string, private bool) string
}

type s3client struct {
	client       *s3.S3
	gopherClient *s3gof3r.S3
}

func NewS3Client(accessKey string, secretKey string, regionName string) (S3Client, error) {
	auth := aws.Auth{
		AccessKey: accessKey,
		SecretKey: secretKey,
	}

	authGopher := s3gof3r.Keys{
		AccessKey: accessKey,
		SecretKey: secretKey,
	}

	region, ok := aws.Regions[regionName]
	if !ok {
		return nil, fmt.Errorf("No such region '%s'", regionName)
	}

	return &s3client{
		client:       s3.New(auth, region),
		gopherClient: s3gof3r.New("", authGopher),
	}, nil
}

func (client *s3client) BucketFiles(bucketName string, prefixHint string) ([]string, error) {
	bucket := client.client.Bucket(bucketName)
	entries, err := client.getBucketContents(bucket, prefixHint)
	if err != nil {
		return []string{}, err
	}

	paths := make([]string, 0, len(*entries))
	for entry := range *entries {
		paths = append(paths, entry)
	}

	return paths, nil
}

func (client *s3client) getBucketContents(bucket *s3.Bucket, prefix string) (*map[string]s3.Key, error) {
	bucketContents := map[string]s3.Key{}
	separator := ""
	marker := ""

	for {
		contents, err := bucket.List(prefix, separator, marker, 1000)
		if err != nil {
			return &bucketContents, err
		}

		lastKey := ""
		for _, key := range contents.Contents {
			bucketContents[key.Key] = key
			lastKey = key.Key
		}

		if contents.IsTruncated {
			marker = contents.NextMarker
			if marker == "" {
				// From the s3 docs: If response does not include the
				// NextMarker and it is truncated, you can use the value of the
				// last Key in the response as the marker in the subsequent
				// request to get the next set of object keys.
				marker = lastKey
			}
		} else {
			break
		}
	}

	return &bucketContents, nil
}

func (client *s3client) UploadFile(bucketName string, remotePath string, localPath string) error {
	bucket := client.gopherClient.Bucket(bucketName)

	localFile, err := os.Open(localPath)
	if err != nil {
		return err
	}

	remoteFile, err := bucket.PutWriter(remotePath, nil, nil)
	if err != nil {
		return err
	}

	if _, err = io.Copy(remoteFile, localFile); err != nil {
		return err
	}

	if err = remoteFile.Close(); err != nil {
		return err
	}

	if err = localFile.Close(); err != nil {
		return err
	}

	return nil
}

func (client *s3client) DownloadFile(bucketName string, remotePath string, localPath string) error {
	bucket := client.client.Bucket(bucketName)

	remoteFile, err := bucket.GetReader(remotePath)
	if err != nil {
		return err
	}
	defer remoteFile.Close()

	localFile, err := os.Create(localPath)
	if err != nil {
		return err
	}
	defer localFile.Close()

	_, err = io.Copy(localFile, remoteFile)
	if err != nil {
		return err
	}

	return nil
}

func (client *s3client) URL(bucketName string, remotePath string, private bool) string {
	bucket := client.client.Bucket(bucketName)

	if private {<commit_msg>use s3gof3r for downloading<commit_after>package s3resource

import (
	"fmt"
	"io"
	"os"
	"time"

	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
	"github.com/rlmcpherson/s3gof3r"
)

//go:generate counterfeiter . S3Client

type S3Client interface {
	BucketFiles(bucketName string, prefixHint string) ([]string, error)

	UploadFile(bucketName string, remotePath string, localPath string) error
	DownloadFile(bucketName string, remotePath string, localPath string) error

	URL(bucketName string, remotePath string, private bool) string
}

type s3client struct {
	client       *s3.S3
	gopherClient *s3gof3r.S3
}

func NewS3Client(accessKey string, secretKey string, regionName string) (S3Client, error) {
	auth := aws.Auth{
		AccessKey: accessKey,
		SecretKey: secretKey,
	}

	authGopher := s3gof3r.Keys{
		AccessKey: accessKey,
		SecretKey: secretKey,
	}

	region, ok := aws.Regions[regionName]
	if !ok {
		return nil, fmt.Errorf("No such region '%s'", regionName)
	}

	return &s3client{
		client:       s3.New(auth, region),
		gopherClient: s3gof3r.New("", authGopher),
	}, nil
}

func (client *s3client) BucketFiles(bucketName string, prefixHint string) ([]string, error) {
	bucket := client.client.Bucket(bucketName)
	entries, err := client.getBucketContents(bucket, prefixHint)
	if err != nil {
		return []string{}, err
	}

	paths := make([]string, 0, len(*entries))
	for entry := range *entries {
		paths = append(paths, entry)
	}

	return paths, nil
}

func (client *s3client) getBucketContents(bucket *s3.Bucket, prefix string) (*map[string]s3.Key, error) {
	bucketContents := map[string]s3.Key{}
	separator := ""
	marker := ""

	for {
		contents, err := bucket.List(prefix, separator, marker, 1000)
		if err != nil {
			return &bucketContents, err
		}

		lastKey := ""
		for _, key := range contents.Contents {
			bucketContents[key.Key] = key
			lastKey = key.Key
		}

		if contents.IsTruncated {
			marker = contents.NextMarker
			if marker == "" {
				// From the s3 docs: If response does not include the
				// NextMarker and it is truncated, you can use the value of the
				// last Key in the response as the marker in the subsequent
				// request to get the next set of object keys.
				marker = lastKey
			}
		} else {
			break
		}
	}

	return &bucketContents, nil
}

func (client *s3client) UploadFile(bucketName string, remotePath string, localPath string) error {
	bucket := client.gopherClient.Bucket(bucketName)

	localFile, err := os.Open(localPath)
	if err != nil {
		return err
	}

	remoteFile, err := bucket.PutWriter(remotePath, nil, nil)
	if err != nil {
		return err
	}

	if _, err = io.Copy(remoteFile, localFile); err != nil {
		return err
	}

	if err = remoteFile.Close(); err != nil {
		return err
	}

	if err = localFile.Close(); err != nil {
		return err
	}

	return nil
}

func (client *s3client) DownloadFile(bucketName string, remotePath string, localPath string) error {
	bucket := client.gopherClient.Bucket(bucketName)

	remoteFile, _, err := bucket.GetReader(remotePath, nil)
	if err != nil {
		return err
	}
	defer remoteFile.Close()

	localFile, err := os.Create(localPath)
	if err != nil {
		return err
	}
	defer localFile.Close()

	_, err = io.Copy(localFile, remoteFile)
	if err != nil {
		return err
	}

	return nil
}

func (client *s3client) URL(bucketName string, remotePath string, private bool) string {
	bucket := client.client.Bucket(bucketName)

	if private {
{'fraction_non_alphanumeric': 0.11907548273844354, 'fraction_numerical': 0.01023990637799883, 'mean_word_length': 3.89756446991404, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4755149', 'n_tokens_mistral': 2481, 'n_tokens_neox': 2313, 'n_words': 658}