s_cherry_e/sr/tool/cleandb/main.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"
)

// Config 配置文件结构 (直接映射环境名称到配置)
type Config map[string]EnvConfig

// EnvConfig 每个环境的配置
type EnvConfig struct {
	SSH   SSHConfig   `json:"ssh"`
	Mongo MongoConfig `json:"mongo"`
	Redis RedisConfig `json:"redis"`
}

// MongoConfig MongoDB 清理配置
type MongoConfig struct {
	Config        DBConfig `json:"config"`
	DropDatabases []string `json:"drop_databases"`
}

// RedisConfig Redis 清理配置
type RedisConfig struct {
	Config     DBConfig `json:"config"`
	DeleteKeys []string `json:"delete_keys"`
}

func main() {
	// 默认配置
	configPath := "config.json"
	var envName string

	args := os.Args[1:] // 获取实际参数 (排除程序名)

	if len(args) >= 2 {
		// 方式: ./cleandb <configPath> <envName>
		configPath = args[0]
		envName = args[1]
	} else if len(args) == 1 {
		// 方式: ./cleandb <envName> (配置路径使用默认值)
		envName = args[0]
	}

	if envName == "" {
		log.Fatal("使用说明:\n  方式1: ./cleandb <configPath> <envName>\n  方式2: ./cleandb <envName> (默认加载当前目录下的 config.json)")
	}

	// 加载配置
	data, err := os.ReadFile(configPath)
	if err != nil {
		log.Fatalf("无法读取配置文件: %v", err)
	}

	var totalConfig Config
	if err := json.Unmarshal(data, &totalConfig); err != nil {
		log.Fatalf("解析配置文件失败: %v", err)
	}

	env, ok := totalConfig[envName]
	if !ok {
		log.Fatalf("未找到环境配置: %s", envName)
	}

	ctx := context.Background()

	// 1. 处理 MongoDB 清理
	if len(env.Mongo.DropDatabases) > 0 {
		cleanMongo(ctx, env.SSH, env.Mongo)
	}

	// 2. 处理 Redis 清理
	if len(env.Redis.DeleteKeys) > 0 {
		cleanRedis(ctx, env.SSH, env.Redis)
	}

	fmt.Println("\n所有任务已完成！")
}

func cleanMongo(ctx context.Context, sshConf SSHConfig, mongoConf MongoConfig) {
	fmt.Printf("\n--- MongoDB 清理 [%v] ---\n", mongoConf.Config.Host)
	fmt.Printf("待删除数据库: %v\n", mongoConf.DropDatabases)

	client, tunnel, err := ConnectMongoViaSSH(sshConf, mongoConf.Config)
	if err != nil {
		log.Fatalf("连接 MongoDB 失败: %v", err)
	}

	defer client.Disconnect(ctx)

	if tunnel != nil {
		defer tunnel.Close()
	}

	for _, dbName := range mongoConf.DropDatabases {
		err := client.Database(dbName).Drop(ctx)
		if err != nil {
			fmt.Printf("警告: 删除数据库 %s 失败: %v\n", dbName, err)
		} else {
			fmt.Printf("成功删除数据库: %s\n", dbName)
		}
	}
}

func cleanRedis(ctx context.Context, sshConf SSHConfig, redisConf RedisConfig) {
	fmt.Printf("\n--- Redis 清理 [%v] ---\n", redisConf.Config.Host)
	fmt.Printf("待删除键: %v\n", redisConf.DeleteKeys)

	rdb, tunnel, err := ConnectRedisViaSSH(sshConf, redisConf.Config)
	if err != nil {
		log.Fatalf("连接 Redis 失败: %v", err)
	}

	defer rdb.Close()

	if tunnel != nil {
		defer tunnel.Close()
	}

	for _, keyPattern := range redisConf.DeleteKeys {
		if strings.Contains(keyPattern, "*") {
			// 处理通配符
			iter := rdb.Scan(ctx, 0, keyPattern, 0).Iterator()
			count := 0
			for iter.Next(ctx) {
				if err := rdb.Del(ctx, iter.Val()).Err(); err != nil {
					fmt.Printf("警告: 删除键 %s 失败: %v\n", iter.Val(), err)
				} else {
					count++
				}
			}
			fmt.Printf("模式 [%s] 已删除 %d 个键\n", keyPattern, count)
		} else {
			// 处理单个键
			if err := rdb.Del(ctx, keyPattern).Err(); err != nil {
				fmt.Printf("警告: 删除键 %s 失败: %v\n", keyPattern, err)
			} else {
				fmt.Printf("成功删除键: %s\n", keyPattern)
			}
		}
	}
}

func confirm(prompt string) bool {
	reader := bufio.NewReader(os.Stdin)
	fmt.Printf("%s (y/n): ", prompt)
	input, _ := reader.ReadString('\n')
	input = strings.TrimSpace(strings.ToLower(input))
	return input == "y" || input == "yes"
}