368 lines
11 KiB
Go
368 lines
11 KiB
Go
package model
|
||
|
||
import (
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"errors"
|
||
"fmt"
|
||
"log"
|
||
|
||
"gorm.io/driver/postgres"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/logger"
|
||
)
|
||
|
||
// DB 全局数据库实例
|
||
var DB *gorm.DB
|
||
|
||
// InitDB 连接 PostgreSQL 并自动迁移
|
||
func InitDB(dsn string) error {
|
||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{
|
||
Logger: logger.Default.LogMode(logger.Warn),
|
||
})
|
||
if err != nil {
|
||
return fmt.Errorf("连接 PostgreSQL 失败: %w", err)
|
||
}
|
||
|
||
// 旧表 refresh_tokens.token（明文）→ token_hash 体系：必须在 AutoMigrate
|
||
// 创建 NOT NULL 列/唯一索引之前完成回填
|
||
if err := prepareRefreshTokenMigration(db); err != nil {
|
||
return fmt.Errorf("refresh token 旧数据迁移失败: %w", err)
|
||
}
|
||
|
||
// 清理前次失败迁移遗留的 FK 约束（constraint:- 标签防止新建，但不删旧的）
|
||
if err := dropStaleChatFKConstraints(db); err != nil {
|
||
return fmt.Errorf("清理遗留 FK 约束失败: %w", err)
|
||
}
|
||
|
||
if err := db.AutoMigrate(
|
||
&User{}, &Board{}, &Post{}, &Comment{}, &CommentEditHistory{}, &RefreshToken{}, &Like{}, &Notification{}, &Checkin{},
|
||
&Announcement{}, &SiteSetting{}, &SiteDailyStats{}, &SiteDailyVisitor{}, &Attachment{}, &UserBoard{}, &LoginLog{},
|
||
&ChatRoom{}, &ChatRoomMember{}, &ChatMessage{},
|
||
&PointLedger{}, &PostContentUnlock{}, &PostAttachment{}, &PostAttachmentUnlock{},
|
||
&PostPollVote{}, &PostLotteryEntry{},
|
||
); err != nil {
|
||
return fmt.Errorf("自动迁移失败: %w", err)
|
||
}
|
||
|
||
// 一次性：用签到累计回填 User.Points（仅余额仍为 0 且有签到积分的用户）
|
||
if err := backfillPointsFromCheckin(db); err != nil {
|
||
return fmt.Errorf("积分余额回填失败: %w", err)
|
||
}
|
||
|
||
// RBAC：把初始管理员（id 最小的 admin，通常即首个注册账号）升级为站长；
|
||
// 已存在 owner 时不动数据，保证幂等
|
||
if err := ensureOwnerRole(db); err != nil {
|
||
return fmt.Errorf("站长角色迁移失败: %w", err)
|
||
}
|
||
|
||
// login_logs.success 早期 default=true 与 GORM 零值省略叠加，
|
||
// 会把失败登录错存为成功；AutoMigrate 不会改列默认值，这里幂等修正
|
||
if err := db.Exec(`ALTER TABLE login_logs ALTER COLUMN success SET DEFAULT false`).Error; err != nil {
|
||
return fmt.Errorf("login_logs 默认值修正失败: %w", err)
|
||
}
|
||
|
||
// 新表结构就位后删除遗留的明文列
|
||
if err := dropLegacyRefreshTokenColumn(db); err != nil {
|
||
return fmt.Errorf("refresh token 旧列清理失败: %w", err)
|
||
}
|
||
|
||
DB = db
|
||
if err := db.Exec(`
|
||
CREATE UNIQUE INDEX IF NOT EXISTS idx_chat_direct_key
|
||
ON chat_rooms (direct_key)
|
||
WHERE room_type = 'direct' AND deleted_at IS NULL AND direct_key <> ''
|
||
`).Error; err != nil {
|
||
return fmt.Errorf("私聊唯一索引创建失败: %w", err)
|
||
}
|
||
seedDefaultBoards(db)
|
||
if err := seedDefaultChatRoom(db); err != nil {
|
||
return fmt.Errorf("默认群聊大厅初始化失败: %w", err)
|
||
}
|
||
if err := ensureDefaultChatMemberships(db); err != nil {
|
||
return fmt.Errorf("默认群聊入群补齐失败: %w", err)
|
||
}
|
||
log.Println("[model] PostgreSQL 数据库初始化完成")
|
||
return nil
|
||
}
|
||
|
||
// PingDB 检测数据库连接
|
||
func PingDB() error {
|
||
if DB == nil {
|
||
return fmt.Errorf("数据库未初始化")
|
||
}
|
||
sqlDB, err := DB.DB()
|
||
if err != nil {
|
||
return err
|
||
}
|
||
return sqlDB.Ping()
|
||
}
|
||
|
||
// dropStaleChatFKConstraints 清理前次迁移失败遗留的 FK 约束。
|
||
// GORM AutoMigrate 先建表再加约束：首次迁移时 Notification.Room 的 FK 添加失败，
|
||
// 但 chat_rooms 表已建好且带 fk_chat_rooms_last_message 约束。
|
||
// 模型已改用 constraint:- 防止新建，但 GORM 不会删除已存在的约束，需手动清理。
|
||
func dropStaleChatFKConstraints(db *gorm.DB) error {
|
||
// 对每个遗留约束名，在可能的两张表上都尝试删除（IF EXISTS 保证幂等）
|
||
pairs := []struct{ table, constraint string }{
|
||
{"chat_rooms", "fk_chat_rooms_last_message"},
|
||
{"notifications", "fk_notifications_room"},
|
||
}
|
||
for _, p := range pairs {
|
||
db.Exec(fmt.Sprintf(`ALTER TABLE IF EXISTS %s DROP CONSTRAINT IF EXISTS %s`, p.table, p.constraint))
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// prepareRefreshTokenMigration 旧版 refresh_tokens 表把明文存在 token 列，
|
||
// 新版改为 token_hash（SHA-256，NOT NULL+唯一索引）。在 AutoMigrate 之前：
|
||
// 1. 新增可带默认值的 token_hash 列（避免对存量行加 NOT NULL 列失败）
|
||
// 2. 用存量明文回填哈希
|
||
// 3. AutoMigrate 随后补唯一索引/其余新列
|
||
//
|
||
// 旧行无法回填 TokenCipher（密钥在 service 层），仅影响该行下一次轮转的
|
||
// 并发重放宽限，属一次性边界；轮转后即完全进入新体系。
|
||
func prepareRefreshTokenMigration(db *gorm.DB) error {
|
||
var tableCount int64
|
||
if err := db.Raw(`SELECT count(1) FROM information_schema.tables WHERE table_name = 'refresh_tokens'`).
|
||
Scan(&tableCount).Error; err != nil {
|
||
return err
|
||
}
|
||
if tableCount == 0 {
|
||
return nil // 全新数据库，AutoMigrate 直接建新表
|
||
}
|
||
|
||
var hasHashCol int64
|
||
if err := db.Raw(`SELECT count(1) FROM information_schema.columns
|
||
WHERE table_name = 'refresh_tokens' AND column_name = 'token_hash'`).
|
||
Scan(&hasHashCol).Error; err != nil {
|
||
return err
|
||
}
|
||
if hasHashCol > 0 {
|
||
return nil // 已是新结构
|
||
}
|
||
|
||
var hasLegacyCol int64
|
||
if err := db.Raw(`SELECT count(1) FROM information_schema.columns
|
||
WHERE table_name = 'refresh_tokens' AND column_name = 'token'`).
|
||
Scan(&hasLegacyCol).Error; err != nil {
|
||
return err
|
||
}
|
||
if hasLegacyCol == 0 {
|
||
return nil
|
||
}
|
||
|
||
if err := db.Exec(`DELETE FROM refresh_tokens WHERE token IS NULL OR token = ''`).Error; err != nil {
|
||
return err
|
||
}
|
||
if err := db.Exec(`ALTER TABLE refresh_tokens ADD COLUMN token_hash varchar(64) NOT NULL DEFAULT ''`).Error; err != nil {
|
||
return err
|
||
}
|
||
|
||
type legacyRow struct {
|
||
ID uint
|
||
Token string
|
||
}
|
||
var rows []legacyRow
|
||
if err := db.Raw(`SELECT id, token FROM refresh_tokens`).Scan(&rows).Error; err != nil {
|
||
return err
|
||
}
|
||
for _, r := range rows {
|
||
sum := sha256.Sum256([]byte(r.Token))
|
||
if err := db.Exec(
|
||
`UPDATE refresh_tokens SET token_hash = ? WHERE id = ?`,
|
||
hex.EncodeToString(sum[:]), r.ID,
|
||
).Error; err != nil {
|
||
return err
|
||
}
|
||
}
|
||
log.Printf("[model] refresh_tokens 已回填 %d 行 token_hash", len(rows))
|
||
return nil
|
||
}
|
||
|
||
// dropLegacyRefreshTokenColumn 新结构就位后删除明文 token 列（PostgreSQL
|
||
// 会连带删除该列上的旧唯一索引）
|
||
func dropLegacyRefreshTokenColumn(db *gorm.DB) error {
|
||
var hasLegacyCol int64
|
||
if err := db.Raw(`SELECT count(1) FROM information_schema.columns
|
||
WHERE table_name = 'refresh_tokens' AND column_name = 'token'`).
|
||
Scan(&hasLegacyCol).Error; err != nil {
|
||
return err
|
||
}
|
||
if hasLegacyCol == 0 {
|
||
return nil
|
||
}
|
||
return db.Exec(`ALTER TABLE refresh_tokens DROP COLUMN token`).Error
|
||
}
|
||
|
||
// ensureOwnerRole 若无站长，则把 id 最小的旧管理员升级为站长；
|
||
// 连管理员都没有的全新库，把 id=1 的初始账号设为站长
|
||
func ensureOwnerRole(db *gorm.DB) error {
|
||
var ownerCount int64
|
||
if err := db.Model(&User{}).Where("role = ?", RoleOwner).Count(&ownerCount).Error; err != nil {
|
||
return err
|
||
}
|
||
if ownerCount > 0 {
|
||
return nil
|
||
}
|
||
res := db.Model(&User{}).Where("role = ?", RoleAdmin).
|
||
Order("id ASC").Limit(1).Update("role", RoleOwner)
|
||
if res.Error != nil {
|
||
return res.Error
|
||
}
|
||
if res.RowsAffected == 0 {
|
||
if err := db.Model(&User{}).Order("id ASC").Limit(1).
|
||
Update("role", RoleOwner).Error; err != nil {
|
||
return err
|
||
}
|
||
}
|
||
log.Println("[model] 已迁移初始账号为站长角色(owner)")
|
||
return nil
|
||
}
|
||
|
||
// backfillPointsFromCheckin 把历史签到积分写入 User.Points（幂等：仅 points=0 且有签到）
|
||
func backfillPointsFromCheckin(db *gorm.DB) error {
|
||
type row struct {
|
||
UserID uint
|
||
Total int
|
||
}
|
||
var rows []row
|
||
if err := db.Raw(`
|
||
SELECT c.user_id, COALESCE(SUM(c.points), 0)::int AS total
|
||
FROM checkins c
|
||
INNER JOIN users u ON u.id = c.user_id AND u.deleted_at IS NULL AND u.points = 0
|
||
GROUP BY c.user_id
|
||
HAVING COALESCE(SUM(c.points), 0) > 0
|
||
`).Scan(&rows).Error; err != nil {
|
||
return err
|
||
}
|
||
for _, r := range rows {
|
||
err := db.Transaction(func(tx *gorm.DB) error {
|
||
if err := tx.Model(&User{}).Where("id = ? AND points = 0", r.UserID).
|
||
Update("points", r.Total).Error; err != nil {
|
||
return err
|
||
}
|
||
return tx.Create(&PointLedger{
|
||
UserID: r.UserID,
|
||
Delta: r.Total,
|
||
Balance: r.Total,
|
||
Reason: PointReasonMigrateCheckin,
|
||
Note: "历史签到积分回填",
|
||
}).Error
|
||
})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
}
|
||
if len(rows) > 0 {
|
||
log.Printf("[model] 已回填 %d 名用户的签到积分余额", len(rows))
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// seedDefaultBoards 写入默认板块
|
||
func seedDefaultBoards(db *gorm.DB) {
|
||
defaults := []Board{
|
||
{Name: "综合讨论", Description: "什么都可以聊", Icon: "message-circle", SortOrder: 1},
|
||
{Name: "技术分享", Description: "分享技术心得与问题", Icon: "code", SortOrder: 2},
|
||
{Name: "问答求助", Description: "提问与解答", Icon: "help-circle", SortOrder: 3},
|
||
{Name: "闲聊灌水", Description: "轻松闲聊", Icon: "coffee", SortOrder: 4},
|
||
}
|
||
for _, b := range defaults {
|
||
var count int64
|
||
db.Model(&Board{}).Where("name = ?", b.Name).Count(&count)
|
||
if count == 0 {
|
||
_ = db.Create(&b).Error
|
||
}
|
||
}
|
||
}
|
||
|
||
// seedDefaultChatRoom 确保存在唯一的全站大厅（公开默认群）
|
||
func seedDefaultChatRoom(db *gorm.DB) error {
|
||
var n int64
|
||
if err := db.Model(&ChatRoom{}).Where("is_default = ?", true).Count(&n).Error; err != nil {
|
||
return err
|
||
}
|
||
if n > 0 {
|
||
return nil
|
||
}
|
||
var owner User
|
||
if err := db.Where("role = ?", RoleOwner).Order("id ASC").First(&owner).Error; err != nil {
|
||
// 尚无用户时跳过，等有站长后再由下次启动补齐
|
||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||
return nil
|
||
}
|
||
// 尝试任意用户
|
||
if err2 := db.Order("id ASC").First(&owner).Error; err2 != nil {
|
||
return nil
|
||
}
|
||
}
|
||
room := ChatRoom{
|
||
Name: "全站大厅",
|
||
Description: "所有成员默认加入的公开群聊",
|
||
OwnerID: owner.ID,
|
||
RoomType: ChatRoomTypeGroup,
|
||
IsPrivate: false,
|
||
IsDefault: true,
|
||
MemberCount: 0,
|
||
}
|
||
if err := db.Create(&room).Error; err != nil {
|
||
return err
|
||
}
|
||
log.Println("[model] 已创建默认群聊「全站大厅」")
|
||
return nil
|
||
}
|
||
|
||
// ensureDefaultChatMemberships 把未封禁用户补进默认大厅（幂等）
|
||
func ensureDefaultChatMemberships(db *gorm.DB) error {
|
||
var room ChatRoom
|
||
if err := db.Where("is_default = ?", true).First(&room).Error; err != nil {
|
||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||
return nil
|
||
}
|
||
return err
|
||
}
|
||
var userIDs []uint
|
||
if err := db.Model(&User{}).
|
||
Where("deleted_at IS NULL AND banned = ?", false).
|
||
Pluck("id", &userIDs).Error; err != nil {
|
||
return err
|
||
}
|
||
if len(userIDs) == 0 {
|
||
return nil
|
||
}
|
||
var existing []uint
|
||
if err := db.Model(&ChatRoomMember{}).
|
||
Where("room_id = ?", room.ID).
|
||
Pluck("user_id", &existing).Error; err != nil {
|
||
return err
|
||
}
|
||
have := make(map[uint]struct{}, len(existing))
|
||
for _, id := range existing {
|
||
have[id] = struct{}{}
|
||
}
|
||
added := 0
|
||
for _, uid := range userIDs {
|
||
if _, ok := have[uid]; ok {
|
||
continue
|
||
}
|
||
role := ChatRoleMember
|
||
if uid == room.OwnerID {
|
||
role = ChatRoleOwner
|
||
}
|
||
m := ChatRoomMember{RoomID: room.ID, UserID: uid, Role: role}
|
||
if err := db.Create(&m).Error; err != nil {
|
||
continue
|
||
}
|
||
added++
|
||
}
|
||
if added > 0 {
|
||
_ = db.Model(&ChatRoom{}).Where("id = ?", room.ID).
|
||
UpdateColumn("member_count", gorm.Expr(
|
||
"(SELECT count(*) FROM chat_room_members WHERE room_id = ?)", room.ID,
|
||
)).Error
|
||
log.Printf("[model] 默认群聊补齐入群 %d 人", added)
|
||
}
|
||
return nil
|
||
}
|