Added support for Badger as an alternative storage backend

internal/datastore/badger_datastore/badger_datastore.go

@@ -0,0 +1,36 @@
+package badgerdatastore
+
+import (
+	"encoding/gob"
+
+	"github.com/dgraph-io/badger/v4"
+	"github.com/pikami/cosmium/internal/logger"
+)
+
+type BadgerDataStore struct {
+	db *badger.DB
+}
+
+func NewBadgerDataStore() *BadgerDataStore {
+	gob.Register([]interface{}{})
+
+	badgerOpts := badger.DefaultOptions("").WithInMemory(true)
+
+	db, err := badger.Open(badgerOpts)
+	if err != nil {
+		panic(err)
+	}
+
+	return &BadgerDataStore{
+		db: db,
+	}
+}
+
+func (r *BadgerDataStore) Close() {
+	r.db.Close()
+}
+
+func (r *BadgerDataStore) DumpToJson() (string, error) {
+	logger.ErrorLn("Badger datastore does not support state export currently.")
+	return "{}", nil
+}

internal/datastore/badger_datastore/collections.go

@@ -0,0 +1,103 @@
+package badgerdatastore
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+	"github.com/pikami/cosmium/internal/resourceid"
+	structhidrators "github.com/pikami/cosmium/internal/struct_hidrators"
+)
+
+func (r *BadgerDataStore) GetAllCollections(databaseId string) ([]datastore.Collection, datastore.DataStoreStatus) {
+	exists, err := keyExists(r.db.NewTransaction(false), generateDatabaseKey(databaseId))
+	if err != nil {
+		logger.ErrorLn("Error while checking if database exists:", err)
+		return nil, datastore.Unknown
+	}
+
+	if !exists {
+		return nil, datastore.StatusNotFound
+	}
+
+	colls, status := listByPrefix[datastore.Collection](r.db, generateKey(resourceid.ResourceTypeCollection, databaseId, "", ""))
+	if status == datastore.StatusOk {
+		return colls, datastore.StatusOk
+	}
+
+	return nil, status
+}
+
+func (r *BadgerDataStore) GetCollection(databaseId string, collectionId string) (datastore.Collection, datastore.DataStoreStatus) {
+	collectionKey := generateCollectionKey(databaseId, collectionId)
+
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	var collection datastore.Collection
+	status := getKey(txn, collectionKey, &collection)
+
+	return collection, status
+}
+
+func (r *BadgerDataStore) DeleteCollection(databaseId string, collectionId string) datastore.DataStoreStatus {
+	collectionKey := generateCollectionKey(databaseId, collectionId)
+
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	prefixes := []string{
+		generateKey(resourceid.ResourceTypeDocument, databaseId, collectionId, ""),
+		generateKey(resourceid.ResourceTypeTrigger, databaseId, collectionId, ""),
+		generateKey(resourceid.ResourceTypeStoredProcedure, databaseId, collectionId, ""),
+		generateKey(resourceid.ResourceTypeUserDefinedFunction, databaseId, collectionId, ""),
+		collectionKey,
+	}
+	for _, prefix := range prefixes {
+		if err := deleteKeysByPrefix(txn, prefix); err != nil {
+			return datastore.Unknown
+		}
+	}
+
+	err := txn.Commit()
+	if err != nil {
+		logger.ErrorLn("Error while committing transaction:", err)
+		return datastore.Unknown
+	}
+
+	return datastore.StatusOk
+}
+
+func (r *BadgerDataStore) CreateCollection(databaseId string, newCollection datastore.Collection) (datastore.Collection, datastore.DataStoreStatus) {
+	collectionKey := generateCollectionKey(databaseId, newCollection.ID)
+
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	collectionExists, err := keyExists(txn, collectionKey)
+	if err != nil || collectionExists {
+		return datastore.Collection{}, datastore.Conflict
+	}
+
+	var database datastore.Database
+	status := getKey(txn, generateDatabaseKey(databaseId), &database)
+	if status != datastore.StatusOk {
+		return datastore.Collection{}, status
+	}
+
+	newCollection = structhidrators.Hidrate(newCollection).(datastore.Collection)
+
+	newCollection.TimeStamp = time.Now().Unix()
+	newCollection.ResourceID = resourceid.NewCombined(database.ResourceID, resourceid.New(resourceid.ResourceTypeCollection))
+	newCollection.ETag = fmt.Sprintf("\"%s\"", uuid.New())
+	newCollection.Self = fmt.Sprintf("dbs/%s/colls/%s/", database.ResourceID, newCollection.ResourceID)
+
+	status = insertKey(txn, collectionKey, newCollection)
+	if status != datastore.StatusOk {
+		return datastore.Collection{}, status
+	}
+
+	return newCollection, datastore.StatusOk
+}

