Fix formatting for grammar file

.github/workflows/compile-shared-libraries.yml

@@ -12,7 +12,7 @@ jobs:
         uses: actions/checkout@v3
 
       - name: Cross-Compile with xgo
-        uses: crazy-max/ghaction-xgo@v3.1.0
+        uses: crazy-max/ghaction-xgo@e22d3c8b089adba750d5a74738b8e95d96f0c991 # v3.1.0
         with:
           xgo_version: latest
           go_version: 1.24.0

.github/workflows/release.yml

@@ -24,7 +24,7 @@ jobs:
           go-version: 1.24.0
 
       - name: Cross-Compile with xgo
-        uses: crazy-max/ghaction-xgo@v3.1.0
+        uses: crazy-max/ghaction-xgo@e22d3c8b089adba750d5a74738b8e95d96f0c991 # v3.1.0
         with:
           xgo_version: latest
           go_version: 1.24.0
@@ -44,7 +44,7 @@ jobs:
           password: ${{ secrets.GITHUB_TOKEN }}
 
       - name: Run GoReleaser
-        uses: goreleaser/goreleaser-action@v5
+        uses: goreleaser/goreleaser-action@5742e2a039330cbb23ebf35f046f814d4c6ff811 # v5
         with:
           distribution: goreleaser
           version: ${{ env.GITHUB_REF_NAME }}

README.md

@@ -86,6 +86,7 @@ To disable SSL and run Cosmium on HTTP instead, you can use the `-DisableTls` fl
 - **-Persist**: Saves data to the given path on application exit (When `-InitialData` argument is not supplied, it will try to load data from path supplied in `-Persist`)
 - **-Port**: Listen port (default 8081)
 - **-LogLevel**: Sets the logging level (one of: debug, info, error, silent) (default info)
