github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/integration/e2ecortex/storage.go (about)

     1  package e2ecortex
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strings"
     7  
     8  	"github.com/go-kit/log"
     9  	"github.com/grafana/dskit/flagext"
    10  	"github.com/oklog/ulid"
    11  	"github.com/thanos-io/thanos/pkg/objstore"
    12  
    13  	"github.com/cortexproject/cortex/integration/e2e"
    14  	e2edb "github.com/cortexproject/cortex/integration/e2e/db"
    15  	"github.com/cortexproject/cortex/pkg/storage/bucket/s3"
    16  )
    17  
    18  type S3Client struct {
    19  	writer objstore.Bucket
    20  	reader objstore.BucketReader
    21  }
    22  
    23  func NewS3Client(cfg s3.Config) (*S3Client, error) {
    24  	writer, err := s3.NewBucketClient(cfg, "test", log.NewNopLogger())
    25  	if err != nil {
    26  		return nil, err
    27  	}
    28  
    29  	reader, err := s3.NewBucketReaderClient(cfg, "test", log.NewNopLogger())
    30  	if err != nil {
    31  		return nil, err
    32  	}
    33  
    34  	return &S3Client{
    35  		writer: writer,
    36  		reader: reader,
    37  	}, nil
    38  }
    39  
    40  func NewS3ClientForMinio(minio *e2e.HTTPService, bucketName string) (*S3Client, error) {
    41  	return NewS3Client(s3.Config{
    42  		Endpoint:        minio.HTTPEndpoint(),
    43  		BucketName:      bucketName,
    44  		SecretAccessKey: flagext.Secret{Value: e2edb.MinioSecretKey},
    45  		AccessKeyID:     e2edb.MinioAccessKey,
    46  		Insecure:        true,
    47  	})
    48  }
    49  
    50  // DeleteBlocks deletes all blocks for a tenant.
    51  func (c *S3Client) DeleteBlocks(userID string) error {
    52  	prefix := fmt.Sprintf("%s/", userID)
    53  
    54  	return c.reader.Iter(context.Background(), prefix, func(entry string) error {
    55  		if !strings.HasPrefix(entry, prefix) {
    56  			return fmt.Errorf("unexpected key in the storage: %s", entry)
    57  		}
    58  
    59  		blockID := strings.TrimPrefix(entry, prefix)
    60  		blockID = strings.TrimSuffix(blockID, "/")
    61  
    62  		// Skip keys which are not block IDs
    63  		if _, err := ulid.Parse(blockID); err != nil {
    64  			return nil
    65  		}
    66  
    67  		return c.DeleteBlock(userID, blockID)
    68  	})
    69  }
    70  
    71  // DeleteBlock deletes a single block.
    72  func (c *S3Client) DeleteBlock(userID, blockID string) error {
    73  	return c.Delete(fmt.Sprintf("%s/%s/", userID, blockID))
    74  }
    75  
    76  // Delete recursively deletes every object within the input prefix.
    77  func (c *S3Client) Delete(prefix string) error {
    78  	return c.reader.Iter(context.Background(), prefix, func(entry string) error {
    79  		if !strings.HasPrefix(entry, prefix) {
    80  			return fmt.Errorf("unexpected key in the storage: %s", entry)
    81  		}
    82  
    83  		// Recursively delete if it's a prefix.
    84  		if strings.HasSuffix(entry, "/") {
    85  			return c.Delete(entry)
    86  		}
    87  
    88  		return c.writer.Delete(context.Background(), entry)
    89  	})
    90  }