优化公共页面&增加自动迁移

2025-12-14 19:41:55 +08:00
1 parent 36eaae9c40
commit c0790e9ced
@@ -1,7 +1,179 @@
 package main

-import "github.com/dushixiang/pika/internal"
+import (
+	"fmt"
+	"log"
+	"os"
+	"time"
+
+	"github.com/dushixiang/pika/internal"
+	"github.com/dushixiang/pika/internal/config"
+	v0_0_13 "github.com/dushixiang/pika/internal/migrate/v0_0_13"
+	"github.com/dushixiang/pika/internal/vmclient"
+	"github.com/go-orz/orz"
+	"github.com/spf13/cobra"
+	"go.uber.org/zap"
+	"gorm.io/gorm"
+)
+
+var (
+	configFile string
+	rootCmd    = &cobra.Command{
+		Use:   "pika-server",
+		Short: "Pika 服务器监控系统",
+		Long:  `Pika 是一个轻量级的服务器监控系统，支持多探针部署和实时监控。`,
+		Run: func(cmd *cobra.Command, args []string) {
+			internal.Run(configFile)
+		},
+	}
+
+	serveCmd = &cobra.Command{
+		Use:   "serve",
+		Short: "启动 Pika 服务器",
+		Long:  `启动 Pika HTTP 服务器，提供 Web 界面和 API 服务。`,
+		Run: func(cmd *cobra.Command, args []string) {
+			internal.Run(configFile)
+		},
+	}
+
+	migrateCmd = &cobra.Command{
+		Use:   "migrate_v0_0_13",
+		Short: "执行数据库迁移",
+		Long:  `将历史 PostgreSQL 指标数据迁移到 VictoriaMetrics（仅需执行一次）。`,
+		Run: func(cmd *cobra.Command, args []string) {
+			runMigration(configFile)
+		},
+	}
+)
+
+func init() {
+	// 添加全局配置文件参数
+	rootCmd.PersistentFlags().StringVarP(&configFile, "config", "c", "./config.yaml", "配置文件路径")
+
+	// 添加子命令
+	rootCmd.AddCommand(serveCmd)
+	rootCmd.AddCommand(migrateCmd)
+}

 func main() {
-	internal.Run("./config.yaml")
+	if err := rootCmd.Execute(); err != nil {
+		fmt.Fprintf(os.Stderr, "错误: %v\n", err)
+		os.Exit(1)
+	}
+}
+
+// runMigration 执行数据迁移
+func runMigration(configPath string) {
+	fmt.Println("=== Pika 数据迁移工具 ===")
+	fmt.Println("配置文件:", configPath)
+	fmt.Println()
+
+	// 初始化 orz 以获取数据库连接和配置
+	var db *gorm.DB
+	var vmClient *vmclient.VMClient
+	var logger *zap.Logger
+
+	err := orz.Quick(configPath, func(app *orz.App) error {
+		logger = app.Logger()
+		db = app.GetDatabase()
+
+		// 读取 VictoriaMetrics 配置
+		var appConfig config.AppConfig
+		_config := app.GetConfig()
+		if _config != nil {
+			if err := _config.App.Unmarshal(&appConfig); err != nil {
+				return fmt.Errorf("读取配置失败: %w", err)
+			}
+		}
+
+		// 初始化 VictoriaMetrics 客户端
+		vmClient = provideVMClient(&appConfig, logger)
+
+		return nil
+	})
+
+	if err != nil {
+		log.Fatalf("初始化失败: %v", err)
+	}
+
+	// 执行迁移
+	fmt.Println("开始迁移历史指标数据到 VictoriaMetrics...")
+	fmt.Println("说明: 数据将从 PostgreSQL 迁移到 VictoriaMetrics。")
+	fmt.Println("警告: 建议在迁移前备份 PostgreSQL 数据！")
+	fmt.Println()
+	fmt.Print("是否在迁移完成后自动删除旧的数据表？(yes/no): ")
+
+	var deleteOldTables bool
+	var answer string
+	fmt.Scanln(&answer)
+	if answer == "yes" || answer == "y" || answer == "Y" {
+		deleteOldTables = true
+		fmt.Println("✓ 将在迁移完成后自动删除旧表")
+	} else {
+		deleteOldTables = false
+		fmt.Println("✓ 将保留旧表，您可以稍后手动删除")
+	}
+
+	fmt.Println()
+	startTime := time.Now()
+
+	// 执行数据迁移
+	if err := v0_0_13.Migrate(db, vmClient); err != nil {
+		logger.Error("迁移失败", zap.Error(err))
+		log.Fatalf("迁移失败: %v", err)
+	}
+
+	elapsed := time.Since(startTime)
+	fmt.Println()
+	fmt.Printf("✓ 数据迁移成功完成！耗时: %s\n", elapsed)
+
+	// 根据用户选择决定是否删除旧表
+	if deleteOldTables {
+		fmt.Println()
+		fmt.Println("正在删除旧的指标数据表...")
+		if err := v0_0_13.DropOldTables(db); err != nil {
+			logger.Error("删除旧表失败", zap.Error(err))
+			log.Printf("警告: 删除旧表失败: %v\n", err)
+			fmt.Println("注意: 旧表删除失败，但数据已成功迁移，您可以手动删除这些表。")
+		} else {
+			fmt.Println("✓ 旧表删除完成")
+			fmt.Println()
+			fmt.Println("✓ 迁移任务全部完成！")
+		}
+	} else {
+		fmt.Println()
+		fmt.Println("提示: 旧数据表已保留，您可以验证数据后手动删除。")
+		fmt.Println("提示: 如需删除旧表，可以重新运行此命令并选择删除。")
+	}
+
+	fmt.Println()
+	fmt.Println("提示: 现在可以启动服务器: ./pika-server serve")
+}
+
+// provideVMClient 提供 VictoriaMetrics 客户端
+func provideVMClient(cfg *config.AppConfig, logger *zap.Logger) *vmclient.VMClient {
+	// 检查配置
+	if cfg.VictoriaMetrics == nil || !cfg.VictoriaMetrics.Enabled {
+		logger.Info("VictoriaMetrics is not enabled, using default configuration")
+		// 返回一个默认配置的客户端（用于本地开发）
+		return vmclient.NewVMClient("http://localhost:8428", 30*time.Second, 60*time.Second)
+	}
+
+	// 使用配置创建客户端
+	writeTimeout := time.Duration(cfg.VictoriaMetrics.WriteTimeout) * time.Second
+	if writeTimeout == 0 {
+		writeTimeout = 30 * time.Second
+	}
+
+	queryTimeout := time.Duration(cfg.VictoriaMetrics.QueryTimeout) * time.Second
+	if queryTimeout == 0 {
+		queryTimeout = 60 * time.Second
+	}
+
+	logger.Info("VictoriaMetrics client initialized",
+		zap.String("url", cfg.VictoriaMetrics.URL),
+		zap.Duration("writeTimeout", writeTimeout),
+		zap.Duration("queryTimeout", queryTimeout))
+
+	return vmclient.NewVMClient(cfg.VictoriaMetrics.URL, writeTimeout, queryTimeout)
 }