internal/datastore/badger_datastore/databases.go

@@ -0,0 +1,80 @@
+package badgerdatastore
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+	"github.com/pikami/cosmium/internal/resourceid"
+)
+
+func (r *BadgerDataStore) GetAllDatabases() ([]datastore.Database, datastore.DataStoreStatus) {
+	dbs, status := listByPrefix[datastore.Database](r.db, DatabaseKeyPrefix)
+	if status == datastore.StatusOk {
+		return dbs, datastore.StatusOk
+	}
+
+	return nil, status
+}
+
+func (r *BadgerDataStore) GetDatabase(id string) (datastore.Database, datastore.DataStoreStatus) {
+	databaseKey := generateDatabaseKey(id)
+
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	var database datastore.Database
+	status := getKey(txn, databaseKey, &database)
+
+	return database, status
+}
+
+func (r *BadgerDataStore) DeleteDatabase(id string) datastore.DataStoreStatus {
+	databaseKey := generateDatabaseKey(id)
+
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	prefixes := []string{
+		generateKey(resourceid.ResourceTypeCollection, id, "", ""),
+		generateKey(resourceid.ResourceTypeDocument, id, "", ""),
+		generateKey(resourceid.ResourceTypeTrigger, id, "", ""),
+		generateKey(resourceid.ResourceTypeStoredProcedure, id, "", ""),
+		generateKey(resourceid.ResourceTypeUserDefinedFunction, id, "", ""),
+		databaseKey,
+	}
+	for _, prefix := range prefixes {
+		if err := deleteKeysByPrefix(txn, prefix); err != nil {
+			return datastore.Unknown
+		}
+	}
+
+	err := txn.Commit()
+	if err != nil {
+		logger.ErrorLn("Error while committing transaction:", err)
+		return datastore.Unknown
+	}
+
+	return datastore.StatusOk
+}
+
+func (r *BadgerDataStore) CreateDatabase(newDatabase datastore.Database) (datastore.Database, datastore.DataStoreStatus) {
+	databaseKey := generateDatabaseKey(newDatabase.ID)
+
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	newDatabase.TimeStamp = time.Now().Unix()
+	newDatabase.ResourceID = resourceid.New(resourceid.ResourceTypeDatabase)
+	newDatabase.ETag = fmt.Sprintf("\"%s\"", uuid.New())
+	newDatabase.Self = fmt.Sprintf("dbs/%s/", newDatabase.ResourceID)
+
+	status := insertKey(txn, databaseKey, newDatabase)
+	if status != datastore.StatusOk {
+		return datastore.Database{}, status
+	}
+
+	return newDatabase, datastore.StatusOk
+}

internal/datastore/badger_datastore/db_abstractions.go