+- **-DataStore**: Allows selecting [storage backend](#data-storage-backends) (default "json")
 
 These arguments allow you to configure various aspects of Cosmium's behavior according to your requirements.
 
@@ -99,6 +100,18 @@ All mentioned arguments can also be set using environment variables:
 - **COSMIUM_PORT** for `-Port`
 - **COSMIUM_LOGLEVEL** for `-LogLevel`
 
+### Data Storage Backends
+
+Cosmium supports multiple storage backends for saving, loading, and managing data at runtime.
+
+| Backend  | Storage Location         | Write Behavior           | Memory Usage         | Supports Initial JSON Load |
+|----------|--------------------------|--------------------------|----------------------|----------------------------|
+| `json` (default) | JSON file on disk 📄 | On application exit ⏳ | 🛑 More than Badger | ✅ Yes                      |
+| `badger`  | BadgerDB database on disk ⚡ | Immediately on write 🚀 | ✅ Less than JSON  | ❌ No                       |
+
+
+The `badger` backend is generally recommended as it uses less memory and writes data to disk immediately. However, if you need to load initial data from a JSON file, use the `json` backend.
+
 # License
 
 This project is [MIT licensed](./LICENSE).

api/config/config.go

@@ -15,6 +15,11 @@ const (
 	ExplorerBaseUrlLocation = "/_explorer"
 )
 
+const (
+	DataStoreJson   = "json"
+	DataStoreBadger = "badger"
+)
+
 func ParseFlags() ServerConfig {
 	host := flag.String("Host", "localhost", "Hostname")
 	port := flag.Int("Port", 8081, "Listen port")
@@ -28,6 +33,8 @@ func ParseFlags() ServerConfig {
 	persistDataPath := flag.String("Persist", "", "Saves data to given path on application exit")
 	logLevel := NewEnumValue("info", []string{"debug", "info", "error", "silent"})
 	flag.Var(logLevel, "LogLevel", fmt.Sprintf("Sets the logging level %s", logLevel.AllowedValuesList()))
+	dataStore := NewEnumValue("json", []string{DataStoreJson, DataStoreBadger})
+	flag.Var(dataStore, "DataStore", fmt.Sprintf("Sets the data store %s", dataStore.AllowedValuesList()))
 
 	flag.Parse()
 	setFlagsFromEnvironment()
@@ -44,6 +51,7 @@ func ParseFlags() ServerConfig {
 	config.DisableTls = *disableTls
 	config.AccountKey = *accountKey
 	config.LogLevel = logLevel.value
+	config.DataStore = dataStore.value
 
 	config.PopulateCalculatedFields()
 
@@ -68,6 +76,29 @@ func (c *ServerConfig) PopulateCalculatedFields() {
 	default:
 		logger.SetLogLevel(logger.LogLevelInfo)
 	}
+
+	fileInfo, err := os.Stat(c.PersistDataFilePath)
+	if c.PersistDataFilePath != "" && !os.IsNotExist(err) {
+		if err != nil {
+			logger.ErrorLn("Failed to get file info for persist path:", err)
+			os.Exit(1)
+		}
+
+		if c.DataStore == DataStoreJson && fileInfo.IsDir() {
+			logger.ErrorLn("--Persist cannot be a directory when using json data store")
+			os.Exit(1)
+		}
+
+		if c.DataStore == DataStoreBadger && !fileInfo.IsDir() {
+			logger.ErrorLn("--Persist must be a directory when using Badger data store")
+			os.Exit(1)
+		}
+	}
+
+	if c.DataStore == DataStoreBadger && c.InitialDataFilePath != "" {
+		logger.ErrorLn("InitialData option is currently not supported with Badger data store")
+		os.Exit(1)
+	}
 }
 
 func (c *ServerConfig) ApplyDefaultsToEmptyFields() {

api/config/models.go

@@ -17,4 +17,6 @@ type ServerConfig struct {
 	DisableTls              bool   `json:"disableTls"`
 	LogLevel                string `json:"logLevel"`
 	ExplorerBaseUrlLocation string `json:"explorerBaseUrlLocation"`
+
+	DataStore string `json:"dataStore"`
 }

api/handlers/collections.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 
 	"github.com/gin-gonic/gin"
+	"github.com/pikami/cosmium/internal/constants"
 	"github.com/pikami/cosmium/internal/datastore"
 )
 
@@ -24,7 +25,7 @@ func (h *Handlers) GetAllCollections(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) GetCollection(c *gin.Context) {
@@ -38,11 +39,11 @@ func (h *Handlers) GetCollection(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) DeleteCollection(c *gin.Context) {
@@ -56,11 +57,11 @@ func (h *Handlers) DeleteCollection(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) CreateCollection(c *gin.Context) {
@@ -73,13 +74,13 @@ func (h *Handlers) CreateCollection(c *gin.Context) {
 	}
 
 	if newCollection.ID == "" {
-		c.JSON(http.StatusBadRequest, gin.H{"message": "BadRequest"})
+		c.JSON(http.StatusBadRequest, constants.BadRequestResponse)
 		return
 	}
 
 	createdCollection, status := h.dataStore.CreateCollection(databaseId, newCollection)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -88,5 +89,5 @@ func (h *Handlers) CreateCollection(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }

api/handlers/databases.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 
 	"github.com/gin-gonic/gin"
+	"github.com/pikami/cosmium/internal/constants"
 	"github.com/pikami/cosmium/internal/datastore"
 )
 
@@ -20,7 +21,7 @@ func (h *Handlers) GetAllDatabases(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) GetDatabase(c *gin.Context) {
@@ -33,11 +34,11 @@ func (h *Handlers) GetDatabase(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) DeleteDatabase(c *gin.Context) {
@@ -50,11 +51,11 @@ func (h *Handlers) DeleteDatabase(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) CreateDatabase(c *gin.Context) {
@@ -66,13 +67,13 @@ func (h *Handlers) CreateDatabase(c *gin.Context) {
 	}
 
 	if newDatabase.ID == "" {
-		c.JSON(http.StatusBadRequest, gin.H{"message": "BadRequest"})
+		c.JSON(http.StatusBadRequest, constants.BadRequestResponse)
 		return
 	}
 
 	createdDatabase, status := h.dataStore.CreateDatabase(newDatabase)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -81,5 +82,5 @@ func (h *Handlers) CreateDatabase(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }

api/handlers/documents.go

@@ -35,7 +35,7 @@ func (h *Handlers) GetAllDocuments(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) GetDocument(c *gin.Context) {
@@ -50,11 +50,11 @@ func (h *Handlers) GetDocument(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) DeleteDocument(c *gin.Context) {
@@ -69,11 +69,11 @@ func (h *Handlers) DeleteDocument(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 // TODO: Maybe move "replace" logic to data store
@@ -90,13 +90,13 @@ func (h *Handlers) ReplaceDocument(c *gin.Context) {
 
 	status := h.dataStore.DeleteDocument(databaseId, collectionId, documentId)
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
 	createdDocument, status := h.dataStore.CreateDocument(databaseId, collectionId, requestBody)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -105,7 +105,7 @@ func (h *Handlers) ReplaceDocument(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) PatchDocument(c *gin.Context) {
@@ -115,7 +115,7 @@ func (h *Handlers) PatchDocument(c *gin.Context) {
 
 	document, status := h.dataStore.GetDocument(databaseId, collectionId, documentId)
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
@@ -166,13 +166,13 @@ func (h *Handlers) PatchDocument(c *gin.Context) {
 
 	status = h.dataStore.DeleteDocument(databaseId, collectionId, documentId)
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
 	createdDocument, status := h.dataStore.CreateDocument(databaseId, collectionId, modifiedDocument)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -181,7 +181,7 @@ func (h *Handlers) PatchDocument(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) DocumentsPost(c *gin.Context) {
@@ -208,7 +208,7 @@ func (h *Handlers) DocumentsPost(c *gin.Context) {
 	}
 
 	if requestBody["id"] == "" {
-		c.JSON(http.StatusBadRequest, gin.H{"message": "BadRequest"})
+		c.JSON(http.StatusBadRequest, constants.BadRequestResponse)
 		return
 	}
 
@@ -219,7 +219,7 @@ func (h *Handlers) DocumentsPost(c *gin.Context) {
 
 	createdDocument, status := h.dataStore.CreateDocument(databaseId, collectionId, requestBody)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -228,7 +228,7 @@ func (h *Handlers) DocumentsPost(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func parametersToMap(pairs []interface{}) map[string]interface{} {
@@ -383,6 +383,7 @@ func (h *Handlers) executeQueryDocuments(databaseId string, collectionId string,
 	if status != datastore.StatusOk {
 		return nil, status
 	}
+	defer allDocumentsIterator.Close()
 
 	rowsIterator := converters.NewDocumentToRowTypeIterator(allDocumentsIterator)
 

api/handlers/partition_key_ranges.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 
 	"github.com/gin-gonic/gin"
+	"github.com/pikami/cosmium/internal/constants"
 	"github.com/pikami/cosmium/internal/datastore"
 	"github.com/pikami/cosmium/internal/resourceid"
 )
@@ -42,9 +43,9 @@ func (h *Handlers) GetPartitionKeyRanges(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }

api/handlers/stored_procedures.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 
 	"github.com/gin-gonic/gin"
+	"github.com/pikami/cosmium/internal/constants"
 	"github.com/pikami/cosmium/internal/datastore"
 )
 
@@ -20,7 +21,7 @@ func (h *Handlers) GetAllStoredProcedures(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) GetStoredProcedure(c *gin.Context) {
@@ -36,11 +37,11 @@ func (h *Handlers) GetStoredProcedure(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) DeleteStoredProcedure(c *gin.Context) {
@@ -55,11 +56,11 @@ func (h *Handlers) DeleteStoredProcedure(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) ReplaceStoredProcedure(c *gin.Context) {
@@ -69,19 +70,19 @@ func (h *Handlers) ReplaceStoredProcedure(c *gin.Context) {
 
 	var sp datastore.StoredProcedure
 	if err := c.BindJSON(&sp); err != nil {
-		c.IndentedJSON(http.StatusBadRequest, gin.H{"message": "Invalid body"})
+		c.IndentedJSON(http.StatusBadRequest, constants.BadRequestResponse)
 		return
 	}
 
 	status := h.dataStore.DeleteStoredProcedure(databaseId, collectionId, spId)
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
 	createdSP, status := h.dataStore.CreateStoredProcedure(databaseId, collectionId, sp)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -90,7 +91,7 @@ func (h *Handlers) ReplaceStoredProcedure(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) CreateStoredProcedure(c *gin.Context) {
@@ -99,13 +100,13 @@ func (h *Handlers) CreateStoredProcedure(c *gin.Context) {
 
 	var sp datastore.StoredProcedure
 	if err := c.BindJSON(&sp); err != nil {
-		c.IndentedJSON(http.StatusBadRequest, gin.H{"message": "Invalid body"})
+		c.IndentedJSON(http.StatusBadRequest, constants.BadRequestResponse)
 		return
 	}
 
 	createdSP, status := h.dataStore.CreateStoredProcedure(databaseId, collectionId, sp)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -114,5 +115,5 @@ func (h *Handlers) CreateStoredProcedure(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }

api/handlers/triggers.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 
 	"github.com/gin-gonic/gin"
+	"github.com/pikami/cosmium/internal/constants"
 	"github.com/pikami/cosmium/internal/datastore"
 )
 
@@ -20,7 +21,7 @@ func (h *Handlers) GetAllTriggers(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) GetTrigger(c *gin.Context) {
@@ -36,11 +37,11 @@ func (h *Handlers) GetTrigger(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) DeleteTrigger(c *gin.Context) {
@@ -55,11 +56,11 @@ func (h *Handlers) DeleteTrigger(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) ReplaceTrigger(c *gin.Context) {
@@ -69,19 +70,19 @@ func (h *Handlers) ReplaceTrigger(c *gin.Context) {
 
 	var trigger datastore.Trigger
 	if err := c.BindJSON(&trigger); err != nil {
-		c.IndentedJSON(http.StatusBadRequest, gin.H{"message": "Invalid body"})
+		c.IndentedJSON(http.StatusBadRequest, constants.BadRequestResponse)
 		return
 	}
 
 	status := h.dataStore.DeleteTrigger(databaseId, collectionId, triggerId)
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
 	createdTrigger, status := h.dataStore.CreateTrigger(databaseId, collectionId, trigger)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -90,7 +91,7 @@ func (h *Handlers) ReplaceTrigger(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) CreateTrigger(c *gin.Context) {
@@ -99,13 +100,13 @@ func (h *Handlers) CreateTrigger(c *gin.Context) {
 
 	var trigger datastore.Trigger
 	if err := c.BindJSON(&trigger); err != nil {
-		c.IndentedJSON(http.StatusBadRequest, gin.H{"message": "Invalid body"})
+		c.IndentedJSON(http.StatusBadRequest, constants.BadRequestResponse)
 		return
 	}
 
 	createdTrigger, status := h.dataStore.CreateTrigger(databaseId, collectionId, trigger)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -114,5 +115,5 @@ func (h *Handlers) CreateTrigger(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }

api/handlers/user_defined_functions.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 
 	"github.com/gin-gonic/gin"
+	"github.com/pikami/cosmium/internal/constants"
 	"github.com/pikami/cosmium/internal/datastore"
 )
 
@@ -20,7 +21,7 @@ func (h *Handlers) GetAllUserDefinedFunctions(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) GetUserDefinedFunction(c *gin.Context) {
@@ -36,11 +37,11 @@ func (h *Handlers) GetUserDefinedFunction(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) DeleteUserDefinedFunction(c *gin.Context) {
@@ -55,11 +56,11 @@ func (h *Handlers) DeleteUserDefinedFunction(c *gin.Context) {
 	}
 
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) ReplaceUserDefinedFunction(c *gin.Context) {
@@ -69,19 +70,19 @@ func (h *Handlers) ReplaceUserDefinedFunction(c *gin.Context) {
 
 	var udf datastore.UserDefinedFunction
 	if err := c.BindJSON(&udf); err != nil {
-		c.IndentedJSON(http.StatusBadRequest, gin.H{"message": "Invalid body"})
+		c.IndentedJSON(http.StatusBadRequest, constants.BadRequestResponse)
 		return
 	}
 
 	status := h.dataStore.DeleteUserDefinedFunction(databaseId, collectionId, udfId)
 	if status == datastore.StatusNotFound {
-		c.IndentedJSON(http.StatusNotFound, gin.H{"message": "NotFound"})
+		c.IndentedJSON(http.StatusNotFound, constants.NotFoundResponse)
 		return
 	}
 
 	createdUdf, status := h.dataStore.CreateUserDefinedFunction(databaseId, collectionId, udf)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -90,7 +91,7 @@ func (h *Handlers) ReplaceUserDefinedFunction(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }
 
 func (h *Handlers) CreateUserDefinedFunction(c *gin.Context) {
@@ -99,13 +100,13 @@ func (h *Handlers) CreateUserDefinedFunction(c *gin.Context) {
 
 	var udf datastore.UserDefinedFunction
 	if err := c.BindJSON(&udf); err != nil {
-		c.IndentedJSON(http.StatusBadRequest, gin.H{"message": "Invalid body"})
+		c.IndentedJSON(http.StatusBadRequest, constants.BadRequestResponse)
 		return
 	}
 
 	createdUdf, status := h.dataStore.CreateUserDefinedFunction(databaseId, collectionId, udf)
 	if status == datastore.Conflict {
-		c.IndentedJSON(http.StatusConflict, gin.H{"message": "Conflict"})
+		c.IndentedJSON(http.StatusConflict, constants.ConflictResponse)
 		return
 	}
 
@@ -114,5 +115,5 @@ func (h *Handlers) CreateUserDefinedFunction(c *gin.Context) {
 		return
 	}
 
-	c.IndentedJSON(http.StatusInternalServerError, gin.H{"message": "Unknown error"})
+	c.IndentedJSON(http.StatusInternalServerError, constants.UnknownErrorResponse)
 }

api/tests/authentication_test.go

@@ -19,7 +19,7 @@ func Test_Authentication(t *testing.T) {
 	t.Run("Should get 200 when correct account key is used", func(t *testing.T) {
 		ts.DataStore.DeleteDatabase(testDatabaseName)
 		client, err := azcosmos.NewClientFromConnectionString(
-			fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, config.DefaultAccountKey),
+			formatConnectionString(ts.URL, config.DefaultAccountKey),
 			&azcosmos.ClientOptions{},
 		)
 		assert.Nil(t, err)
@@ -35,7 +35,7 @@ func Test_Authentication(t *testing.T) {
 	t.Run("Should get 401 when wrong account key is used", func(t *testing.T) {
 		ts.DataStore.DeleteDatabase(testDatabaseName)
 		client, err := azcosmos.NewClientFromConnectionString(
-			fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, "AAAA"),
+			formatConnectionString(ts.URL, "AAAA"),
 			&azcosmos.ClientOptions{},
 		)
 		assert.Nil(t, err)
@@ -72,7 +72,7 @@ func Test_Authentication_Disabled(t *testing.T) {
 	t.Run("Should get 200 when wrong account key is used, but authentication is dissabled", func(t *testing.T) {
 		ts.DataStore.DeleteDatabase(testDatabaseName)
 		client, err := azcosmos.NewClientFromConnectionString(
-			fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, "AAAA"),
+			formatConnectionString(ts.URL, "AAAA"),
 			&azcosmos.ClientOptions{},
 		)
 		assert.Nil(t, err)
@@ -85,3 +85,7 @@ func Test_Authentication_Disabled(t *testing.T) {
 		assert.Equal(t, createResponse.DatabaseProperties.ID, testDatabaseName)
 	})
 }
+
+func formatConnectionString(endpoint, key string) string {
+	return fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", endpoint, key)
+}

api/tests/collections_test.go

@@ -3,32 +3,29 @@ package tests_test
 import (
 	"context"
 	"errors"
-	"fmt"
 	"net/http"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
-	"github.com/pikami/cosmium/api/config"
 	"github.com/pikami/cosmium/internal/datastore"
 	"github.com/stretchr/testify/assert"
 )
 
 func Test_Collections(t *testing.T) {
-	ts := runTestServer()
-	defer ts.Server.Close()
+	presets := []testPreset{PresetJsonStore, PresetBadgerStore}
 
-	client, err := azcosmos.NewClientFromConnectionString(
-		fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, config.DefaultAccountKey),
-		&azcosmos.ClientOptions{},
-	)
-	assert.Nil(t, err)
+	setUp := func(ts *TestServer, client *azcosmos.Client) *azcosmos.DatabaseClient {
+		ts.DataStore.CreateDatabase(datastore.Database{ID: testDatabaseName})
+		databaseClient, err := client.NewDatabase(testDatabaseName)
+		assert.Nil(t, err)
 
-	ts.DataStore.CreateDatabase(datastore.Database{ID: testDatabaseName})
-	databaseClient, err := client.NewDatabase(testDatabaseName)
-	assert.Nil(t, err)
+		return databaseClient
+	}
+
+	runTestsWithPresets(t, "Collection Create", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
+		databaseClient := setUp(ts, client)
 
-	t.Run("Collection Create", func(t *testing.T) {
 		t.Run("Should create collection", func(t *testing.T) {
 			createResponse, err := databaseClient.CreateContainer(context.TODO(), azcosmos.ContainerProperties{
 				ID: testCollectionName,
@@ -57,7 +54,9 @@ func Test_Collections(t *testing.T) {
 		})
 	})
 
-	t.Run("Collection Read", func(t *testing.T) {
+	runTestsWithPresets(t, "Collection Read", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
+		databaseClient := setUp(ts, client)
+
 		t.Run("Should read collection", func(t *testing.T) {
 			ts.DataStore.CreateCollection(testDatabaseName, datastore.Collection{
 				ID: testCollectionName,
@@ -90,7 +89,9 @@ func Test_Collections(t *testing.T) {
 		})
 	})
 
-	t.Run("Collection Delete", func(t *testing.T) {
+	runTestsWithPresets(t, "Collection Delete", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
+		databaseClient := setUp(ts, client)
+
 		t.Run("Should delete collection", func(t *testing.T) {
 			ts.DataStore.CreateCollection(testDatabaseName, datastore.Collection{
 				ID: testCollectionName,

api/tests/config_test.go

@@ -1,13 +1,18 @@
 package tests_test
 
 import (
+	"fmt"
 	"net/http/httptest"
+	"testing"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
 	"github.com/pikami/cosmium/api"
 	"github.com/pikami/cosmium/api/config"
 	"github.com/pikami/cosmium/internal/datastore"
-	mapdatastore "github.com/pikami/cosmium/internal/datastore/map_datastore"
+	badgerdatastore "github.com/pikami/cosmium/internal/datastore/badger_datastore"
+	jsondatastore "github.com/pikami/cosmium/internal/datastore/json_datastore"
 	"github.com/pikami/cosmium/internal/logger"
+	"github.com/stretchr/testify/assert"
 )
 
 type TestServer struct {
@@ -16,14 +21,29 @@ type TestServer struct {
 	URL       string
 }
 
-func runTestServerCustomConfig(config *config.ServerConfig) *TestServer {
-	dataStore := mapdatastore.NewMapDataStore(mapdatastore.MapDataStoreOptions{})
+func getDefaultTestServerConfig() *config.ServerConfig {
+	return &config.ServerConfig{
+		AccountKey:              config.DefaultAccountKey,
+		ExplorerPath:            "/tmp/nothing",
+		ExplorerBaseUrlLocation: config.ExplorerBaseUrlLocation,
+		DataStore:               "json",
+	}
+}
 
-	api := api.NewApiServer(dataStore, config)
+func runTestServerCustomConfig(configuration *config.ServerConfig) *TestServer {
+	var dataStore datastore.DataStore
+	switch configuration.DataStore {
+	case config.DataStoreBadger:
+		dataStore = badgerdatastore.NewBadgerDataStore(badgerdatastore.BadgerDataStoreOptions{})
+	default:
+		dataStore = jsondatastore.NewJsonDataStore(jsondatastore.JsonDataStoreOptions{})
+	}
+
+	api := api.NewApiServer(dataStore, configuration)
 
 	server := httptest.NewServer(api.GetRouter())
 
-	config.DatabaseEndpoint = server.URL
+	configuration.DatabaseEndpoint = server.URL
 
 	return &TestServer{
 		Server:    server,
@@ -33,11 +53,7 @@ func runTestServerCustomConfig(config *config.ServerConfig) *TestServer {
 }
 
 func runTestServer() *TestServer {
-	config := &config.ServerConfig{
-		AccountKey:              config.DefaultAccountKey,
-		ExplorerPath:            "/tmp/nothing",
-		ExplorerBaseUrlLocation: config.ExplorerBaseUrlLocation,
-	}
+	config := getDefaultTestServerConfig()
 
 	config.LogLevel = "debug"
 	logger.SetLogLevel(logger.LogLevelDebug)
@@ -50,3 +66,47 @@ const (
 	testDatabaseName   = "test-db"
 	testCollectionName = "test-coll"
 )
+
+type testFunc func(t *testing.T, ts *TestServer, cosmosClient *azcosmos.Client)
+type testPreset string
+
+const (
+	PresetJsonStore   testPreset = "JsonDS"
+	PresetBadgerStore testPreset = "BadgerDS"
+)
+
+func runTestsWithPreset(t *testing.T, name string, testPreset testPreset, f testFunc) {
+	serverConfig := getDefaultTestServerConfig()
+
+	serverConfig.LogLevel = "debug"
+	logger.SetLogLevel(logger.LogLevelDebug)
+
+	switch testPreset {
+	case PresetBadgerStore:
+		serverConfig.DataStore = config.DataStoreBadger
+	case PresetJsonStore:
+		serverConfig.DataStore = config.DataStoreJson
+	}
+
+	ts := runTestServerCustomConfig(serverConfig)
+	defer ts.Server.Close()
+	defer ts.DataStore.Close()
+
+	client, err := azcosmos.NewClientFromConnectionString(
+		fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, config.DefaultAccountKey),
+		&azcosmos.ClientOptions{},
+	)
+	assert.Nil(t, err)
+
+	testName := fmt.Sprintf("%s_%s", testPreset, name)
+
+	t.Run(testName, func(t *testing.T) {
+		f(t, ts, client)
+	})
+}
+
+func runTestsWithPresets(t *testing.T, name string, testPresets []testPreset, f testFunc) {
+	for _, testPreset := range testPresets {
+		runTestsWithPreset(t, name, testPreset, f)
+	}
+}

api/tests/databases_test.go

@@ -3,28 +3,19 @@ package tests_test
 import (
 	"context"
 	"errors"
-	"fmt"
 	"net/http"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
-	"github.com/pikami/cosmium/api/config"
 	"github.com/pikami/cosmium/internal/datastore"
 	"github.com/stretchr/testify/assert"
 )
 
 func Test_Databases(t *testing.T) {
-	ts := runTestServer()
-	defer ts.Server.Close()
+	presets := []testPreset{PresetJsonStore, PresetBadgerStore}
 
-	client, err := azcosmos.NewClientFromConnectionString(
-		fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, config.DefaultAccountKey),
-		&azcosmos.ClientOptions{},
-	)
-	assert.Nil(t, err)
-
-	t.Run("Database Create", func(t *testing.T) {
+	runTestsWithPresets(t, "Database Create", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
 		t.Run("Should create database", func(t *testing.T) {
 			ts.DataStore.DeleteDatabase(testDatabaseName)
 
@@ -55,7 +46,7 @@ func Test_Databases(t *testing.T) {
 		})
 	})
 
-	t.Run("Database Read", func(t *testing.T) {
+	runTestsWithPresets(t, "Database Read", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
 		t.Run("Should read database", func(t *testing.T) {
 			ts.DataStore.CreateDatabase(datastore.Database{
 				ID: testDatabaseName,
@@ -88,7 +79,7 @@ func Test_Databases(t *testing.T) {
 		})
 	})
 
-	t.Run("Database Delete", func(t *testing.T) {
+	runTestsWithPresets(t, "Database Delete", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
 		t.Run("Should delete database", func(t *testing.T) {
 			ts.DataStore.CreateDatabase(datastore.Database{
 				ID: testDatabaseName,

api/tests/documents_test.go

@@ -53,9 +53,7 @@ func testCosmosQuery(t *testing.T,
 	}
 }
 
-func documents_InitializeDb(t *testing.T) (*TestServer, *azcosmos.ContainerClient) {
-	ts := runTestServer()
-
+func documents_InitializeDb(t *testing.T, ts *TestServer) *azcosmos.ContainerClient {
 	ts.DataStore.CreateDatabase(datastore.Database{ID: testDatabaseName})
 	ts.DataStore.CreateCollection(testDatabaseName, datastore.Collection{
 		ID: testCollectionName,
@@ -79,438 +77,439 @@ func documents_InitializeDb(t *testing.T) (*TestServer, *azcosmos.ContainerClien
 	collectionClient, err := client.NewContainer(testDatabaseName, testCollectionName)
 	assert.Nil(t, err)
 
-	return ts, collectionClient
+	return collectionClient
 }
 
 func Test_Documents(t *testing.T) {
-	ts, collectionClient := documents_InitializeDb(t)
-	defer ts.Server.Close()
+	presets := []testPreset{PresetJsonStore, PresetBadgerStore}
 
-	t.Run("Should query document", func(t *testing.T) {
-		testCosmosQuery(t, collectionClient,
-			"SELECT c.id, c[\"pk\"] FROM c ORDER BY c.id",
-			nil,
-			[]interface{}{
-				map[string]interface{}{"id": "12345", "pk": "123"},
-				map[string]interface{}{"id": "67890", "pk": "456"},
-			},
-		)
+	runTestsWithPresets(t, "Test_Documents", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
+		collectionClient := documents_InitializeDb(t, ts)
+
+		t.Run("Should query document", func(t *testing.T) {
+			testCosmosQuery(t, collectionClient,
+				"SELECT c.id, c[\"pk\"] FROM c ORDER BY c.id",
+				nil,
+				[]interface{}{
+					map[string]interface{}{"id": "12345", "pk": "123"},
+					map[string]interface{}{"id": "67890", "pk": "456"},
+				},
+			)
+		})
+
+		t.Run("Should query VALUE array", func(t *testing.T) {
+			testCosmosQuery(t, collectionClient,
+				"SELECT VALUE [c.id, c[\"pk\"]] FROM c ORDER BY c.id",
+				nil,
+				[]interface{}{
+					[]interface{}{"12345", "123"},
+					[]interface{}{"67890", "456"},
+				},
+			)
+		})
+
+		t.Run("Should query VALUE object", func(t *testing.T) {
+			testCosmosQuery(t, collectionClient,
+				"SELECT VALUE { id: c.id, _pk: c.pk } FROM c ORDER BY c.id",
+				nil,
+				[]interface{}{
+					map[string]interface{}{"id": "12345", "_pk": "123"},
+					map[string]interface{}{"id": "67890", "_pk": "456"},
+				},
+			)
+		})
+
+		t.Run("Should query document with single WHERE condition", func(t *testing.T) {
+			testCosmosQuery(t, collectionClient,
+				`select c.id
+				FROM c
+				WHERE c.isCool=true
+				ORDER BY c.id`,
+				nil,
+				[]interface{}{
+					map[string]interface{}{"id": "67890"},
+				},
+			)
+		})
+
+		t.Run("Should query document with query parameters", func(t *testing.T) {
+			testCosmosQuery(t, collectionClient,
+				`select c.id
+				FROM c
+				WHERE c.id=@param_id
+				ORDER BY c.id`,
+				[]azcosmos.QueryParameter{
+					{Name: "@param_id", Value: "67890"},
+				},
+				[]interface{}{
+					map[string]interface{}{"id": "67890"},
+				},
+			)
+		})
+
+		t.Run("Should query document with query parameters as accessor", func(t *testing.T) {
+			testCosmosQuery(t, collectionClient,
+				`select c.id
+				FROM c
+				WHERE c[@param]="67890"
+				ORDER BY c.id`,
+				[]azcosmos.QueryParameter{
+					{Name: "@param", Value: "id"},
+				},
+				[]interface{}{
+					map[string]interface{}{"id": "67890"},
+				},
+			)
+		})
+
+		t.Run("Should query array accessor", func(t *testing.T) {
+			testCosmosQuery(t, collectionClient,
+				`SELECT c.id,
+					c["arr"][0] AS arr0,
+					c["arr"][1] AS arr1,
+					c["arr"][2] AS arr2,
+					c["arr"][3] AS arr3
+				FROM c ORDER BY c.id`,
+				nil,
+				[]interface{}{
+					map[string]interface{}{"id": "12345", "arr0": 1.0, "arr1": 2.0, "arr2": 3.0, "arr3": nil},
+					map[string]interface{}{"id": "67890", "arr0": 6.0, "arr1": 7.0, "arr2": 8.0, "arr3": nil},
+				},
+			)
+		})
+
+		t.Run("Should handle parallel writes", func(t *testing.T) {
+			var wg sync.WaitGroup
+			rutineCount := 100
+			results := make(chan error, rutineCount)
+
+			createCall := func(i int) {
+				defer wg.Done()
+				item := map[string]interface{}{
+					"id":  fmt.Sprintf("id-%d", i),
+					"pk":  fmt.Sprintf("pk-%d", i),
+					"val": i,
+				}
+				bytes, err := json.Marshal(item)
+				if err != nil {
+					results <- err
+					return
+				}
+
+				ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+				defer cancel()
+
+				_, err = collectionClient.CreateItem(
+					ctx,
+					azcosmos.PartitionKey{},
+					bytes,
+					&azcosmos.ItemOptions{
+						EnableContentResponseOnWrite: false,
+					},
+				)
+				results <- err
+
+				collectionClient.ReadItem(ctx, azcosmos.PartitionKey{}, fmt.Sprintf("id-%d", i), nil)
+				collectionClient.DeleteItem(ctx, azcosmos.PartitionKey{}, fmt.Sprintf("id-%d", i), nil)
+			}
+
+			for i := 0; i < rutineCount; i++ {
+				wg.Add(1)
+				go createCall(i)
+			}
+
+			wg.Wait()
+			close(results)
+
+			for err := range results {
+				if err != nil {
+					t.Errorf("Error creating item: %v", err)
+				}
+			}
+		})
 	})
 
-	t.Run("Should query VALUE array", func(t *testing.T) {
-		testCosmosQuery(t, collectionClient,
-			"SELECT VALUE [c.id, c[\"pk\"]] FROM c ORDER BY c.id",
-			nil,
-			[]interface{}{
-				[]interface{}{"12345", "123"},
-				[]interface{}{"67890", "456"},
-			},
-		)
-	})
+	runTestsWithPresets(t, "Test_Documents_Patch", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
+		collectionClient := documents_InitializeDb(t, ts)
 
-	t.Run("Should query VALUE object", func(t *testing.T) {
-		testCosmosQuery(t, collectionClient,
-			"SELECT VALUE { id: c.id, _pk: c.pk } FROM c ORDER BY c.id",
-			nil,
-			[]interface{}{
-				map[string]interface{}{"id": "12345", "_pk": "123"},
-				map[string]interface{}{"id": "67890", "_pk": "456"},
-			},
-		)
-	})
+		t.Run("Should PATCH document", func(t *testing.T) {
+			context := context.TODO()
+			expectedData := map[string]interface{}{"id": "67890", "pk": "666", "newField": "newValue", "incr": 15., "setted": "isSet"}
 
-	t.Run("Should query document with single WHERE condition", func(t *testing.T) {
-		testCosmosQuery(t, collectionClient,
-			`select c.id
-			FROM c
-			WHERE c.isCool=true
-			ORDER BY c.id`,
-			nil,
-			[]interface{}{
-				map[string]interface{}{"id": "67890"},
-			},
-		)
-	})
+			patch := azcosmos.PatchOperations{}
+			patch.AppendAdd("/newField", "newValue")
+			patch.AppendIncrement("/incr", 15)
+			patch.AppendRemove("/isCool")
+			patch.AppendReplace("/pk", "666")
+			patch.AppendSet("/setted", "isSet")
 
-	t.Run("Should query document with query parameters", func(t *testing.T) {
-		testCosmosQuery(t, collectionClient,
-			`select c.id
-			FROM c
-			WHERE c.id=@param_id
-			ORDER BY c.id`,
-			[]azcosmos.QueryParameter{
-				{Name: "@param_id", Value: "67890"},
-			},
-			[]interface{}{
-				map[string]interface{}{"id": "67890"},
-			},
-		)
-	})
+			itemResponse, err := collectionClient.PatchItem(
+				context,
+				azcosmos.PartitionKey{},
+				"67890",
+				patch,
+				&azcosmos.ItemOptions{
+					EnableContentResponseOnWrite: false,
+				},
+			)
+			assert.Nil(t, err)
 
-	t.Run("Should query document with query parameters as accessor", func(t *testing.T) {
-		testCosmosQuery(t, collectionClient,
-			`select c.id
-			FROM c
-			WHERE c[@param]="67890"
-			ORDER BY c.id`,
-			[]azcosmos.QueryParameter{
-				{Name: "@param", Value: "id"},
-			},
-			[]interface{}{
-				map[string]interface{}{"id": "67890"},
-			},
-		)
-	})
+			var itemResponseBody map[string]interface{}
+			json.Unmarshal(itemResponse.Value, &itemResponseBody)
 
-	t.Run("Should query array accessor", func(t *testing.T) {
-		testCosmosQuery(t, collectionClient,
-			`SELECT c.id,
-				c["arr"][0] AS arr0,
-				c["arr"][1] AS arr1,
-				c["arr"][2] AS arr2,
-				c["arr"][3] AS arr3
-			FROM c ORDER BY c.id`,
-			nil,
-			[]interface{}{
-				map[string]interface{}{"id": "12345", "arr0": 1.0, "arr1": 2.0, "arr2": 3.0, "arr3": nil},
-				map[string]interface{}{"id": "67890", "arr0": 6.0, "arr1": 7.0, "arr2": 8.0, "arr3": nil},
-			},
-		)
-	})
+			assert.Equal(t, expectedData["id"], itemResponseBody["id"])
+			assert.Equal(t, expectedData["pk"], itemResponseBody["pk"])
+			assert.Empty(t, itemResponseBody["isCool"])
+			assert.Equal(t, expectedData["newField"], itemResponseBody["newField"])
+			assert.Equal(t, expectedData["incr"], itemResponseBody["incr"])
+			assert.Equal(t, expectedData["setted"], itemResponseBody["setted"])
+		})
 
-	t.Run("Should handle parallel writes", func(t *testing.T) {
-		var wg sync.WaitGroup
-		rutineCount := 100
-		results := make(chan error, rutineCount)
+		t.Run("Should not allow to PATCH document ID", func(t *testing.T) {
+			context := context.TODO()
+
+			patch := azcosmos.PatchOperations{}
+			patch.AppendReplace("/id", "newValue")
+
+			_, err := collectionClient.PatchItem(
+				context,
+				azcosmos.PartitionKey{},
+				"67890",
+				patch,
+				&azcosmos.ItemOptions{
+					EnableContentResponseOnWrite: false,
+				},
+			)
+			assert.NotNil(t, err)
+
+			var respErr *azcore.ResponseError
+			if errors.As(err, &respErr) {
+				assert.Equal(t, http.StatusUnprocessableEntity, respErr.StatusCode)
+			} else {
+				panic(err)
+			}
+		})
+
+		t.Run("CreateItem", func(t *testing.T) {
+			context := context.TODO()
 
-		createCall := func(i int) {
-			defer wg.Done()
 			item := map[string]interface{}{
-				"id":  fmt.Sprintf("id-%d", i),
-				"pk":  fmt.Sprintf("pk-%d", i),
-				"val": i,
+				"Id":       "6789011",
+				"pk":       "456",
+				"newField": "newValue2",
 			}
 			bytes, err := json.Marshal(item)
-			if err != nil {
-				results <- err
-				return
-			}
+			assert.Nil(t, err)
 
-			ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
-			defer cancel()
-
-			_, err = collectionClient.CreateItem(
-				ctx,
+			r, err2 := collectionClient.CreateItem(
+				context,
 				azcosmos.PartitionKey{},
 				bytes,
 				&azcosmos.ItemOptions{
 					EnableContentResponseOnWrite: false,
 				},
 			)
-			results <- err
+			assert.NotNil(t, r)
+			assert.Nil(t, err2)
+		})
 
-			collectionClient.ReadItem(ctx, azcosmos.PartitionKey{}, fmt.Sprintf("id-%d", i), nil)
-			collectionClient.DeleteItem(ctx, azcosmos.PartitionKey{}, fmt.Sprintf("id-%d", i), nil)
-		}
+		t.Run("CreateItem that already exists", func(t *testing.T) {
+			context := context.TODO()
 
-		for i := 0; i < rutineCount; i++ {
-			wg.Add(1)
-			go createCall(i)
-		}
+			item := map[string]interface{}{"id": "12345", "pk": "123", "isCool": false, "arr": []int{1, 2, 3}}
+			bytes, err := json.Marshal(item)
+			assert.Nil(t, err)
 
-		wg.Wait()
-		close(results)
+			r, err := collectionClient.CreateItem(
+				context,
+				azcosmos.PartitionKey{},
+				bytes,
+				&azcosmos.ItemOptions{
+					EnableContentResponseOnWrite: false,
+				},
+			)
+			assert.NotNil(t, r)
+			assert.NotNil(t, err)
 
-		for err := range results {
-			if err != nil {
-				t.Errorf("Error creating item: %v", err)
+			var respErr *azcore.ResponseError
+			if errors.As(err, &respErr) {
+				assert.Equal(t, http.StatusConflict, respErr.StatusCode)
+			} else {
+				panic(err)
 			}
-		}
-	})
-}
-
-func Test_Documents_Patch(t *testing.T) {
-	ts, collectionClient := documents_InitializeDb(t)
-	defer ts.Server.Close()
-
-	t.Run("Should PATCH document", func(t *testing.T) {
-		context := context.TODO()
-		expectedData := map[string]interface{}{"id": "67890", "pk": "666", "newField": "newValue", "incr": 15., "setted": "isSet"}
-
-		patch := azcosmos.PatchOperations{}
-		patch.AppendAdd("/newField", "newValue")
-		patch.AppendIncrement("/incr", 15)
-		patch.AppendRemove("/isCool")
-		patch.AppendReplace("/pk", "666")
-		patch.AppendSet("/setted", "isSet")
-
-		itemResponse, err := collectionClient.PatchItem(
-			context,
-			azcosmos.PartitionKey{},
-			"67890",
-			patch,
-			&azcosmos.ItemOptions{
-				EnableContentResponseOnWrite: false,
-			},
-		)
-		assert.Nil(t, err)
-
-		var itemResponseBody map[string]interface{}
-		json.Unmarshal(itemResponse.Value, &itemResponseBody)
-
-		assert.Equal(t, expectedData["id"], itemResponseBody["id"])
-		assert.Equal(t, expectedData["pk"], itemResponseBody["pk"])
-		assert.Empty(t, itemResponseBody["isCool"])
-		assert.Equal(t, expectedData["newField"], itemResponseBody["newField"])
-		assert.Equal(t, expectedData["incr"], itemResponseBody["incr"])
-		assert.Equal(t, expectedData["setted"], itemResponseBody["setted"])
-	})
-
-	t.Run("Should not allow to PATCH document ID", func(t *testing.T) {
-		context := context.TODO()
-
-		patch := azcosmos.PatchOperations{}
-		patch.AppendReplace("/id", "newValue")
-
-		_, err := collectionClient.PatchItem(
-			context,
-			azcosmos.PartitionKey{},
-			"67890",
-			patch,
-			&azcosmos.ItemOptions{
-				EnableContentResponseOnWrite: false,
-			},
-		)
-		assert.NotNil(t, err)
-
-		var respErr *azcore.ResponseError
-		if errors.As(err, &respErr) {
-			assert.Equal(t, http.StatusUnprocessableEntity, respErr.StatusCode)
-		} else {
-			panic(err)
-		}
-	})
-
-	t.Run("CreateItem", func(t *testing.T) {
-		context := context.TODO()
-
-		item := map[string]interface{}{
-			"Id":       "6789011",
-			"pk":       "456",
-			"newField": "newValue2",
-		}
-		bytes, err := json.Marshal(item)
-		assert.Nil(t, err)
-
-		r, err2 := collectionClient.CreateItem(
-			context,
-			azcosmos.PartitionKey{},
-			bytes,
-			&azcosmos.ItemOptions{
-				EnableContentResponseOnWrite: false,
-			},
-		)
-		assert.NotNil(t, r)
-		assert.Nil(t, err2)
-	})
-
-	t.Run("CreateItem that already exists", func(t *testing.T) {
-		context := context.TODO()
-
-		item := map[string]interface{}{"id": "12345", "pk": "123", "isCool": false, "arr": []int{1, 2, 3}}
-		bytes, err := json.Marshal(item)
-		assert.Nil(t, err)
-
-		r, err := collectionClient.CreateItem(
-			context,
-			azcosmos.PartitionKey{},
-			bytes,
-			&azcosmos.ItemOptions{
-				EnableContentResponseOnWrite: false,
-			},
-		)
-		assert.NotNil(t, r)
-		assert.NotNil(t, err)
-
-		var respErr *azcore.ResponseError
-		if errors.As(err, &respErr) {
-			assert.Equal(t, http.StatusConflict, respErr.StatusCode)
-		} else {
-			panic(err)
-		}
-	})
-
-	t.Run("UpsertItem new", func(t *testing.T) {
-		context := context.TODO()
-
-		item := map[string]interface{}{"id": "123456", "pk": "1234", "isCool": false, "arr": []int{1, 2, 3}}
-		bytes, err := json.Marshal(item)
-		assert.Nil(t, err)
-
-		r, err2 := collectionClient.UpsertItem(
-			context,
-			azcosmos.PartitionKey{},
-			bytes,
-			&azcosmos.ItemOptions{
-				EnableContentResponseOnWrite: false,
-			},
-		)
-		assert.NotNil(t, r)
-		assert.Nil(t, err2)
-	})
-
-	t.Run("UpsertItem that already exists", func(t *testing.T) {
-		context := context.TODO()
-
-		item := map[string]interface{}{"id": "12345", "pk": "123", "isCool": false, "arr": []int{1, 2, 3, 4}}
-		bytes, err := json.Marshal(item)
-		assert.Nil(t, err)
-
-		r, err2 := collectionClient.UpsertItem(
-			context,
-			azcosmos.PartitionKey{},
-			bytes,
-			&azcosmos.ItemOptions{
-				EnableContentResponseOnWrite: false,
-			},
-		)
-		assert.NotNil(t, r)
-		assert.Nil(t, err2)
-	})
-}
-
-func Test_Documents_TransactionalBatch(t *testing.T) {
-	ts, collectionClient := documents_InitializeDb(t)
-	defer ts.Server.Close()
-
-	t.Run("Should execute CREATE transactional batch", func(t *testing.T) {
-		context := context.TODO()
-		batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
-
-		newItem := map[string]interface{}{
-			"id": "678901",
-		}
-		bytes, err := json.Marshal(newItem)
-		assert.Nil(t, err)
-
-		batch.CreateItem(bytes, nil)
-		response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
-		assert.Nil(t, err)
-		assert.True(t, response.Success)
-		assert.Equal(t, 1, len(response.OperationResults))
-
-		operationResponse := response.OperationResults[0]
-		assert.NotNil(t, operationResponse)
-		assert.NotNil(t, operationResponse.ResourceBody)
-		assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
-
-		var itemResponseBody map[string]interface{}
-		json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
-		assert.Equal(t, newItem["id"], itemResponseBody["id"])
-
-		createdDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
-		assert.Equal(t, newItem["id"], createdDoc["id"])
-	})
-
-	t.Run("Should execute DELETE transactional batch", func(t *testing.T) {
-		context := context.TODO()
-		batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
-
-		batch.DeleteItem("12345", nil)
-		response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
-		assert.Nil(t, err)
-		assert.True(t, response.Success)
-		assert.Equal(t, 1, len(response.OperationResults))
-
-		operationResponse := response.OperationResults[0]
-		assert.NotNil(t, operationResponse)
-		assert.Equal(t, int32(http.StatusNoContent), operationResponse.StatusCode)
-
-		_, status := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, "12345")
-		assert.Equal(t, datastore.StatusNotFound, int(status))
-	})
-
-	t.Run("Should execute REPLACE transactional batch", func(t *testing.T) {
-		context := context.TODO()
-		batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
-
-		newItem := map[string]interface{}{
-			"id": "67890",
-			"pk": "666",
-		}
-		bytes, err := json.Marshal(newItem)
-		assert.Nil(t, err)
-
-		batch.ReplaceItem("67890", bytes, nil)
-		response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
-		assert.Nil(t, err)
-		assert.True(t, response.Success)
-		assert.Equal(t, 1, len(response.OperationResults))
-
-		operationResponse := response.OperationResults[0]
-		assert.NotNil(t, operationResponse)
-		assert.NotNil(t, operationResponse.ResourceBody)
-		assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
-
-		var itemResponseBody map[string]interface{}
-		json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
-		assert.Equal(t, newItem["id"], itemResponseBody["id"])
-		assert.Equal(t, newItem["pk"], itemResponseBody["pk"])
-
-		updatedDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
-		assert.Equal(t, newItem["id"], updatedDoc["id"])
-		assert.Equal(t, newItem["pk"], updatedDoc["pk"])
-	})
-
-	t.Run("Should execute UPSERT transactional batch", func(t *testing.T) {
-		context := context.TODO()
-		batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
-
-		newItem := map[string]interface{}{
-			"id": "678901",
-			"pk": "666",
-		}
-		bytes, err := json.Marshal(newItem)
-		assert.Nil(t, err)
-
-		batch.UpsertItem(bytes, nil)
-		response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
-		assert.Nil(t, err)
-		assert.True(t, response.Success)
-		assert.Equal(t, 1, len(response.OperationResults))
-
-		operationResponse := response.OperationResults[0]
-		assert.NotNil(t, operationResponse)
-		assert.NotNil(t, operationResponse.ResourceBody)
-		assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
-
-		var itemResponseBody map[string]interface{}
-		json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
-		assert.Equal(t, newItem["id"], itemResponseBody["id"])
-		assert.Equal(t, newItem["pk"], itemResponseBody["pk"])
-
-		updatedDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
-		assert.Equal(t, newItem["id"], updatedDoc["id"])
-		assert.Equal(t, newItem["pk"], updatedDoc["pk"])
-	})
-
-	t.Run("Should execute READ transactional batch", func(t *testing.T) {
-		context := context.TODO()
-		batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
-
-		batch.ReadItem("67890", nil)
-		response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
-		assert.Nil(t, err)
-		assert.True(t, response.Success)
-		assert.Equal(t, 1, len(response.OperationResults))
-
-		operationResponse := response.OperationResults[0]
-		assert.NotNil(t, operationResponse)
-		assert.NotNil(t, operationResponse.ResourceBody)
-		assert.Equal(t, int32(http.StatusOK), operationResponse.StatusCode)
-
-		var itemResponseBody map[string]interface{}
-		json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
-		assert.Equal(t, "67890", itemResponseBody["id"])
+		})
+
+		t.Run("UpsertItem new", func(t *testing.T) {
+			context := context.TODO()
+
+			item := map[string]interface{}{"id": "123456", "pk": "1234", "isCool": false, "arr": []int{1, 2, 3}}
+			bytes, err := json.Marshal(item)
+			assert.Nil(t, err)
+
+			r, err2 := collectionClient.UpsertItem(
+				context,
+				azcosmos.PartitionKey{},
+				bytes,
+				&azcosmos.ItemOptions{
+					EnableContentResponseOnWrite: false,
+				},
+			)
+			assert.NotNil(t, r)
+			assert.Nil(t, err2)
+		})
+
+		t.Run("UpsertItem that already exists", func(t *testing.T) {
+			context := context.TODO()
+
+			item := map[string]interface{}{"id": "12345", "pk": "123", "isCool": false, "arr": []int{1, 2, 3, 4}}
+			bytes, err := json.Marshal(item)
+			assert.Nil(t, err)
+
+			r, err2 := collectionClient.UpsertItem(
+				context,
+				azcosmos.PartitionKey{},
+				bytes,
+				&azcosmos.ItemOptions{
+					EnableContentResponseOnWrite: false,
+				},
+			)
+			assert.NotNil(t, r)
+			assert.Nil(t, err2)
+		})
+	})
+
+	runTestsWithPresets(t, "Test_Documents_TransactionalBatch", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
+		collectionClient := documents_InitializeDb(t, ts)
+
+		t.Run("Should execute CREATE transactional batch", func(t *testing.T) {
+			context := context.TODO()
+			batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
+
+			newItem := map[string]interface{}{
+				"id": "678901",
+			}
+			bytes, err := json.Marshal(newItem)
+			assert.Nil(t, err)
+
+			batch.CreateItem(bytes, nil)
+			response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
+			assert.Nil(t, err)
+			assert.True(t, response.Success)
+			assert.Equal(t, 1, len(response.OperationResults))
+
+			operationResponse := response.OperationResults[0]
+			assert.NotNil(t, operationResponse)
+			assert.NotNil(t, operationResponse.ResourceBody)
+			assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
+
+			var itemResponseBody map[string]interface{}
+			json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
+			assert.Equal(t, newItem["id"], itemResponseBody["id"])
+
+			createdDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
+			assert.Equal(t, newItem["id"], createdDoc["id"])
+		})
+
+		t.Run("Should execute DELETE transactional batch", func(t *testing.T) {
+			context := context.TODO()
+			batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
+
+			batch.DeleteItem("12345", nil)
+			response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
+			assert.Nil(t, err)
+			assert.True(t, response.Success)
+			assert.Equal(t, 1, len(response.OperationResults))
+
+			operationResponse := response.OperationResults[0]
+			assert.NotNil(t, operationResponse)
+			assert.Equal(t, int32(http.StatusNoContent), operationResponse.StatusCode)
+
+			_, status := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, "12345")
+			assert.Equal(t, datastore.StatusNotFound, int(status))
+		})
+
+		t.Run("Should execute REPLACE transactional batch", func(t *testing.T) {
+			context := context.TODO()
+			batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
+
+			newItem := map[string]interface{}{
+				"id": "67890",
+				"pk": "666",
+			}
+			bytes, err := json.Marshal(newItem)
+			assert.Nil(t, err)
+
+			batch.ReplaceItem("67890", bytes, nil)
+			response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
+			assert.Nil(t, err)
+			assert.True(t, response.Success)
+			assert.Equal(t, 1, len(response.OperationResults))
+
+			operationResponse := response.OperationResults[0]
+			assert.NotNil(t, operationResponse)
+			assert.NotNil(t, operationResponse.ResourceBody)
+			assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
+
+			var itemResponseBody map[string]interface{}
+			json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
+			assert.Equal(t, newItem["id"], itemResponseBody["id"])
+			assert.Equal(t, newItem["pk"], itemResponseBody["pk"])
+
+			updatedDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
+			assert.Equal(t, newItem["id"], updatedDoc["id"])
+			assert.Equal(t, newItem["pk"], updatedDoc["pk"])
+		})
+
+		t.Run("Should execute UPSERT transactional batch", func(t *testing.T) {
+			context := context.TODO()
+			batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
+
+			newItem := map[string]interface{}{
+				"id": "678901",
+				"pk": "666",
+			}
+			bytes, err := json.Marshal(newItem)
+			assert.Nil(t, err)
+
+			batch.UpsertItem(bytes, nil)
+			response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
+			assert.Nil(t, err)
+			assert.True(t, response.Success)
+			assert.Equal(t, 1, len(response.OperationResults))
+
+			operationResponse := response.OperationResults[0]
+			assert.NotNil(t, operationResponse)
+			assert.NotNil(t, operationResponse.ResourceBody)
+			assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
+
+			var itemResponseBody map[string]interface{}
+			json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
+			assert.Equal(t, newItem["id"], itemResponseBody["id"])
+			assert.Equal(t, newItem["pk"], itemResponseBody["pk"])
+
+			updatedDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
+			assert.Equal(t, newItem["id"], updatedDoc["id"])
+			assert.Equal(t, newItem["pk"], updatedDoc["pk"])
+		})
+
+		t.Run("Should execute READ transactional batch", func(t *testing.T) {
+			context := context.TODO()
+			batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
+
+			batch.ReadItem("67890", nil)
+			response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
+			assert.Nil(t, err)
+			assert.True(t, response.Success)
+			assert.Equal(t, 1, len(response.OperationResults))
+
+			operationResponse := response.OperationResults[0]
+			assert.NotNil(t, operationResponse)
+			assert.NotNil(t, operationResponse.ResourceBody)
+			assert.Equal(t, int32(http.StatusOK), operationResponse.StatusCode)
+
+			var itemResponseBody map[string]interface{}
+			json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
+			assert.Equal(t, "67890", itemResponseBody["id"])
+		})
 	})
 }

api/tests/documents_trailingslash_test.go

@@ -14,7 +14,8 @@ import (
 
 // Request document with trailing slash like python cosmosdb client does.
 func Test_Documents_Read_Trailing_Slash(t *testing.T) {
-	ts, _ := documents_InitializeDb(t)
+	ts := runTestServer()
+	documents_InitializeDb(t, ts)
 	defer ts.Server.Close()
 
 	t.Run("Read doc with client that appends slash to path", func(t *testing.T) {

cmd/server/server.go

@@ -8,16 +8,28 @@ import (
 	"github.com/pikami/cosmium/api"
 	"github.com/pikami/cosmium/api/config"
 	"github.com/pikami/cosmium/internal/datastore"
-	mapdatastore "github.com/pikami/cosmium/internal/datastore/map_datastore"
+	badgerdatastore "github.com/pikami/cosmium/internal/datastore/badger_datastore"
+	jsondatastore "github.com/pikami/cosmium/internal/datastore/json_datastore"
+	"github.com/pikami/cosmium/internal/logger"
 )
 
 func main() {
 	configuration := config.ParseFlags()
 
-	var dataStore datastore.DataStore = mapdatastore.NewMapDataStore(mapdatastore.MapDataStoreOptions{
-		InitialDataFilePath: configuration.InitialDataFilePath,
-		PersistDataFilePath: configuration.PersistDataFilePath,
-	})
+	var dataStore datastore.DataStore
+	switch configuration.DataStore {
+	case config.DataStoreBadger:
+		dataStore = badgerdatastore.NewBadgerDataStore(badgerdatastore.BadgerDataStoreOptions{
+			PersistDataFilePath: configuration.PersistDataFilePath,
+		})
+		logger.InfoLn("Using Badger data store")
+	default:
+		dataStore = jsondatastore.NewJsonDataStore(jsondatastore.JsonDataStoreOptions{
+			InitialDataFilePath: configuration.InitialDataFilePath,
+			PersistDataFilePath: configuration.PersistDataFilePath,
+		})
+		logger.InfoLn("Using in-memory data store")
+	}
 
 	server := api.NewApiServer(dataStore, &configuration)
 	err := server.Start()
@@ -25,10 +37,10 @@ func main() {
 		panic(err)
 	}
 
-	waitForExit(server, dataStore, configuration)
+	waitForExit(server, dataStore)
 }
 
-func waitForExit(server *api.ApiServer, dataStore datastore.DataStore, config config.ServerConfig) {
+func waitForExit(server *api.ApiServer, dataStore datastore.DataStore) {
 	sigs := make(chan os.Signal, 1)
 	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
 

docs/COMPATIBILITY.md

@@ -79,7 +79,7 @@ Cosmium strives to support the core features of Cosmos DB, including:
 
 | Function | Implemented |
 | -------- | ----------- |
-| IIF      | No          |
+| IIF      | Yes         |
 
 ### Date and time Functions
 

go.mod

@@ -3,44 +3,58 @@ module github.com/pikami/cosmium
 go 1.24.0
 
 require (
-	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0
-	github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.3.0
-	github.com/cosmiumdev/json-patch/v5 v5.9.3
+	github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0
+	github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.4.0
+	github.com/cosmiumdev/json-patch/v5 v5.9.11
+	github.com/dgraph-io/badger/v4 v4.7.0
 	github.com/gin-gonic/gin v1.10.0
 	github.com/google/uuid v1.6.0
 	github.com/stretchr/testify v1.10.0
-	golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa
+	github.com/vmihailenco/msgpack/v5 v5.4.1
+	golang.org/x/exp v0.0.0-20250506013437-ce4c2cf36ca6
 )
 
 require (
 	github.com/Azure/azure-sdk-for-go v68.0.0+incompatible // indirect
-	github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0 // indirect
-	github.com/bytedance/sonic v1.12.9 // indirect
-	github.com/bytedance/sonic/loader v0.2.3 // indirect
+	github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1 // indirect
+	github.com/bytedance/sonic v1.13.2 // indirect
+	github.com/bytedance/sonic/loader v0.2.4 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cloudwego/base64x v0.1.5 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/gabriel-vasile/mimetype v1.4.8 // indirect
-	github.com/gin-contrib/sse v1.0.0 // indirect
+	github.com/dgraph-io/ristretto/v2 v2.2.0 // indirect
+	github.com/dustin/go-humanize v1.0.1 // indirect
+	github.com/gabriel-vasile/mimetype v1.4.9 // indirect
+	github.com/gin-contrib/sse v1.1.0 // indirect
+	github.com/go-logr/logr v1.4.2 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect
 	github.com/go-playground/universal-translator v0.18.1 // indirect
-	github.com/go-playground/validator/v10 v10.25.0 // indirect
+	github.com/go-playground/validator/v10 v10.26.0 // indirect
 	github.com/goccy/go-json v0.10.5 // indirect
+	github.com/google/flatbuffers v25.2.10+incompatible // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/klauspost/compress v1.18.0 // indirect
 	github.com/klauspost/cpuid/v2 v2.2.10 // indirect
 	github.com/leodido/go-urn v1.4.0 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
-	github.com/pelletier/go-toml/v2 v2.2.3 // indirect
+	github.com/pelletier/go-toml/v2 v2.2.4 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.2.12 // indirect
-	golang.org/x/arch v0.14.0 // indirect
-	golang.org/x/crypto v0.35.0 // indirect
-	golang.org/x/net v0.35.0 // indirect
-	golang.org/x/sys v0.30.0 // indirect
-	golang.org/x/text v0.22.0 // indirect
-	google.golang.org/protobuf v1.36.5 // indirect
+	github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
+	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
+	go.opentelemetry.io/otel v1.35.0 // indirect
+	go.opentelemetry.io/otel/metric v1.35.0 // indirect
+	go.opentelemetry.io/otel/trace v1.35.0 // indirect
+	golang.org/x/arch v0.17.0 // indirect
+	golang.org/x/crypto v0.38.0 // indirect
+	golang.org/x/net v0.40.0 // indirect
+	golang.org/x/sys v0.33.0 // indirect
+	golang.org/x/text v0.25.0 // indirect
+	google.golang.org/protobuf v1.36.6 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -2,33 +2,67 @@ github.com/Azure/azure-sdk-for-go v68.0.0+incompatible h1:fcYLmCpyNYRnvJbPerq7U0
 github.com/Azure/azure-sdk-for-go v68.0.0+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0 h1:g0EZJwz7xkXQiZAI5xi9f3WWFYBlX1CPTrR+NDToRkQ=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v1.17.0/go.mod h1:XCW7KnZet0Opnr7HccfUw1PLc4CjHqpcaxW8DHklNkQ=
+github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0 h1:Gt0j3wceWMwPmiazCa8MzMA0MfhmPIz0Qp0FJ6qcM0U=
+github.com/Azure/azure-sdk-for-go/sdk/azcore v1.18.0/go.mod h1:Ot/6aikWnKWi4l9QB7qVSwa8iMphQNqkWALMoNT3rzM=
 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.8.0 h1:B/dfvscEQtew9dVuoxqxrUKKv8Ih2f55PydknDamU+g=
 github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.8.0/go.mod h1:fiPSssYvltE08HJchL04dOy+RD4hgrjph0cwGGMntdI=
+github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.8.2 h1:F0gBpfdPLGsw+nsgk6aqqkZS1jiixa5WwFe3fk/T3Ys=
 github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.3.0 h1:RGcdpSElvcXCwxydI0xzOBu1Gvp88OoiTGfbtO/z1m0=
 github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.3.0/go.mod h1:YwUyrNUtcZcibA99JcfCP6UUp95VVQKO2MJfBzgJDwA=
+github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.4.0 h1:TSaH6Lj0m8bDr4vX1+LC1KLQTnLzZb3tOxrx/PLqw+c=
+github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos v1.4.0/go.mod h1:Krtog/7tz27z75TwM5cIS8bxEH4dcBUezcq+kGVeZEo=
 github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0 h1:ywEEhmNahHBihViHepv3xPBn1663uRv2t2q/ESv9seY=
 github.com/Azure/azure-sdk-for-go/sdk/internal v1.10.0/go.mod h1:iZDifYGJTIgIIkYRNWPENUnqx6bJ2xnSDFI2tjwZNuY=
+github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1 h1:FPKJS1T+clwv+OLGt13a8UjqeRuh0O4SJ3lUriThc+4=
+github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.1/go.mod h1:j2chePtV91HrC22tGoRX3sGY42uF13WzmmV80/OdVAA=
 github.com/AzureAD/microsoft-authentication-library-for-go v1.3.2 h1:kYRSnvJju5gYVyhkij+RTJ/VR6QIUaCfWeaFm2ycsjQ=
 github.com/AzureAD/microsoft-authentication-library-for-go v1.3.2/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI=
-github.com/bytedance/sonic v1.12.9 h1:Od1BvK55NnewtGaJsTDeAOSnLVO2BTSLOe0+ooKokmQ=
-github.com/bytedance/sonic v1.12.9/go.mod h1:uVvFidNmlt9+wa31S1urfwwthTWteBgG0hWuoKAXTx8=
+github.com/AzureAD/microsoft-authentication-library-for-go v1.4.2 h1:oygO0locgZJe7PpYPXT5A29ZkwJaPqcva7BVeemZOZs=
+github.com/bytedance/sonic v1.13.1 h1:Jyd5CIvdFnkOWuKXr+wm4Nyk2h0yAFsr8ucJgEasO3g=
+github.com/bytedance/sonic v1.13.1/go.mod h1:o68xyaF9u2gvVBuGHPlUVCy+ZfmNNO5ETf1+KgkJhz4=
+github.com/bytedance/sonic v1.13.2 h1:8/H1FempDZqC4VqjptGo14QQlJx8VdZJegxs6wwfqpQ=
+github.com/bytedance/sonic v1.13.2/go.mod h1:o68xyaF9u2gvVBuGHPlUVCy+ZfmNNO5ETf1+KgkJhz4=
 github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
-github.com/bytedance/sonic/loader v0.2.3 h1:yctD0Q3v2NOGfSWPLPvG2ggA2kV6TS6s4wioyEqssH0=
-github.com/bytedance/sonic/loader v0.2.3/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI=
+github.com/bytedance/sonic/loader v0.2.4 h1:ZWCw4stuXUsn1/+zQDqeE7JKP+QO47tz7QCNan80NzY=
+github.com/bytedance/sonic/loader v0.2.4/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI=
+github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
+github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4=
 github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
 github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
-github.com/cosmiumdev/json-patch/v5 v5.9.3 h1:l+Og3+5edqV2NHDo58sz72eS733lbXVYP61seYK43Do=
-github.com/cosmiumdev/json-patch/v5 v5.9.3/go.mod h1:WzSTCdia0WrlZtjnL19P4RiwWtfdyArm/E7stgEeP5g=
+github.com/cosmiumdev/json-patch/v5 v5.9.11 h1:WD2Wqaz/vO987z2FFdqgkj15HgYZ/Y5TpqE3I4T/iOQ=
+github.com/cosmiumdev/json-patch/v5 v5.9.11/go.mod h1:YPZmckmv4ZY+oxKIOjgq3sIudHVB6VEMcicCS9LtVLM=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dgraph-io/badger/v4 v4.6.0 h1:acOwfOOZ4p1dPRnYzvkVm7rUk2Y21TgPVepCy5dJdFQ=
+github.com/dgraph-io/badger/v4 v4.6.0/go.mod h1:KSJ5VTuZNC3Sd+YhvVjk2nYua9UZnnTr/SkXvdtiPgI=
+github.com/dgraph-io/badger/v4 v4.7.0 h1:Q+J8HApYAY7UMpL8d9owqiB+odzEc0zn/aqOD9jhc6Y=
+github.com/dgraph-io/badger/v4 v4.7.0/go.mod h1:He7TzG3YBy3j4f5baj5B7Zl2XyfNe5bl4Udl0aPemVA=
+github.com/dgraph-io/ristretto/v2 v2.1.0 h1:59LjpOJLNDULHh8MC4UaegN52lC4JnO2dITsie/Pa8I=
+github.com/dgraph-io/ristretto/v2 v2.1.0/go.mod h1:uejeqfYXpUomfse0+lO+13ATz4TypQYLJZzBSAemuB4=
+github.com/dgraph-io/ristretto/v2 v2.2.0 h1:bkY3XzJcXoMuELV8F+vS8kzNgicwQFAaGINAEJdWGOM=
+github.com/dgraph-io/ristretto/v2 v2.2.0/go.mod h1:RZrm63UmcBAaYWC1DotLYBmTvgkrs0+XhBd7Npn7/zI=
+github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13 h1:fAjc9m62+UWV/WAFKLNi6ZS0675eEUC9y3AlwSbQu1Y=
+github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
+github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da h1:aIftn67I1fkbMa512G+w+Pxci9hJPB8oMnkcP3iZF38=
+github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
+github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/gabriel-vasile/mimetype v1.4.8 h1:FfZ3gj38NjllZIeJAmMhr+qKL8Wu+nOoI3GqacKw1NM=
 github.com/gabriel-vasile/mimetype v1.4.8/go.mod h1:ByKUIKGjh1ODkGM1asKUbQZOLGrPjydw3hYPU2YU9t8=
+github.com/gabriel-vasile/mimetype v1.4.9 h1:5k+WDwEsD9eTLL8Tz3L0VnmVh9QxGjRmjBvAG7U/oYY=
+github.com/gabriel-vasile/mimetype v1.4.9/go.mod h1:WnSQhFKJuBlRyLiKohA/2DtIlPFAbguNaG7QCHcyGok=
 github.com/gin-contrib/sse v1.0.0 h1:y3bT1mUWUxDpW4JLQg/HnTqV4rozuW4tC9eFKTxYI9E=
 github.com/gin-contrib/sse v1.0.0/go.mod h1:zNuFdwarAygJBht0NTKiSi3jRf6RbqeILZ9Sp6Slhe0=
+github.com/gin-contrib/sse v1.1.0 h1:n0w2GMuUpWDVp7qSpvze6fAu9iRxJY4Hmj6AmBOU05w=
+github.com/gin-contrib/sse v1.1.0/go.mod h1:hxRZ5gVpWMT7Z0B0gSNYqqsSCNIJMjzvm6fqCz9vjwM=
 github.com/gin-gonic/gin v1.10.0 h1:nTuyha1TYqgedzytsKYqna+DfLos46nTv2ygFy86HFU=
 github.com/gin-gonic/gin v1.10.0/go.mod h1:4PMNQiOhvDRa013RKVbsiNwoyezlm2rm0uX/T7kzp5Y=
+github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
+github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
+github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
+github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
 github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s=
 github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4=
 github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA=
@@ -37,17 +71,24 @@ github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJn
 github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY=
 github.com/go-playground/validator/v10 v10.25.0 h1:5Dh7cjvzR7BRZadnsVOzPhWsrwUr0nmsZJxEAnFLNO8=
 github.com/go-playground/validator/v10 v10.25.0/go.mod h1:GGzBIJMuE98Ic/kJsBXbz1x/7cByt++cQ+YOuDM5wus=
+github.com/go-playground/validator/v10 v10.26.0 h1:SP05Nqhjcvz81uJaRfEV0YBSSSGMc/iMaVtFbr3Sw2k=
+github.com/go-playground/validator/v10 v10.26.0/go.mod h1:I5QpIEbmr8On7W0TktmJAumgzX4CA1XNl4ZmDuVHKKo=
 github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
 github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
 github.com/golang-jwt/jwt/v5 v5.2.1 h1:OuVbFODueb089Lh128TAcimifWaLhJwVflnrgM17wHk=
 github.com/golang-jwt/jwt/v5 v5.2.1/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
-github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
-github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8=
+github.com/google/flatbuffers v25.2.10+incompatible h1:F3vclr7C3HpB1k9mxCGRMXq6FdUalZ6H/pNX4FP1v0Q=
+github.com/google/flatbuffers v25.2.10+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8=
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
+github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
+github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
 github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
 github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE=
 github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0=
@@ -69,14 +110,16 @@ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9G
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
 github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M=
 github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc=
+github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4=
+github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
-github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
+github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
+github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
@@ -93,21 +136,47 @@ github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS
 github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
 github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE=
 github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
-golang.org/x/arch v0.14.0 h1:z9JUEZWr8x4rR0OU6c4/4t6E6jOZ8/QBS2bBYBm4tx4=
-golang.org/x/arch v0.14.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys=
-golang.org/x/crypto v0.35.0 h1:b15kiHdrGCHrP6LvwaQ3c03kgNhhiMgvlhxHQhmg2Xs=
-golang.org/x/crypto v0.35.0/go.mod h1:dy7dXNW32cAb/6/PRuTNsix8T+vJAqvuIy5Bli/x0YQ=
-golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa h1:t2QcU6V556bFjYgu4L6C+6VrCPyJZ+eyRsABUPs1mz4=
-golang.org/x/exp v0.0.0-20250218142911-aa4b98e5adaa/go.mod h1:BHOTPb3L19zxehTsLoJXVaTktb06DFgmdW6Wb9s8jqk=
-golang.org/x/net v0.35.0 h1:T5GQRQb2y08kTAByq9L4/bz8cipCdA8FbRTXewonqY8=
-golang.org/x/net v0.35.0/go.mod h1:EglIi67kWsHKlRzzVMUD93VMSWGFOMSZgxFjparz1Qk=
+github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8=
+github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok=
+github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g=
+github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds=
+go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
+go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
+go.opentelemetry.io/otel v1.35.0 h1:xKWKPxrxB6OtMCbmMY021CqC45J+3Onta9MqjhnusiQ=
+go.opentelemetry.io/otel v1.35.0/go.mod h1:UEqy8Zp11hpkUrL73gSlELM0DupHoiq72dR+Zqel/+Y=
+go.opentelemetry.io/otel/metric v1.35.0 h1:0znxYu2SNyuMSQT4Y9WDWej0VpcsxkuklLa4/siN90M=
+go.opentelemetry.io/otel/metric v1.35.0/go.mod h1:nKVFgxBZ2fReX6IlyW28MgZojkoAkJGaE8CpgeAU3oE=
+go.opentelemetry.io/otel/trace v1.35.0 h1:dPpEfJu1sDIqruz7BHFG3c7528f6ddfSWfFDVt/xgMs=
+go.opentelemetry.io/otel/trace v1.35.0/go.mod h1:WUk7DtFp1Aw2MkvqGdwiXYDZZNvA/1J8o6xRXLrIkyc=
+golang.org/x/arch v0.15.0 h1:QtOrQd0bTUnhNVNndMpLHNWrDmYzZ2KDqSrEymqInZw=
+golang.org/x/arch v0.15.0/go.mod h1:JmwW7aLIoRUKgaTzhkiEFxvcEiQGyOg9BMonBJUS7EE=
+golang.org/x/arch v0.17.0 h1:4O3dfLzd+lQewptAHqjewQZQDyEdejz3VwgeYwkZneU=
+golang.org/x/arch v0.17.0/go.mod h1:bdwinDaKcfZUGpH09BB7ZmOfhalA8lQdzl62l8gGWsk=
+golang.org/x/crypto v0.36.0 h1:AnAEvhDddvBdpY+uR+MyHmuZzzNqXSe/GvuDeob5L34=
+golang.org/x/crypto v0.36.0/go.mod h1:Y4J0ReaxCR1IMaabaSMugxJES1EpwhBHhv2bDHklZvc=
+golang.org/x/crypto v0.38.0 h1:jt+WWG8IZlBnVbomuhg2Mdq0+BBQaHbtqHEFEigjUV8=
+golang.org/x/crypto v0.38.0/go.mod h1:MvrbAqul58NNYPKnOra203SB9vpuZW0e+RRZV+Ggqjw=
+golang.org/x/exp v0.0.0-20250305212735-054e65f0b394 h1:nDVHiLt8aIbd/VzvPWN6kSOPE7+F/fNFDSXLVYkE/Iw=
+golang.org/x/exp v0.0.0-20250305212735-054e65f0b394/go.mod h1:sIifuuw/Yco/y6yb6+bDNfyeQ/MdPUy/hKEMYQV17cM=
+golang.org/x/exp v0.0.0-20250506013437-ce4c2cf36ca6 h1:y5zboxd6LQAqYIhHnB48p0ByQ/GnQx2BE33L8BOHQkI=
+golang.org/x/exp v0.0.0-20250506013437-ce4c2cf36ca6/go.mod h1:U6Lno4MTRCDY+Ba7aCcauB9T60gsv5s4ralQzP72ZoQ=
+golang.org/x/net v0.37.0 h1:1zLorHbz+LYj7MQlSf1+2tPIIgibq2eL5xkrGk6f+2c=
+golang.org/x/net v0.37.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8=
+golang.org/x/net v0.40.0 h1:79Xs7wF06Gbdcg4kdCCIQArK11Z1hr5POQ6+fIYHNuY=
+golang.org/x/net v0.40.0/go.mod h1:y0hY0exeL2Pku80/zKK7tpntoX23cqL3Oa6njdgRtds=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.30.0 h1:QjkSwP/36a20jFYWkSue1YwXzLmsV5Gfq7Eiy72C1uc=
-golang.org/x/sys v0.30.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/text v0.22.0 h1:bofq7m3/HAFvbF51jz3Q9wLg3jkvSPuiZu/pD1XwgtM=
-golang.org/x/text v0.22.0/go.mod h1:YRoo4H8PVmsu+E3Ou7cqLVH8oXWIHVoX0jqUWALQhfY=
+golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik=
+golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
+golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/text v0.23.0 h1:D71I7dUrlY+VX0gQShAThNGHFxZ13dGLBHQLVl1mJlY=
+golang.org/x/text v0.23.0/go.mod h1:/BLNzu4aZCJ1+kcD0DNRotWKage4q2rGVAg4o22unh4=
+golang.org/x/text v0.25.0 h1:qVyWApTSYLk/drJRO5mDlNYskwQznZmkpV2c8q9zls4=
+golang.org/x/text v0.25.0/go.mod h1:WEdwpYrmk1qmdHvhkSTNPm3app7v4rsT8F2UD6+VHIA=
 google.golang.org/protobuf v1.36.5 h1:tPhr+woSbjfYvY6/GPufUoYizxw1cF/yFoxJ2fmpwlM=
 google.golang.org/protobuf v1.36.5/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
+google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
+google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=

internal/constants/responses.go

@@ -30,3 +30,8 @@ var QueryPlanResponse = gin.H{
 		},
 	},
 }
+
+var UnknownErrorResponse = gin.H{"message": "Unknown error"}
+var NotFoundResponse = gin.H{"message": "NotFound"}
+var ConflictResponse = gin.H{"message": "Conflict"}
+var BadRequestResponse = gin.H{"message": "BadRequest"}

internal/datastore/badger_datastore/badger_datastore.go

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

internal/datastore/badger_datastore/badger_logger.go

@@ -0,0 +1,28 @@
+package badgerdatastore
+
+import (
+	"github.com/dgraph-io/badger/v4"
+	"github.com/pikami/cosmium/internal/logger"
+)
+
+type badgerLogger struct{}
+
+func newBadgerLogger() badger.Logger {
+	return &badgerLogger{}
+}
+
+func (l *badgerLogger) Errorf(format string, v ...interface{}) {
+	logger.Errorf(format, v...)
+}
+
+func (l *badgerLogger) Warningf(format string, v ...interface{}) {
+	logger.Infof(format, v...)
+}
+
+func (l *badgerLogger) Infof(format string, v ...interface{}) {
+	logger.Infof(format, v...)
+}
+
+func (l *badgerLogger) Debugf(format string, v ...interface{}) {
+	logger.Debugf(format, v...)
+}

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,204 @@
+package badgerdatastore
+
+import (
+	"github.com/dgraph-io/badger/v4"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+	"github.com/pikami/cosmium/internal/resourceid"
+	"github.com/vmihailenco/msgpack/v5"
+)
+
+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
+	}
+
+	buf, err := msgpack.Marshal(value)
+	if err != nil {
+		logger.ErrorLn("Error while encoding value:", err)
+		return datastore.Unknown
+	}
+
+	err = txn.Set([]byte(key), buf)
+	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 = msgpack.Unmarshal(val, &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) {
+	results := make([]T, 0)
+
+	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,58 @@
+package badgerdatastore
+
+import (
+	"github.com/dgraph-io/badger/v4"
+	"github.com/pikami/cosmium/internal/datastore"
+	"github.com/pikami/cosmium/internal/logger"
+	"github.com/vmihailenco/msgpack/v5"
+)
+
+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 = msgpack.Unmarshal(val, &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/json_datastore/array_document_iterator.go

@@ -1,4 +1,4 @@
-package mapdatastore
+package jsondatastore
 
 import "github.com/pikami/cosmium/internal/datastore"
 
@@ -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/json_datastore/collections.go

@@ -1,4 +1,4 @@
-package mapdatastore
+package jsondatastore
 
 import (
 	"fmt"
@@ -11,7 +11,7 @@ import (
 	"golang.org/x/exp/maps"
 )
 
-func (r *MapDataStore) GetAllCollections(databaseId string) ([]datastore.Collection, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetAllCollections(databaseId string) ([]datastore.Collection, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -22,7 +22,7 @@ func (r *MapDataStore) GetAllCollections(databaseId string) ([]datastore.Collect
 	return maps.Values(r.storeState.Collections[databaseId]), datastore.StatusOk
 }
 
-func (r *MapDataStore) GetCollection(databaseId string, collectionId string) (datastore.Collection, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetCollection(databaseId string, collectionId string) (datastore.Collection, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -37,7 +37,7 @@ func (r *MapDataStore) GetCollection(databaseId string, collectionId string) (da
 	return r.storeState.Collections[databaseId][collectionId], datastore.StatusOk
 }
 
-func (r *MapDataStore) DeleteCollection(databaseId string, collectionId string) datastore.DataStoreStatus {
+func (r *JsonDataStore) DeleteCollection(databaseId string, collectionId string) datastore.DataStoreStatus {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 
@@ -58,7 +58,7 @@ func (r *MapDataStore) DeleteCollection(databaseId string, collectionId string)
 	return datastore.StatusOk
 }
 
-func (r *MapDataStore) CreateCollection(databaseId string, newCollection datastore.Collection) (datastore.Collection, datastore.DataStoreStatus) {
+func (r *JsonDataStore) CreateCollection(databaseId string, newCollection datastore.Collection) (datastore.Collection, datastore.DataStoreStatus) {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 

internal/datastore/json_datastore/databases.go

@@ -1,4 +1,4 @@
-package mapdatastore
+package jsondatastore
 
 import (
 	"fmt"
@@ -10,14 +10,14 @@ import (
 	"golang.org/x/exp/maps"
 )
 
-func (r *MapDataStore) GetAllDatabases() ([]datastore.Database, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetAllDatabases() ([]datastore.Database, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
 	return maps.Values(r.storeState.Databases), datastore.StatusOk
 }
 
-func (r *MapDataStore) GetDatabase(id string) (datastore.Database, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetDatabase(id string) (datastore.Database, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -28,7 +28,7 @@ func (r *MapDataStore) GetDatabase(id string) (datastore.Database, datastore.Dat
 	return datastore.Database{}, datastore.StatusNotFound
 }
 
-func (r *MapDataStore) DeleteDatabase(id string) datastore.DataStoreStatus {
+func (r *JsonDataStore) DeleteDatabase(id string) datastore.DataStoreStatus {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 
@@ -46,7 +46,7 @@ func (r *MapDataStore) DeleteDatabase(id string) datastore.DataStoreStatus {
 	return datastore.StatusOk
 }
 
-func (r *MapDataStore) CreateDatabase(newDatabase datastore.Database) (datastore.Database, datastore.DataStoreStatus) {
+func (r *JsonDataStore) CreateDatabase(newDatabase datastore.Database) (datastore.Database, datastore.DataStoreStatus) {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 

internal/datastore/json_datastore/documents.go

@@ -1,4 +1,4 @@
-package mapdatastore
+package jsondatastore
 
 import (
 	"fmt"
@@ -10,7 +10,7 @@ import (
 	"golang.org/x/exp/maps"
 )
 
-func (r *MapDataStore) GetAllDocuments(databaseId string, collectionId string) ([]datastore.Document, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetAllDocuments(databaseId string, collectionId string) ([]datastore.Document, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -25,7 +25,7 @@ func (r *MapDataStore) GetAllDocuments(databaseId string, collectionId string) (
 	return maps.Values(r.storeState.Documents[databaseId][collectionId]), datastore.StatusOk
 }
 
-func (r *MapDataStore) GetDocument(databaseId string, collectionId string, documentId string) (datastore.Document, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetDocument(databaseId string, collectionId string, documentId string) (datastore.Document, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -44,7 +44,7 @@ func (r *MapDataStore) GetDocument(databaseId string, collectionId string, docum
 	return r.storeState.Documents[databaseId][collectionId][documentId], datastore.StatusOk
 }
 
-func (r *MapDataStore) DeleteDocument(databaseId string, collectionId string, documentId string) datastore.DataStoreStatus {
+func (r *JsonDataStore) DeleteDocument(databaseId string, collectionId string, documentId string) datastore.DataStoreStatus {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 
@@ -65,7 +65,7 @@ func (r *MapDataStore) DeleteDocument(databaseId string, collectionId string, do
 	return datastore.StatusOk
 }
 
-func (r *MapDataStore) CreateDocument(databaseId string, collectionId string, document map[string]interface{}) (datastore.Document, datastore.DataStoreStatus) {
+func (r *JsonDataStore) CreateDocument(databaseId string, collectionId string, document map[string]interface{}) (datastore.Document, datastore.DataStoreStatus) {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 
@@ -100,7 +100,7 @@ func (r *MapDataStore) CreateDocument(databaseId string, collectionId string, do
 	return document, datastore.StatusOk
 }
 
-func (r *MapDataStore) GetDocumentIterator(databaseId string, collectionId string) (datastore.DocumentIterator, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetDocumentIterator(databaseId string, collectionId string) (datastore.DocumentIterator, datastore.DataStoreStatus) {
 	documents, status := r.GetAllDocuments(databaseId, collectionId)
 	if status != datastore.StatusOk {
 		return nil, status

internal/datastore/json_datastore/map_datastore.go

@@ -1,21 +1,21 @@
-package mapdatastore
+package jsondatastore
 
 import "github.com/pikami/cosmium/internal/datastore"
 
-type MapDataStore struct {
+type JsonDataStore struct {
 	storeState State
 
 	initialDataFilePath string
 	persistDataFilePath string
 }
 
-type MapDataStoreOptions struct {
+type JsonDataStoreOptions struct {
 	InitialDataFilePath string
 	PersistDataFilePath string
 }
 
-func NewMapDataStore(options MapDataStoreOptions) *MapDataStore {
-	dataStore := &MapDataStore{
+func NewJsonDataStore(options JsonDataStoreOptions) *JsonDataStore {
+	dataStore := &JsonDataStore{
 		storeState: State{
 			Databases:            make(map[string]datastore.Database),
 			Collections:          make(map[string]map[string]datastore.Collection),

internal/datastore/json_datastore/partition_key_ranges.go

@@ -1,4 +1,4 @@
-package mapdatastore
+package jsondatastore
 
 import (
 	"fmt"
@@ -9,7 +9,7 @@ import (
 )
 
 // I have no idea what this is tbh
-func (r *MapDataStore) GetPartitionKeyRanges(databaseId string, collectionId string) ([]datastore.PartitionKeyRange, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetPartitionKeyRanges(databaseId string, collectionId string) ([]datastore.PartitionKeyRange, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 

internal/datastore/json_datastore/state.go

@@ -1,4 +1,4 @@
-package mapdatastore
+package jsondatastore
 
 import (
 	"encoding/json"
@@ -33,7 +33,7 @@ type State struct {
 	UserDefinedFunctions map[string]map[string]map[string]datastore.UserDefinedFunction `json:"udfs"`
 }
 
-func (r *MapDataStore) InitializeDataStore() {
+func (r *JsonDataStore) InitializeDataStore() {
 	if r.initialDataFilePath != "" {
 		r.LoadStateFS(r.initialDataFilePath)
 		return
@@ -55,7 +55,7 @@ func (r *MapDataStore) InitializeDataStore() {
 	}
 }
 
-func (r *MapDataStore) LoadStateFS(filePath string) {
+func (r *JsonDataStore) LoadStateFS(filePath string) {
 	data, err := os.ReadFile(filePath)
 	if err != nil {
 		log.Fatalf("Error reading state JSON file: %v", err)
@@ -68,7 +68,7 @@ func (r *MapDataStore) LoadStateFS(filePath string) {
 	}
 }
 
-func (r *MapDataStore) LoadStateJSON(jsonData string) error {
+func (r *JsonDataStore) LoadStateJSON(jsonData string) error {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 
@@ -94,7 +94,7 @@ func (r *MapDataStore) LoadStateJSON(jsonData string) error {
 	return nil
 }
 
-func (r *MapDataStore) SaveStateFS(filePath string) {
+func (r *JsonDataStore) SaveStateFS(filePath string) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -115,7 +115,7 @@ func (r *MapDataStore) SaveStateFS(filePath string) {
 	logger.Infof("User defined functions: %d\n", getLength(r.storeState.UserDefinedFunctions))
 }
 
-func (r *MapDataStore) DumpToJson() (string, error) {
+func (r *JsonDataStore) DumpToJson() (string, error) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -129,7 +129,7 @@ func (r *MapDataStore) DumpToJson() (string, error) {
 
 }
 
-func (r *MapDataStore) Close() {
+func (r *JsonDataStore) Close() {
 	if r.persistDataFilePath != "" {
 		r.SaveStateFS(r.persistDataFilePath)
 	}
@@ -163,7 +163,7 @@ func getLength(v interface{}) int {
 	return count
 }
 
-func (r *MapDataStore) ensureStoreStateNoNullReferences() {
+func (r *JsonDataStore) ensureStoreStateNoNullReferences() {
 	if r.storeState.Databases == nil {
 		r.storeState.Databases = make(map[string]datastore.Database)
 	}

internal/datastore/json_datastore/stored_procedures.go

@@ -1,4 +1,4 @@
-package mapdatastore
+package jsondatastore
 
 import (
 	"fmt"
@@ -10,14 +10,14 @@ import (
 	"golang.org/x/exp/maps"
 )
 
-func (r *MapDataStore) GetAllStoredProcedures(databaseId string, collectionId string) ([]datastore.StoredProcedure, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetAllStoredProcedures(databaseId string, collectionId string) ([]datastore.StoredProcedure, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
 	return maps.Values(r.storeState.StoredProcedures[databaseId][collectionId]), datastore.StatusOk
 }
 
-func (r *MapDataStore) GetStoredProcedure(databaseId string, collectionId string, spId string) (datastore.StoredProcedure, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetStoredProcedure(databaseId string, collectionId string, spId string) (datastore.StoredProcedure, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -36,7 +36,7 @@ func (r *MapDataStore) GetStoredProcedure(databaseId string, collectionId string
 	return datastore.StoredProcedure{}, datastore.StatusNotFound
 }
 
-func (r *MapDataStore) DeleteStoredProcedure(databaseId string, collectionId string, spId string) datastore.DataStoreStatus {
+func (r *JsonDataStore) DeleteStoredProcedure(databaseId string, collectionId string, spId string) datastore.DataStoreStatus {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 
@@ -57,7 +57,7 @@ func (r *MapDataStore) DeleteStoredProcedure(databaseId string, collectionId str
 	return datastore.StatusOk
 }
 
-func (r *MapDataStore) CreateStoredProcedure(databaseId string, collectionId string, sp datastore.StoredProcedure) (datastore.StoredProcedure, datastore.DataStoreStatus) {
+func (r *JsonDataStore) CreateStoredProcedure(databaseId string, collectionId string, sp datastore.StoredProcedure) (datastore.StoredProcedure, datastore.DataStoreStatus) {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 

internal/datastore/json_datastore/triggers.go

@@ -1,4 +1,4 @@
-package mapdatastore
+package jsondatastore
 
 import (
 	"fmt"
@@ -10,14 +10,14 @@ import (
 	"golang.org/x/exp/maps"
 )
 
-func (r *MapDataStore) GetAllTriggers(databaseId string, collectionId string) ([]datastore.Trigger, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetAllTriggers(databaseId string, collectionId string) ([]datastore.Trigger, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
 	return maps.Values(r.storeState.Triggers[databaseId][collectionId]), datastore.StatusOk
 }
 
-func (r *MapDataStore) GetTrigger(databaseId string, collectionId string, triggerId string) (datastore.Trigger, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetTrigger(databaseId string, collectionId string, triggerId string) (datastore.Trigger, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -36,7 +36,7 @@ func (r *MapDataStore) GetTrigger(databaseId string, collectionId string, trigge
 	return datastore.Trigger{}, datastore.StatusNotFound
 }
 
-func (r *MapDataStore) DeleteTrigger(databaseId string, collectionId string, triggerId string) datastore.DataStoreStatus {
+func (r *JsonDataStore) DeleteTrigger(databaseId string, collectionId string, triggerId string) datastore.DataStoreStatus {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 
@@ -57,7 +57,7 @@ func (r *MapDataStore) DeleteTrigger(databaseId string, collectionId string, tri
 	return datastore.StatusOk
 }
 
-func (r *MapDataStore) CreateTrigger(databaseId string, collectionId string, trigger datastore.Trigger) (datastore.Trigger, datastore.DataStoreStatus) {
+func (r *JsonDataStore) CreateTrigger(databaseId string, collectionId string, trigger datastore.Trigger) (datastore.Trigger, datastore.DataStoreStatus) {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 

internal/datastore/json_datastore/user_defined_functions.go

@@ -1,4 +1,4 @@
-package mapdatastore
+package jsondatastore
 
 import (
 	"fmt"
@@ -10,14 +10,14 @@ import (
 	"golang.org/x/exp/maps"
 )
 
-func (r *MapDataStore) GetAllUserDefinedFunctions(databaseId string, collectionId string) ([]datastore.UserDefinedFunction, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetAllUserDefinedFunctions(databaseId string, collectionId string) ([]datastore.UserDefinedFunction, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
 	return maps.Values(r.storeState.UserDefinedFunctions[databaseId][collectionId]), datastore.StatusOk
 }
 
-func (r *MapDataStore) GetUserDefinedFunction(databaseId string, collectionId string, udfId string) (datastore.UserDefinedFunction, datastore.DataStoreStatus) {
+func (r *JsonDataStore) GetUserDefinedFunction(databaseId string, collectionId string, udfId string) (datastore.UserDefinedFunction, datastore.DataStoreStatus) {
 	r.storeState.RLock()
 	defer r.storeState.RUnlock()
 
@@ -36,7 +36,7 @@ func (r *MapDataStore) GetUserDefinedFunction(databaseId string, collectionId st
 	return datastore.UserDefinedFunction{}, datastore.StatusNotFound
 }
 
-func (r *MapDataStore) DeleteUserDefinedFunction(databaseId string, collectionId string, udfId string) datastore.DataStoreStatus {
+func (r *JsonDataStore) DeleteUserDefinedFunction(databaseId string, collectionId string, udfId string) datastore.DataStoreStatus {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 
@@ -57,7 +57,7 @@ func (r *MapDataStore) DeleteUserDefinedFunction(databaseId string, collectionId
 	return datastore.StatusOk
 }
 
-func (r *MapDataStore) CreateUserDefinedFunction(databaseId string, collectionId string, udf datastore.UserDefinedFunction) (datastore.UserDefinedFunction, datastore.DataStoreStatus) {
+func (r *JsonDataStore) CreateUserDefinedFunction(databaseId string, collectionId string, udf datastore.UserDefinedFunction) (datastore.UserDefinedFunction, datastore.DataStoreStatus) {
 	r.storeState.Lock()
 	defer r.storeState.Unlock()
 

internal/datastore/models.go

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

parsers/models.go

@@ -34,6 +34,7 @@ const (
 	SelectItemTypeConstant
 	SelectItemTypeFunctionCall
 	SelectItemTypeSubQuery
+	SelectItemTypeExpression
 )
 
 type SelectItem struct {
@@ -134,6 +135,8 @@ const (
 	FunctionCallSetIntersect     FunctionCallType = "SetIntersect"
 	FunctionCallSetUnion         FunctionCallType = "SetUnion"
 
+	FunctionCallIif FunctionCallType = "Iif"
+
 	FunctionCallMathAbs              FunctionCallType = "MathAbs"
 	FunctionCallMathAcos             FunctionCallType = "MathAcos"
 	FunctionCallMathAsin             FunctionCallType = "MathAsin"

parsers/nosql/misc_test.go

@@ -163,4 +163,27 @@ func Test_Parse(t *testing.T) {
 			},
 		)
 	})
+
+	t.Run("Should parse IIF function", func(t *testing.T) {
+		testQueryParse(
+			t,
+			`SELECT IIF(true, c.pk, c.id) FROM c`,
+			parsers.SelectStmt{
+				SelectItems: []parsers.SelectItem{
+					{
+						Type: parsers.SelectItemTypeFunctionCall,
+						Value: parsers.FunctionCall{
+							Type: parsers.FunctionCallIif,
+							Arguments: []interface{}{
+								testutils.SelectItem_Constant_Bool(true),
+								testutils.SelectItem_Path("c", "pk"),
+								testutils.SelectItem_Path("c", "id"),
+							},
+						},
+					},
+				},
+				Table: parsers.Table{SelectItem: testutils.SelectItem_Path("c")},
+			},
+		)
+	})
 }

parsers/nosql/nosql.peg

@@ -4,137 +4,137 @@ package nosql
 import "github.com/pikami/cosmium/parsers"
 
 func makeSelectStmt(
-	columns, fromClause, joinItems,
-	whereClause interface{}, distinctClause interface{},
-	count interface{}, groupByClause interface{}, orderList interface{},
-	offsetClause interface{},
+    columns, fromClause, joinItems,
+    whereClause interface{}, distinctClause interface{},
+    count interface{}, groupByClause interface{}, orderList interface{},
+    offsetClause interface{},
 ) (parsers.SelectStmt, error) {
-	selectStmt := parsers.SelectStmt{
-		SelectItems: columns.([]parsers.SelectItem),
-	}
+    selectStmt := parsers.SelectStmt{
+        SelectItems: columns.([]parsers.SelectItem),
+    }
 
     if fromTable, ok := fromClause.(parsers.Table); ok {
         selectStmt.Table = fromTable
     }
 
     if joinItemsArray, ok := joinItems.([]interface{}); ok && len(joinItemsArray) > 0 {
-		selectStmt.JoinItems = make([]parsers.JoinItem, len(joinItemsArray))
-		for i, joinItem := range joinItemsArray {
-			selectStmt.JoinItems[i] = joinItem.(parsers.JoinItem)
-		}
-	}
+        selectStmt.JoinItems = make([]parsers.JoinItem, len(joinItemsArray))
+        for i, joinItem := range joinItemsArray {
+            selectStmt.JoinItems[i] = joinItem.(parsers.JoinItem)
+        }
+    }
 
-	switch v := whereClause.(type) {
-	case parsers.ComparisonExpression, parsers.LogicalExpression, parsers.Constant, parsers.SelectItem:
-		selectStmt.Filters = v
-	}
+    switch v := whereClause.(type) {
+    case parsers.ComparisonExpression, parsers.LogicalExpression, parsers.Constant, parsers.SelectItem:
+        selectStmt.Filters = v
+    }
 
-	if distinctClause != nil {
-		selectStmt.Distinct = true
-	}
+    if distinctClause != nil {
+        selectStmt.Distinct = true
+    }
 
-	if n, ok := count.(int); ok {
-		selectStmt.Count = n
-	}
+    if n, ok := count.(int); ok {
+        selectStmt.Count = n
+    }
 
-	if offsetArr, ok := offsetClause.([]interface{}); ok && len(offsetArr) == 2 {
-		if n, ok := offsetArr[0].(int); ok {
-			selectStmt.Offset = n
-		}
+    if offsetArr, ok := offsetClause.([]interface{}); ok && len(offsetArr) == 2 {
+        if n, ok := offsetArr[0].(int); ok {
+            selectStmt.Offset = n
+        }
 
-		if n, ok := offsetArr[1].(int); ok {
-			selectStmt.Count = n
-		}
-	}
+        if n, ok := offsetArr[1].(int); ok {
+            selectStmt.Count = n
+        }
+    }
 
-	if orderExpressions, ok := orderList.([]parsers.OrderExpression); ok {
-		selectStmt.OrderExpressions = orderExpressions
-	}
+    if orderExpressions, ok := orderList.([]parsers.OrderExpression); ok {
+        selectStmt.OrderExpressions = orderExpressions
+    }
 
-	if groupByClause != nil {
-		selectStmt.GroupBy = groupByClause.([]parsers.SelectItem)
-	}
+    if groupByClause != nil {
+        selectStmt.GroupBy = groupByClause.([]parsers.SelectItem)
+    }
 
-	return selectStmt, nil
+    return selectStmt, nil
 }
 
 func makeJoin(table interface{}, column interface{}) (parsers.JoinItem, error) {
-	joinItem := parsers.JoinItem{}
+    joinItem := parsers.JoinItem{}
 
-	if selectItem, isSelectItem := column.(parsers.SelectItem); isSelectItem {
-		joinItem.SelectItem = selectItem
-		joinItem.Table.Value = selectItem.Alias
-	}
+    if selectItem, isSelectItem := column.(parsers.SelectItem); isSelectItem {
+        joinItem.SelectItem = selectItem
+        joinItem.Table.Value = selectItem.Alias
+    }
 
-	if tableTyped, isTable := table.(parsers.Table); isTable {
-		joinItem.Table = tableTyped
-	}
+    if tableTyped, isTable := table.(parsers.Table); isTable {
+        joinItem.Table = tableTyped
+    }
 
-	return joinItem, nil
+    return joinItem, nil
 }
 
 func makeSelectItem(name interface{}, path interface{}, selectItemType parsers.SelectItemType) (parsers.SelectItem, error) {
-	ps := path.([]interface{})
+    ps := path.([]interface{})
 
-	paths := make([]string, 1)
+    paths := make([]string, 1)
     paths[0] = name.(string)
-	for _, p := range ps {
+    for _, p := range ps {
         paths = append(paths, p.(string))
-	}
+    }
 
-	return parsers.SelectItem{Path: paths, Type: selectItemType}, nil
+    return parsers.SelectItem{Path: paths, Type: selectItemType}, nil
 }
 
 func makeColumnList(column interface{}, other_columns interface{}) ([]parsers.SelectItem, error) {
     collsAsArray := other_columns.([]interface{})
-	columnList := make([]parsers.SelectItem, len(collsAsArray) + 1)
-	columnList[0] = column.(parsers.SelectItem)
+    columnList := make([]parsers.SelectItem, len(collsAsArray) + 1)
+    columnList[0] = column.(parsers.SelectItem)
 
-	for i, v := range collsAsArray {
+    for i, v := range collsAsArray {
         if col, ok := v.(parsers.SelectItem); ok {
-			columnList[i+1] = col
+            columnList[i+1] = col
         }
     }
 
-	return columnList, nil
+    return columnList, nil
 }
 
 func makeSelectArray(columns interface{}) (parsers.SelectItem, error) {
-	return parsers.SelectItem{
-		SelectItems: columns.([]parsers.SelectItem),
-		Type:        parsers.SelectItemTypeArray,
-	}, nil
+    return parsers.SelectItem{
+        SelectItems: columns.([]parsers.SelectItem),
+        Type:        parsers.SelectItemTypeArray,
+    }, nil
 }
 
 func makeSelectObject(field interface{}, other_fields interface{}) (parsers.SelectItem, error) {
     fieldsAsArray := other_fields.([]interface{})
-	fieldsList := make([]parsers.SelectItem, len(fieldsAsArray)+1)
-	fieldsList[0] = field.(parsers.SelectItem)
+    fieldsList := make([]parsers.SelectItem, len(fieldsAsArray)+1)
+    fieldsList[0] = field.(parsers.SelectItem)
 
-	for i, v := range fieldsAsArray {
-		if col, ok := v.(parsers.SelectItem); ok {
-			fieldsList[i+1] = col
-		}
-	}
+    for i, v := range fieldsAsArray {
+        if col, ok := v.(parsers.SelectItem); ok {
+            fieldsList[i+1] = col
+        }
+    }
 
-	return parsers.SelectItem{
-		SelectItems: fieldsList,
-		Type:        parsers.SelectItemTypeObject,
-	}, nil
+    return parsers.SelectItem{
+        SelectItems: fieldsList,
+        Type:        parsers.SelectItemTypeObject,
+    }, nil
 }
 
 func makeOrderByClause(ex1 interface{}, others interface{}) ([]parsers.OrderExpression, error) {
     othersArray := others.([]interface{})
-	orderList := make([]parsers.OrderExpression, len(othersArray)+1)
-	orderList[0] = ex1.(parsers.OrderExpression)
+    orderList := make([]parsers.OrderExpression, len(othersArray)+1)
+    orderList[0] = ex1.(parsers.OrderExpression)
 
-	for i, v := range othersArray {
-		if col, ok := v.(parsers.OrderExpression); ok {
-			orderList[i+1] = col
-		}
-	}
+    for i, v := range othersArray {
+        if col, ok := v.(parsers.OrderExpression); ok {
+            orderList[i+1] = col
+        }
+    }
 
-	return orderList, nil
+    return orderList, nil
 }
 
 func makeOrderExpression(field interface{}, order interface{}) (parsers.OrderExpression, error) {
@@ -144,8 +144,8 @@ func makeOrderExpression(field interface{}, order interface{}) (parsers.OrderExp
     }
 
     if orderValue, ok := order.(parsers.OrderDirection); ok {
-		value.Direction = orderValue
-	}
+        value.Direction = orderValue
+    }
 
     return value, nil
 }
@@ -169,13 +169,13 @@ func joinStrings(array []interface{}) string {
 
 func combineExpressions(ex1 interface{}, exs interface{}, operation parsers.LogicalExpressionType) (interface{}, error) {
     if exs == nil || len(exs.([]interface{})) < 1 {
-		return ex1, nil
-	}
+        return ex1, nil
+    }
 
-	return parsers.LogicalExpression{
-		Expressions: append([]interface{}{ex1}, exs.([]interface{})...),
-		Operation:   operation,
-	}, nil
+    return parsers.LogicalExpression{
+        Expressions: append([]interface{}{ex1}, exs.([]interface{})...),
+        Operation:   operation,
+    }, nil
 }
 
 }
@@ -204,16 +204,16 @@ TopClause <- Top ws count:Integer {
     return count, nil
 }
 
-FromClause <- From ws table:TableName selectItem:(ws In ws column:SelectItem { return column, nil }) {
+FromClause <- From ws table:TableName selectItem:(ws In ws column:SelectItemWithAlias { return column, nil }) {
     tableTyped := table.(parsers.Table)
 
     if selectItem != nil {
-	    tableTyped.SelectItem = selectItem.(parsers.SelectItem)
+        tableTyped.SelectItem = selectItem.(parsers.SelectItem)
         tableTyped.IsInSelect = true
-	}
+    }
 
-	return tableTyped, nil
-} / From ws column:SelectItem {
+    return tableTyped, nil
+} / From ws column:SelectItemWithAlias {
     tableSelectItem := column.(parsers.SelectItem)
     table := parsers.Table{
         Value: tableSelectItem.Alias,
@@ -222,11 +222,11 @@ FromClause <- From ws table:TableName selectItem:(ws In ws column:SelectItem { r
     return table, nil
 } / From ws subQuery:SubQuerySelectItem {
     subQueryTyped := subQuery.(parsers.SelectItem)
-	table := parsers.Table{
-		Value: subQueryTyped.Alias,
-		SelectItem: subQueryTyped,
-	}
-	return table, nil
+    table := parsers.Table{
+        Value: subQueryTyped.Alias,
+        SelectItem: subQueryTyped,
+    }
+    return table, nil
 }
 
 SubQuery <- exists:(exists:Exists ws { return exists, nil })? "(" ws selectStmt:SelectStmt ws ")" {
@@ -251,7 +251,7 @@ SubQuerySelectItem <- subQuery:SubQuery asClause:(ws alias:AsClause { return ali
     return selectItem, nil
 }
 
-JoinClause <- Join ws table:TableName ws In ws column:SelectItem {
+JoinClause <- Join ws table:TableName ws In ws column:SelectItemWithAlias {
     return makeJoin(table, column)
 } / Join ws subQuery:SubQuerySelectItem {
     return makeJoin(nil, subQuery)
@@ -265,17 +265,40 @@ Selection <- SelectValueSpec / ColumnList / SelectAsterisk
 
 SelectAsterisk <- "*" {
     selectItem, _ := makeSelectItem("c", make([]interface{}, 0), parsers.SelectItemTypeField)
-	selectItem.IsTopLevel = true
+    selectItem.IsTopLevel = true
     return makeColumnList(selectItem, make([]interface{}, 0))
 }
 
-ColumnList <- column:SelectItem other_columns:(ws "," ws coll:SelectItem {return coll, nil })* {
+ColumnList <- column:ExpressionOrSelectItem other_columns:(ws "," ws coll:ExpressionOrSelectItem {return coll, nil })* {
     return makeColumnList(column, other_columns)
 }
 
-SelectValueSpec <- "VALUE"i ws column:SelectItem {
+ExpressionOrSelectItem <- expression:OrExpression asClause:AsClause? {
+    switch typedValue := expression.(type) {
+    case parsers.ComparisonExpression, parsers.LogicalExpression:
+        selectItem := parsers.SelectItem{
+            Type: parsers.SelectItemTypeExpression,
+            Value: typedValue,
+        }
+
+        if aliasValue, ok := asClause.(string); ok {
+            selectItem.Alias = aliasValue
+        }
+
+        return selectItem, nil
+    case parsers.SelectItem:
+        if aliasValue, ok := asClause.(string); ok {
+            typedValue.Alias = aliasValue
+        }
+        return typedValue, nil
+    default:
+        return typedValue, nil
+    }
+} / item:SelectItemWithAlias { return item, nil }
+
+SelectValueSpec <- "VALUE"i ws column:SelectItemWithAlias {
     selectItem := column.(parsers.SelectItem)
-	selectItem.IsTopLevel = true
+    selectItem.IsTopLevel = true
     return makeColumnList(selectItem, make([]interface{}, 0))
 }
 
@@ -289,19 +312,32 @@ SelectArray <- "[" ws columns:ColumnList ws "]" {
 
 SelectObject <- "{" ws field:SelectObjectField ws other_fields:(ws "," ws coll:SelectObjectField {return coll, nil })* ws "}" {
     return makeSelectObject(field, other_fields)
+} / "{" ws "}" {
+    return parsers.SelectItem{
+        SelectItems: []parsers.SelectItem{},
+        Type:        parsers.SelectItemTypeObject,
+    }, nil
 }
 
 SelectObjectField <- name:(Identifier / "\"" key:Identifier "\"" { return key, nil }) ws ":" ws selectItem:SelectItem {
     item := selectItem.(parsers.SelectItem)
-	item.Alias = name.(string)
-	return item, nil
+    item.Alias = name.(string)
+    return item, nil
 }
 
 SelectProperty <- name:Identifier path:(DotFieldAccess / ArrayFieldAccess)* {
     return makeSelectItem(name, path, parsers.SelectItemTypeField)
 }
 
-SelectItem <- selectItem:(SubQuerySelectItem / Literal / FunctionCall / SelectArray / SelectObject / SelectProperty) asClause:AsClause? {
+SelectItemWithAlias <- selectItem:SelectItem asClause:AsClause? {
+    item := selectItem.(parsers.SelectItem)
+    if aliasValue, ok := asClause.(string); ok {
+        item.Alias = aliasValue
+    }
+    return item, nil
+}
+
+SelectItem <- selectItem:(SubQuerySelectItem / Literal / FunctionCall / SelectArray / SelectObject / SelectProperty) {
     var itemResult parsers.SelectItem
     switch typedValue := selectItem.(type) {
     case parsers.SelectItem:
@@ -318,11 +354,7 @@ SelectItem <- selectItem:(SubQuerySelectItem / Literal / FunctionCall / SelectAr
         }
     }
 
-	if aliasValue, ok := asClause.(string); ok {
-		itemResult.Alias = aliasValue
-	}
-
-	return itemResult, nil
+    return itemResult, nil
 }
 
 AsClause <- (ws As)? ws !ExcludedKeywords alias:Identifier {
@@ -360,10 +392,10 @@ ComparisonExpression <- "(" ws ex:OrExpression ws ")" { return ex, nil }
     return parsers.ComparisonExpression{Left:left,Right:right,Operation:op.(string)}, nil
 } / inv:(Not ws)? ex:SelectItem {
     if inv != nil {
-		ex1 := ex.(parsers.SelectItem)
-		ex1.Invert = true
-		return ex1, nil
-	}
+        ex1 := ex.(parsers.SelectItem)
+        ex1.Invert = true
+        return ex1, nil
+    }
     return ex, nil
 } / ex:BooleanLiteral { return ex, nil }
 
@@ -377,10 +409,10 @@ OrderExpression <- field:SelectProperty ws order:OrderDirection? {
 
 OrderDirection <- ("ASC"i / "DESC"i) {
     if strings.EqualFold(string(c.text), "DESC") {
-		return parsers.OrderDirectionDesc, nil
+        return parsers.OrderDirectionDesc, nil
     }
 
-	return parsers.OrderDirectionAsc, nil
+    return parsers.OrderDirectionAsc, nil
 }
 
 Select <- "SELECT"i
@@ -442,6 +474,7 @@ BooleanLiteral <- ("true"i / "false"i) {
 FunctionCall <- StringFunctions
     / TypeCheckingFunctions
     / ArrayFunctions
+    / ConditionalFunctions
     / InFunction
     / AggregateFunctions
     / MathFunctions
@@ -489,6 +522,8 @@ ArrayFunctions <- ArrayConcatExpression
     / SetIntersectExpression
     / SetUnionExpression
 
+ConditionalFunctions <- IifExpression
+
 MathFunctions <- MathAbsExpression
     / MathAcosExpression
     / MathAsinExpression
@@ -681,6 +716,10 @@ SetUnionExpression <- "SetUnion"i ws "(" ws set1:SelectItem ws "," ws set2:Selec
     return createFunctionCall(parsers.FunctionCallSetUnion, []interface{}{set1, set2})
 }
 
+IifExpression <- "IIF"i ws "(" ws condition:SelectItem ws "," ws trueValue:SelectItem ws "," ws falseValue:SelectItem ws ")" {
+    return createFunctionCall(parsers.FunctionCallIif, []interface{}{condition, trueValue, falseValue})
+}
+
 MathAbsExpression <- "ABS"i ws "(" ws ex:SelectItem ws ")" { return createFunctionCall(parsers.FunctionCallMathAbs, []interface{}{ex}) }
 MathAcosExpression <- "ACOS"i ws "(" ws ex:SelectItem ws ")" { return createFunctionCall(parsers.FunctionCallMathAcos, []interface{}{ex}) }
 MathAsinExpression <- "ASIN"i ws "(" ws ex:SelectItem ws ")" { return createFunctionCall(parsers.FunctionCallMathAsin, []interface{}{ex}) }

parsers/nosql/select_test.go

@@ -178,4 +178,90 @@ func Test_Parse_Select(t *testing.T) {
 			},
 		)
 	})
+
+	t.Run("Should parse SELECT empty object", func(t *testing.T) {
+		testQueryParse(
+			t,
+			`SELECT {} AS obj FROM c`,
+			parsers.SelectStmt{
+				SelectItems: []parsers.SelectItem{
+					{
+						Alias:       "obj",
+						Type:        parsers.SelectItemTypeObject,
+						SelectItems: []parsers.SelectItem{},
+					},
+				},
+				Table: parsers.Table{SelectItem: testutils.SelectItem_Path("c")},
+			},
+		)
+	})
+
+	t.Run("Should parse comparison expressions in SELECT", func(t *testing.T) {
+		testQueryParse(
+			t,
+			`SELECT c["id"] = "123", c["pk"] > 456 FROM c`,
+			parsers.SelectStmt{
+				SelectItems: []parsers.SelectItem{
+					{
+						Type: parsers.SelectItemTypeExpression,
+						Value: parsers.ComparisonExpression{
+							Operation: "=",
+							Left:      testutils.SelectItem_Path("c", "id"),
+							Right:     testutils.SelectItem_Constant_String("123"),
+						},
+					},
+					{
+						Type: parsers.SelectItemTypeExpression,
+						Value: parsers.ComparisonExpression{
+							Operation: ">",
+							Left:      testutils.SelectItem_Path("c", "pk"),
+							Right:     testutils.SelectItem_Constant_Int(456),
+						},
+					},
+				},
+				Table: parsers.Table{SelectItem: testutils.SelectItem_Path("c")},
+			},
+		)
+	})
+
+	t.Run("Should parse logical expressions in SELECT", func(t *testing.T) {
+		testQueryParse(
+			t,
+			`SELECT c["id"] = "123" OR c["pk"] > 456, c["isCool"] AND c["hasRizz"] AS isRizzler FROM c`,
+			parsers.SelectStmt{
+				SelectItems: []parsers.SelectItem{
+					{
+						Type: parsers.SelectItemTypeExpression,
+						Value: parsers.LogicalExpression{
+							Operation: parsers.LogicalExpressionTypeOr,
+							Expressions: []interface{}{
+								parsers.ComparisonExpression{
+									Operation: "=",
+									Left:      testutils.SelectItem_Path("c", "id"),
+									Right:     testutils.SelectItem_Constant_String("123"),
+								},
+								parsers.ComparisonExpression{
+									Operation: ">",
+									Left:      testutils.SelectItem_Path("c", "pk"),
+									Right:     testutils.SelectItem_Constant_Int(456),
+								},
+							},
+						},
+					},
+					{
+						Type:  parsers.SelectItemTypeExpression,
+						Alias: "isRizzler",
+						Value: parsers.LogicalExpression{
+							Operation: parsers.LogicalExpressionTypeAnd,
+							Expressions: []interface{}{
+								testutils.SelectItem_Path("c", "isCool"),
+								testutils.SelectItem_Path("c", "hasRizz"),
+							},
+						},
+					},
+				},
+				Table: parsers.Table{SelectItem: testutils.SelectItem_Path("c")},
+			},
+		)
+	})
 }

query_executors/memory_executor/common.go

@@ -69,6 +69,19 @@ func (r rowContext) resolveSelectItem(selectItem parsers.SelectItem) interface{}
 		return nil
 	}
 
+	if selectItem.Type == parsers.SelectItemTypeExpression {
+		if typedExpression, ok := selectItem.Value.(parsers.ComparisonExpression); ok {
+			return r.filters_ComparisonExpression(typedExpression)
+		}
+
+		if typedExpression, ok := selectItem.Value.(parsers.LogicalExpression); ok {
+			return r.filters_LogicalExpression(typedExpression)
+		}
+
+		logger.ErrorLn("parsers.SelectItem has incorrect Value type (expected parsers.ComparisonExpression)")
+		return nil
+	}
+
 	return r.selectItem_SelectItemTypeField(selectItem)
 }
 
@@ -209,6 +222,9 @@ func (r rowContext) selectItem_SelectItemTypeFunctionCall(functionCall parsers.F
 	case parsers.FunctionCallSetUnion:
 		return r.set_Union(functionCall.Arguments)
 
+	case parsers.FunctionCallIif:
+		return r.misc_Iif(functionCall.Arguments)
+
 	case parsers.FunctionCallMathAbs:
 		return r.math_Abs(functionCall.Arguments)
 	case parsers.FunctionCallMathAcos:

query_executors/memory_executor/comparison_expression_test.go

@@ -0,0 +1,92 @@
+package memoryexecutor_test
+
+import (
+	"testing"
+
+	"github.com/pikami/cosmium/parsers"
+	memoryexecutor "github.com/pikami/cosmium/query_executors/memory_executor"
+	testutils "github.com/pikami/cosmium/test_utils"
+)
+
+func Test_Execute_Expressions(t *testing.T) {
+	mockData := []memoryexecutor.RowType{
+		map[string]interface{}{"id": "123", "age": 10, "isCool": true},
+		map[string]interface{}{"id": "456", "age": 20, "isCool": false},
+		map[string]interface{}{"id": "789", "age": 30, "isCool": true},
+	}
+
+	t.Run("Should execute comparison expressions in SELECT", func(t *testing.T) {
+		testQueryExecute(
+			t,
+			parsers.SelectStmt{
+				SelectItems: []parsers.SelectItem{
+					{
+						Path: []string{"c", "id"},
+						Type: parsers.SelectItemTypeField,
+					},
+					{
+						Alias: "isAdult",
+						Type:  parsers.SelectItemTypeExpression,
+						Value: parsers.ComparisonExpression{
+							Operation: ">=",
+							Left:      testutils.SelectItem_Path("c", "age"),
+							Right:     testutils.SelectItem_Constant_Int(18),
+						},
+					},
+					{
+						Alias: "isNotCool",
+						Type:  parsers.SelectItemTypeExpression,
+						Value: parsers.ComparisonExpression{
+							Operation: "!=",
+							Left:      testutils.SelectItem_Path("c", "isCool"),
+							Right:     testutils.SelectItem_Constant_Bool(true),
+						},
+					},
+				},
+				Table: parsers.Table{SelectItem: testutils.SelectItem_Path("c")},
+			},
+			mockData,
+			[]memoryexecutor.RowType{
+				map[string]interface{}{"id": "123", "isAdult": false, "isNotCool": false},
+				map[string]interface{}{"id": "456", "isAdult": true, "isNotCool": true},
+				map[string]interface{}{"id": "789", "isAdult": true, "isNotCool": false},
+			},
+		)
+	})
+
+	t.Run("Should execute logical expressions in SELECT", func(t *testing.T) {
+		testQueryExecute(
+			t,
+			parsers.SelectStmt{
+				SelectItems: []parsers.SelectItem{
+					{
+						Path: []string{"c", "id"},
+						Type: parsers.SelectItemTypeField,
+					},
+					{
+						Alias: "isCoolAndAdult",
+						Type:  parsers.SelectItemTypeExpression,
+						Value: parsers.LogicalExpression{
+							Operation: parsers.LogicalExpressionTypeAnd,
+							Expressions: []interface{}{
+								testutils.SelectItem_Path("c", "isCool"),
+								parsers.ComparisonExpression{
+									Operation: ">=",
+									Left:      testutils.SelectItem_Path("c", "age"),
+									Right:     testutils.SelectItem_Constant_Int(18),
+								},
+							},
+						},
+					},
+				},
+				Table: parsers.Table{SelectItem: testutils.SelectItem_Path("c")},
+			},
+			mockData,
+			[]memoryexecutor.RowType{
+				map[string]interface{}{"id": "123", "isCoolAndAdult": false},
+				map[string]interface{}{"id": "456", "isCoolAndAdult": false},
+				map[string]interface{}{"id": "789", "isCoolAndAdult": true},
+			},
+		)
+	})
+}

query_executors/memory_executor/misc_functions.go

@@ -16,3 +16,16 @@ func (r rowContext) misc_In(arguments []interface{}) bool {
 
 	return false
 }
+
+func (r rowContext) misc_Iif(arguments []interface{}) interface{} {
+	if len(arguments) != 3 {
+		return nil
+	}
+
+	condition := r.resolveSelectItem(arguments[0].(parsers.SelectItem))
+	if condition != nil && condition == true {
+		return r.resolveSelectItem(arguments[1].(parsers.SelectItem))
+	}
+
+	return r.resolveSelectItem(arguments[2].(parsers.SelectItem))
+}

query_executors/memory_executor/misc_test.go

@@ -210,4 +210,35 @@ func Test_Execute(t *testing.T) {
 			},
 		)
 	})
+
+	t.Run("Should execute function IIF()", func(t *testing.T) {
+		testQueryExecute(
+			t,
+			parsers.SelectStmt{
+				SelectItems: []parsers.SelectItem{
+					testutils.SelectItem_Path("c", "id"),
+					{
+						Alias: "coolness",
+						Type:  parsers.SelectItemTypeFunctionCall,
+						Value: parsers.FunctionCall{
+							Type: parsers.FunctionCallIif,
+							Arguments: []interface{}{
+								testutils.SelectItem_Path("c", "isCool"),
+								testutils.SelectItem_Constant_String("real cool"),
+								testutils.SelectItem_Constant_String("not cool"),
+							},
+						},
+					},
+				},
+				Table: parsers.Table{SelectItem: testutils.SelectItem_Path("c")},
+			},
+			mockData,
+			[]memoryexecutor.RowType{
+				map[string]interface{}{"id": "12345", "coolness": "not cool"},
+				map[string]interface{}{"id": "67890", "coolness": "real cool"},
+				map[string]interface{}{"id": "456", "coolness": "real cool"},
+				map[string]interface{}{"id": "123", "coolness": "real cool"},
+			},
+		)
+	})
 }

query_executors/memory_executor/select_test.go

@@ -205,4 +205,27 @@ func Test_Execute_Select(t *testing.T) {
 			},
 		)
 	})
+
+	t.Run("Should execute SELECT empty object", func(t *testing.T) {
+		testQueryExecute(
+			t,
+			parsers.SelectStmt{
+				SelectItems: []parsers.SelectItem{
+					{
+						Alias:       "obj",
+						Type:        parsers.SelectItemTypeObject,
+						SelectItems: []parsers.SelectItem{},
+					},
+				},
+				Table: parsers.Table{SelectItem: testutils.SelectItem_Path("c")},
+			},
+			mockData,
+			[]memoryexecutor.RowType{
+				map[string]interface{}{"obj": map[string]interface{}{}},
+				map[string]interface{}{"obj": map[string]interface{}{}},
+				map[string]interface{}{"obj": map[string]interface{}{}},
+				map[string]interface{}{"obj": map[string]interface{}{}},
+			},
+		)
+	})
 }

sharedlibrary/shared.go

@@ -20,14 +20,14 @@ var (
 const (
 	ResponseSuccess = 0
 
-	ResponseUnknown                                   = 100
-	ResponseFailedToParseConfiguration                = 101
-	ResponseFailedToLoadState                         = 102
-	ResponseFailedToParseRequest                      = 103
-	ResponseServerInstanceAlreadyExists               = 104
-	ResponseServerInstanceNotFound                    = 105
-	ResponseFailedToStartServer                       = 106
-	ResponseCurentDataStoreDoesNotSupportStateLoading = 107
+	ResponseUnknown                                    = 100
+	ResponseFailedToParseConfiguration                 = 101
+	ResponseFailedToLoadState                          = 102
+	ResponseFailedToParseRequest                       = 103
+	ResponseServerInstanceAlreadyExists                = 104
+	ResponseServerInstanceNotFound                     = 105
+	ResponseFailedToStartServer                        = 106
+	ResponseCurrentDataStoreDoesNotSupportStateLoading = 107
 
 	ResponseDataStoreNotFound   = 200
 	ResponseDataStoreConflict   = 201

sharedlibrary/sharedlibrary.go

@@ -11,7 +11,9 @@ import (
 
 	"github.com/pikami/cosmium/api"
 	"github.com/pikami/cosmium/api/config"
-	mapdatastore "github.com/pikami/cosmium/internal/datastore/map_datastore"
+	"github.com/pikami/cosmium/internal/datastore"
+	badgerdatastore "github.com/pikami/cosmium/internal/datastore/badger_datastore"
+	jsondatastore "github.com/pikami/cosmium/internal/datastore/json_datastore"
 )
 
 //export CreateServerInstance
@@ -32,10 +34,18 @@ func CreateServerInstance(serverName *C.char, configurationJSON *C.char) int {
 	configuration.ApplyDefaultsToEmptyFields()
 	configuration.PopulateCalculatedFields()
 
-	dataStore := mapdatastore.NewMapDataStore(mapdatastore.MapDataStoreOptions{
-		InitialDataFilePath: configuration.InitialDataFilePath,
-		PersistDataFilePath: configuration.PersistDataFilePath,
-	})
+	var dataStore datastore.DataStore
+	switch configuration.DataStore {
+	case config.DataStoreBadger:
+		dataStore = badgerdatastore.NewBadgerDataStore(badgerdatastore.BadgerDataStoreOptions{
+			PersistDataFilePath: configuration.PersistDataFilePath,
+		})
+	default:
+		dataStore = jsondatastore.NewJsonDataStore(jsondatastore.JsonDataStoreOptions{
+			InitialDataFilePath: configuration.InitialDataFilePath,
+			PersistDataFilePath: configuration.PersistDataFilePath,
+		})
+	}
 
 	server := api.NewApiServer(dataStore, &configuration)
 	err = server.Start()
@@ -57,6 +67,7 @@ func StopServerInstance(serverName *C.char) int {
 
 	if serverInstance, ok := getInstance(serverNameStr); ok {
 		serverInstance.server.Stop()
+		serverInstance.dataStore.Close()
 		removeInstance(serverNameStr)
 		return ResponseSuccess
 	}
@@ -85,14 +96,14 @@ func LoadServerInstanceState(serverName *C.char, stateJSON *C.char) int {
 	stateJSONStr := C.GoString(stateJSON)
 
 	if serverInstance, ok := getInstance(serverNameStr); ok {
-		if mapDS, ok := serverInstance.dataStore.(*mapdatastore.MapDataStore); ok {
-			err := mapDS.LoadStateJSON(stateJSONStr)
+		if jsonDS, ok := serverInstance.dataStore.(*jsondatastore.JsonDataStore); ok {
+			err := jsonDS.LoadStateJSON(stateJSONStr)
 			if err != nil {
 				return ResponseFailedToLoadState
 			}
 			return ResponseSuccess
 		}
-		return ResponseCurentDataStoreDoesNotSupportStateLoading
+		return ResponseCurrentDataStoreDoesNotSupportStateLoading
 	}
 
 	return ResponseServerInstanceNotFound