From af5d97794ad229b12e0e85922c3610b212075309 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Thu, 18 May 2023 15:28:22 +0800
Subject: [PATCH] fix setup bugs

---
 connector/connector.go | 16 ++++++++++++++++
 engine/load.go         | 17 ++++++++++-------
 query/query.go         |  8 ++++++++
 setup/handler.go       |  7 +++++++
 setup/install.go       |  8 ++++++++
 setup/setup.go         | 11 +++--------
 6 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/connector/connector.go b/connector/connector.go
index 864e974d..eca6261e 100644
--- a/connector/connector.go
+++ b/connector/connector.go
@@ -35,6 +35,22 @@ func Load(cfg config.Config) error {
 	return nil
 }
 
+// Unload Connector
+func Unload() error {
+	messages := []string{}
+	for id, conn := range connector.Connectors {
+		err := conn.Close()
+		if err != nil {
+			messages = append(messages, err.Error())
+		}
+		delete(connector.Connectors, id)
+	}
+	if len(messages) > 0 {
+		return fmt.Errorf("%s", strings.Join(messages, ";\n"))
+	}
+	return nil
+}
+
 // Close close connector
 func Close() error {
 	messages := []string{}
diff --git a/engine/load.go b/engine/load.go
index f97bb74d..dde7c14b 100644
--- a/engine/load.go
+++ b/engine/load.go
@@ -56,6 +56,12 @@ func Load(cfg config.Config) (err error) {
 		printErr(cfg.Mode, "Load Application", err)
 	}
 
+	// Make Database connections
+	err = share.DBConnect(cfg.DB)
+	if err != nil {
+		printErr(cfg.Mode, "DB", err)
+	}
+
 	// Load Certs
 	err = cert.Load(cfg)
 	if err != nil {
@@ -86,12 +92,6 @@ func Load(cfg config.Config) (err error) {
 		printErr(cfg.Mode, "Runtime", err)
 	}
 
-	// Make Database connections
-	err = share.DBConnect(cfg.DB)
-	if err != nil {
-		printErr(cfg.Mode, "DB", err)
-	}
-
 	// Load Query Engine
 	err = query.Load(cfg)
 	if err != nil {
@@ -203,8 +203,11 @@ func Unload() (err error) {
 	// Close DB
 	err = share.DBClose()
 
+	// Close Query Engine
+	err = query.Unload()
+
 	// Close Connectors
-	err = connector.Close()
+	err = connector.Unload()
 
 	// Recycle
 	// api
diff --git a/query/query.go b/query/query.go
index 485b9c42..706dbc3a 100644
--- a/query/query.go
+++ b/query/query.go
@@ -40,6 +40,14 @@ func Load(cfg config.Config) error {
 	return nil
 }
 
+// Unload Query Engine
+func Unload() error {
+	for id := range query.Engines {
+		query.Unregister(id)
+	}
+	return nil
+}
+
 // registerDefaultQuery register the default engine
 func registerDefault() {
 	if capsule.Global != nil {
diff --git a/setup/handler.go b/setup/handler.go
index 720e6232..d8c94138 100644
--- a/setup/handler.go
+++ b/setup/handler.go
@@ -10,6 +10,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/yaoapp/kun/exception"
 	"github.com/yaoapp/xun"
+	"github.com/yaoapp/yao/engine"
 	"github.com/yaoapp/yao/widgets/app"
 )
 
@@ -85,6 +86,12 @@ func runSetup(c *gin.Context) {
 		return
 	}
 
+	err = engine.Load(cfg)
+	if err != nil {
+		c.JSON(500, gin.H{"code": 500, "message": err.Error()})
+		return
+	}
+
 	// Return
 	urls, err := AdminURL(cfg)
 	if err != nil {
diff --git a/setup/install.go b/setup/install.go
index 1113211e..fc79e23f 100644
--- a/setup/install.go
+++ b/setup/install.go
@@ -15,6 +15,7 @@ import (
 	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/data"
 	"github.com/yaoapp/yao/engine"
+	"github.com/yaoapp/yao/studio"
 	"github.com/yaoapp/yao/widgets/app"
 )
 
@@ -112,11 +113,18 @@ func Install(payload map[string]map[string]string) error {
 	if err != nil {
 		return err
 	}
+
 	defer func() {
 		engine.Unload()
 		time.Sleep(time.Millisecond * 200)
 	}()
 
+	// Load Studio
+	err = studio.Load(cfg)
+	if err != nil {
+		return err
+	}
+
 	//  Migrage & Setup
 	err = makeMigrate(root, cfg)
 	if err != nil {
diff --git a/setup/setup.go b/setup/setup.go
index b33fb36f..ee947036 100644
--- a/setup/setup.go
+++ b/setup/setup.go
@@ -132,24 +132,23 @@ func Start() (err error) {
 
 // Complete stop the studio api server
 func Complete() {
+	engine.Unload()
 	Done <- true
 }
 
 // Stop stop the studio api server
 func Stop() {
-	engine.Unload()
 	shutdown <- true
 }
 
 // AdminURL get admin url
 func AdminURL(cfg config.Config) ([]string, error) {
-	err := engine.Load(cfg)
+
+	urls, err := URLs(cfg)
 	if err != nil {
 		return nil, err
 	}
 
-	urls, err := URLs(cfg)
-
 	adminRoot := "yao"
 	if app.Setting.AdminRoot != "" {
 		adminRoot = app.Setting.AdminRoot
@@ -164,10 +163,6 @@ func AdminURL(cfg config.Config) ([]string, error) {
 
 // URLs get admin url
 func URLs(cfg config.Config) ([]string, error) {
-	err := engine.Load(cfg)
-	if err != nil {
-		return nil, err
-	}
 
 	ips, err := Ips()
 	if err != nil {