@@ -0,0 +1,207 @@
+package badgerdatastore
+
+import (
+	"bytes"
+	"encoding/gob"
+
+	"github.com/dgraph-io/badger/v4"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+	"github.com/pikami/cosmium/internal/resourceid"
+)
+
+const (
+	DatabaseKeyPrefix            = "DB:"
+	CollectionKeyPrefix          = "COL:"
+	DocumentKeyPrefix            = "DOC:"
+	TriggerKeyPrefix             = "TRG:"
+	StoredProcedureKeyPrefix     = "SP:"
+	UserDefinedFunctionKeyPrefix = "UDF:"
+)
+
+func generateKey(
+	resourceType resourceid.ResourceType,
+	databaseId string,
+	collectionId string,
+	resourceId string,
+) string {
+	result := ""
+
+	switch resourceType {
+	case resourceid.ResourceTypeDatabase:
+		result += DatabaseKeyPrefix
+	case resourceid.ResourceTypeCollection:
+		result += CollectionKeyPrefix
+	case resourceid.ResourceTypeDocument:
+		result += DocumentKeyPrefix
+	case resourceid.ResourceTypeTrigger:
+		result += TriggerKeyPrefix
+	case resourceid.ResourceTypeStoredProcedure:
+		result += StoredProcedureKeyPrefix
+	case resourceid.ResourceTypeUserDefinedFunction:
+		result += UserDefinedFunctionKeyPrefix
+	}
+
+	if databaseId != "" {
+		result += databaseId
+	}
+
+	if collectionId != "" {
+		result += "/colls/" + collectionId
+	}
+
+	if resourceId != "" {
+		result += "/" + resourceId
+	}
+
+	return result
+}
+
+func generateDatabaseKey(databaseId string) string {
+	return generateKey(resourceid.ResourceTypeDatabase, databaseId, "", "")
+}
+
+func generateCollectionKey(databaseId string, collectionId string) string {
+	return generateKey(resourceid.ResourceTypeCollection, databaseId, collectionId, "")
+}
+
+func generateDocumentKey(databaseId string, collectionId string, documentId string) string {
+	return generateKey(resourceid.ResourceTypeDocument, databaseId, collectionId, documentId)
+}
+
+func generateTriggerKey(databaseId string, collectionId string, triggerId string) string {
+	return generateKey(resourceid.ResourceTypeTrigger, databaseId, collectionId, triggerId)
+}
+
+func generateStoredProcedureKey(databaseId string, collectionId string, storedProcedureId string) string {
+	return generateKey(resourceid.ResourceTypeStoredProcedure, databaseId, collectionId, storedProcedureId)
+}
+
+func generateUserDefinedFunctionKey(databaseId string, collectionId string, udfId string) string {
+	return generateKey(resourceid.ResourceTypeUserDefinedFunction, databaseId, collectionId, udfId)
+}
+
+func insertKey(txn *badger.Txn, key string, value interface{}) datastore.DataStoreStatus {
+	_, err := txn.Get([]byte(key))
+	if err == nil {
+		return datastore.Conflict
+	}
+
+	if err != badger.ErrKeyNotFound {
+		logger.ErrorLn("Error while checking if key exists:", err)
+		return datastore.Unknown
+	}
+
+	var buf bytes.Buffer
+	err = gob.NewEncoder(&buf).Encode(value)
+	if err != nil {
+		logger.ErrorLn("Error while encoding value:", err)
+		return datastore.Unknown
+	}
+
+	err = txn.Set([]byte(key), buf.Bytes())
+	if err != nil {
+		logger.ErrorLn("Error while setting key:", err)
+		return datastore.Unknown
+	}
+
+	err = txn.Commit()
+	if err != nil {
+		logger.ErrorLn("Error while committing transaction:", err)
+		return datastore.Unknown
+	}
+
+	return datastore.StatusOk
+}
+
+func getKey(txn *badger.Txn, key string, value interface{}) datastore.DataStoreStatus {
+	item, err := txn.Get([]byte(key))
+	if err != nil {
+		if err == badger.ErrKeyNotFound {
+			return datastore.StatusNotFound
+		}
+		logger.ErrorLn("Error while getting key:", err)
+		return datastore.Unknown
+	}
+
+	val, err := item.ValueCopy(nil)
+	if err != nil {
+		logger.ErrorLn("Error while copying value:", err)
+		return datastore.Unknown
+	}
+
+	if value == nil {
+		logger.ErrorLn("getKey called with nil value")
+		return datastore.Unknown
+	}
+
+	err = gob.NewDecoder(bytes.NewReader(val)).Decode(value)
+	if err != nil {
+		logger.ErrorLn("Error while decoding value:", err)
+		return datastore.Unknown
+	}
+
+	return datastore.StatusOk
+}
+
+func keyExists(txn *badger.Txn, key string) (bool, error) {
+	_, err := txn.Get([]byte(key))
+	if err == nil {
+		return true, nil
+	}
+
+	if err == badger.ErrKeyNotFound {
+		return false, nil
+	}
+
+	return false, err
+}
+
+func listByPrefix[T any](db *badger.DB, prefix string) ([]T, datastore.DataStoreStatus) {
+	var results []T
+
+	err := db.View(func(txn *badger.Txn) error {
+		opts := badger.DefaultIteratorOptions
+		opts.Prefix = []byte(prefix)
+		it := txn.NewIterator(opts)
+		defer it.Close()
+
+		for it.Rewind(); it.Valid(); it.Next() {
+			item := it.Item()
+			var entry T
+
+			status := getKey(txn, string(item.Key()), &entry)
+			if status != datastore.StatusOk {
+				logger.ErrorLn("Failed to retrieve entry:", string(item.Key()))
+				continue
+			}
+
+			results = append(results, entry)
+		}
+		return nil
+	})
+
+	if err != nil {
+		logger.ErrorLn("Error while listing entries:", err)
+		return nil, datastore.Unknown
+	}
+
+	return results, datastore.StatusOk
+}
+
+func deleteKeysByPrefix(txn *badger.Txn, prefix string) error {
+	opts := badger.DefaultIteratorOptions
+	opts.Prefix = []byte(prefix)
+	it := txn.NewIterator(opts)
+	defer it.Close()
+
+	for it.Rewind(); it.Valid(); it.Next() {
+		key := it.Item().KeyCopy(nil)
+		if err := txn.Delete(key); err != nil {
+			logger.ErrorLn("Failed to delete key:", string(key), "Error:", err)
+			return err
+		}
+	}
+
+	return nil
+}

