diff --git a/arangodb/README.md b/arangodb/README.md new file mode 100644 index 00000000..f6eb8bd4 --- /dev/null +++ b/arangodb/README.md @@ -0,0 +1,107 @@ +# ArangoDB +A ArangoDB storage driver using `arangodb/go-driver` and [arangodb/go-driver](https://github.com/arangodb/go-driver). + +### Table of Contents +- [Signatures](#signatures) +- [Installation](#installation) +- [Examples](#examples) +- [Config](#config) +- [Default Config](#default-config) + +### Signatures +```go +func New(config ...Config) Storage +func (s *Storage) Get(key string) ([]byte, error) +func (s *Storage) Set(key string, val []byte, exp time.Duration) error +func (s *Storage) Delete(key string) error +func (s *Storage) Reset() error +func (s *Storage) Close() error +``` +### Installation +ArangoDB is tested on the 2 last (1.14/1.15) [Go versions](https://golang.org/dl/) with support for modules. So make sure to initialize one first if you didn't do that yet: +```bash +go mod init github.com// +``` +And then install the mysql implementation: +```bash +go get github.com/gofiber/storage/arangodb +``` + +### Examples +Import the storage package. +```go +import "github.com/gofiber/storage/arangodb" +``` + +You can use the following possibilities to create a storage: +```go +// Initialize custom config +// *http* is mandatory +store := arangodb.New(arangodb.Config{ + Host: "http://127.0.0.1", + Port: "8529", + Username: "username", + Password: "password" + Database: "fiber", + Collection: "fiber_storage", + Reset: false, + GCInterval: 10 * time.Second, +}) +``` + +### Config +```go +type Config struct { + // Host name where the DB is hosted + // + // Optional. Default is "http://127.0.0.1" + Host string + + // Port where the DB is listening on + // + // Optional. Default is 8529 + Port string + + // Server username + // + // Mandatory + Username string + + // Server password + // + // Mandatory + Password string + + // Database name + // + // Optional. Default is "fiber" + Database string + + // Collection name + // + // Optional. Default is "fiber_storage" + Collection string + + // Reset clears any existing keys in existing collection + // + // Optional. Default is false + Reset bool + // Time before deleting expired keys + // + // Optional. Default is 10 * time.Second + GCInterval time.Duration +} +``` + +### Default Config +Used only for optional fields +```go +var ConfigDefault = Config{ + Host: "http://127.0.0.1", + Port: "8529", + Database: "fiber", + Collection: "fiber_storage", + Reset: false, + GCInterval: 10 * time.Second, +} +``` diff --git a/arangodb/arangodb.go b/arangodb/arangodb.go new file mode 100644 index 00000000..2136ca62 --- /dev/null +++ b/arangodb/arangodb.go @@ -0,0 +1,247 @@ +package arangodb + +import ( + "context" + "fmt" + "github.com/arangodb/go-driver" + "github.com/arangodb/go-driver/http" + "github.com/gofiber/utils" + "time" +) + +// 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{cfg.hostComposed()}, + }) + 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) + } + } +} diff --git a/arangodb/arangodb_test.go b/arangodb/arangodb_test.go new file mode 100644 index 00000000..c65b5751 --- /dev/null +++ b/arangodb/arangodb_test.go @@ -0,0 +1,126 @@ +package arangodb + +import ( + "github.com/gofiber/utils" + "os" + "testing" + "time" +) + +var testStore = New(Config{ + Host: os.Getenv("ARANGODB_HOST"), + Username: os.Getenv("ARANGODB_USERNAME"), + Password: os.Getenv("ARANGODB_PASSWORD"), + Reset: true, +}) + +func Test_ARANGODB_Set(t *testing.T) { + var ( + key = "john" + val = []byte("doe") + ) + + err := testStore.Set(key, val, 0) + utils.AssertEqual(t, nil, err) +} + +func Test_ARANGODB_Upsert(t *testing.T) { + var ( + key = "john" + val = []byte("doe") + ) + + err := testStore.Set(key, val, 0) + utils.AssertEqual(t, nil, err) + + err = testStore.Set(key, val, 0) + utils.AssertEqual(t, nil, err) +} + +func Test_ARANGODB_Get(t *testing.T) { + var ( + key = "john" + val = []byte("doe") + ) + + err := testStore.Set(key, val, 0) + utils.AssertEqual(t, nil, err) + + result, err := testStore.Get(key) + utils.AssertEqual(t, nil, err) + utils.AssertEqual(t, val, result) +} + +func Test_ARANGODB_Set_Expiration(t *testing.T) { + var ( + key = "john" + val = []byte("doe") + exp = 1 * time.Second + ) + + err := testStore.Set(key, val, exp) + utils.AssertEqual(t, nil, err) + + time.Sleep(1100 * time.Millisecond) +} + +func Test_ARANGODB_Get_Expired(t *testing.T) { + var ( + key = "john" + ) + + result, err := testStore.Get(key) + utils.AssertEqual(t, nil, err) + utils.AssertEqual(t, true, len(result) == 0) +} + +func Test_ARANGODB_Get_NotExist(t *testing.T) { + + result, err := testStore.Get("notexist") + utils.AssertEqual(t, nil, err) + utils.AssertEqual(t, true, len(result) == 0) +} + +func Test_ARANGODB_Delete(t *testing.T) { + var ( + key = "john" + val = []byte("doe") + ) + + err := testStore.Set(key, val, 0) + utils.AssertEqual(t, nil, err) + + err = testStore.Delete(key) + utils.AssertEqual(t, nil, err) + + result, err := testStore.Get(key) + utils.AssertEqual(t, nil, err) + utils.AssertEqual(t, true, len(result) == 0) +} + +func Test_ARANGODB_Reset(t *testing.T) { + var ( + val = []byte("doe") + ) + + err := testStore.Set("john1", val, 0) + utils.AssertEqual(t, nil, err) + + err = testStore.Set("john2", val, 0) + utils.AssertEqual(t, nil, err) + + err = testStore.Reset() + utils.AssertEqual(t, nil, err) + + result, err := testStore.Get("john1") + utils.AssertEqual(t, nil, err) + utils.AssertEqual(t, true, len(result) == 0) + + result, err = testStore.Get("john2") + utils.AssertEqual(t, nil, err) + utils.AssertEqual(t, true, len(result) == 0) +} + +func Test_ARANGODB_Close(t *testing.T) { + utils.AssertEqual(t, nil, testStore.Close()) +} diff --git a/arangodb/config.go b/arangodb/config.go new file mode 100644 index 00000000..80b65af4 --- /dev/null +++ b/arangodb/config.go @@ -0,0 +1,92 @@ +package arangodb + +import ( + "fmt" + "strings" + "time" +) + +// Config defines the config for storage. +type Config struct { + // Host name where the DB is hosted + // + // Optional. Default is "http://127.0.0.1" + Host string + + // Port where the DB is listening on + // + // Optional. Default is 8529 + Port string + + // Server username + // + // Mandatory + Username string + + // Server password + // + // Mandatory + Password string + + // Database name + // + // Optional. Default is "fiber" + Database string + + // Collection name + // + // Optional. Default is "fiber_storage" + Collection string + + // Reset clears any existing keys in existing collection + // + // Optional. Default is false + Reset bool + // Time before deleting expired keys + // + // Optional. Default is 10 * time.Second + GCInterval time.Duration +} + +// ConfigDefault is the default config +var ConfigDefault = Config{ + Host: "http://127.0.0.1", + Port: "8529", + Database: "fiber", + Collection: "fiber_storage", + Reset: false, + GCInterval: 10 * time.Second, +} + +func (c Config) hostComposed() string { + return fmt.Sprintf("%s:%s", c.Host, c.Port) +} + +// Helper function to set default values +func configDefault(cfg Config) Config { + if cfg.Username == "" || cfg.Password == "" { + panic("username and password are mandatory") + } + // Set default values + if cfg.Host == "" { + cfg.Host = ConfigDefault.Host + } else { + if !strings.HasPrefix(cfg.Host, "http") { + panic("the host should start with http:// or https://") + } + } + if len(cfg.Port) <= 0 { + cfg.Port = ConfigDefault.Port + } + if cfg.Database == "" { + cfg.Database = ConfigDefault.Database + } + if cfg.Collection == "" { + cfg.Collection = ConfigDefault.Collection + } + + if int(cfg.GCInterval.Seconds()) <= 0 { + cfg.GCInterval = ConfigDefault.GCInterval + } + return cfg +} diff --git a/arangodb/go.mod b/arangodb/go.mod new file mode 100644 index 00000000..d06bf599 --- /dev/null +++ b/arangodb/go.mod @@ -0,0 +1,8 @@ +module github.com/gofiber/storage/arangodb + +go 1.14 + +require ( + github.com/arangodb/go-driver v0.0.0-20201106193344-56ae8fd24510 + github.com/gofiber/utils v0.1.2 +)