From 1f6922c9e597ca2e94c8b767df9d3c2fa2ba4188 Mon Sep 17 00:00:00 2001
From: MaysWind <i@mayswind.net>
Date: Thu, 3 Dec 2020 23:22:01 +0800
Subject: [PATCH] support sync database structure when starting web server

---
 cmd/database.go         | 64 ++++++++++++++++++++++++++++++++++++++---
 cmd/webserver.go        |  9 ++++++
 conf/lab.ini            |  3 ++
 pkg/settings/setting.go |  6 ++--
 4 files changed, 76 insertions(+), 6 deletions(-)

diff --git a/cmd/database.go b/cmd/database.go
index ad0c6101..d8ef2244 100644
--- a/cmd/database.go
+++ b/cmd/database.go
@@ -29,11 +29,67 @@ func updateDatabaseStructure(c *cli.Context) error {
 
 	log.BootInfof("[database.updateDatabaseStructure] starting maintaining")
 
-	_ = datastore.Container.UserStore.SyncStructs(new(models.User), new(models.TwoFactor), new(models.TwoFactorRecoveryCode))
-	_ = datastore.Container.TokenStore.SyncStructs(new(models.TokenRecord))
-	_ = datastore.Container.UserDataStore.SyncStructs(new(models.Account), new(models.TransactionCategory))
+	err = updateAllDatabaseTablesStructure()
 
-	log.BootInfof("[database.updateDatabaseStructure] maintained successfully")
+	if err != nil {
+		log.BootErrorf("[database.updateDatabaseStructure] update database table structure failed, because %s", err.Error())
+		return err
+	}
+
+	log.BootInfof("[database.updateDatabaseStructure] all tables maintained successfully")
+	return nil
+}
+
+func updateAllDatabaseTablesStructure() error {
+	var err error
+
+	err = datastore.Container.UserStore.SyncStructs(new(models.User))
+
+	if err != nil {
+		return err
+	} else {
+		log.BootInfof("[database.updateAllDatabaseTablesStructure] user table maintained successfully")
+	}
+
+	err = datastore.Container.UserStore.SyncStructs(new(models.TwoFactor))
+
+	if err != nil {
+		return err
+	} else {
+		log.BootInfof("[database.updateAllDatabaseTablesStructure] two factor table maintained successfully")
+	}
+
+	err = datastore.Container.UserStore.SyncStructs(new(models.TwoFactorRecoveryCode))
+
+	if err != nil {
+		return err
+	} else {
+		log.BootInfof("[database.updateAllDatabaseTablesStructure] two factor recovery code table maintained successfully")
+	}
+
+	err = datastore.Container.TokenStore.SyncStructs(new(models.TokenRecord))
+
+	if err != nil {
+		return err
+	} else {
+		log.BootInfof("[database.updateAllDatabaseTablesStructure] token record table maintained successfully")
+	}
+
+	err = datastore.Container.UserDataStore.SyncStructs(new(models.Account))
+
+	if err != nil {
+		return err
+	} else {
+		log.BootInfof("[database.updateAllDatabaseTablesStructure] account table maintained successfully")
+	}
+
+	err = datastore.Container.UserDataStore.SyncStructs(new(models.TransactionCategory))
+
+	if err != nil {
+		return err
+	} else {
+		log.BootInfof("[database.updateAllDatabaseTablesStructure] transaction category table maintained successfully")
+	}
 
 	return nil
 }
diff --git a/cmd/webserver.go b/cmd/webserver.go
index a5ec3740..d5a38388 100644
--- a/cmd/webserver.go
+++ b/cmd/webserver.go
@@ -44,6 +44,15 @@ func startWebServer(c *cli.Context) error {
 
 	log.BootInfof("[server.startWebServer] static root path is %s", config.StaticRootPath)
 
+	if config.AutoUpdateDatabase {
+		err = updateAllDatabaseTablesStructure()
+
+		if err != nil {
+			log.BootErrorf("[server.startWebServer] update database table structure failed, because %s", err.Error())
+			return err
+		}
+	}
+
 	err = requestid.InitializeRequestIdGenerator(config)
 
 	if err != nil {
diff --git a/conf/lab.ini b/conf/lab.ini
index 223b40a7..9488373d 100644
--- a/conf/lab.ini
+++ b/conf/lab.ini
@@ -65,6 +65,9 @@ conn_max_lifetime = 14400
 # Set to true to log each sql statement and execution times
 log_query = false
 
+# Set to true to automatically update database structure when starting web server
+auto_update_database = true
+
 [log]
 # Either "console", "file", default is "console"
 # Use space to separate multiple modes, e.g. "console file"
diff --git a/pkg/settings/setting.go b/pkg/settings/setting.go
index 893974ec..d5a80037 100644
--- a/pkg/settings/setting.go
+++ b/pkg/settings/setting.go
@@ -116,8 +116,9 @@ type Config struct {
 	EnableRequestLog bool
 
 	// Database
-	DatabaseConfig *DatabaseConfig
-	EnableQueryLog bool
+	DatabaseConfig     *DatabaseConfig
+	EnableQueryLog     bool
+	AutoUpdateDatabase bool
 
 	// Log
 	LogModes         []string
@@ -300,6 +301,7 @@ func loadDatabaseConfiguration(config *Config, configFile *ini.File, sectionName
 
 	config.DatabaseConfig = dbConfig
 	config.EnableQueryLog = getConfigItemBoolValue(configFile, sectionName, "log_query", false)
+	config.AutoUpdateDatabase = getConfigItemBoolValue(configFile, sectionName, "auto_update_database", true)
 
 	return nil
 }