internal/datastore/badger_datastore/document_iterator.go

@@ -0,0 +1,60 @@
+package badgerdatastore
+
+import (
+	"bytes"
+	"encoding/gob"
+
+	"github.com/dgraph-io/badger/v4"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+)
+
+type BadgerDocumentIterator struct {
+	txn    *badger.Txn
+	it     *badger.Iterator
+	prefix string
+}
+
+func NewBadgerDocumentIterator(txn *badger.Txn, prefix string) *BadgerDocumentIterator {
+	opts := badger.DefaultIteratorOptions
+	opts.Prefix = []byte(prefix)
+
+	it := txn.NewIterator(opts)
+	it.Rewind()
+
+	return &BadgerDocumentIterator{
+		txn:    txn,
+		it:     it,
+		prefix: prefix,
+	}
+}
+
+func (i *BadgerDocumentIterator) Next() (datastore.Document, datastore.DataStoreStatus) {
+	if !i.it.Valid() {
+		i.it.Close()
+		return datastore.Document{}, datastore.IterEOF
+	}
+
+	item := i.it.Item()
+	val, err := item.ValueCopy(nil)
+	if err != nil {
+		logger.ErrorLn("Error while copying value:", err)
+		return datastore.Document{}, datastore.Unknown
+	}
+
+	current := &datastore.Document{}
+	err = gob.NewDecoder(bytes.NewReader(val)).Decode(current)
+	if err != nil {
+		logger.ErrorLn("Error while decoding value:", err)
+		return datastore.Document{}, datastore.Unknown
+	}
+
+	i.it.Next()
+
+	return *current, datastore.StatusOk
+}
+
+func (i *BadgerDocumentIterator) Close() {
+	i.it.Close()
+	i.txn.Discard()
+}

internal/datastore/badger_datastore/documents.go

