xiubuzhe/sunhpc-go
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

优化数据库,日志,命令模块

Browse Source
This commit is contained in:
Qichao
2026-02-23 18:54:54 +08:00
parent 47a2dfeda1
commit 3a5f5ddd5d
6 changed files with 330 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files

219
pkg/config/config.go
View File

@@ -4,70 +4,145 @@ import (
"fmt"
"os"
"path/filepath"
"sunhpc/pkg/logger"
"sunhpc/pkg/utils"
"sync"
"github.com/spf13/viper"
"go.yaml.in/yaml/v3"
)
// ==============================================================
// 全局变量
// ==============================================================
var (
GlobalConfig *Config
configOnce sync.Once // 确保配置只加载一次
configMutex sync.RWMutex // 读写锁保护 GlobalConfig
)
// ==============================================================
// 目录常量 (可在 main.go 中通过 flag 覆盖默认值)
// ==============================================================
var (
BaseDir string = utils.DefaultBaseDir
TmplDir string = utils.DefaultTmplDir
LogDir string = utils.DefaultLogDir
)
// ==============================================================
// 配置结构体
// ==============================================================
type Config struct {
Database DatabaseConfig `yaml:"database"`
Log LogConfig `yaml:"log"`
Log logger.LogConfig `mapstructure:"log" yaml:"log"`
Database DatabaseConfig `mapstructure:"database" yaml:"database"`
}
type DatabaseConfig struct {
DSN string `yaml:"dsn"` // 数据库连接字符串
Path string `yaml:"path"` // SQLite: 目录路径
Name string `yaml:"name"` // SQLite: 文件名
DSN string `mapstructure:"dsn" yaml:"dsn"` // 数据库连接字符串
Path string `mapstructure:"path" yaml:"path"` // SQLite: 目录路径
Name string `mapstructure:"name" yaml:"name"` // SQLite: 文件名
Args string `mapstructure:"args" yaml:"args"` // 数据库连接参数
}
type LogConfig struct {
Level string `yaml:"level"`
Format string `yaml:"format"`
Output string `yaml:"output"`
Verbose bool `yaml:"verbose"`
LogFile string `yaml:"log_file"`
ShowColor bool `yaml:"show_color"`
type CLIParamsType struct {
Verbose bool // -v/--verbose
NoColor bool // --no-color
Config string // -c/--config
}
// --------------------------------- 全局单例配置(核心) ---------------------------------
var (
// GlobalConfig 全局配置单例实例
GlobalConfig *Config
// 命令行参数配置(全局、由root命令绑定)
CLIParams = struct {
Verbose bool // -v/--verbose
NoColor bool // --no-color
Config string // -c/--config
}{}
BaseDir string = "/etc/sunhpc"
LogDir string = "/var/log/sunhpc"
TmplDir string = BaseDir + "/tmpl.d"
appName string = "sunhpc"
defaultDBPath string = "/var/lib/sunhpc"
defaultDBName string = "sunhpc.db"
)
var CLIParams CLIParamsType // 命令行参数
// InitConfigs 初始化所有配置目录等数据
func InitConfigs() error {
dirs := []string{
BaseDir,
TmplDir,
LogDir,
}
for _, d := range dirs {
if err := os.MkdirAll(d, 0755); err != nil {
return fmt.Errorf("创建目录 %s 失败: %w", d, err)
}
}
// 检查配置文件是否存在，不存在则创建默认配置
configPath := filepath.Join(BaseDir, "sunhpc.yaml")
if _, err := os.Stat(configPath); os.IsNotExist(err) {
if err := createDefaultConfig(configPath); err != nil {
return fmt.Errorf("创建默认配置文件失败: %w", err)
}
}
return nil
}
func createDefaultConfig(configPath string) error {
defaultConfig := &Config{
Database: DatabaseConfig{
Path: utils.DefaultDBPath,
Name: utils.DefaultDBName,
Args: utils.DefaultDBArgs,
},
Log: logger.LogConfig{
Level: "info",
Format: "text",
Output: "stdout",
Verbose: false,
LogFile: utils.DefaultLogFile,
ShowColor: true,
},
}
// 确保数据库目录存在
if err := os.MkdirAll(defaultConfig.Database.Path, 0755); err != nil {
return fmt.Errorf("创建数据库目录失败: %w", err)
}
// 序列号并写入
data, err := yaml.Marshal(defaultConfig)
if err != nil {
return fmt.Errorf("序列化配置失败: %w", err)
}
return os.WriteFile(configPath, data, 0644)
}
// ----------------------------------- 配置加载(只加载一次) -----------------------------------
func LoadConfig() (*Config, error) {
// 如果已经加载过，直接返回
configMutex.RLock()
if GlobalConfig != nil {
// 如果已经加载过，直接返回
configMutex.RUnlock()
return GlobalConfig, nil
}
configMutex.RUnlock()
configMutex.Lock()
defer configMutex.Unlock()
// 双重检查
if GlobalConfig != nil {
// 如果已经加载过，直接返回
return GlobalConfig, nil
}
viper.SetConfigName("sunhpc")
viper.SetConfigType("yaml")
viper.AddConfigPath(BaseDir)
viper.AddConfigPath(".")
viper.AddConfigPath(filepath.Join(os.Getenv("HOME"), "."))
// 配置文件路径
if CLIParams.Config != "" {
viper.SetConfigFile(CLIParams.Config)
} else {
viper.SetConfigName("sunhpc")
viper.SetConfigType("yaml")
viper.AddConfigPath(BaseDir)
viper.AddConfigPath(".")
viper.AddConfigPath(filepath.Join(os.Getenv("HOME"), "."))
}
// Step 1: 设置默认值(最低优先级)
viper.SetDefault("log.level", "info")
viper.SetDefault("log.format", "text")
viper.SetDefault("log.output", "stdout")
viper.SetDefault("log.verbose", false)
viper.SetDefault("log.log_file", filepath.Join(LogDir, "sunhpc.log"))
viper.SetDefault("database.name", "sunhpc.db")
viper.SetDefault("database.path", "/var/lib/sunhpc")
viper.SetDefault("log.log_file", utils.DefaultLogFile)
viper.SetDefault("database.name", utils.DefaultDBName)
viper.SetDefault("database.path", utils.DefaultDBPath)
if err := viper.ReadInConfig(); err != nil {
// 配置文件不存在时，使用默认值
@@ -87,10 +162,12 @@ func LoadConfig() (*Config, error) {
viper.Set("log.show_color", false)
}
fullPath := filepath.Join(
viper.GetString("database.path"), viper.GetString("database.name"))
dsn := fmt.Sprintf(
"%s?_foreign_keys=on&_journal_mode=WAL&_timeout=5000", fullPath)
// 计算派生配置 (如数据库 DSN)
dbPath := viper.GetString("database.path")
dbName := viper.GetString("database.name")
dbArgs := viper.GetString("database.args")
fullPath := filepath.Join(dbPath, dbName)
dsn := fmt.Sprintf("%s?%s", fullPath, dbArgs)
viper.Set("database.dsn", dsn)
// 解码到结构体
@@ -104,57 +181,23 @@ func LoadConfig() (*Config, error) {
return GlobalConfig, nil
}
// InitDirs 创建所有必需目录
func InitDirs() error {
dirs := []string{
BaseDir,
TmplDir,
LogDir,
}
for _, d := range dirs {
if err := os.MkdirAll(d, 0755); err != nil {
return fmt.Errorf("创建目录 %s 失败: %w", d, err)
}
}
return nil
}
// ==============================================================
// SaveConfig - 保存全局配置到文件、运行时配置
// ==============================================================
func SaveConfig() error {
configMutex.RLock()
defer configMutex.RUnlock()
func (c *Config) WriteDefaultConfig(path string) error {
// 确保目录存在
dir := filepath.Dir(path)
if err := os.MkdirAll(dir, 0755); err != nil {
return fmt.Errorf("创建目录失败: %w", err)
if GlobalConfig == nil {
return fmt.Errorf("全局配置为空")
}
// 生成默认配置
cfg := DefaultConfig(path)
// 序列化为 YAML
data, err := yaml.Marshal(cfg)
configPath := filepath.Join(BaseDir, "config.yaml")
data, err := yaml.Marshal(GlobalConfig)
if err != nil {
return fmt.Errorf("序列化配置失败: %w", err)
}
// 写入文件（0644 权限）
return os.WriteFile(path, data, 0644)
}
func DefaultConfig(path string) *Config {
return &Config{
Database: DatabaseConfig{
DSN: fmt.Sprintf("%s?_foreign_keys=on&_journal_mode=WAL&_timeout=5000",
filepath.Join(filepath.Dir(path), defaultDBName)),
Path: filepath.Dir(path),
Name: defaultDBName,
},
Log: LogConfig{
Level: "info",
Format: "text",
Output: "stdout",
LogFile: filepath.Join(filepath.Dir(path), "sunhpc.log"),
Verbose: false,
},
return err
}
return os.WriteFile(configPath, data, 0644)
}
// ResetConfig 重置全局配置为默认值

156
pkg/database/database.go
View File

@@ -5,7 +5,6 @@ import (
"database/sql"
"fmt"
"os"
"path/filepath"
"strings"
"sync"
@@ -16,65 +15,57 @@ import (
_ "github.com/mattn/go-sqlite3"
)
type DB struct {
db *sql.DB
logger logger.Logger
}
// =========================================================
// 全局变量
// =========================================================
var (
dbInstance *DB
dbInstance *sql.DB
dbOnce sync.Once
dbMutex sync.RWMutex
dbErr error
)
func GetInstance(dbConfig *config.DatabaseConfig, log logger.Logger) (*DB, error) {
// =========================================================
// GetDB - 获取数据库连接(单例模式)
// =========================================================
func GetDB() (*sql.DB, error) {
dbOnce.Do(func() {
// 兜底: 未注入则使用全局默认日志实例
if log == nil {
log = logger.DefaultLogger
}
log.Debugf("开始初始化数据库,路径: %s", dbConfig.Path)
// 确认数据库目录存在
if err := os.MkdirAll(dbConfig.Path, 0755); err != nil {
log.Errorf("创建数据库目录失败: %v", err)
dbErr = err
if dbInstance != nil {
return
}
fullPath := filepath.Join(dbConfig.Path, dbConfig.Name)
log.Debugf("数据库路径: %s", fullPath)
// 确保配置已加载
cfg, err := config.LoadConfig()
if err != nil {
dbErr = fmt.Errorf("加载配置失败: %w", err)
return
}
// 构建DSN
dsn := fmt.Sprintf("%s?_foreign_keys=on&_journal_mode=WAL&_timeout=5000&cache=shared",
fullPath)
log.Debugf("DSN: %s", dsn)
logger.Debugf("DSN: %s", cfg.Database.DSN)
// 打开SQLite 连接
sqlDB, err := sql.Open("sqlite3", dsn)
sqlDB, err := sql.Open("sqlite3", cfg.Database.DSN)
if err != nil {
log.Errorf("数据库打开失败: %v", err)
dbErr = err
dbErr = fmt.Errorf("数据库打开失败: %w", err)
return
}
// 设置连接池参数
sqlDB.SetMaxOpenConns(1) // SQLite 只支持单连接
sqlDB.SetMaxIdleConns(1) // 保持一个空闲连接
sqlDB.SetMaxOpenConns(10) // 最大打开连接数
sqlDB.SetMaxIdleConns(5) // 保持空闲连接
sqlDB.SetConnMaxLifetime(0) // 禁用连接生命周期超时
sqlDB.SetConnMaxIdleTime(0) // 禁用空闲连接超时
// 测试数据库连接
if err := sqlDB.Ping(); err != nil {
sqlDB.Close()
log.Errorf("数据库连接失败: %v", err)
dbErr = err
dbErr = fmt.Errorf("数据库连接失败: %w", err)
return
}
// 赋值 *DB 类型的单例(而非直接复制 *sql.DB)
log.Info("数据库连接成功")
dbInstance = &DB{sqlDB, log}
logger.Debug("数据库连接成功")
dbInstance = sqlDB
})
if dbErr != nil {
@@ -84,21 +75,6 @@ func GetInstance(dbConfig *config.DatabaseConfig, log logger.Logger) (*DB, error
return dbInstance, nil
}
// Close 关闭数据库连接
func (d *DB) Close() error {
if d.db != nil {
d.logger.Debug("关闭数据库连接")
err := d.db.Close()
if err != nil {
d.logger.Errorf("数据库连接关闭失败: %v", err)
return err
}
d.logger.Info("数据库连接关闭成功")
return nil
}
return nil
}
func confirmAction(prompt string) bool {
reader := bufio.NewReader(os.Stdin)
@@ -112,38 +88,38 @@ func confirmAction(prompt string) bool {
return response == "y" || response == "yes"
}
func (d *DB) InitTables(force bool) error {
d.logger.Info("开始初始化数据库表...")
func InitTables(db *sql.DB, force bool) error {
if force {
// 确认是否强制删除
if !confirmAction("确认强制删除所有表和触发器？") {
d.logger.Info("操作已取消")
logger.Info("操作已取消")
db.Close()
os.Exit(0)
return nil
}
// 强制删除所有表和触发器
d.logger.Debug("强制删除所有表和触发器...")
if err := dropTables(d.db); err != nil {
logger.Debug("强制删除所有表和触发器...")
if err := dropTables(db); err != nil {
return fmt.Errorf("删除表失败: %w", err)
}
d.logger.Debug("删除所有表和触发器成功")
logger.Debug("删除所有表和触发器成功")
if err := dropTriggers(d.db); err != nil {
if err := dropTriggers(db); err != nil {
return fmt.Errorf("删除触发器失败: %w", err)
}
d.logger.Debug("删除所有触发器成功")
logger.Debug("删除所有触发器成功")
}
// ✅ 调用 schema.go 中的函数
for _, ddl := range CreateTableStatements() {
d.logger.Debugf("执行: %s", ddl)
if _, err := d.db.Exec(ddl); err != nil {
logger.Debugf("执行: %s", ddl)
if _, err := db.Exec(ddl); err != nil {
return fmt.Errorf("数据表创建失败: %w", err)
}
}
d.logger.Info("数据库表创建成功")
logger.Info("数据库表创建成功")
/*
使用sqlite3命令 测试数据库是否存在表
✅ 查询所有表
@@ -174,3 +150,65 @@ func dropTriggers(db *sql.DB) error {
}
return nil
}
func CloseDB() error {
dbMutex.Lock()
defer dbMutex.Unlock()
if dbInstance == nil {
if err := dbInstance.Close(); err != nil {
return err
}
dbInstance = nil
}
return nil
}
// 使用事务回滚测试
func RunTestWithRollback(db *sql.DB, testFunc func(*sql.Tx) error) error {
tx, err := db.Begin()
if err != nil {
return err
}
// 执行测试
if err := testFunc(tx); err != nil {
tx.Rollback()
return err
}
// 回滚事务，所有更改（包括 ID 递增）都会撤销
return tx.Rollback()
}
// 使用示例
func TestNodeInsert(db *sql.DB) error {
logger.Debug("测试数据插入...")
return RunTestWithRollback(db, func(tx *sql.Tx) error {
// 插入测试数据
logger.Debug("执行插入测试数据...")
_, err := tx.Exec(`
INSERT INTO nodes (name, cpus, rack, rank)
VALUES (?, ?, ?, ?)
`, "test-node", 64, 1, 1)
if err != nil {
return err
}
// 验证插入
var count int
logger.Debug("执行查询测试数据...")
err = tx.QueryRow(`
SELECT COUNT(*) FROM nodes WHERE name = ?
`, "test-node").Scan(&count)
if err != nil {
return err
}
logger.Infof("测试数据插入成功，共 %d 条", count)
// 不需要手动删除，回滚会自动撤销
return nil
})
}

54
pkg/logger/logger.go
View File

@@ -67,21 +67,25 @@ var (
// DefaultLogger 全局默认日志实例（所有模块可直接用，也可注入自定义实现）
DefaultLogger Logger
// once 保证日志只初始化一次
once sync.Once
initOnce sync.Once
)
// LogConfig 日志配置结构体（和项目的config包对齐）
type LogConfig struct {
Verbose bool // 是否开启详细模式（Debug级别）
Level string // 日志级别：debug/info/warn/error
ShowColor bool // 是否显示彩色输出
LogFile string // 日志文件路径（可选，空则只输出到控制台）
Level string `mapstructure:"level" yaml:"level"`
Format string `mapstructure:"format" yaml:"format"`
Output string `mapstructure:"output" yaml:"output"`
Verbose bool `mapstructure:"verbose" yaml:"verbose"`
LogFile string `mapstructure:"log_file" yaml:"log_file"`
ShowColor bool `mapstructure:"show_color" yaml:"show_color"`
}
// 默认配置
var defaultConfig = LogConfig{
Verbose: false,
Level: "info",
Format: "text",
Output: "stdout",
Verbose: false,
ShowColor: true,
LogFile: "/var/log/sunhpc/sunhpc.log",
}
@@ -113,6 +117,9 @@ func (f *CustomFormatter) Format(entry *logrus.Entry) ([]byte, error) {
colorReset = ""
}
// Debug,打印原始字节,用于调试
// fmt.Printf("%q\n", entry.Message)
// 拼接格式：
// 灰色日期 + 空格 + 带颜色的[级别] + 空格 + 日志内容 + 空格 + 灰色文件行号 + 重置
fmt.Fprintf(&buf, "%s%s%s %s[%s]%s %s %s%s:%d%s\n",
@@ -150,27 +157,6 @@ func getLevelInfo(level logrus.Level) (string, string) {
}
}
// getCallerInfo 获取调用日志的文件和行号（跳过logrus内部调用）
func _getCallerInfo() (string, int) {
// 跳过的调用栈深度：根据实际情况调整（这里跳过logrus和logger包的调用）
skip := 6
pc, file, line, ok := runtime.Caller(skip)
if !ok {
return "unknown.go", 0
}
// 只保留文件名（如 db.go），去掉完整路径
fileName := filepath.Base(file)
// 过滤logrus内部调用（可选）
funcName := runtime.FuncForPC(pc).Name()
if funcName == "" || filepath.Base(funcName) == "logrus" {
return getCallerInfoWithSkip(skip + 1)
}
return fileName, line
}
func getCallerInfo() (string, int) {
// 从当前调用开始，逐层向上查找
for i := 2; i < 15; i++ { // i从2开始（跳过getCallerInfo自身）
@@ -203,8 +189,7 @@ func shouldSkipPackage(funcName, file string) bool {
}
// 跳过logger包（你自己的包装包）
if strings.Contains(funcName, "your/package/logger") ||
strings.Contains(file, "logger") {
if strings.Contains(file, "/logger/") {
return true
}
@@ -227,21 +212,14 @@ func getCallerInfoWithSkip(skip int) (string, int) {
// Init 初始化全局默认日志实例（全局只执行一次）
func Init(cfg LogConfig) {
once.Do(func() {
// 合并配置：传入的配置为空则用默认值
if cfg.Level == "" {
cfg.Level = defaultConfig.Level
}
if cfg.LogFile == "" {
cfg.LogFile = defaultConfig.LogFile
}
initOnce.Do(func() {
// 1. 创建logrus实例
logrusInst := logrus.New()
// 2. 配置输出（控制台 + 文件，可选）
var outputs []io.Writer
outputs = append(outputs, os.Stdout) // 控制台输出
// 如果配置了日志文件，添加文件输出
if cfg.LogFile != "" {
// 确保日志目录存在

15
pkg/utils/utils.go
View File

@@ -36,3 +36,18 @@ func GetTimestamp() string {
const (
NoAvailableNetworkInterfaces = "No available network interfaces"
)
// 定义目录
const (
DefaultBaseDir string = "/etc/sunhpc"
DefaultTmplDir string = DefaultBaseDir + "/tmpl.d"
DefaultLogDir string = "/var/log/sunhpc"
DefaultLogFile string = DefaultLogDir + "/sunhpc.log"
)
// 定义数据库
const (
DefaultDBName string = "sunhpc.db"
DefaultDBPath string = "/var/lib/sunhpc"
DefaultDBArgs string = "_foreign_keys=on&_journal_mode=WAL&_timeout=5000"
)