From 150612a1d9f5cd06c9adc8d59b992784b87affb3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8D=96=E5=A5=B3=E5=AD=A9=E7=9A=84=E5=B0=8F=E7=81=AB?=
 =?UTF-8?q?=E6=9F=B4?= <44471469+Erope@users.noreply.github.com>
Date: Thu, 2 Nov 2023 21:06:34 +0800
Subject: [PATCH] Feat: Add some parameters (#294)

---
 cmd/dashboard/main.go | 33 +++++++++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/cmd/dashboard/main.go b/cmd/dashboard/main.go
index 3e69f37..db5dce8 100644
--- a/cmd/dashboard/main.go
+++ b/cmd/dashboard/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"fmt"
 	"log"
 
 	"github.com/naiba/nezha/cmd/dashboard/controller"
@@ -9,13 +10,30 @@ import (
 	"github.com/naiba/nezha/model"
 	"github.com/naiba/nezha/service/singleton"
 	"github.com/ory/graceful"
+	flag "github.com/spf13/pflag"
+)
+
+type DashboardCliParam struct {
+	Version          bool   // 当前版本号
+	ConfigFile       string // 配置文件路径
+	DatebaseLocation string // Sqlite3 数据库文件路径
+}
+
+var (
+	dashboardCliParam DashboardCliParam
 )
 
 func init() {
+	flag.CommandLine.ParseErrorsWhitelist.UnknownFlags = true
+	flag.BoolVarP(&dashboardCliParam.Version, "version", "v", false, "查看当前版本号")
+	flag.StringVarP(&dashboardCliParam.ConfigFile, "config", "c", "data/config.yaml", "配置文件路径")
+	flag.StringVar(&dashboardCliParam.DatebaseLocation, "db", "data/sqlite.db", "Sqlite3数据库文件路径")
+	flag.Parse()
+
 	// 初始化 dao 包
-	singleton.InitConfigFromPath("data/config.yaml")
+	singleton.InitConfigFromPath(dashboardCliParam.ConfigFile)
 	singleton.InitTimezoneAndCache()
-	singleton.InitDBFromPath("data/sqlite.db")
+	singleton.InitDBFromPath(dashboardCliParam.DatebaseLocation)
 	singleton.InitLocalizer()
 	initSystem()
 }
@@ -36,6 +54,11 @@ func initSystem() {
 }
 
 func main() {
+	if dashboardCliParam.Version {
+		fmt.Println(singleton.Version)
+		return
+	}
+
 	singleton.CleanMonitorHistory()
 	go rpc.ServeRPC(singleton.Conf.GRPCPort)
 	serviceSentinelDispatchBus := make(chan model.Monitor) // 用于传递服务监控任务信息的channel
@@ -44,7 +67,7 @@ func main() {
 	go singleton.AlertSentinelStart()
 	singleton.NewServiceSentinel(serviceSentinelDispatchBus)
 	srv := controller.ServeWeb(singleton.Conf.HTTPPort)
-	graceful.Graceful(func() error {
+	if err := graceful.Graceful(func() error {
 		return srv.ListenAndServe()
 	}, func(c context.Context) error {
 		log.Println("NEZHA>> Graceful::START")
@@ -52,5 +75,7 @@ func main() {
 		log.Println("NEZHA>> Graceful::END")
 		srv.Shutdown(c)
 		return nil
-	})
+	}); err != nil {
+		log.Printf("NEZHA>> ERROR: %v", err)
+	}
 }