cosmium/internal/datastore/badger_datastore/badger_datastore.go

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"
)

type BadgerDataStore struct {
	db       *badger.DB
	gcTicker *time.Ticker
}

type BadgerDataStoreOptions struct {
	InitialDataFilePath string
	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)
	}

	gcTicker := time.NewTicker(5 * time.Minute)

	ds := &BadgerDataStore{
		db:       db,
		gcTicker: gcTicker,
	}

	ds.initializeDataStore(options.InitialDataFilePath)

	go ds.runGarbageCollector()

	return ds
}

func (r *BadgerDataStore) Close() {
	if r.gcTicker != nil {
		r.gcTicker.Stop()
		r.gcTicker = nil
	}

	r.db.Close()
	r.db = nil
}

func (r *BadgerDataStore) DumpToJson() (string, error) {
	logger.ErrorLn("Badger datastore does not support state export currently.")
	return "{}", nil
}

func (r *BadgerDataStore) runGarbageCollector() {
	for range r.gcTicker.C {
	again:
		err := r.db.RunValueLogGC(0.7)
		if err == nil {
			goto again
		}
	}
}

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)
			}
		}
	}
}