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/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{}, &LLMModel{}, &KB{}, &Doc{}, &Agent{}, &DocLink{}, &Pricing{}); 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
}

// 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
}

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