@@ -0,0 +1,127 @@
+package badgerdatastore
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+	"github.com/pikami/cosmium/internal/resourceid"
+)
+
+func (r *BadgerDataStore) GetAllDocuments(databaseId string, collectionId string) ([]datastore.Document, datastore.DataStoreStatus) {
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	dbExists, err := keyExists(txn, generateDatabaseKey(databaseId))
+	if err != nil || !dbExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	collExists, err := keyExists(txn, generateCollectionKey(databaseId, collectionId))
+	if err != nil || !collExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	docs, status := listByPrefix[datastore.Document](r.db, generateKey(resourceid.ResourceTypeDocument, databaseId, collectionId, ""))
+	if status == datastore.StatusOk {
+		return docs, datastore.StatusOk
+	}
+
+	return nil, status
+}
+
+func (r *BadgerDataStore) GetDocumentIterator(databaseId string, collectionId string) (datastore.DocumentIterator, datastore.DataStoreStatus) {
+	txn := r.db.NewTransaction(false)
+
+	dbExists, err := keyExists(txn, generateDatabaseKey(databaseId))
+	if err != nil || !dbExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	collExists, err := keyExists(txn, generateCollectionKey(databaseId, collectionId))
+	if err != nil || !collExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	iter := NewBadgerDocumentIterator(txn, generateKey(resourceid.ResourceTypeDocument, databaseId, collectionId, ""))
+	return iter, datastore.StatusOk
+}
+
+func (r *BadgerDataStore) GetDocument(databaseId string, collectionId string, documentId string) (datastore.Document, datastore.DataStoreStatus) {
+	documentKey := generateDocumentKey(databaseId, collectionId, documentId)
+
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	var document datastore.Document
+	status := getKey(txn, documentKey, &document)
+
+	return document, status
+}
+
+func (r *BadgerDataStore) DeleteDocument(databaseId string, collectionId string, documentId string) datastore.DataStoreStatus {
+	documentKey := generateDocumentKey(databaseId, collectionId, documentId)
+
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	exists, err := keyExists(txn, documentKey)
+	if err != nil {
+		return datastore.Unknown
+	}
+	if !exists {
+		return datastore.StatusNotFound
+	}
+
+	err = txn.Delete([]byte(documentKey))
+	if err != nil {
+		logger.ErrorLn("Error while deleting document:", err)
+		return datastore.Unknown
+	}
+
+	err = txn.Commit()
+	if err != nil {
+		logger.ErrorLn("Error while committing transaction:", err)
+		return datastore.Unknown
+	}
+
+	return datastore.StatusOk
+}
+
+func (r *BadgerDataStore) CreateDocument(databaseId string, collectionId string, document map[string]interface{}) (datastore.Document, datastore.DataStoreStatus) {
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	var database datastore.Database
+	status := getKey(txn, generateDatabaseKey(databaseId), &database)
+	if status != datastore.StatusOk {
+		return datastore.Document{}, status
+	}
+
+	var collection datastore.Collection
+	status = getKey(txn, generateCollectionKey(databaseId, collectionId), &collection)
+	if status != datastore.StatusOk {
+		return datastore.Document{}, status
+	}
+
+	var ok bool
+	var documentId string
+	if documentId, ok = document["id"].(string); !ok || documentId == "" {
+		documentId = fmt.Sprint(uuid.New())
+		document["id"] = documentId
+	}
+
+	document["_ts"] = time.Now().Unix()
+	document["_rid"] = resourceid.NewCombined(collection.ResourceID, resourceid.New(resourceid.ResourceTypeDocument))
+	document["_etag"] = fmt.Sprintf("\"%s\"", uuid.New())
+	document["_self"] = fmt.Sprintf("dbs/%s/colls/%s/docs/%s/", database.ResourceID, collection.ResourceID, document["_rid"])
+
+	status = insertKey(txn, generateDocumentKey(databaseId, collectionId, documentId), document)
+	if status != datastore.StatusOk {
+		return datastore.Document{}, status
+	}
+
+	return document, datastore.StatusOk
+}

internal/datastore/badger_datastore/partition_key_ranges.go

@@ -0,0 +1,53 @@
+package badgerdatastore
+
+import (
+	"fmt"
+
+	"github.com/google/uuid"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/resourceid"
+)
+
+// I have no idea what this is tbh
+func (r *BadgerDataStore) GetPartitionKeyRanges(databaseId string, collectionId string) ([]datastore.PartitionKeyRange, datastore.DataStoreStatus) {
+	databaseRid := databaseId
+	collectionRid := collectionId
+	var timestamp int64 = 0
+
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	var database datastore.Database
+	status := getKey(txn, generateDatabaseKey(databaseId), &database)
+	if status != datastore.StatusOk {
+		databaseRid = database.ResourceID
+	}
+
+	var collection datastore.Collection
+	status = getKey(txn, generateCollectionKey(databaseId, collectionId), &collection)
+	if status != datastore.StatusOk {
+		collectionRid = collection.ResourceID
+		timestamp = collection.TimeStamp
+	}
+
+	pkrResourceId := resourceid.NewCombined(collectionRid, resourceid.New(resourceid.ResourceTypePartitionKeyRange))
+	pkrSelf := fmt.Sprintf("dbs/%s/colls/%s/pkranges/%s/", databaseRid, collectionRid, pkrResourceId)
+	etag := fmt.Sprintf("\"%s\"", uuid.New())
+
+	return []datastore.PartitionKeyRange{
+		{
+			ResourceID:         pkrResourceId,
+			ID:                 "0",
+			Etag:               etag,
+			MinInclusive:       "",
+			MaxExclusive:       "FF",
+			RidPrefix:          0,
+			Self:               pkrSelf,
+			ThroughputFraction: 1,
+			Status:             "online",
+			Parents:            []interface{}{},
+			TimeStamp:          timestamp,
+			Lsn:                17,
+		},
+	}, datastore.StatusOk
+}

