sundynix-agentix/sundynix-gateway/internal/store/pgsql.go

// 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{}); err != nil {
		log.Printf("[store] postgres AutoMigrate 失败，降级运行: %v", err)
		return &Postgres{}
	}
	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, id, graph string) error {
	if p.db == nil {
		return nil
	}
	return p.db.WithContext(ctx).Create(&Task{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
	}
	return p.db.WithContext(ctx).Clauses(clause.OnConflict{
		Columns: []clause.Column{{Name: "task_id"}},
		DoUpdates: clause.AssignmentColumns([]string{
			"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）。
func (p *Postgres) RecentTasks(ctx context.Context, n int) []Task {
	if p.db == nil {
		return nil
	}
	var out []Task
	p.db.WithContext(ctx).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 条运行（含评测分级，供「运行历史」列表）。
func (p *Postgres) RecentRuns(ctx context.Context, n int) []RunRow {
	if p.db == nil {
		return nil
	}
	var out []RunRow
	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").
		Order("t.created_at desc").Limit(n).Scan(&out)
	return out
}

// Close 释放底层连接。
func (p *Postgres) Close() {
	if p.db == nil {
		return
	}
	if sqlDB, err := p.db.DB(); err == nil {
		_ = sqlDB.Close()
	}
}