// Package store 封装 MainDB(PgSQL) 与 CacheDB(Redis) 的访问。 package store import ( "context" "errors" "log" "os" "strconv" "time" "gorm.io/driver/postgres" "gorm.io/gorm" "gorm.io/gorm/clause" "gorm.io/gorm/schema" "github.com/sundynix/sundynix-shared/contract" ) // envInt 读正整数环境变量,缺省回退 def。 func envInt(key string, def int) int { if v := os.Getenv(key); v != "" { if n, err := strconv.Atoi(v); err == nil && n > 0 { return n } } return def } // tunePool 给连接池设上限:高并发下不至于无限开连接打爆 PG(max_connections 默认 100)。 // 各服务默认 25,可经 DB_MAX_OPEN_CONNS / DB_MAX_IDLE_CONNS 调整。 func tunePool(db *gorm.DB) { if sqlDB, err := db.DB(); err == nil { sqlDB.SetMaxOpenConns(envInt("DB_MAX_OPEN_CONNS", 25)) sqlDB.SetMaxIdleConns(envInt("DB_MAX_IDLE_CONNS", 5)) sqlDB.SetConnMaxLifetime(time.Hour) } } // errStoreDisabled 表示 Postgres 处于降级(未连接)模式,写操作无法进行。 var errStoreDisabled = errors.New("postgres store disabled") // Postgres 持有 MainDB 连接(Users / Billing / DSL)。 // db 为 nil 表示降级模式(连接失败时仍允许网关启动)。 type Postgres struct { db *gorm.DB } // OpenPostgres 用 GORM 连接 MainDB 并自动迁移表结构。 // 表名统一 sundynix_ 前缀 + 单数(User→sundynix_user, Task→sundynix_task)。 // 连接失败不 fatal:返回降级实例,网关仍可启动(无 Docker 跑 demo 时即此路径)。 func OpenPostgres(dsn string) *Postgres { db, err := gorm.Open(postgres.New(postgres.Config{DSN: dsn}), &gorm.Config{ NamingStrategy: schema.NamingStrategy{ TablePrefix: "sundynix_", // 所有表加前缀 SingularTable: true, // 单数表名 }, }) if err != nil { log.Printf("[store] postgres 不可用,降级运行(不持久化): %v", err) return &Postgres{} } tunePool(db) // 连接池上限,防高并发打爆 PG // 一次性迁移:旧表用整型自增 id,与新雪花字符串 id 不兼容(AutoMigrate 不改主键类型)。 // 备份模型密钥(唯一不可再生的数据) → 重建全部表 → 回灌模型。其余为可重建的测试数据。 migrateLegacyIntIDs(db) migrateDocLinkToID(db) if err := db.AutoMigrate(&User{}, &Task{}, &Eval{}, &LLMModel{}, &KB{}, &Doc{}, &Agent{}, &DocLink{}, &Pricing{}, &Prompt{}, &AuditLog{}, &GuardrailEvent{}, &Tenant{}, &TenantMember{}, &Space{}, &SpaceMember{}, &UsageEvent{}, &CreditLedger{}, &UsageRollup{}, &Setting{}); err != nil { log.Printf("[store] postgres AutoMigrate 失败,降级运行: %v", err) return &Postgres{} } registerTenantScope(db) // 多租户:受租户模型的查询/创建自动按上下文注入 tenant_id(统一强制隔离) log.Println("[store] postgres connected & migrated (雪花 id + 软删 规约)") return &Postgres{db: db} } // migrateLegacyIntIDs 检测到旧整型 id 表则备份模型密钥、删旧表(AutoMigrate 随后按新规约重建)。 func migrateLegacyIntIDs(db *gorm.DB) { var dt string db.Raw(`SELECT data_type FROM information_schema.columns WHERE table_name='sundynix_model' AND column_name='id'`).Scan(&dt) if dt != "bigint" && dt != "integer" { return // 全新库或已是新规约 } log.Println("[store] 检测到旧整型 id 表,执行雪花 id 迁移(保模型密钥,重置其它测试表)") var saved []map[string]any db.Table("sundynix_model").Find(&saved) for _, t := range []string{"sundynix_doc_link", "sundynix_doc", "sundynix_agent", "sundynix_kb", "sundynix_model", "sundynix_task", "sundynix_user"} { db.Exec("DROP TABLE IF EXISTS " + t + " CASCADE") } _ = db.AutoMigrate(&LLMModel{}) // 先建模型表以回灌 for _, r := range saved { s := func(k string) string { v, _ := r[k].(string); return v } b, _ := r["active"].(bool) _ = db.Create(&LLMModel{ Kind: s("kind"), Provider: s("provider"), BaseURL: s("base_url"), APIKey: s("api_key"), Model: s("model"), Active: b, }).Error } log.Printf("[store] 已回灌 %d 条模型配置(新雪花 id)", len(saved)) } // migrateDocLinkToID 把旧的按名双链表(from_name/to_name)迁到按 Doc.ID 关联的新表。 // 旧表无 from_id 列即判定为旧 schema:直接删表,由 AutoMigrate 重建;链接随文档再入库/编辑重建。 func migrateDocLinkToID(db *gorm.DB) { if !db.Migrator().HasTable("sundynix_doc_link") { return } if db.Migrator().HasColumn(&DocLink{}, "from_id") { return // 已是按 ID 关联的新 schema } log.Println("[store] 双链表升级为按文件 ID 关联,重建 sundynix_doc_link(链接随文档再入库重建)") db.Exec("DROP TABLE IF EXISTS sundynix_doc_link CASCADE") } // Enabled 报告是否处于真实持久化模式。 func (p *Postgres) Enabled() bool { return p.db != nil } // SaveTask 持久化一次任务提交(best-effort:降级模式下静默跳过)。 func (p *Postgres) SaveTask(ctx context.Context, owner, id, graph string) error { if p.db == nil { return nil } // TenantID 由 tenant 插件按请求 ctx 自动填;Owner 显式记录提交者(供个人工作台过滤)。 return p.db.WithContext(ctx).Create(&Task{Owner: owner, TaskID: id, Graph: graph, Status: contract.TaskSubmitted}).Error } // UpdateTaskStatus 流转任务状态(running/done/failed/timeout),由 dispatcher 经 NATS 回写驱动。 func (p *Postgres) UpdateTaskStatus(ctx context.Context, id, status, detail string) error { if p.db == nil { return nil } return p.db.WithContext(ctx).Model(&Task{}). Where("task_id = ?", id). Updates(map[string]any{"status": status, "detail": detail}).Error } // GetTaskStatus 取一条任务的当前状态(供 UI 轮询;不存在返回空串)。 func (p *Postgres) GetTaskStatus(ctx context.Context, id string) (status, detail string) { if p.db == nil { return "", "" } var t Task if err := p.db.WithContext(ctx).Select("status", "detail").Where("task_id = ?", id).First(&t).Error; err != nil { return "", "" } return t.Status, t.Detail } // SaveTaskOutput 收尾时落库最终模型输出(供历史复盘,best-effort)。 func (p *Postgres) SaveTaskOutput(ctx context.Context, id, output string) error { if p.db == nil { return nil } return p.db.WithContext(ctx).Model(&Task{}).Where("task_id = ?", id).Update("output", output).Error } // SaveTaskTrace 收尾时落库执行轨迹 JSON(供历史复盘,best-effort)。 func (p *Postgres) SaveTaskTrace(ctx context.Context, id, traceJSON string) error { if p.db == nil { return nil } return p.db.WithContext(ctx).Model(&Task{}).Where("task_id = ?", id).Update("trace", traceJSON).Error } // GetRunDetail 取一条任务持久化的输出 + 轨迹(历史复盘读库,不依赖 Redis TTL)。 func (p *Postgres) GetRunDetail(ctx context.Context, id string) (output, trace string) { if p.db == nil { return "", "" } var t Task if err := p.db.WithContext(ctx).Select("output", "trace").Where("task_id = ?", id).First(&t).Error; err != nil { return "", "" } return t.Output, t.Trace } // SaveEval 落库一条评测结果(按 task_id upsert:重评覆盖)。 func (p *Postgres) SaveEval(ctx context.Context, e *Eval) error { if p.db == nil { return nil } // 评测由 dispatcher 经 NATS 回写(background ctx,无请求租户)→ 从对应 task 复制 owner+tenant。 var t Task if err := p.db.WithContext(ctx).Select("owner", "tenant_id").Where("task_id = ?", e.TaskID).First(&t).Error; err == nil { e.Owner, e.TenantID = t.Owner, t.TenantID } return p.db.WithContext(ctx).Clauses(clause.OnConflict{ Columns: []clause.Column{{Name: "task_id"}}, DoUpdates: clause.AssignmentColumns([]string{ "owner", "tenant_id", "overall", "rule", "llm", "faithful", "level", "flags", "reason", "sources", "corrected", "updated_at", }), }).Create(e).Error } // GetEval 取一条任务的评测结果(不存在返回 nil)。 func (p *Postgres) GetEval(ctx context.Context, taskID string) *Eval { if p.db == nil { return nil } var e Eval if err := p.db.WithContext(ctx).Where("task_id = ?", taskID).First(&e).Error; err != nil { return nil } return &e } // CountTasks 返回已提交任务数(降级模式返回 0)。 func (p *Postgres) CountTasks(ctx context.Context) (int64, error) { if p.db == nil { return 0, nil } var n int64 err := p.db.WithContext(ctx).Model(&Task{}).Count(&n).Error return n, err } // DayCount 是「某天 / 某状态 → 计数」的一行(工作台趋势/分布用)。 type DayCount struct { Key string `json:"key"` Count int64 `json:"count"` } // Overview 是工作台概览的聚合数据。任务/评测为实例级(Task 无 owner,单租户部署即全量), // 知识库为 owner 级。降级模式(db==nil)返回零值。 type Overview struct { TasksToday int64 `json:"tasks_today"` TasksTotal int64 `json:"tasks_total"` StatusCount []DayCount `json:"status_count"` // 近 7 天各终态分布 TaskTrend []DayCount `json:"task_trend"` // 近 7 天每日任务数(MM-DD) EvalAvg float64 `json:"eval_avg"` // 综合分均值 FaithfulAvg float64 `json:"faithful_avg"` // 忠实度均值(仅有来源的) EvalCount int64 `json:"eval_count"` KBDocs int64 `json:"kb_docs"` // owner 文档数 KBCount int64 `json:"kb_count"` // owner 知识库数 } // StatsOverview 聚合工作台概览(几条轻量查询)。owner 用于知识库口径。 func (p *Postgres) StatsOverview(ctx context.Context, owner string) *Overview { o := &Overview{StatusCount: []DayCount{}, TaskTrend: []DayCount{}} if p.db == nil { return o } db := p.db.WithContext(ctx) startOfDay := time.Now().Truncate(24 * time.Hour) db.Model(&Task{}).Count(&o.TasksTotal) db.Model(&Task{}).Where("created_at >= ?", startOfDay).Count(&o.TasksToday) // 近 7 天每日任务数(按日期分组,缺的天补 0 在前端/此处处理)。 db.Model(&Task{}). Select("to_char(created_at, 'MM-DD') as key, count(*) as count"). Where("created_at >= ?", time.Now().AddDate(0, 0, -6).Truncate(24*time.Hour)). Group("key").Order("key").Scan(&o.TaskTrend) // 近 7 天终态分布。 db.Model(&Task{}). Select("status as key, count(*) as count"). Where("created_at >= ?", time.Now().AddDate(0, 0, -6)). Group("status").Scan(&o.StatusCount) // 评测均值(综合 + 忠实度仅算有来源的)。 var ev struct { Avg float64 Faithful float64 N int64 } db.Model(&Eval{}).Select("coalesce(avg(overall),0) as avg, coalesce(avg(nullif(faithful,0)),0) as faithful, count(*) as n").Scan(&ev) o.EvalAvg, o.FaithfulAvg, o.EvalCount = ev.Avg, ev.Faithful, ev.N if owner != "" { db.Model(&Doc{}).Where("owner = ?", owner).Count(&o.KBDocs) db.Model(&KB{}).Where("owner = ?", owner).Count(&o.KBCount) } return o } // RecentTasks 返回最近 n 条任务(工作台「近期运行」feed)。 // RecentTasks 返回某用户最近 n 条任务(个人工作台「近期运行」feed)。 // owner 过滤"我的运行";tenant 由插件自动叠加(双保险:跨用户/跨租户都隔离)。 func (p *Postgres) RecentTasks(ctx context.Context, owner string, n int) []Task { if p.db == nil { return nil } var out []Task p.db.WithContext(ctx).Where("owner = ?", owner).Order("created_at desc").Limit(n).Find(&out) return out } // RunRow 是「运行历史」一行:任务 + 其评测(LEFT JOIN,未评则 level 空)。 type RunRow struct { TaskID string `json:"task_id"` Status string `json:"status"` Detail string `json:"detail"` At time.Time `json:"at"` EvalLevel string `json:"eval_level"` EvalOverall float64 `json:"eval_overall"` } // RecentRuns 返回某用户最近 n 条运行(含评测分级,供「运行历史」列表)。 // 注:raw Table 查询绕过 gorm 模型回调 → 租户插件不生效,故此处**手动**按 owner(+ctx 租户) 过滤。 func (p *Postgres) RecentRuns(ctx context.Context, owner string, n int) []RunRow { if p.db == nil { return nil } var out []RunRow q := p.db.WithContext(ctx).Table("sundynix_task as t"). Select("t.task_id, t.status, t.detail, t.created_at as at, " + "coalesce(e.level,'') as eval_level, coalesce(e.overall,0) as eval_overall"). Joins("left join sundynix_eval e on e.task_id = t.task_id"). Where("t.deleted_at is null AND t.owner = ?", owner) if tid := tenantFromCtx(ctx); tid != "" && !isSystemCtx(ctx) { q = q.Where("t.tenant_id = ?", tid) } q.Order("t.created_at desc").Limit(n).Scan(&out) return out } // SystemCounts 返回系统级计数(管理端 overview 用:全平台口径,非 owner 隔离)。 // 软删行由 GORM DeletedAt 作用域自动排除。降级模式返回零值。 func (p *Postgres) SystemCounts(ctx context.Context) (users, kbs, docs int64) { if p.db == nil { return 0, 0, 0 } db := p.db.WithContext(ctx) db.Model(&User{}).Count(&users) db.Model(&KB{}).Count(&kbs) db.Model(&Doc{}).Count(&docs) return } // Close 释放底层连接。 func (p *Postgres) Close() { if p.db == nil { return } if sqlDB, err := p.db.DB(); err == nil { _ = sqlDB.Close() } }