238 lines
9.3 KiB
Go
238 lines
9.3 KiB
Go
package db
|
||
|
||
import (
|
||
"fmt"
|
||
"time"
|
||
|
||
"github.com/pkg/errors"
|
||
"github.com/sirupsen/logrus"
|
||
"github.com/uptrace/opentelemetry-go-extra/otelgorm"
|
||
"gorm.io/driver/mysql"
|
||
"gorm.io/driver/postgres"
|
||
"gorm.io/gorm"
|
||
|
||
"github.com/tamakyi/TamaBox/internal/conf"
|
||
)
|
||
|
||
var AllTables = []interface{}{
|
||
&User{}, &Question{}, &FollowUp{}, &CensorLog{}, &UploadImage{}, &UploadImageQuestion{},
|
||
&Conversation{}, &Message{}, &MessageImage{}, &UserBlock{},
|
||
// 匿名投票
|
||
&Vote{}, &VoteOption{}, &VoteFormField{}, &VoteSubmission{},
|
||
// 功能开关（提问 / 私聊 / 投票 的开关、定时、定量收集）
|
||
&FeatureControl{},
|
||
// 站方处置与举报（无管理后台，处置经 CLI 写入、运行时读表生效）
|
||
&UserRestriction{}, &ContentReport{},
|
||
}
|
||
|
||
func Init(typ, dsn string) (*gorm.DB, error) {
|
||
var dialector gorm.Dialector
|
||
|
||
switch typ {
|
||
case "mysql", "":
|
||
dialector = mysql.Open(dsn)
|
||
case "postgres":
|
||
dialector = postgres.Open(dsn)
|
||
default:
|
||
return nil, errors.Errorf("unknown database type: %q", typ)
|
||
}
|
||
|
||
db, err := gorm.Open(dialector, &gorm.Config{
|
||
SkipDefaultTransaction: true,
|
||
})
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "connect to database")
|
||
}
|
||
|
||
sqlDB, err := db.DB()
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "get sql.DB from gorm")
|
||
}
|
||
|
||
// 1C1G 低配环境：调低连接池避免单连接占用过多内存。提问箱低并发，10 个打开连接绰绰有余。
|
||
sqlDB.SetMaxOpenConns(10)
|
||
sqlDB.SetMaxIdleConns(4)
|
||
sqlDB.SetConnMaxLifetime(5 * time.Minute) // 避免长连接被数据库服务端断开
|
||
sqlDB.SetConnMaxIdleTime(3 * time.Minute)
|
||
|
||
// 仅在明确开启迁移清理开关时处理 users 表历史重复值，避免服务启动时静默改写账号数据。
|
||
// 正常部署应先离线检查并清理重复 email/domain，再执行 AutoMigrate。
|
||
if conf.Database.DedupeUniqueFields {
|
||
if err := dedupeUniqueFields(db); err != nil {
|
||
return nil, errors.Wrap(err, "dedupe unique fields before migrate")
|
||
}
|
||
}
|
||
|
||
if err := db.AutoMigrate(AllTables...); err != nil {
|
||
return nil, errors.Wrap(err, "auto migrate")
|
||
}
|
||
|
||
// 功能开关初始化：为历史账号补齐 feature_controls 行（私聊沿用 users.allow_chat）。
|
||
// 幂等，已初始化的账号不会重复写入；失败不阻断启动，仅告警，避免依赖表抖动导致服务起不来。
|
||
if err := backfillFeatureControls(db); err != nil {
|
||
logrus.WithError(err).Warn("初始化功能开关配置时异常（可在功能管理页手动保存一次以生成配置）")
|
||
}
|
||
|
||
// 私聊开关纠偏：旧版 backfill 曾把未开启私聊的账号写为开启（GORM 零值陷阱），
|
||
// 启动期幂等复位；同样不阻断启动，失败仅告警。
|
||
if err := reconcileLegacyChatEnabled(db); err != nil {
|
||
logrus.WithError(err).Warn("私聊开关历史数据纠偏时异常（可在功能管理页手动调整私聊开关）")
|
||
}
|
||
|
||
// 私聊收集上限下线：私聊不再支持收集上限（UI 已隐藏，保存时强制清零），
|
||
// 启动期把历史残留的上限/计数一并清零，避免私聊「莫名关闭」。幂等，失败仅告警。
|
||
if err := normalizeChatQuota(db); err != nil {
|
||
logrus.WithError(err).Warn("私聊收集上限历史数据清理时异常（可在功能管理页保存一次私聊设置）")
|
||
}
|
||
|
||
// 防骚扰设置归一：历史账号 harassment_setting 为空串（建表时无默认值），与 none 语义等价
|
||
// 但取值不同，会导致按 == none 判断时漏判。幂等，失败仅告警。
|
||
if err := normalizeHarassmentSetting(db); err != nil {
|
||
logrus.WithError(err).Warn("防骚扰设置历史数据归一时异常（可在防骚扰设置页保存一次）")
|
||
}
|
||
|
||
// 启动期补充高频查询复合索引（幂等，适合既有库）。
|
||
// 仅 PostgreSQL 需要：这些索引针对服务端真实查询模式（按归属用户+私密+时间、
|
||
// 按 asker、按展示/投票状态+可见性+时间、按会话/发送者+时间）。
|
||
// 本项目实际使用 PostgreSQL（详见 code-review-2026-08-28.md P1-4）。
|
||
if typ == "postgres" {
|
||
if err := ensureFrequentlyUsedIndexes(db); err != nil {
|
||
// 索引创建失败不应阻断启动，仅告警，便于在运维侧排查。
|
||
logrus.WithError(err).Warn("创建高频查询复合索引时异常（可能部分索引未建立，不影响服务启动）")
|
||
}
|
||
}
|
||
|
||
Users = NewUsersStore(db)
|
||
Questions = NewQuestionsStore(db)
|
||
FollowUps = NewFollowUpsStore(db)
|
||
CensorLogs = NewCensorLogsStore(db)
|
||
UploadImages = NewUploadImagesStore(db)
|
||
Conversations = NewConversationStore(db)
|
||
Messages = NewMessageStore(db)
|
||
MessageImages = NewMessageImageStore(db)
|
||
Votes = NewVotesStore(db)
|
||
VoteOptions = NewVoteOptionsStore(db)
|
||
VoteFormFields = NewVoteFormFieldsStore(db)
|
||
VoteSubmissions = NewVoteSubmissionsStore(db)
|
||
FeatureControls = NewFeatureControlsStore(db)
|
||
Restrictions = NewRestrictionsStore(db)
|
||
|
||
if err := db.Use(otelgorm.NewPlugin(
|
||
otelgorm.WithDBName(conf.Database.Name),
|
||
)); err != nil {
|
||
return nil, errors.Wrap(err, "register otelgorm plugin")
|
||
}
|
||
|
||
if typ == "postgres" {
|
||
if err := fixPostgresSequences(db); err != nil {
|
||
return nil, errors.Wrap(err, "fix postgres sequences")
|
||
}
|
||
}
|
||
|
||
return db, nil
|
||
}
|
||
|
||
// dedupeUniqueFields 在 AutoMigrate 之前清理 users 表中重复的 email / domain：
|
||
// 每个重复值仅保留 id 最小（最早注册）的一行，其余行该字段置为 NULL。
|
||
// NULL 在 MySQL / PostgreSQL / SQLite 的唯一索引中互不冲突，因此置空即可让唯一索引成功建立。
|
||
// 幂等：重复数据被清理后再次启动时无任何操作。跨数据库兼容（仅使用标准 SQL，无目标表子查询）。
|
||
func dedupeUniqueFields(db *gorm.DB) error {
|
||
for _, field := range []string{"email", "domain"} {
|
||
var keepIDs []uint
|
||
var dupVals []string
|
||
|
||
rows, err := db.Table("users").
|
||
Select("MIN(id) AS keep_id, " + field + " AS val").
|
||
Where(field + " IS NOT NULL").
|
||
Group(field).
|
||
Having("COUNT(*) > 1").
|
||
Rows()
|
||
if err != nil {
|
||
return errors.Wrapf(err, "query duplicate %s", field)
|
||
}
|
||
for rows.Next() {
|
||
var keepID uint
|
||
var val string
|
||
if err := rows.Scan(&keepID, &val); err != nil {
|
||
rows.Close()
|
||
return errors.Wrapf(err, "scan duplicate %s", field)
|
||
}
|
||
keepIDs = append(keepIDs, keepID)
|
||
dupVals = append(dupVals, val)
|
||
}
|
||
if err := rows.Close(); err != nil {
|
||
return errors.Wrapf(err, "close rows for %s", field)
|
||
}
|
||
if rows.Err() != nil {
|
||
return errors.Wrapf(rows.Err(), "iterate duplicate %s", field)
|
||
}
|
||
if len(keepIDs) == 0 {
|
||
continue
|
||
}
|
||
|
||
res := db.Table("users").
|
||
Where(field+" IS NOT NULL").
|
||
Where(field+" IN ?", dupVals).
|
||
Where("id NOT IN ?", keepIDs).
|
||
Update(field, nil)
|
||
if res.Error != nil {
|
||
return errors.Wrapf(res.Error, "clean duplicate %s", field)
|
||
}
|
||
if res.RowsAffected > 0 {
|
||
logrus.WithField("field", field).
|
||
Warnf("迁移前置清理：%d 行重复的 %s 已置空（保留最早注册的账号）", res.RowsAffected, field)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// ensureFrequentlyUsedIndexes 在 AutoMigrate 之后创建高频查询所需的复合索引。
|
||
// 全部使用 CREATE INDEX IF NOT EXISTS，幂等、可在既有库上安全运行，重复执行无副作用。
|
||
// 索引设计依据服务端真实查询模式（详见 code-review-2026-08-28.md P1-4）：
|
||
// - questions: 按 user_id + is_private + created_at 过滤排序（提问箱列表）；
|
||
// 按 asker_user_id 查询（我发起的提问）。
|
||
// - users: 按 show_in_home + created_at（首页展示箱列表）。
|
||
// - votes: 按 status + visibility + created_at（探索页/提问箱页公开投票）。
|
||
// - messages: 按 conversation_id + created_at、sender_id + created_at（私信列表/会话消息）。
|
||
func ensureFrequentlyUsedIndexes(db *gorm.DB) error {
|
||
indexes := []string{
|
||
`CREATE INDEX IF NOT EXISTS idx_questions_asker_user_id ON questions(asker_user_id)`,
|
||
`CREATE INDEX IF NOT EXISTS idx_questions_user_private_created ON questions(user_id, is_private, created_at)`,
|
||
`CREATE INDEX IF NOT EXISTS idx_users_show_in_home_created ON users(show_in_home, created_at)`,
|
||
`CREATE INDEX IF NOT EXISTS idx_votes_status_visibility_created ON votes(status, visibility, created_at)`,
|
||
`CREATE INDEX IF NOT EXISTS idx_messages_sender_created ON messages(sender_id, created_at)`,
|
||
`CREATE INDEX IF NOT EXISTS idx_messages_conversation_created ON messages(conversation_id, created_at)`,
|
||
}
|
||
for _, ddl := range indexes {
|
||
if err := db.Exec(ddl).Error; err != nil {
|
||
return errors.Wrapf(err, "create index: %s", ddl)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func fixPostgresSequences(db *gorm.DB) error {
|
||
// 使用配置的 schema（未配置时回退 public），避免硬编码 nekobox。
|
||
schema := conf.Database.Schema
|
||
if schema == "" {
|
||
schema = "public"
|
||
}
|
||
|
||
tablesWithID := []string{"users", "questions", "follow_ups", "censor_logs", "upload_images", "conversations", "messages", "message_images", "votes", "vote_options", "vote_form_fields", "vote_submissions", "feature_controls"}
|
||
for _, table := range tablesWithID {
|
||
seqName := table + "_id_seq"
|
||
resetQuery := fmt.Sprintf(`
|
||
CREATE SEQUENCE IF NOT EXISTS %s.%s;
|
||
ALTER TABLE %s.%s
|
||
ALTER COLUMN id SET DEFAULT nextval('%s.%s');
|
||
SELECT setval('%s.%s',
|
||
COALESCE((SELECT MAX(id) FROM %s.%s), 1)::bigint,
|
||
true);
|
||
`, schema, seqName, schema, table, schema, seqName, schema, seqName, schema, table)
|
||
if err := db.Exec(resetQuery).Error; err != nil {
|
||
return fmt.Errorf("修复表 %s 序列时出错: %w", table, err)
|
||
}
|
||
}
|
||
return nil
|
||
}
|