ai-agent/cmd/testdata/main.go

package main

import (
	"code.tczkiot.com/wlw/ai-agent/cmd/testdata/aiagent"
	"code.tczkiot.com/wlw/ai-agent/cmd/testdata/aiconfig"
	"code.tczkiot.com/wlw/ai-agent/cmd/testdata/channel"
	"code.tczkiot.com/wlw/ai-agent/cmd/testdata/kb"
	"code.tczkiot.com/wlw/ai-agent/cmd/testdata/quickreply"
	"code.tczkiot.com/wlw/ai-agent/cmd/testdata/seedlang"
	"code.tczkiot.com/wlw/ai-agent/internal/bootstrap"
	"code.tczkiot.com/wlw/ai-agent/internal/models"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/config"
	"flag"
	"fmt"
	"log/slog"
	"os"
	"strings"

	"github.com/mlogclub/simple/sqls"
	"gorm.io/gorm"
	gormlogger "gorm.io/gorm/logger"
)

func main() {
	if err := run(); err != nil {
		slog.Error("init testdata failed", "error", err)
		os.Exit(1)
	}
}

func run() error {
	configPath := flag.String("config", "config/config.yaml", "path to config file")
	autoConfirm := flag.Bool("yes", false, "skip confirmation prompt")
	langValue := flag.String("lang", string(seedlang.Chinese), "testdata language: zh or en")
	flag.Parse()

	lang, err := seedlang.Parse(*langValue)
	if err != nil {
		return err
	}

	if err := confirmDestructiveAction(*autoConfirm); err != nil {
		return err
	}

	cfg, err := config.Load(*configPath)
	if err != nil {
		return fmt.Errorf("load config failed: %w", err)
	}

	db, err := bootstrap.InitDB(cfg.DB)
	if err != nil {
		return fmt.Errorf("init db failed: %w", err)
	}
	db = withSilentSQLLogger(db)
	sqls.SetDB(db)
	slog.Info("connected database success")

	droppedTableCount, err := resetAllTables(db, cfg.DB.Type)
	if err != nil {
		return fmt.Errorf("reset all tables failed: %w", err)
	}
	slog.Info("reset all tables success", slog.Int("droppedTableCount", droppedTableCount))

	if err := db.AutoMigrate(models.Models...); err != nil {
		return fmt.Errorf("create testdata schema failed: %w", err)
	}
	slog.Info("create testdata schema success")

	aiConfigResult, err := aiconfig.Init()
	if err != nil {
		return fmt.Errorf("init ai config failed: %w", err)
	}
	slog.Info("ai config init success", slog.Bool("skipped", aiConfigResult.Skipped),
		slog.String("filePath", aiConfigResult.FilePath),
		slog.Int("created", aiConfigResult.Created),
		slog.Int("updated", aiConfigResult.Updated))

	kbResult, err := kb.Init(lang)
	if err != nil {
		return fmt.Errorf("init knowledge base failed: %w", err)
	}
	slog.Info("knowledge base init success",
		slog.Int64("faqKnowledgeBaseID", kbResult.FAQKnowledgeBaseID),
		slog.Int("createdFAQs", kbResult.CreatedFAQs),
		slog.Int("updatedFAQs", kbResult.UpdatedFAQs),
	)

	aiAgentResult, err := aiagent.Init(lang)
	if err != nil {
		return fmt.Errorf("init ai agent failed: %w", err)
	}
	slog.Info("ai agent init success", slog.Int("created", aiAgentResult.Created), slog.Int("updated", aiAgentResult.Updated))

	channelResult, err := channel.Init(lang)
	if err != nil {
		return fmt.Errorf("init channel failed: %w", err)
	}
	slog.Info("channel init success", slog.Int("created", channelResult.Created), slog.Int("updated", channelResult.Updated))

	if err := quickreply.Init(lang); err != nil {
		return fmt.Errorf("init quick reply failed: %w", err)
	}
	slog.Info("quick reply init success")

	slog.Info("testdata initialization completed")
	return nil
}

func withSilentSQLLogger(db *gorm.DB) *gorm.DB {
	if db == nil {
		return nil
	}
	return db.Session(&gorm.Session{Logger: db.Logger.LogMode(gormlogger.Silent)})
}

func confirmDestructiveAction(autoConfirm bool) error {
	if autoConfirm {
		return nil
	}

	fmt.Println("警告：该操作会清空当前数据库中的所有表和数据。")
	fmt.Print("请输入 INIT 继续，输入其他任意内容取消：")

	var input string
	if _, err := fmt.Scanln(&input); err != nil {
		return fmt.Errorf("read confirmation failed: %w", err)
	}

	if strings.TrimSpace(input) != "INIT" {
		return fmt.Errorf("initialization cancelled")
	}
	return nil
}

func resetAllTables(db *gorm.DB, dbType string) (int, error) {
	tables, err := db.Migrator().GetTables()
	if err != nil {
		return 0, err
	}

	filterSystemTables := func(tables []string) []string {
		ret := make([]string, 0, len(tables))
		for _, table := range tables {
			if strings.HasPrefix(table, "sqlite_") {
				continue
			}
			ret = append(ret, table)
		}
		return ret
	}

	filtered := filterSystemTables(tables)
	if len(filtered) == 0 {
		return 0, nil
	}

	err = withForeignKeyChecksDisabled(db, dbType, func() error {
		items := make([]any, 0, len(filtered))
		for _, table := range filtered {
			items = append(items, table)
		}
		return db.Migrator().DropTable(items...)
	})
	if err != nil {
		return 0, err
	}

	return len(filtered), nil
}

func withForeignKeyChecksDisabled(db *gorm.DB, dbType string, fn func() error) error {
	var disableSQL string
	var enableSQL string

	switch dbType {
	case "sqlite":
		disableSQL = "PRAGMA foreign_keys = OFF"
		enableSQL = "PRAGMA foreign_keys = ON"
	case "mysql":
		disableSQL = "SET FOREIGN_KEY_CHECKS = 0"
		enableSQL = "SET FOREIGN_KEY_CHECKS = 1"
	default:
		return fn()
	}

	if err := db.Exec(disableSQL).Error; err != nil {
		return err
	}
	defer func() {
		_ = db.Exec(enableSQL).Error
	}()

	return fn()
}