302 lines
12 KiB
Go
302 lines
12 KiB
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{}, &R2DeleteTask{},
|
||
// 匿名投票
|
||
&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{
|
||
Logger: newSafeLogger(),
|
||
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)
|
||
R2DeleteTasks = NewR2DeleteTaskStore(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),
|
||
otelgorm.WithoutQueryVariables(),
|
||
)); 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
|
||
}
|