package gcs
import (
"context"
"fmt"
"io"
"runtime"
"strings"
"testing"
"cloud.google.com/go/storage"
"github.com/go-kit/log"
"github.com/pkg/errors"
"github.com/prometheus/common/version"
"golang.org/x/oauth2/google"
"google.golang.org/api/iterator"
"google.golang.org/api/option"
"gopkg.in/yaml.v2"
"github.com/thanos-io/objstore"
)
const DirDelim = "/"
type Config struct {
Bucket string `yaml:"bucket"`
ServiceAccount string `yaml:"service_account"`
}
type Bucket struct {
logger log.Logger
bkt *storage.BucketHandle
name string
closer io.Closer
}
func NewBucket(ctx context.Context, logger log.Logger, conf []byte, component string) (*Bucket, error) {
var gc Config
if err := yaml.Unmarshal(conf, &gc); err != nil {
return nil, err
}
return NewBucketWithConfig(ctx, logger, gc, component)
}
func NewBucketWithConfig(ctx context.Context, logger log.Logger, gc Config, component string) (*Bucket, error) {
if gc.Bucket == "" {
return nil, errors.New("missing Google Cloud Storage bucket name for stored blocks")
}
var opts []option.ClientOption
if gc.ServiceAccount != "" {
credentials, err := google.CredentialsFromJSON(ctx, []byte(gc.ServiceAccount), storage.ScopeFullControl)
if err != nil {
return nil, errors.Wrap(err, "failed to create credentials from JSON")
}
opts = append(opts, option.WithCredentials(credentials))
}
opts = append(opts,
option.WithUserAgent(fmt.Sprintf("thanos-%s/%s (%s)", component, version.Version, runtime.Version())),
)
gcsClient, err := storage.NewClient(ctx, opts...)
if err != nil {
return nil, err
}
bkt := &Bucket{
logger: logger,
bkt: gcsClient.Bucket(gc.Bucket),
closer: gcsClient,
name: gc.Bucket,
}
return bkt, nil
}
func (b *Bucket) Name() string {
return b.name
}
func (b *Bucket) Iter(ctx context.Context, dir string, f func(string) error, options ...objstore.IterOption) error {
if dir != "" {
dir = strings.TrimSuffix(dir, DirDelim) + DirDelim
}
delimiter := DirDelim
if objstore.ApplyIterOptions(options...).Recursive {
delimiter = ""
}
it := b.bkt.Objects(ctx, &storage.Query{
Prefix: dir,
Delimiter: delimiter,
})
for {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
attrs, err := it.Next()
if err == iterator.Done {
return nil
}
if err != nil {
return err
}
if err := f(attrs.Prefix + attrs.Name); err != nil {
return err
}
}
}
func (b *Bucket) Get(ctx context.Context, name string) (io.ReadCloser, error) {
return b.bkt.Object(name).NewReader(ctx)
}
func (b *Bucket) GetRange(ctx context.Context, name string, off, length int64) (io.ReadCloser, error) {
return b.bkt.Object(name).NewRangeReader(ctx, off, length)
}
func (b *Bucket) Attributes(ctx context.Context, name string) (objstore.ObjectAttributes, error) {
attrs, err := b.bkt.Object(name).Attrs(ctx)
if err != nil {
return objstore.ObjectAttributes{}, err
}
return objstore.ObjectAttributes{
Size: attrs.Size,
LastModified: attrs.Updated,
}, nil
}
func (b *Bucket) Handle() *storage.BucketHandle {
return b.bkt
}
func (b *Bucket) Exists(ctx context.Context, name string) (bool, error) {
if _, err := b.bkt.Object(name).Attrs(ctx); err == nil {
return true, nil
} else if err != storage.ErrObjectNotExist {
return false, err
}
return false, nil
}
func (b *Bucket) Upload(ctx context.Context, name string, r io.Reader) error {
w := b.bkt.Object(name).NewWriter(ctx)
if _, err := io.Copy(w, r); err != nil {
return err
}
return w.Close()
}
func (b *Bucket) Delete(ctx context.Context, name string) error {
return b.bkt.Object(name).Delete(ctx)
}
func (b *Bucket) IsObjNotFoundErr(err error) bool {
return errors.Is(err, storage.ErrObjectNotExist)
}
func (b *Bucket) Close() error {
return b.closer.Close()
}
func NewTestBucket(t testing.TB, project string) (objstore.Bucket, func(), error) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
gTestConfig := Config{
Bucket: objstore.CreateTemporaryTestBucketName(t),
}
bc, err := yaml.Marshal(gTestConfig)
if err != nil {
return nil, nil, err
}
b, err := NewBucket(ctx, log.NewNopLogger(), bc, "thanos-e2e-test")
if err != nil {
return nil, nil, err
}
if err = b.bkt.Create(ctx, project, nil); err != nil {
_ = b.Close()
return nil, nil, err
}
t.Log("created temporary GCS bucket for GCS tests with name", b.name, "in project", project)
return b, func() {
objstore.EmptyBucket(t, ctx, b)
if err := b.bkt.Delete(ctx); err != nil {
t.Logf("deleting bucket failed: %s", err)
}
if err := b.Close(); err != nil {
t.Logf("closing bucket failed: %s", err)
}
}, nil
}