storage/azureblob/azureblob.go

package azureblob

import (
	"context"
	"fmt"
	"io"
	"time"

	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/bloberror"
)

// Storage interface that is implemented by storage providers
type Storage struct {
	client         *azblob.Client
	container      string
	requestTimeout time.Duration
}

// New creates a new storage
func New(config ...Config) *Storage {
	// Set default config
	cfg := configure(config...)
	// Set the azure credentials
	cred, err := azblob.NewSharedKeyCredential(cfg.Credentials.Account, cfg.Credentials.Key)
	handleError(err)
	client, err := azblob.NewClientWithSharedKeyCredential(cfg.Endpoint, cred, nil)
	handleError(err)
	_, err = client.CreateContainer(context.TODO(), cfg.Container, nil)
	if err != nil {
		if !bloberror.HasCode(err, bloberror.ContainerAlreadyExists) {
			panic(fmt.Sprintf("invalid config:, %v", err))
		}
	}
	storage := &Storage{
		client:         client,
		container:      cfg.Container,
		requestTimeout: cfg.RequestTimeout,
	}

	// Reset all entries if set to true
	if cfg.Reset {
		if err := storage.Reset(); err != nil {
			panic(err)
		}
	}

	return storage
}

// Get value by key
func (s *Storage) Get(key string) ([]byte, error) {
	if len(key) <= 0 {
		return nil, nil
	}
	ctx, cancel := s.requestContext()
	defer cancel()
	resp, err := s.client.DownloadStream(ctx, s.container, key, nil)
	if err != nil {
		return []byte{}, err
	}
	data, err := io.ReadAll(resp.Body)
	return data, err
}

// Set key with value
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
	if len(key) <= 0 {
		return nil
	}
	ctx, cancel := s.requestContext()
	defer cancel()
	_, err := s.client.UploadBuffer(ctx, s.container, key, val, nil)
	return err
}

// Delete entry by key
func (s *Storage) Delete(key string) error {
	if len(key) <= 0 {
		return nil
	}

	ctx, cancel := s.requestContext()
	defer cancel()
	_, err := s.client.DeleteBlob(ctx, s.container, key, nil)
	return err
}

// Reset all entries
func (s *Storage) Reset() error {
	ctx, cancel := s.requestContext()
	defer cancel()
	//_, err := s.client.DeleteContainer(ctx, s.container, nil)
	//return err
	pager := s.client.NewListBlobsFlatPager(s.container, nil)
	errCounter := 0
	for pager.More() {
		resp, err := pager.NextPage(ctx)
		if err != nil {
			errCounter = errCounter + 1
		}
		for _, v := range resp.Segment.BlobItems {
			_, err = s.client.DeleteBlob(ctx, s.container, *v.Name, nil)
			if err != nil {
				errCounter = errCounter + 1
			}
		}
	}
	if errCounter > 0 {
		return fmt.Errorf("%d errors occured while resetting", errCounter)
	}
	return nil
}

// Conn returns storage client
func (s *Storage) Conn() *azblob.Client {
	return s.client
}

// Close the storage connextion
func (s *Storage) Close() error {
	return nil
}

// Context for making requests will timeout if a non-zero timeout is configured
func (s *Storage) requestContext() (context.Context, context.CancelFunc) {
	if s.requestTimeout > 0 {
		return context.WithTimeout(context.Background(), s.requestTimeout)
	}
	return context.Background(), func() {}
}

// handleError is a helper to panic on error
func handleError(err error) {
	if err != nil {
		panic(fmt.Sprintf("invalid config:, %v", err))
	}
}