From 62ac9ddf6b714f496a08e18cf5903224112ce436 Mon Sep 17 00:00:00 2001
From: Pijus Kamandulis <pkpjuklas@gmail.com>
Date: Wed, 14 Feb 2024 20:46:40 +0200
Subject: [PATCH] Added ability to load initial storage data

---
 api/config/config.go                 |  2 ++
 api/config/models.go                 |  1 +
 internal/repositories/state.go       | 31 ++++++++++++++++++++++++++++
 internal/repository_models/models.go |  6 ++++++
 main.go                              |  5 +++++
 5 files changed, 45 insertions(+)
 create mode 100644 internal/repositories/state.go

diff --git a/api/config/config.go b/api/config/config.go
index 52a9917..1b114e2 100644
--- a/api/config/config.go
+++ b/api/config/config.go
@@ -13,6 +13,7 @@ func ParseFlags() {
 	explorerPath := flag.String("ExplorerDir", "/home/pk/pro/cosmos-explorer/dist", "Path to cosmos-explorer files")
 	tlsCertificatePath := flag.String("Cert", "../example.crt", "Hostname")
 	tlsCertificateKey := flag.String("CertKey", "../example.key", "Hostname")
+	initialDataPath := flag.String("InitialData", "", "Path to JSON containing initial state")
 
 	flag.Parse()
 
@@ -21,6 +22,7 @@ func ParseFlags() {
 	Config.ExplorerPath = *explorerPath
 	Config.TLS_CertificatePath = *tlsCertificatePath
 	Config.TLS_CertificateKey = *tlsCertificateKey
+	Config.DataFilePath = *initialDataPath
 
 	Config.DatabaseAccount = Config.Host
 	Config.DatabaseDomain = Config.Host
diff --git a/api/config/models.go b/api/config/models.go
index e7b98f1..f1bfc16 100644
--- a/api/config/models.go
+++ b/api/config/models.go
@@ -10,4 +10,5 @@ type ServerConfig struct {
 	Host                string
 	TLS_CertificatePath string
 	TLS_CertificateKey  string
+	DataFilePath        string
 }
diff --git a/internal/repositories/state.go b/internal/repositories/state.go
new file mode 100644
index 0000000..80345c5
--- /dev/null
+++ b/internal/repositories/state.go
@@ -0,0 +1,31 @@
+package repositories
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"os"
+
+	repositorymodels "github.com/pikami/cosmium/internal/repository_models"
+)
+
+func LoadStateFS(filePath string) {
+	data, err := os.ReadFile(filePath)
+	if err != nil {
+		log.Fatalf("Error reading state JSON file: %v", err)
+	}
+
+	var state repositorymodels.State
+	if err := json.Unmarshal(data, &state); err != nil {
+		log.Fatalf("Error unmarshalling state JSON: %v", err)
+	}
+
+	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))
+
+	databases = state.Databases
+	collections = state.Collections
+	documents = state.Documents
+}
diff --git a/internal/repository_models/models.go b/internal/repository_models/models.go
index ac16270..1021301 100644
--- a/internal/repository_models/models.go
+++ b/internal/repository_models/models.go
@@ -110,3 +110,9 @@ type PartitionKeyRange struct {
 	Ts                 int    `json:"_ts"`
 	Lsn                int    `json:"lsn"`
 }
+
+type State struct {
+	Databases   []Database   `json:"databases"`
+	Collections []Collection `json:"collections"`
+	Documents   []Document   `json:"documents"`
+}
diff --git a/main.go b/main.go
index 6669780..cfa32cf 100644
--- a/main.go
+++ b/main.go
@@ -5,11 +5,16 @@ import (
 
 	"github.com/pikami/cosmium/api"
 	"github.com/pikami/cosmium/api/config"
+	"github.com/pikami/cosmium/internal/repositories"
 )
 
 func main() {
 	config.ParseFlags()
 
+	if config.Config.DataFilePath != "" {
+		repositories.LoadStateFS(config.Config.DataFilePath)
+	}
+
 	router := api.CreateRouter()
 	router.RunTLS(
 		fmt.Sprintf(":%d", config.Config.Port),