internal/datastore/badger_datastore/stored_procedures.go

@@ -0,0 +1,107 @@
+package badgerdatastore
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+	"github.com/pikami/cosmium/internal/resourceid"
+)
+
+func (r *BadgerDataStore) GetAllStoredProcedures(databaseId string, collectionId string) ([]datastore.StoredProcedure, datastore.DataStoreStatus) {
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	dbExists, err := keyExists(txn, generateDatabaseKey(databaseId))
+	if err != nil || !dbExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	collExists, err := keyExists(txn, generateCollectionKey(databaseId, collectionId))
+	if err != nil || !collExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	storedProcedures, status := listByPrefix[datastore.StoredProcedure](r.db, generateKey(resourceid.ResourceTypeStoredProcedure, databaseId, collectionId, ""))
+	if status == datastore.StatusOk {
+		return storedProcedures, datastore.StatusOk
+	}
+
+	return nil, status
+}
+
+func (r *BadgerDataStore) GetStoredProcedure(databaseId string, collectionId string, storedProcedureId string) (datastore.StoredProcedure, datastore.DataStoreStatus) {
+	storedProcedureKey := generateStoredProcedureKey(databaseId, collectionId, storedProcedureId)
+
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	var storedProcedure datastore.StoredProcedure
+	status := getKey(txn, storedProcedureKey, &storedProcedure)
+
+	return storedProcedure, status
+}
+
+func (r *BadgerDataStore) DeleteStoredProcedure(databaseId string, collectionId string, storedProcedureId string) datastore.DataStoreStatus {
+	storedProcedureKey := generateStoredProcedureKey(databaseId, collectionId, storedProcedureId)
+
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	exists, err := keyExists(txn, storedProcedureKey)
+	if err != nil {
+		return datastore.Unknown
+	}
+	if !exists {
+		return datastore.StatusNotFound
+	}
+
+	err = txn.Delete([]byte(storedProcedureKey))
+	if err != nil {
+		logger.ErrorLn("Error while deleting stored procedure:", err)
+		return datastore.Unknown
+	}
+
+	err = txn.Commit()
+	if err != nil {
+		logger.ErrorLn("Error while committing transaction:", err)
+		return datastore.Unknown
+	}
+
+	return datastore.StatusOk
+}
+
+func (r *BadgerDataStore) CreateStoredProcedure(databaseId string, collectionId string, storedProcedure datastore.StoredProcedure) (datastore.StoredProcedure, datastore.DataStoreStatus) {
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	if storedProcedure.ID == "" {
+		return datastore.StoredProcedure{}, datastore.BadRequest
+	}
+
+	var database datastore.Database
+	status := getKey(txn, generateDatabaseKey(databaseId), &database)
+	if status != datastore.StatusOk {
+		return datastore.StoredProcedure{}, status
+	}
+
+	var collection datastore.Collection
+	status = getKey(txn, generateCollectionKey(databaseId, collectionId), &collection)
+	if status != datastore.StatusOk {
+		return datastore.StoredProcedure{}, status
+	}
+
+	storedProcedure.TimeStamp = time.Now().Unix()
+	storedProcedure.ResourceID = resourceid.NewCombined(collection.ResourceID, resourceid.New(resourceid.ResourceTypeStoredProcedure))
+	storedProcedure.ETag = fmt.Sprintf("\"%s\"", uuid.New())
+	storedProcedure.Self = fmt.Sprintf("dbs/%s/colls/%s/sprocs/%s/", database.ResourceID, collection.ResourceID, storedProcedure.ResourceID)
+
+	status = insertKey(txn, generateStoredProcedureKey(databaseId, collectionId, storedProcedure.ID), storedProcedure)
+	if status != datastore.StatusOk {
+		return datastore.StoredProcedure{}, status
+	}
+
+	return storedProcedure, datastore.StatusOk
+}

