github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/cmd/bootstrap/gcs/gcs.go (about)

     1  package gcs
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"path/filepath"
     9  
    10  	"cloud.google.com/go/storage"
    11  	"google.golang.org/api/iterator"
    12  	"google.golang.org/api/option"
    13  )
    14  
    15  // GoogleBucket ...
    16  type googleBucket struct {
    17  	Name string
    18  }
    19  
    20  // NewGoogleBucket ...
    21  func NewGoogleBucket(bucketName string) *googleBucket {
    22  	return &googleBucket{
    23  		Name: bucketName,
    24  	}
    25  }
    26  
    27  // NewClient ...
    28  func (g *googleBucket) NewClient(ctx context.Context) (*storage.Client, error) {
    29  	client, err := storage.NewClient(ctx, option.WithoutAuthentication())
    30  	if err != nil {
    31  		return nil, err
    32  	}
    33  	return client, nil
    34  }
    35  
    36  // GetFiles returns a list of file names within the Google bucket
    37  func (g *googleBucket) GetFiles(ctx context.Context, client *storage.Client, prefix, delimiter string) ([]string, error) {
    38  	it := client.Bucket(g.Name).Objects(ctx, &storage.Query{
    39  		Prefix:    prefix,
    40  		Delimiter: delimiter,
    41  	})
    42  
    43  	var files []string
    44  	for {
    45  		attrs, err := it.Next()
    46  		if err == iterator.Done {
    47  			break
    48  		}
    49  		if err != nil {
    50  			return nil, err
    51  		}
    52  
    53  		files = append(files, attrs.Name)
    54  	}
    55  
    56  	return files, nil
    57  }
    58  
    59  // DownloadFile downloads a file from the bucket to a desination folder
    60  func (g *googleBucket) DownloadFile(ctx context.Context, client *storage.Client, destination, source string) error {
    61  
    62  	// create dir of destination
    63  	dir := filepath.Dir(destination)
    64  	err := os.MkdirAll(dir, os.ModePerm)
    65  	if err != nil {
    66  		return fmt.Errorf("error creating destination directory: %w", err)
    67  	}
    68  
    69  	download, err := client.Bucket(g.Name).Object(source).NewReader(ctx)
    70  	if err != nil {
    71  		return fmt.Errorf("error creating GCS object reader: %w", err)
    72  	}
    73  	defer download.Close()
    74  
    75  	file, err := os.Create(destination)
    76  	if err != nil {
    77  		return fmt.Errorf("error creating download file: %w", err)
    78  	}
    79  	defer file.Close()
    80  
    81  	_, err = io.Copy(file, download)
    82  	if err != nil {
    83  		return fmt.Errorf("error downloading file: %w", err)
    84  	}
    85  	return nil
    86  }
    87  
    88  // UploadFile uploads a file to the google bucket
    89  func (g *googleBucket) UploadFile(ctx context.Context, client *storage.Client, destination, source string) error {
    90  
    91  	upload := client.Bucket(g.Name).Object(destination).NewWriter(ctx)
    92  	defer upload.Close()
    93  
    94  	file, err := os.Open(source)
    95  	if err != nil {
    96  		return fmt.Errorf("Error opening upload file: %w", err)
    97  	}
    98  	defer file.Close()
    99  
   100  	_, err = io.Copy(upload, file)
   101  	if err != nil {
   102  		return fmt.Errorf("Error uploading file: %w", err)
   103  	}
   104  
   105  	return nil
   106  }