From 838490e8419b7cdd69d2b8c953f1f7359f32ead5 Mon Sep 17 00:00:00 2001
From: Evan Chen <evanchen@kumoly.io>
Date: Sun, 9 Jan 2022 16:31:12 +0800
Subject: [PATCH] update

---
 attribute/attribute.go      | 25 +++++++++++++++++++++----
 attribute/attribute_test.go |  8 ++++----
 attribute/service.go        | 18 +++++++++++++-----
 3 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/attribute/attribute.go b/attribute/attribute.go
index 12db61c..36e3ee4 100644
--- a/attribute/attribute.go
+++ b/attribute/attribute.go
@@ -1,6 +1,7 @@
 package attribute
 
 import (
+	"sync"
 	"time"
 
 	"github.com/rs/zerolog"
@@ -16,6 +17,7 @@ type Attribute struct {
 	Description string
 	Value       string
 	Default     string
+	NeedRestart bool
 
 	CreatedAt time.Time
 	UpdatedAt time.Time
@@ -34,7 +36,7 @@ func Init() {
 	}
 }
 
-func Add(Key, Name, Description, Default string) error {
+func Add(Key, Name, Description, Default string, NeedRestart bool) error {
 	ctr := 0
 	store.DB.Raw(`select count(*) from "attributes" where "key" = ?`, Key).
 		Scan(&ctr)
@@ -54,9 +56,7 @@ func Add(Key, Name, Description, Default string) error {
 }
 
 func RestoreDefault(key string) error {
-	result := store.DB.Exec(`update "attributes" set value = (
-	select "default" from "attributes" where "key" = ?
-) where "key" = ?`, key, key)
+	result := store.DB.Exec(`update "attributes" set value = "default" where "key" = ?`, key)
 	if result.RowsAffected == 0 {
 		return ErrorAttributeNotFound
 	}
@@ -79,3 +79,20 @@ func Set(key, value string) error {
 	}
 	return result.Error
 }
+
+var load []*func()
+var loadLock sync.Mutex
+
+func AddSoftLoad(f *func()) {
+	loadLock.Lock()
+	load = append(load, f)
+	loadLock.Unlock()
+}
+
+func Load() {
+	for _, f := range load {
+		if f != nil {
+			(*f)()
+		}
+	}
+}
diff --git a/attribute/attribute_test.go b/attribute/attribute_test.go
index b59fec7..de6e8d5 100644
--- a/attribute/attribute_test.go
+++ b/attribute/attribute_test.go
@@ -13,10 +13,10 @@ func TestAttr(t *testing.T) {
 	viper.Set("db.type", "sqlite")
 	store.Setup()
 	Init()
-	fmt.Println(Add("test", "test", "", "1"))
-	fmt.Println(Add("test1", "test1", "", "2"))
-	fmt.Println(Add("test1", "test2", "", "2"))
-	fmt.Println(Add("test1", "test3", "", "2"))
+	fmt.Println(Add("test", "test", "", "1", false))
+	fmt.Println(Add("test1", "test1", "", "2", false))
+	fmt.Println(Add("test1", "test2", "", "2", false))
+	fmt.Println(Add("test1", "test3", "", "2", false))
 	fmt.Println(Set("test", "reset"))
 	fmt.Println(Get("test"))
 }
diff --git a/attribute/service.go b/attribute/service.go
index 42d3eb7..cfc31ea 100644
--- a/attribute/service.go
+++ b/attribute/service.go
@@ -12,7 +12,6 @@ type Service struct{}
 func (srv Service) GetName() string           { return "attribute.Service" }
 func (srv Service) GetDependencies() []string { return []string{} }
 func (srv Service) IsService() bool           { return false }
-func (srv Service) Main() error               { return nil }
 func (srv Service) Del()                      {}
 func (srv Service) Health() error             { return nil }
 
@@ -59,11 +58,10 @@ func (srv Service) Load() error {
 		if err := c.ShouldBindJSON(&data); err != nil {
 			panic(err)
 		}
-		result := store.DB.Exec(`update "attributes" set value = "default" where "key" = ?`, data.Key)
-		if result.Error != nil {
-			panic(result.Error)
+		if err := RestoreDefault(data.Key); err != nil {
+			panic(err)
 		}
-		server.OK(c, result.RowsAffected)
+		server.OK(c, "ok")
 	})
 
 	api.POST("restore_all", func(c *gin.Context) {
@@ -74,5 +72,15 @@ func (srv Service) Load() error {
 		server.OK(c, result.RowsAffected)
 	})
 
+	api.POST("reload", func(c *gin.Context) {
+		Load()
+		server.OK(c, "ok")
+	})
+
+	return nil
+}
+
+func (srv Service) Main() error {
+	Load()
 	return nil
 }