github.com/stellar/stellar-etl@v1.0.1-0.20240312145900-4874b6bf2b89/cmd/upload_to_gcs.go (about)

     1  package cmd
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"os"
     8  	"time"
     9  
    10  	"cloud.google.com/go/storage"
    11  )
    12  
    13  type GCS struct {
    14  	gcsCredentialsPath string
    15  	gcsBucket          string
    16  }
    17  
    18  func newGCS(gcsCredentialsPath, gcsBucket string) CloudStorage {
    19  	return &GCS{
    20  		gcsCredentialsPath: gcsCredentialsPath,
    21  		gcsBucket:          gcsBucket,
    22  	}
    23  }
    24  
    25  func (g *GCS) UploadTo(credentialsPath, bucket, path string) error {
    26  	// Use credentials file in dev/local runs. Otherwise, derive credentials from the service account.
    27  	if len(credentialsPath) > 0 {
    28  		os.Setenv("GOOGLE_APPLICATION_CREDENTIALS", credentialsPath)
    29  		cmdLogger.Infof("Using credentials found at: %s", credentialsPath)
    30  	}
    31  
    32  	reader, err := os.Open(path)
    33  	if err != nil {
    34  		return fmt.Errorf("failed to open file %s: %v", path, err)
    35  	}
    36  
    37  	ctx := context.Background()
    38  	client, err := storage.NewClient(ctx)
    39  	if err != nil {
    40  		return fmt.Errorf("failed to create client: %v", err)
    41  	}
    42  	defer client.Close()
    43  
    44  	ctx, cancel := context.WithTimeout(ctx, time.Hour)
    45  	defer cancel()
    46  
    47  	wc := client.Bucket(bucket).Object(path).NewWriter(ctx)
    48  
    49  	uploadLocation := fmt.Sprintf("gs://%s/%s", bucket, path)
    50  	cmdLogger.Infof("Uploading %s to %s", path, uploadLocation)
    51  
    52  	var written int64
    53  	if written, err = io.Copy(wc, reader); err != nil {
    54  		return fmt.Errorf("unable to copy: %v", err)
    55  	}
    56  	err = wc.Close()
    57  	if err != nil {
    58  		return err
    59  	}
    60  
    61  	cmdLogger.Infof("Successfully uploaded %d bytes to gs://%s/%s", written, bucket, path)
    62  
    63  	deleteLocalFiles(path)
    64  
    65  	return nil
    66  }