internal/datastore/badger_datastore/triggers.go

@@ -0,0 +1,107 @@
+package badgerdatastore
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+	"github.com/pikami/cosmium/internal/resourceid"
+)
+
+func (r *BadgerDataStore) GetAllTriggers(databaseId string, collectionId string) ([]datastore.Trigger, datastore.DataStoreStatus) {
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	dbExists, err := keyExists(txn, generateDatabaseKey(databaseId))
+	if err != nil || !dbExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	collExists, err := keyExists(txn, generateCollectionKey(databaseId, collectionId))
+	if err != nil || !collExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	triggers, status := listByPrefix[datastore.Trigger](r.db, generateKey(resourceid.ResourceTypeTrigger, databaseId, collectionId, ""))
+	if status == datastore.StatusOk {
+		return triggers, datastore.StatusOk
+	}
+
+	return nil, status
+}
+
+func (r *BadgerDataStore) GetTrigger(databaseId string, collectionId string, triggerId string) (datastore.Trigger, datastore.DataStoreStatus) {
+	triggerKey := generateTriggerKey(databaseId, collectionId, triggerId)
+
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	var trigger datastore.Trigger
+	status := getKey(txn, triggerKey, &trigger)
+
+	return trigger, status
+}
+
+func (r *BadgerDataStore) DeleteTrigger(databaseId string, collectionId string, triggerId string) datastore.DataStoreStatus {
+	triggerKey := generateTriggerKey(databaseId, collectionId, triggerId)
+
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	exists, err := keyExists(txn, triggerKey)
+	if err != nil {
+		return datastore.Unknown
+	}
+	if !exists {
+		return datastore.StatusNotFound
+	}
+
+	err = txn.Delete([]byte(triggerKey))
+	if err != nil {
+		logger.ErrorLn("Error while deleting trigger:", err)
+		return datastore.Unknown
+	}
+
+	err = txn.Commit()
+	if err != nil {
+		logger.ErrorLn("Error while committing transaction:", err)
+		return datastore.Unknown
+	}
+
+	return datastore.StatusOk
+}
+
+func (r *BadgerDataStore) CreateTrigger(databaseId string, collectionId string, trigger datastore.Trigger) (datastore.Trigger, datastore.DataStoreStatus) {
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	if trigger.ID == "" {
+		return datastore.Trigger{}, datastore.BadRequest
+	}
+
+	var database datastore.Database
+	status := getKey(txn, generateDatabaseKey(databaseId), &database)
+	if status != datastore.StatusOk {
+		return datastore.Trigger{}, status
+	}
+
+	var collection datastore.Collection
+	status = getKey(txn, generateCollectionKey(databaseId, collectionId), &collection)
+	if status != datastore.StatusOk {
+		return datastore.Trigger{}, status
+	}
+
+	trigger.TimeStamp = time.Now().Unix()
+	trigger.ResourceID = resourceid.NewCombined(collection.ResourceID, resourceid.New(resourceid.ResourceTypeTrigger))
+	trigger.ETag = fmt.Sprintf("\"%s\"", uuid.New())
+	trigger.Self = fmt.Sprintf("dbs/%s/colls/%s/triggers/%s/", database.ResourceID, collection.ResourceID, trigger.ResourceID)
+
+	status = insertKey(txn, generateTriggerKey(databaseId, collectionId, trigger.ID), trigger)
+	if status != datastore.StatusOk {
+		return datastore.Trigger{}, status
+	}
+
+	return trigger, datastore.StatusOk
+}

internal/datastore/badger_datastore/user_defined_functions.go

