tamabox/internal/db/db.go

package db

import (
	"context"
	"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"
)

// handle 保存 Init 建立的连接，供「不需要具体 store」的场景使用（如就绪探针）。
// 与下面那些业务 store 的包级变量区分开：那些是接口，这个是连接本身。
var handle *gorm.DB

// Ping 检查数据库连接是否可用。供 /readyz 就绪探针调用。
//
// 为什么需要它：原来的「探活」只能 curl /，而首页只渲染模板、不碰数据库 ——
// DB 挂了首页照样 200，反过来 Redis 挂了首页又可能 500。于是探针既会漏报也会误报。
func Ping(ctx context.Context) error {
	if handle == nil {
		return errors.New("database not initialized")
	}
	sqlDB, err := handle.DB()
	if err != nil {
		return errors.Wrap(err, "get sql.DB")
	}
	return sqlDB.PingContext(ctx)
}

var AllTables = []interface{}{
	&User{}, &Question{}, &FollowUp{}, &CensorLog{}, &UploadImage{}, &UploadImageQuestion{},
	&Conversation{}, &Message{}, &MessageImage{}, &UserBlock{},
	// 匿名投票
	&Vote{}, &VoteOption{}, &VoteFormField{}, &VoteSubmission{},
	// 功能开关（提问 / 私聊 / 投票 的开关、定时、定量收集）
	&FeatureControl{},
	// 站方处置与举报（无管理后台，处置经 CLI 写入、运行时读表生效）
	&UserRestriction{}, &ContentReport{},
	// 蜜罐命中记录（审计证据，只追加不软删）
	&HoneypotHit{},
	// 来源 IP 的地理定位缓存（安全汇总报告用，纯 KV 缓存）
	&IPGeoRecord{},
	// 周期安全汇总的报告台账（period_key 唯一，保证重启不会重复发信）
	&SecurityReportLog{},
	// 自动封禁记录（内存态封禁的落库账本，用于事后解释与取证）
	&BanRecord{},
	// 箱主数据面板的访问流水（每 访客×天×箱 一行，Unscoped 语义，见 BoxView 注释）
	&BoxView{},
}

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")
	}
	handle = db

	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")
		}
	}

	// AutoMigrate 只负责声明式的「加表 / 加列 / 加索引」——这些变更幂等、无数据丢失风险。
	if err := db.AutoMigrate(AllTables...); err != nil {
		return nil, errors.Wrap(err, "auto migrate")
	}

	// 版本化迁移：AutoMigrate 做不了的破坏性变更（改列名/删列/改约束）与一次性数据回填
	// 都收编进这里，靠 schema_migrations 账本保证「只跑一次、顺序固定、可审计」。
	// 迁移失败会阻断启动 —— 半途而废的迁移比静默跳过更危险（见 runMigrations 注释）。
	if err := runMigrations(db); err != nil {
		return nil, errors.Wrap(err, "run versioned migrations")
	}

	// 启动期补充高频查询复合索引（幂等，适合既有库）。
	// 仅 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)
	BanRecords = NewBanRecordsStore(db)
	Analytics = NewSiteAnalyticsStore(db)
	BoxViews = NewBoxViewsStore(db)
	BoxAnalytics = NewBoxAnalyticsStore(db)

	// 旧版本把消息/密码明文存在 honeypot_hits 上；结构体已删掉这些字段，
	// 但 AutoMigrate 不会删除既有列 —— 老库升级上来后历史明文仍在。
	// 这里只告警不自动删：DROP COLUMN 不可逆，必须由运维确认后执行
	// `TamaBox honeypot purge-creds`。
	if cols, err := Honeypots.LegacyCredentialColumns(context.Background()); err == nil && len(cols) > 0 {
		logrus.WithField("columns", cols).
			Warn("honeypot_hits 上仍残留旧版明文凭据列，请执行 TamaBox honeypot purge-creds 清除")
	}

	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)`,
		// 全站数据页（/stats）：按 created_at 做范围扫描后按天/小时分组。
		// 没有这个索引时，趋势与热力图会退化成全表扫描 + 排序；加上之后是范围扫描 + 聚合。
		`CREATE INDEX IF NOT EXISTS idx_questions_created_at ON questions(created_at)`,
		`CREATE INDEX IF NOT EXISTS idx_questions_user_private_created ON questions(user_id, is_private, created_at)`,
		// 「下一条未读」：user_id + is_read 过滤后按 created_at 取最早一条。
		// 上面那条中间隔着 is_private，帮不上这个排序。
		`CREATE INDEX IF NOT EXISTS idx_questions_user_unread_created ON questions(user_id, is_read, 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
}