storage/arangodb/arangodb.go

package arangodb

import (
	"context"
	"fmt"
	"time"

	"github.com/arangodb/go-driver"
	"github.com/arangodb/go-driver/http"
	"github.com/gofiber/utils"
)

// Storage interface that is implemented by storage providers
type Storage struct {
	db         driver.Database
	gcInterval time.Duration
	done       chan struct{}

	// Arango mandatory fields
	connection    driver.Connection
	client        driver.Client
	collection    driver.Collection
	bindingParams map[string]interface{}
	config        Config
	// AQL query used to remove expired keys
	aqlRemoveGC string
}

type model struct {
	// respect key format field name for ArangoDB
	Key string `json:"_key"`
	Val string `json:"val"`
	Exp int64  `json:"exp"`
}

// New creates a new storage
func New(config ...Config) *Storage {
	// Set default config
	cfg := configDefault(config...)

	// create connection object to arango
	conn, err := http.NewConnection(http.ConnectionConfig{
		Endpoints: []string{fmt.Sprintf("%s:%d", cfg.Host, cfg.Port)},
	})
	if err != nil {
		panic(err)
	}

	// instantiate client after the connection is started
	client, err := driver.NewClient(driver.ClientConfig{
		Connection:     conn,
		Authentication: driver.BasicAuthentication(cfg.Username, cfg.Password),
	})
	if err != nil {
		panic(err)
	}

	// check if the database exists
	// if not create it
	// (it works only with admin privilege user)
	exists, err := client.DatabaseExists(context.Background(), cfg.Database)
	if err != nil {
		panic(err)
	}
	if !exists {
		_, err = client.CreateDatabase(context.Background(), cfg.Database, nil)
		if err != nil {
			panic(err)
		}
	}
	database, err := client.Database(context.Background(), cfg.Database)
	if err != nil {
		panic(err)
	}
	found, _ := database.CollectionExists(context.Background(), cfg.Collection)

	// Create the collection if not exists
	var collection driver.Collection
	if !found {
		// Create
		collection, err = database.CreateCollection(context.Background(), cfg.Collection, &driver.CreateCollectionOptions{})
		if err != nil {
			panic(err)
		}
	} else {
		// Get the collection
		collection, err = database.Collection(context.Background(), cfg.Collection)
		if err != nil {
			panic(err)
		}
	}

	// Truncate collection if Reset set to true
	if cfg.Reset {
		err = collection.Truncate(context.Background())
		if err != nil {
			panic(err)
		}
	}

	// Create storage
	store := &Storage{
		gcInterval:  cfg.GCInterval,
		db:          database,
		collection:  collection,
		client:      client,
		connection:  conn,
		config:      cfg,
		done:        make(chan struct{}),
		aqlRemoveGC: fmt.Sprintf("FOR doc IN %s\n  FILTER doc.exp <= @exp \n REMOVE { _key: doc._key } IN %s", collection.Name(), collection.Name()),
	}

	// Start garbage collector
	go store.gc()

	return store
}

// Get value by key
func (s *Storage) Get(key string) ([]byte, error) {
	if len(key) <= 0 {
		return nil, nil
	}

	ctx := context.Background()

	// Check if the document exists
	// to avoid errors later
	exists, err := s.collection.DocumentExists(ctx, key)
	if err != nil {
		return nil, err
	}

	// instead of returning an error if not exists
	// return nil
	if !exists {
		return nil, nil
	}

	// result model
	var model model
	_, err = s.collection.ReadDocument(ctx, key, &model)
	if err != nil {
		return nil, err
	}
	// If the expiration time has already passed, then return nil
	if model.Exp != 0 && model.Exp <= time.Now().Unix() {
		return nil, nil
	}

	return utils.UnsafeBytes(model.Val), nil
}

// Set key with value
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
	// Ain't Nobody Got Time For That
	if len(key) <= 0 || len(val) <= 0 {
		return nil
	}
	var expireAt int64
	if exp != 0 {
		expireAt = time.Now().Add(exp).Unix()
	}
	valStr := utils.UnsafeString(val)

	// create the structure for the storage
	data := model{
		Key: key,
		Val: valStr,
		Exp: expireAt,
	}
	ctx := context.Background()

	// Arango does not support documents with the same key
	// So we need to check if the document exists
	exists, err := s.collection.DocumentExists(ctx, key)
	if err != nil {
		return err
	}
	// Update the document if exists
	if exists {
		_, err = s.collection.UpdateDocument(ctx, key, data)
		return err
	}
	// Otherwise create it
	_, err = s.collection.CreateDocument(ctx, data)

	return err
}

// Delete value by key
func (s *Storage) Delete(key string) error {
	// Ain't Nobody Got Time For That
	if len(key) <= 0 {
		return nil
	}
	_, err := s.collection.RemoveDocument(context.Background(), key)
	return err
}

// Reset all keys
// truncate the collection
func (s *Storage) Reset() error {
	return s.collection.Truncate(context.Background())
}

// Close the database
// Arango does not provide a method to close the connection
// more info @https://github.com/arangodb/go-driver/issues/43
func (s *Storage) Close() error {
	// Stop gc
	s.done <- struct{}{}
	// reset connection params
	s.db = nil
	s.collection = nil
	s.connection = nil
	s.bindingParams = nil

	return nil
}

// execute query
func (s *Storage) exec(query string) error {
	// execute query
	_, err := s.db.Query(context.Background(), query, s.bindingParams)
	if err != nil {
		return err
	}
	// reset binding params
	s.bindingParams = map[string]interface{}{}
	return nil
}

// Garbage collector to delete expired keys
func (s *Storage) gc() {
	ticker := time.NewTicker(s.gcInterval)
	defer ticker.Stop()
	for {
		select {
		case <-s.done:
			return
		case t := <-ticker.C:
			// set the expiration
			s.bindingParams["exp"] = t.Unix()
			_ = s.exec(s.aqlRemoveGC)
		}
	}
}

// Return database client
func (s *Storage) Conn() driver.Client {
	return s.client
}