go/gcs/storage.go

package gcs

import (
	"fmt"
	"io"
	"io/ioutil"

	"cloud.google.com/go/storage"
	"go.skia.org/infra/go/util"
	"golang.org/x/net/context"
	"google.golang.org/api/iterator"
)

// GCSClient is an interface for interacting with Google Cloud Storage (GCS). Introducing
// the interface allows for easier mocking and testing for unit (small) tests. GCSClient
// should have common, general functionality. Users should feel free to create a
// instance-specific GCSClient that creates an abstraction for more instance-specific method calls
// (see fuzzer for an example).
// One intentional thing missing from these method calls is bucket name. The bucket name
// is given at creation time, so as to simplify the method signatures.
// In all methods, context.Background() is a safe value for ctx if you don't want to use
// the context of the web request, for example.
// See also mock_gcs_client.New() for mocking this for unit tests.
type GCSClient interface {
	// FileReader returns an io.ReadCloser pointing to path on GCS, using the provided
	// context. storage.ErrObjectNotExist will be returned if the file is not found.
	// The caller must call Close on the returned Reader when done reading.
	FileReader(ctx context.Context, path string) (io.ReadCloser, error)
	// FileReader returns an io.WriteCloser that writes to the GCS file given by path
	// using the provided context. A new GCS file will be created if it doesn't already exist.
	// Otherwise, the existing file will be overwritten. The caller must call Close on
	// the returned Writer to flush the writes.
	FileWriter(ctx context.Context, path string, opts FileWriteOptions) io.WriteCloser
	// GetFileContents returns the []byte represented by the GCS file at path. This is a
	// convenience wrapper around FileReader. storage.ErrObjectNotExist will be returned
	// if the file is not found.
	GetFileContents(ctx context.Context, path string) ([]byte, error)
	// SetFileContents writes the []byte to the GCS file at path. This is a
	// convenience wrapper around FileWriter. The GCS file will be created if it doesn't exist.
	SetFileContents(ctx context.Context, path string, opts FileWriteOptions, contents []byte) error
	// AllFilesInDirectory executes the callback on all GCS files with the given prefix,
	// i.e. in the directory prefix. It returns an error if it fails to read any of the
	// ObjectAttrs belonging to files.
	AllFilesInDirectory(ctx context.Context, prefix string, callback func(item *storage.ObjectAttrs)) error
	// DeleteFile deletes the given file, returning any error.
	DeleteFile(ctx context.Context, path string) error
	// Bucket() returns the bucket name of this client
	Bucket() string
}

// FileWriteOptions represents the metadata for a GCS file.  See storage.ObjectAttrs
// for a more detailed description of what these are.
type FileWriteOptions struct {
	ContentEncoding    string
	ContentType        string
	ContentLanguage    string
	ContentDisposition string
	Metadata           map[string]string
}

var FILE_WRITE_OPTS_TEXT = FileWriteOptions{ContentEncoding: "text/plain"}

// gcsclient holds the information needed to talk to cloud storage.
type gcsclient struct {
	client *storage.Client
	bucket string
}

// NewGCSClient returns a GCSClient. See the interface for more information.
func NewGCSClient(s *storage.Client, bucket string) GCSClient {
	return &gcsclient{
		client: s,
		bucket: bucket,
	}
}

// See the GCSClient interface for more information about FileReader.
func (g *gcsclient) FileReader(ctx context.Context, path string) (io.ReadCloser, error) {
	return g.client.Bucket(g.bucket).Object(path).NewReader(ctx)
}

// See the GCSClient interface for more information about FileWriter.
func (g *gcsclient) FileWriter(ctx context.Context, path string, opts FileWriteOptions) io.WriteCloser {
	w := g.client.Bucket(g.bucket).Object(path).NewWriter(ctx)
	w.ObjectAttrs.ContentEncoding = opts.ContentEncoding
	w.ObjectAttrs.ContentType = opts.ContentType
	w.ObjectAttrs.ContentLanguage = opts.ContentLanguage
	w.ObjectAttrs.ContentDisposition = opts.ContentDisposition
	w.ObjectAttrs.Metadata = opts.Metadata

	return w
}

// See the GCSClient interface for more information about GetFileContents.
func (g *gcsclient) GetFileContents(ctx context.Context, path string) ([]byte, error) {
	response, err := g.client.Bucket(g.bucket).Object(path).NewReader(ctx)
	if err != nil {
		return nil, err
	}
	defer util.Close(response)
	return ioutil.ReadAll(response)
}

// See the GCSClient interface for more information about SetFileContents.
func (g *gcsclient) SetFileContents(ctx context.Context, path string, opts FileWriteOptions, contents []byte) error {
	w := g.FileWriter(ctx, path, opts)
	defer util.Close(w)
	if n, err := w.Write(contents); err != nil {
		return fmt.Errorf("There was a problem uploading %s.  Only uploaded %d bytes: %s", path, n, err)
	}
	return nil
}

// See the GCSClient interface for more information about AllFilesInDirectory.
func (g *gcsclient) AllFilesInDirectory(ctx context.Context, prefix string, callback func(item *storage.ObjectAttrs)) error {
	total := 0
	q := &storage.Query{Prefix: prefix, Versions: false}
	it := g.client.Bucket(g.bucket).Objects(ctx, q)
	for obj, err := it.Next(); err != iterator.Done; obj, err = it.Next() {
		if err != nil {
			return fmt.Errorf("Problem reading from Google Storage: %v", err)
		}
		total++
		callback(obj)
	}
	return nil
}

// See the GCSClient interface for more information about DeleteFile.
func (g *gcsclient) DeleteFile(ctx context.Context, path string) error {
	return g.client.Bucket(g.bucket).Object(path).Delete(ctx)
}

// See the GCSClient interface for more information about Bucket.
func (g *gcsclient) Bucket() string {
	return g.bucket
}