Added support for initial data loading when using badger db

README.md

@@ -105,10 +105,10 @@ All mentioned arguments can also be set using environment variables:
 
 Cosmium supports multiple storage backends for saving, loading, and managing data at runtime.
 
-| Backend  | Storage Location         | Write Behavior           | Memory Usage         | Supports Initial JSON Load |
+| Backend  | Storage Location         | Write Behavior           | Memory Usage         |
-|----------|--------------------------|--------------------------|----------------------|----------------------------|
+|----------|--------------------------|--------------------------|----------------------|
-| `json` (default) | JSON file on disk 📄 | On application exit ⏳ | 🛑 More than Badger | ✅ Yes                      |
+| `json` (default) | JSON file on disk 📄 | On application exit ⏳ | 🛑 More than Badger |
-| `badger`  | BadgerDB database on disk ⚡ | Immediately on write 🚀 | ✅ Less than JSON  | ❌ No                       |
+| `badger`  | BadgerDB database on disk ⚡ | Immediately on write 🚀 | ✅ Less than JSON  |
 
 
 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.

api/config/config.go

@@ -94,11 +94,6 @@ func (c *ServerConfig) PopulateCalculatedFields() {
 			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() {

cmd/server/server.go

@@ -20,6 +20,7 @@ func main() {
 	switch configuration.DataStore {
 	case config.DataStoreBadger:
 		dataStore = badgerdatastore.NewBadgerDataStore(badgerdatastore.BadgerDataStoreOptions{
+			InitialDataFilePath: configuration.InitialDataFilePath,
 			PersistDataFilePath: configuration.PersistDataFilePath,
 		})
 		logger.InfoLn("Using Badger data store")

internal/datastore/badger_datastore/badger_datastore.go

@@ -1,9 +1,13 @@
 package badgerdatastore
 
 import (
+	"encoding/json"
+	"log"
+	"os"
 	"time"
 
 	"github.com/dgraph-io/badger/v4"
+	"github.com/pikami/cosmium/internal/datastore"
 	"github.com/pikami/cosmium/internal/logger"
 )
 
@@ -13,6 +17,7 @@ type BadgerDataStore struct {
 }
 
 type BadgerDataStoreOptions struct {
+	InitialDataFilePath string
 	PersistDataFilePath string
 }
 
@@ -35,6 +40,8 @@ func NewBadgerDataStore(options BadgerDataStoreOptions) *BadgerDataStore {
 		gcTicker: gcTicker,
 	}
 
+	ds.initializeDataStore(options.InitialDataFilePath)
+
 	go ds.runGarbageCollector()
 
 	return ds
@@ -64,3 +71,53 @@ func (r *BadgerDataStore) runGarbageCollector() {
 		}
 	}
 }
+
+func (r *BadgerDataStore) initializeDataStore(initialDataFilePath string) {
+	if initialDataFilePath == "" {
+		return
+	}
+
+	stat, err := os.Stat(initialDataFilePath)
+	if err != nil {
+		panic(err)
+	}
+
+	if stat.IsDir() {
+		logger.ErrorLn("Argument '-Persist' must be a path to file, not a directory.")
+		os.Exit(1)
+	}
+
+	jsonData, err := os.ReadFile(initialDataFilePath)
+	if err != nil {
+		log.Fatalf("Error reading state JSON file: %v", err)
+		return
+	}
+
+	var state datastore.InitialDataModel
+	if err := json.Unmarshal([]byte(jsonData), &state); err != nil {
+		log.Fatalf("Error parsing state JSON file: %v", err)
+		return
+	}
+
+	for dbName, dbModel := range state.Databases {
+		r.CreateDatabase(dbModel)
+		for colName, colModel := range state.Collections[dbName] {
+			r.CreateCollection(dbName, colModel)
+			for _, docModel := range state.Documents[dbName][colName] {
+				r.CreateDocument(dbName, colName, docModel)
+			}
+
+			for _, triggerModel := range state.Triggers[dbName][colName] {
+				r.CreateTrigger(dbName, colName, triggerModel)
+			}
+
+			for _, spModel := range state.StoredProcedures[dbName][colName] {
+				r.CreateStoredProcedure(dbName, colName, spModel)
+			}
+
+			for _, udfModel := range state.UserDefinedFunctions[dbName][colName] {
+				r.CreateUserDefinedFunction(dbName, colName, udfModel)
+			}
+		}
+	}
+}

internal/datastore/initial_data.go

@@ -0,0 +1,21 @@
+package datastore
+
+type InitialDataModel struct {
+	// Map databaseId -> Database
+	Databases map[string]Database `json:"databases"`
+
+	// Map databaseId -> collectionId -> Collection
+	Collections map[string]map[string]Collection `json:"collections"`
+
+	// Map databaseId -> collectionId -> documentId -> Documents
+	Documents map[string]map[string]map[string]Document `json:"documents"`
+
+	// Map databaseId -> collectionId -> triggerId -> Trigger
+	Triggers map[string]map[string]map[string]Trigger `json:"triggers"`
+
+	// Map databaseId -> collectionId -> spId -> StoredProcedure
+	StoredProcedures map[string]map[string]map[string]StoredProcedure `json:"sprocs"`
+
+	// Map databaseId -> collectionId -> udfId -> UserDefinedFunction
+	UserDefinedFunctions map[string]map[string]map[string]UserDefinedFunction `json:"udfs"`
+}

sharedlibrary/sharedlibrary.go

@@ -38,6 +38,7 @@ func CreateServerInstance(serverName *C.char, configurationJSON *C.char) int {
 	switch configuration.DataStore {
 	case config.DataStoreBadger:
 		dataStore = badgerdatastore.NewBadgerDataStore(badgerdatastore.BadgerDataStoreOptions{
+			InitialDataFilePath: configuration.InitialDataFilePath,
 			PersistDataFilePath: configuration.PersistDataFilePath,
 		})
 	default: