Code cleanup; Implement persistant storage; Use maps for storage

internal/repositories/collections.go

@@ -7,48 +7,50 @@ import (
 	"github.com/google/uuid"
 	repositorymodels "github.com/pikami/cosmium/internal/repository_models"
 	structhidrators "github.com/pikami/cosmium/internal/struct_hidrators"
+	"golang.org/x/exp/maps"
 )
 
-var collections = []repositorymodels.Collection{}
-
 func GetAllCollections(databaseId string) ([]repositorymodels.Collection, repositorymodels.RepositoryStatus) {
-	dbCollections := make([]repositorymodels.Collection, 0)
-
-	for _, coll := range collections {
-		if coll.Internals.DatabaseId == databaseId {
-			dbCollections = append(dbCollections, coll)
-		}
+	if _, ok := storeState.Databases[databaseId]; !ok {
+		return make([]repositorymodels.Collection, 0), repositorymodels.StatusNotFound
 	}
 
-	return dbCollections, repositorymodels.StatusOk
+	return maps.Values(storeState.Collections[databaseId]), repositorymodels.StatusOk
 }
 
-func GetCollection(databaseId string, id string) (repositorymodels.Collection, repositorymodels.RepositoryStatus) {
-	for _, coll := range collections {
-		if coll.Internals.DatabaseId == databaseId && coll.ID == id {
-			return coll, repositorymodels.StatusOk
-		}
+func GetCollection(databaseId string, collectionId string) (repositorymodels.Collection, repositorymodels.RepositoryStatus) {
+	if _, ok := storeState.Databases[databaseId]; !ok {
+		return repositorymodels.Collection{}, repositorymodels.StatusNotFound
 	}
 
-	return repositorymodels.Collection{}, repositorymodels.StatusNotFound
+	if _, ok := storeState.Collections[databaseId][collectionId]; !ok {
+		return repositorymodels.Collection{}, repositorymodels.StatusNotFound
+	}
+
+	return storeState.Collections[databaseId][collectionId], repositorymodels.StatusOk
 }
 
-func DeleteCollection(databaseId string, id string) repositorymodels.RepositoryStatus {
-	for index, coll := range collections {
-		if coll.Internals.DatabaseId == databaseId && coll.ID == id {
-			collections = append(collections[:index], collections[index+1:]...)
-			return repositorymodels.StatusOk
-		}
+func DeleteCollection(databaseId string, collectionId string) repositorymodels.RepositoryStatus {
+	if _, ok := storeState.Databases[databaseId]; !ok {
+		return repositorymodels.StatusNotFound
 	}
 
-	return repositorymodels.StatusNotFound
+	if _, ok := storeState.Collections[databaseId][collectionId]; !ok {
+		return repositorymodels.StatusNotFound
+	}
+
+	delete(storeState.Collections[databaseId], collectionId)
+
+	return repositorymodels.StatusOk
 }
 
 func CreateCollection(databaseId string, newCollection repositorymodels.Collection) (repositorymodels.Collection, repositorymodels.RepositoryStatus) {
-	for _, coll := range collections {
-		if coll.Internals.DatabaseId == databaseId && coll.ID == newCollection.ID {
-			return repositorymodels.Collection{}, repositorymodels.Conflict
-		}
+	if _, ok := storeState.Databases[databaseId]; !ok {
+		return repositorymodels.Collection{}, repositorymodels.StatusNotFound
+	}
+
+	if _, ok := storeState.Collections[databaseId][newCollection.ID]; ok {
+		return repositorymodels.Collection{}, repositorymodels.Conflict
 	}
 
 	newCollection = structhidrators.Hidrate(newCollection).(repositorymodels.Collection)
@@ -56,9 +58,9 @@ func CreateCollection(databaseId string, newCollection repositorymodels.Collecti
 	newCollection.TimeStamp = time.Now().Unix()
 	newCollection.UniqueID = uuid.New().String()
 	newCollection.ETag = fmt.Sprintf("\"%s\"", newCollection.UniqueID)
-	newCollection.Internals = struct{ DatabaseId string }{
-		DatabaseId: databaseId,
-	}
-	collections = append(collections, newCollection)
+
+	storeState.Collections[databaseId][newCollection.ID] = newCollection
+	storeState.Documents[databaseId][newCollection.ID] = make(map[string]repositorymodels.Document)
+
 	return newCollection, repositorymodels.StatusOk
 }

internal/repositories/databases.go

@@ -6,48 +6,42 @@ import (
 
 	"github.com/google/uuid"
 	repositorymodels "github.com/pikami/cosmium/internal/repository_models"
+	"golang.org/x/exp/maps"
 )
 
-var databases = []repositorymodels.Database{
-	{ID: "db1"},
-	{ID: "db2"},
-}
-
 func GetAllDatabases() ([]repositorymodels.Database, repositorymodels.RepositoryStatus) {
-	return databases, repositorymodels.StatusOk
+	return maps.Values(storeState.Databases), repositorymodels.StatusOk
 }
 
 func GetDatabase(id string) (repositorymodels.Database, repositorymodels.RepositoryStatus) {
-	for _, db := range databases {
-		if db.ID == id {
-			return db, repositorymodels.StatusOk
-		}
+	if database, ok := storeState.Databases[id]; ok {
+		return database, repositorymodels.StatusOk
 	}
 
 	return repositorymodels.Database{}, repositorymodels.StatusNotFound
 }
 
 func DeleteDatabase(id string) repositorymodels.RepositoryStatus {
-	for index, db := range databases {
-		if db.ID == id {
-			databases = append(databases[:index], databases[index+1:]...)
-			return repositorymodels.StatusOk
-		}
+	if _, ok := storeState.Databases[id]; !ok {
+		return repositorymodels.StatusNotFound
 	}
 
-	return repositorymodels.StatusNotFound
+	delete(storeState.Databases, id)
+
+	return repositorymodels.StatusOk
 }
 
-func CreateDatabase(newDatabase repositorymodels.Database) repositorymodels.RepositoryStatus {
-	for _, db := range databases {
-		if db.ID == newDatabase.ID {
-			return repositorymodels.Conflict
-		}
+func CreateDatabase(newDatabase repositorymodels.Database) (repositorymodels.Database, repositorymodels.RepositoryStatus) {
+	if _, ok := storeState.Databases[newDatabase.ID]; ok {
+		return repositorymodels.Database{}, repositorymodels.Conflict
 	}
 
 	newDatabase.TimeStamp = time.Now().Unix()
 	newDatabase.UniqueID = uuid.New().String()
 	newDatabase.ETag = fmt.Sprintf("\"%s\"", newDatabase.UniqueID)
-	databases = append(databases, newDatabase)
-	return repositorymodels.StatusOk
+	storeState.Databases[newDatabase.ID] = newDatabase
+	storeState.Collections[newDatabase.ID] = make(map[string]repositorymodels.Collection)
+	storeState.Documents[newDatabase.ID] = make(map[string]map[string]repositorymodels.Document)
+
+	return newDatabase, repositorymodels.StatusOk
 }

internal/repositories/documents.go

@@ -3,7 +3,6 @@ package repositories
 import (
 	"fmt"
 	"log"
-	"strings"
 	"time"
 
 	"github.com/google/uuid"
@@ -11,102 +10,81 @@ import (
 	"github.com/pikami/cosmium/parsers"
 	"github.com/pikami/cosmium/parsers/nosql"
 	memoryexecutor "github.com/pikami/cosmium/query_executors/memory_executor"
+	"golang.org/x/exp/maps"
 )
 
-var documents = []repositorymodels.Document{}
-
 func GetAllDocuments(databaseId string, collectionId string) ([]repositorymodels.Document, repositorymodels.RepositoryStatus) {
-	filteredDocuments := make([]repositorymodels.Document, 0)
-
-	for _, doc := range documents {
-		docDbId := doc["_internal"].(map[string]interface{})["databaseId"]
-		docCollId := doc["_internal"].(map[string]interface{})["collectionId"]
-
-		if docDbId == databaseId && docCollId == collectionId {
-			doc["_partitionKeyValue"] = doc["_internal"].(map[string]interface{})["partitionKeyValue"]
-			filteredDocuments = append(filteredDocuments, doc)
-		}
+	if _, ok := storeState.Databases[databaseId]; !ok {
+		return make([]repositorymodels.Document, 0), repositorymodels.StatusNotFound
 	}
 
-	return filteredDocuments, repositorymodels.StatusOk
+	if _, ok := storeState.Collections[databaseId][collectionId]; !ok {
+		return make([]repositorymodels.Document, 0), repositorymodels.StatusNotFound
+	}
+
+	return maps.Values(storeState.Documents[databaseId][collectionId]), repositorymodels.StatusOk
 }
 
 func GetDocument(databaseId string, collectionId string, documentId string) (repositorymodels.Document, repositorymodels.RepositoryStatus) {
-	for _, doc := range documents {
-		docDbId := doc["_internal"].(map[string]interface{})["databaseId"]
-		docCollId := doc["_internal"].(map[string]interface{})["collectionId"]
-		docId := doc["id"]
-
-		if docDbId == databaseId && docCollId == collectionId && docId == documentId {
-			doc["_partitionKeyValue"] = doc["_internal"].(map[string]interface{})["partitionKeyValue"]
-			return doc, repositorymodels.StatusOk
-		}
+	if _, ok := storeState.Databases[databaseId]; !ok {
+		return repositorymodels.Document{}, repositorymodels.StatusNotFound
 	}
 
-	return repositorymodels.Document{}, repositorymodels.StatusNotFound
+	if _, ok := storeState.Collections[databaseId][collectionId]; !ok {
+		return repositorymodels.Document{}, repositorymodels.StatusNotFound
+	}
+
+	if _, ok := storeState.Documents[databaseId][collectionId][documentId]; !ok {
+		return repositorymodels.Document{}, repositorymodels.StatusNotFound
+	}
+
+	return storeState.Documents[databaseId][collectionId][documentId], repositorymodels.StatusOk
 }
 
 func DeleteDocument(databaseId string, collectionId string, documentId string) repositorymodels.RepositoryStatus {
-	for index, doc := range documents {
-		docDbId := doc["_internal"].(map[string]interface{})["databaseId"]
-		docCollId := doc["_internal"].(map[string]interface{})["collectionId"]
-		docId := doc["id"]
-
-		if docDbId == databaseId && docCollId == collectionId && docId == documentId {
-			documents = append(documents[:index], documents[index+1:]...)
-			return repositorymodels.StatusOk
-		}
-	}
-
-	return repositorymodels.StatusNotFound
-}
-
-func CreateDocument(databaseId string, collectionId string, document map[string]interface{}) repositorymodels.RepositoryStatus {
-	if document["id"] == "" {
-		return repositorymodels.BadRequest
-	}
-
-	collection, status := GetCollection(databaseId, collectionId)
-	if status != repositorymodels.StatusOk {
+	if _, ok := storeState.Databases[databaseId]; !ok {
 		return repositorymodels.StatusNotFound
 	}
 
-	for _, doc := range documents {
-		docDbId := doc["_internal"].(map[string]interface{})["databaseId"]
-		docCollId := doc["_internal"].(map[string]interface{})["collectionId"]
-		docId := doc["id"]
-
-		if docDbId == databaseId && docCollId == collectionId && docId == document["id"] {
-			return repositorymodels.Conflict
-		}
+	if _, ok := storeState.Collections[databaseId][collectionId]; !ok {
+		return repositorymodels.StatusNotFound
 	}
 
-	partitionKeyValue := make([]string, 0)
-	for _, path := range collection.PartitionKey.Paths {
-		var val interface{}
-		for _, part := range strings.Split(path, "/") {
-			val = document[part]
-		}
+	if _, ok := storeState.Documents[databaseId][collectionId][documentId]; !ok {
+		return repositorymodels.StatusNotFound
+	}
 
-		if val == nil {
-			val = ""
-		}
+	delete(storeState.Documents[databaseId][collectionId], documentId)
 
-		// TODO: handle non-string partition keys
-		partitionKeyValue = append(partitionKeyValue, val.(string))
+	return repositorymodels.StatusOk
+}
+
+func CreateDocument(databaseId string, collectionId string, document map[string]interface{}) (repositorymodels.Document, repositorymodels.RepositoryStatus) {
+	var documentId string
+	var ok bool
+	if documentId, ok = document["id"].(string); !ok || documentId == "" {
+		return repositorymodels.Document{}, repositorymodels.BadRequest
+	}
+
+	if _, ok := storeState.Databases[databaseId]; !ok {
+		return repositorymodels.Document{}, repositorymodels.StatusNotFound
+	}
+
+	if _, ok = storeState.Collections[databaseId][collectionId]; !ok {
+		return repositorymodels.Document{}, repositorymodels.StatusNotFound
+	}
+
+	if _, ok := storeState.Documents[databaseId][collectionId][documentId]; ok {
+		return repositorymodels.Document{}, repositorymodels.Conflict
 	}
 
 	document["_ts"] = time.Now().Unix()
 	document["_rid"] = uuid.New().String()
 	document["_etag"] = fmt.Sprintf("\"%s\"", document["_rid"])
-	document["_internal"] = map[string]interface{}{
-		"databaseId":        databaseId,
-		"collectionId":      collectionId,
-		"partitionKeyValue": partitionKeyValue,
-	}
-	documents = append(documents, document)
 
-	return repositorymodels.StatusOk
+	storeState.Documents[databaseId][collectionId][documentId] = document
+
+	return document, repositorymodels.StatusOk
 }
 
 func ExecuteQueryDocuments(databaseId string, collectionId string, query string, queryParameters map[string]interface{}) ([]memoryexecutor.RowType, repositorymodels.RepositoryStatus) {

internal/repositories/state.go

@@ -5,10 +5,20 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"reflect"
 
 	repositorymodels "github.com/pikami/cosmium/internal/repository_models"
 )
 
+var storedProcedures = []repositorymodels.StoredProcedure{}
+var triggers = []repositorymodels.Trigger{}
+var userDefinedFunctions = []repositorymodels.UserDefinedFunction{}
+var storeState = repositorymodels.State{
+	Databases:   make(map[string]repositorymodels.Database),
+	Collections: make(map[string]map[string]repositorymodels.Collection),
+	Documents:   make(map[string]map[string]map[string]repositorymodels.Document),
+}
+
 func LoadStateFS(filePath string) {
 	data, err := os.ReadFile(filePath)
 	if err != nil {
@@ -21,19 +31,91 @@ func LoadStateFS(filePath string) {
 	}
 
 	fmt.Println("Loaded state:")
-	fmt.Printf("Databases: %d\n", len(state.Databases))
-	fmt.Printf("Collections: %d\n", len(state.Collections))
-	fmt.Printf("Documents: %d\n", len(state.Documents))
+	fmt.Printf("Databases: %d\n", getLength(state.Databases))
+	fmt.Printf("Collections: %d\n", getLength(state.Collections))
+	fmt.Printf("Documents: %d\n", getLength(state.Documents))
 
-	databases = state.Databases
-	collections = state.Collections
-	documents = state.Documents
+	storeState = state
+
+	ensureStoreStateNoNullReferences()
 }
 
-func GetState() map[string]interface{} {
-	return map[string]interface{}{
-		"databases":   databases,
-		"collections": collections,
-		"documents":   documents,
+func SaveStateFS(filePath string) {
+	data, err := json.MarshalIndent(storeState, "", "\t")
+	if err != nil {
+		fmt.Printf("Failed to save state: %v\n", err)
+		return
+	}
+
+	os.WriteFile(filePath, data, os.ModePerm)
+
+	fmt.Println("Saved state:")
+	fmt.Printf("Databases: %d\n", getLength(storeState.Databases))
+	fmt.Printf("Collections: %d\n", getLength(storeState.Collections))
+	fmt.Printf("Documents: %d\n", getLength(storeState.Documents))
+}
+
+func GetState() repositorymodels.State {
+	return storeState
+}
+
+func getLength(v interface{}) int {
+	switch v.(type) {
+	case repositorymodels.Database,
+		repositorymodels.Collection,
+		repositorymodels.Document:
+		return 1
+	}
+
+	rv := reflect.ValueOf(v)
+	if rv.Kind() != reflect.Map {
+		return -1
+	}
+
+	count := 0
+	for _, key := range rv.MapKeys() {
+		if rv.MapIndex(key).Kind() == reflect.Map {
+			count += getLength(rv.MapIndex(key).Interface())
+		} else {
+			count++
+		}
+	}
+
+	return count
+}
+
+func ensureStoreStateNoNullReferences() {
+	if storeState.Databases == nil {
+		storeState.Databases = make(map[string]repositorymodels.Database)
+	}
+
+	if storeState.Collections == nil {
+		storeState.Collections = make(map[string]map[string]repositorymodels.Collection)
+	}
+
+	if storeState.Documents == nil {
+		storeState.Documents = make(map[string]map[string]map[string]repositorymodels.Document)
+	}
+
+	for database := range storeState.Databases {
+		if storeState.Collections[database] == nil {
+			storeState.Collections[database] = make(map[string]repositorymodels.Collection)
+		}
+
+		if storeState.Documents[database] == nil {
+			storeState.Documents[database] = make(map[string]map[string]repositorymodels.Document)
+		}
+
+		for collection := range storeState.Collections[database] {
+			if storeState.Documents[database][collection] == nil {
+				storeState.Documents[database][collection] = make(map[string]repositorymodels.Document)
+			}
+
+			for document := range storeState.Documents[database][collection] {
+				if storeState.Documents[database][collection][document] == nil {
+					delete(storeState.Documents[database][collection], document)
+				}
+			}
+		}
 	}
 }

internal/repositories/stored_procedures.go

@@ -2,16 +2,6 @@ package repositories
 
 import repositorymodels "github.com/pikami/cosmium/internal/repository_models"
 
-var storedProcedures = []repositorymodels.StoredProcedure{}
-
 func GetAllStoredProcedures(databaseId string, collectionId string) ([]repositorymodels.StoredProcedure, repositorymodels.RepositoryStatus) {
-	sps := make([]repositorymodels.StoredProcedure, 0)
-
-	for _, coll := range storedProcedures {
-		if coll.Internals.DatabaseId == databaseId && coll.Internals.CollectionId == collectionId {
-			sps = append(sps, coll)
-		}
-	}
-
-	return sps, repositorymodels.StatusOk
+	return storedProcedures, repositorymodels.StatusOk
 }

internal/repositories/triggers.go

@@ -2,16 +2,6 @@ package repositories
 
 import repositorymodels "github.com/pikami/cosmium/internal/repository_models"
 
-var triggers = []repositorymodels.Trigger{}
-
 func GetAllTriggers(databaseId string, collectionId string) ([]repositorymodels.Trigger, repositorymodels.RepositoryStatus) {
-	filteredTriggers := make([]repositorymodels.Trigger, 0)
-
-	for _, coll := range triggers {
-		if coll.Internals.DatabaseId == databaseId && coll.Internals.CollectionId == collectionId {
-			filteredTriggers = append(filteredTriggers, coll)
-		}
-	}
-
-	return filteredTriggers, repositorymodels.StatusOk
+	return triggers, repositorymodels.StatusOk
 }

internal/repositories/user_defined_functions.go

@@ -2,16 +2,6 @@ package repositories
 
 import repositorymodels "github.com/pikami/cosmium/internal/repository_models"
 
-var userDefinedFunctions = []repositorymodels.UserDefinedFunction{}
-
 func GetAllUserDefinedFunctions(databaseId string, collectionId string) ([]repositorymodels.UserDefinedFunction, repositorymodels.RepositoryStatus) {
-	udfs := make([]repositorymodels.UserDefinedFunction, 0)
-
-	for _, coll := range userDefinedFunctions {
-		if coll.Internals.DatabaseId == databaseId && coll.Internals.CollectionId == collectionId {
-			udfs = append(udfs, coll)
-		}
-	}
-
-	return udfs, repositorymodels.StatusOk
+	return userDefinedFunctions, repositorymodels.StatusOk
 }

internal/repository_models/models.go

@@ -29,9 +29,6 @@ type Collection struct {
 	Triggers       string                   `json:"_triggers"`
 	Udfs           string                   `json:"_udfs"`
 	Conflicts      string                   `json:"_conflicts"`
-	Internals      struct {
-		DatabaseId string
-	}
 }
 
 type CollectionIndexingPolicy struct {
@@ -57,29 +54,21 @@ type CollectionPartitionKey struct {
 }
 
 type UserDefinedFunction struct {
-	Body      string `json:"body"`
-	ID        string `json:"id"`
-	Rid       string `json:"_rid"`
-	Ts        int    `json:"_ts"`
-	Self      string `json:"_self"`
-	Etag      string `json:"_etag"`
-	Internals struct {
-		DatabaseId   string
-		CollectionId string
-	}
+	Body string `json:"body"`
+	ID   string `json:"id"`
+	Rid  string `json:"_rid"`
+	Ts   int    `json:"_ts"`
+	Self string `json:"_self"`
+	Etag string `json:"_etag"`
 }
 
 type StoredProcedure struct {
-	Body      string `json:"body"`
-	ID        string `json:"id"`
-	Rid       string `json:"_rid"`
-	Ts        int    `json:"_ts"`
-	Self      string `json:"_self"`
-	Etag      string `json:"_etag"`
-	Internals struct {
-		DatabaseId   string
-		CollectionId string
-	}
+	Body string `json:"body"`
+	ID   string `json:"id"`
+	Rid  string `json:"_rid"`
+	Ts   int    `json:"_ts"`
+	Self string `json:"_self"`
+	Etag string `json:"_etag"`
 }
 
 type Trigger struct {
@@ -91,10 +80,6 @@ type Trigger struct {
 	Ts               int    `json:"_ts"`
 	Self             string `json:"_self"`
 	Etag             string `json:"_etag"`
-	Internals        struct {
-		DatabaseId   string
-		CollectionId string
-	}
 }
 
 type Document map[string]interface{}
@@ -115,7 +100,12 @@ type PartitionKeyRange struct {
 }
 
 type State struct {
-	Databases   []Database   `json:"databases"`
-	Collections []Collection `json:"collections"`
-	Documents   []Document   `json:"documents"`
+	// 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"`
 }