@@ -0,0 +1,107 @@
+package badgerdatastore
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+	"github.com/pikami/cosmium/internal/resourceid"
+)
+
+func (r *BadgerDataStore) GetAllUserDefinedFunctions(databaseId string, collectionId string) ([]datastore.UserDefinedFunction, datastore.DataStoreStatus) {
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	dbExists, err := keyExists(txn, generateDatabaseKey(databaseId))
+	if err != nil || !dbExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	collExists, err := keyExists(txn, generateCollectionKey(databaseId, collectionId))
+	if err != nil || !collExists {
+		return nil, datastore.StatusNotFound
+	}
+
+	udfs, status := listByPrefix[datastore.UserDefinedFunction](r.db, generateKey(resourceid.ResourceTypeUserDefinedFunction, databaseId, collectionId, ""))
+	if status == datastore.StatusOk {
+		return udfs, datastore.StatusOk
+	}
+
+	return nil, status
+}
+
+func (r *BadgerDataStore) GetUserDefinedFunction(databaseId string, collectionId string, udfId string) (datastore.UserDefinedFunction, datastore.DataStoreStatus) {
+	udfKey := generateUserDefinedFunctionKey(databaseId, collectionId, udfId)
+
+	txn := r.db.NewTransaction(false)
+	defer txn.Discard()
+
+	var udf datastore.UserDefinedFunction
+	status := getKey(txn, udfKey, &udf)
+
+	return udf, status
+}
+
+func (r *BadgerDataStore) DeleteUserDefinedFunction(databaseId string, collectionId string, udfId string) datastore.DataStoreStatus {
+	udfKey := generateUserDefinedFunctionKey(databaseId, collectionId, udfId)
+
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	exists, err := keyExists(txn, udfKey)
+	if err != nil {
+		return datastore.Unknown
+	}
+	if !exists {
+		return datastore.StatusNotFound
+	}
+
+	err = txn.Delete([]byte(udfKey))
+	if err != nil {
+		logger.ErrorLn("Error while deleting user defined function:", err)
+		return datastore.Unknown
+	}
+
+	err = txn.Commit()
+	if err != nil {
+		logger.ErrorLn("Error while committing transaction:", err)
+		return datastore.Unknown
+	}
+
+	return datastore.StatusOk
+}
+
+func (r *BadgerDataStore) CreateUserDefinedFunction(databaseId string, collectionId string, udf datastore.UserDefinedFunction) (datastore.UserDefinedFunction, datastore.DataStoreStatus) {
+	txn := r.db.NewTransaction(true)
+	defer txn.Discard()
+
+	if udf.ID == "" {
+		return datastore.UserDefinedFunction{}, datastore.BadRequest
+	}
+
+	var database datastore.Database
+	status := getKey(txn, generateDatabaseKey(databaseId), &database)
+	if status != datastore.StatusOk {
+		return datastore.UserDefinedFunction{}, status
+	}
+
+	var collection datastore.Collection
+	status = getKey(txn, generateCollectionKey(databaseId, collectionId), &collection)
+	if status != datastore.StatusOk {
+		return datastore.UserDefinedFunction{}, status
+	}
+
+	udf.TimeStamp = time.Now().Unix()
+	udf.ResourceID = resourceid.NewCombined(collection.ResourceID, resourceid.New(resourceid.ResourceTypeUserDefinedFunction))
+	udf.ETag = fmt.Sprintf("\"%s\"", uuid.New())
+	udf.Self = fmt.Sprintf("dbs/%s/colls/%s/udfs/%s/", database.ResourceID, collection.ResourceID, udf.ResourceID)
+
+	status = insertKey(txn, generateUserDefinedFunctionKey(databaseId, collectionId, udf.ID), udf)
+	if status != datastore.StatusOk {
+		return datastore.UserDefinedFunction{}, status
+	}
+
+	return udf, datastore.StatusOk
+}

internal/datastore/datastore.go

@@ -40,4 +40,5 @@ type DataStore interface {
 
 type DocumentIterator interface {
 	Next() (Document, DataStoreStatus)
+	Close()
 }

internal/datastore/map_datastore/array_document_iterator.go

@@ -15,3 +15,7 @@ func (i *ArrayDocumentIterator) Next() (datastore.Document, datastore.DataStoreS
 
 	return i.documents[i.index], datastore.StatusOk
 }
+
+func (i *ArrayDocumentIterator) Close() {
+	i.documents = []datastore.Document{}
+}

internal/datastore/models.go

@@ -16,6 +16,7 @@ const (
 	Conflict       = 3
 	BadRequest     = 4
 	IterEOF        = 5
+	Unknown        = 6
 )
 
 type TriggerOperation string