tamabox/internal/db/db.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{},
	// 蜜罐命中记录（审计证据，只追加不软删）
	&HoneypotHit{},
}

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("创建高频查询复合索引时异常（可能部分索引未建立，不影响服务启动）")
		}
	}

	// 蜜罐表上无查询使用的索引（provider）：GORM 去掉 tag 不会自动删既有索引，
	// 这里显式清理一次。失败仅告警。
	if err := dropUnusedHoneypotIndexes(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)
	Reports = NewReportsStore(db)
	Honeypots = NewHoneypotStore(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
}

// dropUnusedHoneypotIndexes 删除蜜罐表上没有任何查询使用的索引。
//
// 背景：provider 字段原先带 gorm:"index"，但全代码库没有一处按 provider
// 过滤或排序（TopSources 用 max(provider) 随 group by 顺带取，不需要独立索引）。
// 实测该索引占 10 万行约 648 kB，对一张「只增不减」的审计表是纯粹的空间浪费
// 与写放大。去掉 GORM tag 只影响新建库，既有库上的旧索引要显式 DROP。
//
// 幂等：索引不存在时不报错。失败仅告警，不阻断启动。
func dropUnusedHoneypotIndexes(db *gorm.DB) error {
	// GORM 生成的索引名规则：idx_<表名>_<列名>。
	if db.Migrator().HasIndex(&HoneypotHit{}, "idx_honeypot_hits_provider") {
		if err := db.Migrator().DropIndex(&HoneypotHit{}, "idx_honeypot_hits_provider"); err != nil {
			return errors.Wrap(err, "drop unused honeypot index")
		}
		logrus.Info("已删除蜜罐表上无查询使用的 provider 索引")
	}
	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
}