125 lines
3.5 KiB
Go
125 lines
3.5 KiB
Go
// Package db 提供基于PostgreSQL的数据存储功能
|
|
// 使用GORM作为ORM库来操作数据库
|
|
// 实现与PostgreSQL数据库的连接和基本操作
|
|
package db
|
|
|
|
import (
|
|
"fmt"
|
|
"time"
|
|
|
|
"git.huangwc.com/pig/pig-farm-controller/internal/logs"
|
|
"git.huangwc.com/pig/pig-farm-controller/internal/model"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
// migrateModels 需要自动迁移的数据库模型列表
|
|
var migrateModels = []interface{}{
|
|
&model.User{},
|
|
&model.OperationHistory{},
|
|
&model.Device{},
|
|
&model.DeviceControl{},
|
|
}
|
|
|
|
// PostgresStorage 代表基于PostgreSQL的存储实现
|
|
// 使用GORM作为ORM库
|
|
type PostgresStorage struct {
|
|
// db GORM数据库实例
|
|
db *gorm.DB
|
|
|
|
// connectionString 数据库连接字符串
|
|
connectionString string
|
|
|
|
// maxOpenConns 最大开放连接数
|
|
maxOpenConns int
|
|
|
|
// maxIdleConns 最大空闲连接数
|
|
maxIdleConns int
|
|
|
|
// connMaxLifetime 连接最大生命周期(秒)
|
|
connMaxLifetime int
|
|
|
|
// logger 日志记录器
|
|
logger *logs.Logger
|
|
}
|
|
|
|
// NewPostgresStorage 创建并返回一个新的PostgreSQL存储实例
|
|
// 初始化数据库连接和相关配置
|
|
func NewPostgresStorage(connectionString string, maxOpenConns, maxIdleConns, connMaxLifetime int) *PostgresStorage {
|
|
return &PostgresStorage{
|
|
connectionString: connectionString,
|
|
maxOpenConns: maxOpenConns,
|
|
maxIdleConns: maxIdleConns,
|
|
connMaxLifetime: connMaxLifetime,
|
|
logger: logs.NewLogger(),
|
|
}
|
|
}
|
|
|
|
// Connect 建立与PostgreSQL数据库的连接
|
|
// 使用GORM建立数据库连接
|
|
func (ps *PostgresStorage) Connect() error {
|
|
ps.logger.Info("正在连接PostgreSQL数据库")
|
|
|
|
var err error
|
|
ps.db, err = gorm.Open(postgres.Open(ps.connectionString), &gorm.Config{})
|
|
if err != nil {
|
|
ps.logger.Error(fmt.Sprintf("数据库连接失败: %v", err))
|
|
return fmt.Errorf("数据库连接失败: %v", err)
|
|
}
|
|
|
|
// 测试连接
|
|
sqlDB, err := ps.db.DB()
|
|
if err != nil {
|
|
ps.logger.Error(fmt.Sprintf("获取数据库实例失败: %v", err))
|
|
return fmt.Errorf("获取数据库实例失败: %v", err)
|
|
}
|
|
|
|
if err = sqlDB.Ping(); err != nil {
|
|
ps.logger.Error(fmt.Sprintf("数据库连接测试失败: %v", err))
|
|
return fmt.Errorf("数据库连接测试失败: %v", err)
|
|
}
|
|
|
|
// 设置连接池参数
|
|
sqlDB.SetMaxOpenConns(ps.maxOpenConns)
|
|
sqlDB.SetMaxIdleConns(ps.maxIdleConns)
|
|
sqlDB.SetConnMaxLifetime(time.Duration(ps.connMaxLifetime) * time.Second)
|
|
|
|
// 自动迁移数据库表结构
|
|
ps.logger.Info("正在自动迁移数据库表结构")
|
|
if err = ps.db.AutoMigrate(migrateModels...); err != nil {
|
|
ps.logger.Error(fmt.Sprintf("数据库表结构迁移失败: %v", err))
|
|
return fmt.Errorf("数据库表结构迁移失败: %v", err)
|
|
}
|
|
ps.logger.Info("数据库表结构迁移完成")
|
|
|
|
ps.logger.Info("PostgreSQL数据库连接成功")
|
|
return nil
|
|
}
|
|
|
|
// Disconnect 断开与PostgreSQL数据库的连接
|
|
// 安全地关闭所有数据库连接
|
|
func (ps *PostgresStorage) Disconnect() error {
|
|
if ps.db != nil {
|
|
ps.logger.Info("正在断开PostgreSQL数据库连接")
|
|
|
|
sqlDB, err := ps.db.DB()
|
|
if err != nil {
|
|
ps.logger.Error(fmt.Sprintf("获取数据库实例失败: %v", err))
|
|
return fmt.Errorf("获取数据库实例失败: %v", err)
|
|
}
|
|
|
|
if err := sqlDB.Close(); err != nil {
|
|
ps.logger.Error(fmt.Sprintf("关闭数据库连接失败: %v", err))
|
|
return fmt.Errorf("关闭数据库连接失败: %v", err)
|
|
}
|
|
ps.logger.Info("PostgreSQL数据库连接已断开")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetDB 获取GORM数据库实例
|
|
// 用于执行具体的数据库操作
|
|
func (ps *PostgresStorage) GetDB() *gorm.DB {
|
|
return ps.db
|
|
}
|