From 550a12a3f7be3964e47f14332f190ed8d64bd4c1 Mon Sep 17 00:00:00 2001
From: ning <710leo@gmail.com>
Date: Thu, 9 Mar 2023 19:07:10 +0800
Subject: [PATCH] update upgrade cli

---
 cli/cli.go             |  4 +--
 cli/upgrade/upgrade.go | 55 ++++++++++++++++++++++++++++++++----------
 cmd/cli/main.go        |  8 +-----
 etc/config.toml        |  9 +++++++
 4 files changed, 54 insertions(+), 22 deletions(-)

diff --git a/cli/cli.go b/cli/cli.go
index 5ea14f176..afd2a0d94 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -4,6 +4,6 @@ import (
 	"github.com/ccfos/nightingale/v6/cli/upgrade"
 )
 
-func Upgrade(configFile string, sqlFile string) error {
-	return upgrade.Upgrade(configFile, sqlFile)
+func Upgrade(configFile string) error {
+	return upgrade.Upgrade(configFile)
 }
diff --git a/cli/upgrade/upgrade.go b/cli/upgrade/upgrade.go
index 588a5cb32..78256c7df 100644
--- a/cli/upgrade/upgrade.go
+++ b/cli/upgrade/upgrade.go
@@ -2,33 +2,62 @@ package upgrade
 
 import (
 	"context"
-	"fmt"
-	"os/exec"
 
 	"github.com/ccfos/nightingale/v6/models"
 	"github.com/ccfos/nightingale/v6/pkg/ctx"
 	"github.com/ccfos/nightingale/v6/storage"
-
-	"github.com/go-sql-driver/mysql"
+	"github.com/toolkits/pkg/logger"
 )
 
-func Upgrade(configFile string, sqlFile string) error {
+func Upgrade(configFile string) error {
 	var config Config
 	Parse(configFile, &config)
-	dsnConf, _ := mysql.ParseDSN(config.DB.DSN)
-	passwd := dsnConf.Passwd
-	user := dsnConf.User
-
-	cmd := exec.Command("mysql", fmt.Sprintf("-u%s", user), fmt.Sprintf("-p%s", passwd), "<", sqlFile)
-	if err := cmd.Run(); err != nil {
-		return err
-	}
 
 	db, err := storage.New(config.DB)
 	if err != nil {
 		return err
 	}
+
 	ctx := ctx.NewContext(context.Background(), db)
+	for _, cluster := range config.Clusters {
+		header := make(map[string]string)
+		headerCount := len(cluster.Headers)
+		if headerCount > 0 && headerCount%2 == 0 {
+			for i := 0; i < len(cluster.Headers); i += 2 {
+				header[cluster.Headers[i]] = cluster.Headers[i+1]
+			}
+		}
+
+		authJosn := models.Auth{
+			BasicAuthUser:     cluster.BasicAuthUser,
+			BasicAuthPassword: cluster.BasicAuthPass,
+		}
+
+		httpJson := models.HTTP{
+			Timeout:     cluster.Timeout,
+			DialTimeout: cluster.DialTimeout,
+			UseTLS: models.TLS{
+				SkipTlsVerify: cluster.UseTLS,
+			},
+			MaxIdleConnsPerHost: cluster.MaxIdleConnsPerHost,
+			Url:                 cluster.Prom,
+			Headers:             header,
+		}
+
+		datasrouce := models.Datasource{
+			PluginId:       1,
+			PluginType:     "prometheus",
+			PluginTypeName: "Prometheus Like",
+			Name:           cluster.Name,
+			HTTPJson:       httpJson,
+			AuthJson:       authJosn,
+			ClusterName:    "default",
+		}
+		err := datasrouce.Add(ctx)
+		if err != nil {
+			logger.Errorf("add datasource %s error: %v", cluster.Name, err)
+		}
+	}
 
 	datasources, err := models.GetDatasources(ctx)
 	if err != nil {
diff --git a/cmd/cli/main.go b/cmd/cli/main.go
index c630319ee..77b0c6c89 100644
--- a/cmd/cli/main.go
+++ b/cmd/cli/main.go
@@ -12,7 +12,6 @@ import (
 var (
 	upgrade     = flag.Bool("upgrade", false, "Upgrade the database.")
 	showVersion = flag.Bool("version", false, "Show version.")
-	sqlFile     = flag.String("sql", "", "Specify the sql file to be executed.")
 	configDir   = flag.String("config", "", "Specify configuration directory.(env:N9E_CONFIGS)")
 )
 
@@ -25,17 +24,12 @@ func main() {
 	}
 
 	if *upgrade {
-		if *sqlFile == "" {
-			fmt.Println("Please specify the sql file to be executed.")
-			os.Exit(1)
-		}
-
 		if *configDir == "" {
 			fmt.Println("Please specify the configuration directory.")
 			os.Exit(1)
 		}
 
-		err := cli.Upgrade(*configDir, *sqlFile)
+		err := cli.Upgrade(*configDir)
 		if err != nil {
 			fmt.Println(err)
 			os.Exit(1)
diff --git a/etc/config.toml b/etc/config.toml
index aba216ff1..bf75b8932 100644
--- a/etc/config.toml
+++ b/etc/config.toml
@@ -90,6 +90,15 @@ ClusterName = "default"
 Timeout=30000
 NotifyConcurrency = 10
 
+[Alert.SMTP]
+Host = "smtp.163.com"
+Port = 994
+User = "username"
+Pass = "password"
+From = "username@163.com"
+InsecureSkipVerify = true
+Batch = 5
+
 [Center]
 MetricsYamlFile = "./etc/metrics.yaml"
 I18NHeaderKey = "X-Language"