GORM 企业级实战：大规模 Go 应用的生产模式与最佳实践

package database

import (
    "time"

    "gorm.io/driver/postgres"
    "gorm.io/gorm"
    "gorm.io/plugin/dbresolver"
)

func NewProductionDB() (*gorm.DB, error) {
    db, err := gorm.Open(postgres.Open("host=primary dbname=app"), &gorm.Config{
        // 生产环境禁用默认事务包装，性能提升约 30%
        SkipDefaultTransaction: true,
        // 全局预编译语句缓存
        PrepareStmt: true,
    })
    if err != nil {
        return nil, err
    }

    // 注册读写分离 + 按模型路由
    err = db.Use(dbresolver.Register(dbresolver.Config{
        Sources:  []gorm.Dialector{postgres.Open("host=primary dbname=app")},
        Replicas: []gorm.Dialector{
            postgres.Open("host=replica1 dbname=app"),
            postgres.Open("host=replica2 dbname=app"),
        },
        Policy:             dbresolver.RandomPolicy{},
        TraceResolverMode:  true, // 日志中显示使用的是 source 还是 replica
    }).Register(dbresolver.Config{
        // 特定模型走独立数据库
        Sources:  []gorm.Dialector{postgres.Open("host=analytics-primary dbname=analytics")},
        Replicas: []gorm.Dialector{postgres.Open("host=analytics-replica dbname=analytics")},
    }, "analytics_events", "analytics_reports").
        // 连接池参数
        SetConnMaxIdleTime(10 * time.Minute).
        SetConnMaxLifetime(1 * time.Hour).
        SetMaxIdleConns(50).
        SetMaxOpenConns(200))
    if err != nil {
        return nil, err
    }

    return db, nil
}

// 手动控制读写路由
func ExampleManualRouting(db *gorm.DB) {
    var user User

    // 强制走主库读（读取刚写入的数据，避免主从延迟）
    db.Clauses(dbresolver.Write).First(&user, 1)

    // 强制走从库读
    db.Clauses(dbresolver.Read).Find(&users)

    // 指定命名数据源
    db.Clauses(dbresolver.Use("analytics")).Find(&events)

    // 事务控制：基于从库的只读事务
    tx := db.Clauses(dbresolver.Read).Begin()
    defer tx.Rollback()
    tx.Find(&reports)
    tx.Commit()
}

package database

import (
    "gorm.io/gorm"
    "gorm.io/sharding"
)

func SetupSharding(db *gorm.DB) error {
    return db.Use(sharding.Register(sharding.Config{
        ShardingKey:         "user_id",
        NumberOfShards:      64,
        PrimaryKeyGenerator: sharding.PKSnowflake,
        // 自定义分片算法（可选）
        // ShardingAlgorithm: func(columnValue interface{}) (suffix string, err error) {
        //     uid := columnValue.(int64)
        //     return fmt.Sprintf("_%04d", uid % 64), nil
        // },
    }, "orders", "order_items", "notifications"))
}

// 使用方式 -- 完全透明
func CreateOrder(db *gorm.DB, order *Order) error {
    // 自动路由到 orders_{hash(user_id) % 64}
    return db.Create(order).Error
    // 生成 SQL: INSERT INTO orders_2 ...
}

func GetUserOrders(db *gorm.DB, userID int64) ([]Order, error) {
    var orders []Order
    // 自动路由到正确的分片表
    err := db.Where("user_id = ?", userID).Find(&orders).Error
    return orders, err
    // 生成 SQL: SELECT * FROM orders_2 WHERE user_id = ?
}

// cmd/gen/main.go -- 代码生成入口
package main

import (
    "gorm.io/gen"
    "gorm.io/driver/postgres"
    "gorm.io/gorm"
    "myapp/internal/model"
)

// 自定义查询接口：支持动态 SQL 模板
type OrderQuerier interface {
    // SELECT * FROM @@table WHERE user_id = @userID
    //   {{if status != ""}} AND status = @status {{end}}
    //   ORDER BY created_at DESC
    //   LIMIT @limit OFFSET @offset
    FindByUserWithStatus(userID int64, status string, limit, offset int) ([]*gen.T, error)

    // UPDATE @@table SET status = @status WHERE id = @id AND version = @version
    UpdateStatusWithVersion(id int64, status string, version int) (gen.RowsAffected, error)
}

func main() {
    db, _ := gorm.Open(postgres.Open("dsn"))

    g := gen.NewGenerator(gen.Config{
        OutPath:      "../internal/query",
        Mode:         gen.WithDefaultQuery | gen.WithQueryInterface,
        ModelPkgPath: "../internal/model",
    })

    g.UseDB(db)

    // 从数据库 Schema 生成模型
    g.ApplyBasic(
        g.GenerateModel("users"),
        g.GenerateModel("orders"),
        g.GenerateModel("products"),
    )

    // 绑定自定义查询接口
    g.ApplyInterface(func(OrderQuerier) {}, model.Order{})

    g.Execute()
}

package service

import "myapp/internal/query"

func (s *OrderService) GetUserOrders(ctx context.Context, userID int64, status string) ([]*model.Order, error) {
    o := query.Order

    // 类型安全：字段名拼写错误会导致编译失败
    orders, err := o.WithContext(ctx).
        Where(o.UserID.Eq(userID)).
        Where(o.Status.Eq(status)).
        Order(o.CreatedAt.Desc()).
        Limit(20).
        Find()

    return orders, err
}

func (s *OrderService) GetOrderStats(ctx context.Context) {
    o := query.Order

    // 动态 SQL 模板查询
    orders, err := o.WithContext(ctx).FindByUserWithStatus(123, "paid", 10, 0)

    // 带版本号的更新（乐观锁场景）
    rowsAffected, err := o.WithContext(ctx).UpdateStatusWithVersion(1, "shipped", 3)
}

package plugin

import (
    "context"
    "encoding/json"
    "time"

    "gorm.io/gorm"
)

// AuditLog 审计日志模型
type AuditLog struct {
    ID        uint      `gorm:"primaryKey"`
    TableName string    `gorm:"index;size:64"`
    Operation string    `gorm:"size:16"` // CREATE, UPDATE, DELETE
    RecordID  string    `gorm:"index;size:64"`
    OldValue  string    `gorm:"type:jsonb"`
    NewValue  string    `gorm:"type:jsonb"`
    UserID    string    `gorm:"index;size:64"`
    IP        string    `gorm:"size:45"`
    CreatedAt time.Time `gorm:"index"`
}

type AuditPlugin struct{}

func (p *AuditPlugin) Name() string { return "audit" }

func (p *AuditPlugin) Initialize(db *gorm.DB) error {
    // 注册 Create 后回调
    db.Callback().Create().After("gorm:create").Register("audit:create", func(db *gorm.DB) {
        if db.Error != nil || db.Statement.Schema == nil {
            return
        }
        p.log(db, "CREATE", nil)
    })

    // 注册 Update 前后回调
    db.Callback().Update().Before("gorm:update").Register("audit:before_update", func(db *gorm.DB) {
        // 快照更新前的值
        if db.Statement.Schema == nil {
            return
        }
        var oldRecord map[string]interface{}
        db.Session(&gorm.Session{NewDB: true}).
            Table(db.Statement.Table).
            Where(db.Statement.Clauses["WHERE"]).
            First(&oldRecord)
        db.Set("audit:old_value", oldRecord)
    })

    db.Callback().Update().After("gorm:update").Register("audit:after_update", func(db *gorm.DB) {
        if db.Error != nil {
            return
        }
        old, _ := db.Get("audit:old_value")
        p.log(db, "UPDATE", old)
    })

    // 注册 Delete 回调
    db.Callback().Delete().After("gorm:delete").Register("audit:delete", func(db *gorm.DB) {
        if db.Error != nil {
            return
        }
        p.log(db, "DELETE", nil)
    })

    return nil
}

func (p *AuditPlugin) log(db *gorm.DB, op string, oldVal interface{}) {
    ctx := db.Statement.Context
    userID, _ := ctx.Value("user_id").(string)
    ip, _ := ctx.Value("client_ip").(string)

    oldJSON, _ := json.Marshal(oldVal)
    newJSON, _ := json.Marshal(db.Statement.Dest)

    audit := AuditLog{
        TableName: db.Statement.Table,
        Operation: op,
        OldValue:  string(oldJSON),
        NewValue:  string(newJSON),
        UserID:    userID,
        IP:        ip,
        CreatedAt: time.Now(),
    }

    // 使用独立会话写入，避免影响主事务
    db.Session(&gorm.Session{NewDB: true}).Create(&audit)
}

// 注册插件
// db.Use(&AuditPlugin{})

package model

import (
    "gorm.io/plugin/optimisticlock"
)

type Product struct {
    ID      uint   `gorm:"primaryKey"`
    Name    string
    Stock   int
    Price   float64
    Version optimisticlock.Version // 自动管理版本号
}

// 乐观锁更新：带重试机制
func DeductStock(db *gorm.DB, productID uint, quantity int) error {
    const maxRetries = 3

    for i := 0; i < maxRetries; i++ {
        var product Product
        if err := db.First(&product, productID).Error; err != nil {
            return err
        }

        if product.Stock < quantity {
            return errors.New("insufficient stock")
        }

        product.Stock -= quantity
        result := db.Save(&product)

        if result.Error != nil {
            return result.Error
        }
        if result.RowsAffected == 0 {
            // 版本冲突，重试
            continue
        }
        return nil // 成功
    }

    return errors.New("optimistic lock: max retries exceeded")
}

// 悲观锁：SELECT ... FOR UPDATE
func TransferBalance(db *gorm.DB, fromID, toID uint, amount float64) error {
    return db.Transaction(func(tx *gorm.DB) error {
        var from, to Account

        // FOR UPDATE 锁定两行（按 ID 排序避免死锁）
        ids := sortIDs(fromID, toID)
        if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
            Where("id IN ?", ids).
            Order("id").
            Find(&[]Account{&from, &to}).Error; err != nil {
            return err
        }

        if from.Balance < amount {
            return errors.New("insufficient balance")
        }

        from.Balance -= amount
        to.Balance += amount

        if err := tx.Save(&from).Error; err != nil {
            return err
        }
        return tx.Save(&to).Error
    })
}

// FOR SHARE -- 共享锁，允许并发读但阻塞写
func GetAccountForRead(db *gorm.DB, id uint) (*Account, error) {
    var account Account
    err := db.Clauses(clause.Locking{
        Strength: "SHARE",
        Options:  "NOWAIT", // 获取不到锁立即返回错误，而非等待
    }).First(&account, id).Error
    return &account, err
}

package saga

import (
    "context"
    "fmt"
)

// Step 定义 Saga 的一个步骤
type Step struct {
    Name       string
    Execute    func(ctx context.Context) error
    Compensate func(ctx context.Context) error
}

// Saga 编排器
type Saga struct {
    steps     []Step
    completed []int // 已完成步骤的索引
}

func New() *Saga {
    return &Saga{}
}

func (s *Saga) AddStep(step Step) *Saga {
    s.steps = append(s.steps, step)
    return s
}

func (s *Saga) Execute(ctx context.Context) error {
    for i, step := range s.steps {
        if err := step.Execute(ctx); err != nil {
            // 执行补偿：逆序回滚已完成的步骤
            compensateErr := s.compensate(ctx)
            if compensateErr != nil {
                return fmt.Errorf("step %q failed: %w; compensation also failed: %v",
                    step.Name, err, compensateErr)
            }
            return fmt.Errorf("step %q failed (compensated): %w", step.Name, err)
        }
        s.completed = append(s.completed, i)
    }
    return nil
}

func (s *Saga) compensate(ctx context.Context) error {
    // 逆序补偿
    for i := len(s.completed) - 1; i >= 0; i-- {
        idx := s.completed[i]
        if err := s.steps[idx].Compensate(ctx); err != nil {
            return fmt.Errorf("compensation for step %q failed: %w",
                s.steps[idx].Name, err)
        }
    }
    return nil
}

// 业务使用示例：创建订单 Saga
func CreateOrderSaga(
    orderSvc *OrderService,
    paymentSvc *PaymentService,
    inventorySvc *InventoryService,
    req CreateOrderRequest,
) error {
    ctx := context.Background()

    s := New().
        AddStep(Step{
            Name: "create_order",
            Execute: func(ctx context.Context) error {
                return orderSvc.Create(ctx, req)
            },
            Compensate: func(ctx context.Context) error {
                return orderSvc.Cancel(ctx, req.OrderID)
            },
        }).
        AddStep(Step{
            Name: "reserve_inventory",
            Execute: func(ctx context.Context) error {
                return inventorySvc.Reserve(ctx, req.Items)
            },
            Compensate: func(ctx context.Context) error {
                return inventorySvc.Release(ctx, req.Items)
            },
        }).
        AddStep(Step{
            Name: "process_payment",
            Execute: func(ctx context.Context) error {
                return paymentSvc.Charge(ctx, req.PaymentInfo)
            },
            Compensate: func(ctx context.Context) error {
                return paymentSvc.Refund(ctx, req.PaymentInfo)
            },
        })

    return s.Execute(ctx)
}

package uow

import "gorm.io/gorm"

// UnitOfWork 封装事务边界
type UnitOfWork struct {
    db *gorm.DB
}

func New(db *gorm.DB) *UnitOfWork {
    return &UnitOfWork{db: db}
}

// Execute 在事务中执行工作单元
func (u *UnitOfWork) Execute(fn func(tx *gorm.DB) error) error {
    return u.db.Transaction(func(tx *gorm.DB) error {
        return fn(tx)
    })
}

// ExecuteWithSavepoint 支持嵌套事务（Savepoint）
func (u *UnitOfWork) ExecuteWithSavepoint(fn func(tx *gorm.DB) error) error {
    return u.db.Transaction(func(tx *gorm.DB) error {
        // 外层操作
        if err := fn(tx); err != nil {
            return err
        }

        // 嵌套事务：可独立回滚而不影响外层
        return tx.Transaction(func(nested *gorm.DB) error {
            // 嵌套操作（使用 SAVEPOINT）
            return nil
        })
    })
}

// 业务使用
func (s *OrderService) PlaceOrder(ctx context.Context, req PlaceOrderReq) error {
    uow := uow.New(s.db)
    return uow.Execute(func(tx *gorm.DB) error {
        order := &Order{UserID: req.UserID, Total: req.Total}
        if err := tx.Create(order).Error; err != nil {
            return err
        }

        for _, item := range req.Items {
            item.OrderID = order.ID
            if err := tx.Create(&item).Error; err != nil {
                return err
            }
        }

        // 扣减库存
        for _, item := range req.Items {
            result := tx.Model(&Product{}).
                Where("id = ? AND stock >= ?", item.ProductID, item.Quantity).
                Update("stock", gorm.Expr("stock - ?", item.Quantity))
            if result.RowsAffected == 0 {
                return fmt.Errorf("insufficient stock for product %d", item.ProductID)
            }
        }

        return nil
    })
}

// 全局启用
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{
    PrepareStmt: true,
    // 禁用默认事务包装，单条操作性能提升约 30%
    SkipDefaultTransaction: true,
})

// 按会话启用
tx := db.Session(&gorm.Session{PrepareStmt: true})
tx.First(&user, 1)  // 首次执行：PREPARE + EXECUTE
tx.Find(&users)      // 后续执行：直接 EXECUTE（命中缓存）

// CreateInBatches：分批插入，控制单次 SQL 大小
users := make([]User, 10000)
// ... 填充数据

// 每批 200 条，自动拆分为 50 个 INSERT 语句
db.CreateInBatches(users, 200)

// 全局配置批量大小
db, _ := gorm.Open(postgres.Open(dsn), &gorm.Config{
    CreateBatchSize: 500, // 所有 Create 操作自动分批
})

// FindInBatches：分批读取大数据集，控制内存使用
var allUsers []User
result := db.Where("active = ?", true).FindInBatches(&allUsers, 1000,
    func(tx *gorm.DB, batch int) error {
        for i := range allUsers {
            // 处理每条记录
            allUsers[i].LastSyncAt = time.Now()
        }
        // 批量更新
        tx.Save(&allUsers)
        return nil
    })
fmt.Printf("总处理 %d 条记录\n", result.RowsAffected)

// 直接指定字段
db.Select("id", "name", "email").Find(&users)

// 使用 API 专用结构体 -- Smart Select
type UserBrief struct {
    ID   uint
    Name string
}

// GORM 自动推断：SELECT `id`, `name` FROM `users`
db.Model(&User{}).Limit(100).Find(&[]UserBrief{})

// Preload：分离查询（适合一对多关系）
// 执行 2 条 SQL：SELECT * FROM users; SELECT * FROM orders WHERE user_id IN (1,2,3...)
db.Preload("Orders").Find(&users)

// 条件 Preload
db.Preload("Orders", "status = ?", "paid").Find(&users)

// 自定义 Preload SQL
db.Preload("Orders", func(db *gorm.DB) *gorm.DB {
    return db.Where("amount > ?", 100).Order("created_at DESC").Limit(5)
}).Find(&users)

// 嵌套 Preload
db.Preload("Orders.OrderItems.Product").
    Preload("CreditCard").
    Find(&users)

// Joins Preloading：SQL JOIN（适合一对一、多对一关系）
// 执行 1 条 SQL：SELECT users.*, companies.* FROM users LEFT JOIN companies ...
db.Joins("Company").Joins("Manager").First(&user, 1)

// 带条件的 Joins
db.Joins("Company", db.Where(&Company{Active: true})).Find(&users)

// 嵌套 Joins
db.Joins("Manager").Joins("Manager.Company").Find(&users)

package cache

import (
    "context"
    "encoding/json"
    "fmt"
    "time"

    "github.com/redis/go-redis/v9"
    "gorm.io/gorm"
)

type CachedRepository[T any] struct {
    db    *gorm.DB
    redis *redis.Client
    ttl   time.Duration
}

func NewCachedRepo[T any](db *gorm.DB, rdb *redis.Client, ttl time.Duration) *CachedRepository[T] {
    return &CachedRepository[T]{db: db, redis: rdb, ttl: ttl}
}

// FindByID 实现 Cache-Aside 模式
func (r *CachedRepository[T]) FindByID(ctx context.Context, id uint) (*T, error) {
    key := fmt.Sprintf("%T:%d", *new(T), id)

    // 1. 先查缓存
    cached, err := r.redis.Get(ctx, key).Result()
    if err == nil {
        var result T
        if err := json.Unmarshal([]byte(cached), &result); err == nil {
            return &result, nil
        }
    }

    // 2. 缓存未命中，查数据库
    var result T
    if err := r.db.WithContext(ctx).First(&result, id).Error; err != nil {
        return nil, err
    }

    // 3. 写入缓存
    data, _ := json.Marshal(result)
    r.redis.Set(ctx, key, data, r.ttl)

    return &result, nil
}

// Invalidate 缓存失效
func (r *CachedRepository[T]) Invalidate(ctx context.Context, id uint) error {
    key := fmt.Sprintf("%T:%d", *new(T), id)
    return r.redis.Del(ctx, key).Err()
}

// Update 更新数据并失效缓存
func (r *CachedRepository[T]) Update(ctx context.Context, id uint, updates map[string]interface{}) error {
    // 先更新数据库
    if err := r.db.WithContext(ctx).Model(new(T)).Where("id = ?", id).Updates(updates).Error; err != nil {
        return err
    }
    // 再失效缓存（Delete 而非 Set，避免缓存与数据库不一致）
    return r.Invalidate(ctx, id)
}

// 使用示例
// repo := cache.NewCachedRepo[User](db, redisClient, 15*time.Minute)
// user, err := repo.FindByID(ctx, 123)

package middleware

import (
    "log"
    "sync"

    "gorm.io/gorm"
)

// N+1 检测插件：开发环境使用
type N1DetectorPlugin struct{}

func (p *N1DetectorPlugin) Name() string { return "n1_detector" }

func (p *N1DetectorPlugin) Initialize(db *gorm.DB) error {
    var (
        mu       sync.Mutex
        queries  = make(map[string]int) // SQL 模式 -> 执行次数
    )

    db.Callback().Query().After("gorm:query").Register("n1:detect", func(db *gorm.DB) {
        sql := db.Statement.SQL.String()
        mu.Lock()
        queries[sql]++
        count := queries[sql]
        mu.Unlock()

        if count > 5 { // 同一 SQL 模式执行超过 5 次
            log.Printf("[N+1 WARNING] Query executed %d times: %s", count, sql)
        }
    })

    return nil
}

// 开发环境注册
// if env == "development" {
//     db.Use(&N1DetectorPlugin{})
// }

package tenant

import (
    "context"

    "gorm.io/gorm"
)

type contextKey string

const TenantIDKey contextKey = "tenant_id"

// TenantScope 自动注入租户条件
func TenantScope(ctx context.Context) func(db *gorm.DB) *gorm.DB {
    return func(db *gorm.DB) *gorm.DB {
        tenantID, ok := ctx.Value(TenantIDKey).(string)
        if !ok || tenantID == "" {
            // 没有租户上下文，阻止查询（安全兜底）
            db.AddError(fmt.Errorf("tenant_id is required"))
            return db
        }
        return db.Where("tenant_id = ?", tenantID)
    }
}

// Gin 中间件：从请求头提取租户 ID
func TenantMiddleware() gin.HandlerFunc {
    return func(c *gin.Context) {
        tenantID := c.GetHeader("X-Tenant-ID")
        if tenantID == "" {
            c.AbortWithStatusJSON(400, gin.H{"error": "X-Tenant-ID header required"})
            return
        }
        ctx := context.WithValue(c.Request.Context(), TenantIDKey, tenantID)
        c.Request = c.Request.WithContext(ctx)
        c.Next()
    }
}

// Repository 层使用
func (r *OrderRepo) List(ctx context.Context, page, size int) ([]Order, error) {
    var orders []Order
    err := r.db.WithContext(ctx).
        Scopes(TenantScope(ctx)).
        Offset((page - 1) * size).
        Limit(size).
        Find(&orders).Error
    return orders, err
}

package tenant

import (
    "fmt"

    "gorm.io/driver/postgres"
    "gorm.io/gorm"
)

// SchemaResolver 根据租户 ID 返回对应的数据库连接
type SchemaResolver struct {
    baseDB *gorm.DB
}

func NewSchemaResolver(baseDB *gorm.DB) *SchemaResolver {
    return &SchemaResolver{baseDB: baseDB}
}

// GetTenantDB 返回设置了正确 search_path 的连接
func (r *SchemaResolver) GetTenantDB(tenantID string) *gorm.DB {
    schema := fmt.Sprintf("tenant_%s", tenantID)
    return r.baseDB.Session(&gorm.Session{}).
        Exec(fmt.Sprintf("SET search_path TO %s, public", schema))
}

// Gin 中间件版本
func SchemaMiddleware(resolver *SchemaResolver) gin.HandlerFunc {
    return func(c *gin.Context) {
        tenantID := c.GetHeader("X-Tenant-ID")
        if tenantID == "" {
            c.AbortWithStatusJSON(400, gin.H{"error": "X-Tenant-ID required"})
            return
        }

        tenantDB := resolver.GetTenantDB(tenantID)
        c.Set("db", tenantDB)
        c.Next()
    }
}

// 租户 Schema 生命周期管理
func CreateTenantSchema(db *gorm.DB, tenantID string) error {
    schema := fmt.Sprintf("tenant_%s", tenantID)

    // 创建 Schema
    if err := db.Exec(fmt.Sprintf("CREATE SCHEMA IF NOT EXISTS %s", schema)).Error; err != nil {
        return err
    }

    // 在租户 Schema 中创建表
    tenantDB := db.Session(&gorm.Session{}).Exec(
        fmt.Sprintf("SET search_path TO %s", schema),
    )
    return tenantDB.AutoMigrate(&User{}, &Order{}, &Product{})
}

// 数据库层面设置 RLS
func SetupRLS(db *gorm.DB) error {
    sqls := []string{
        "ALTER TABLE orders ENABLE ROW LEVEL SECURITY",
        "ALTER TABLE orders FORCE ROW LEVEL SECURITY",
        `CREATE POLICY tenant_isolation ON orders
            USING (tenant_id = current_setting('app.tenant_id')::text)`,
    }
    for _, sql := range sqls {
        if err := db.Exec(sql).Error; err != nil {
            return err
        }
    }
    return nil
}

// 中间件设置当前租户
func RLSMiddleware(db *gorm.DB) gin.HandlerFunc {
    return func(c *gin.Context) {
        tenantID := c.GetHeader("X-Tenant-ID")
        // 设置 PostgreSQL 会话变量
        tenantDB := db.Session(&gorm.Session{}).
            Exec("SET app.tenant_id = ?", tenantID)
        c.Set("db", tenantDB)
        c.Next()
    }
}

package observability

import (
    "gorm.io/gorm"
    "gorm.io/plugin/opentelemetry/tracing"
)

func SetupTracing(db *gorm.DB) error {
    // 启用追踪 + 指标
    return db.Use(tracing.NewPlugin(
        // 可选：禁用指标，只保留追踪
        // tracing.WithoutMetrics(),
    ))
}

// 搭配 Gin 的 OpenTelemetry 中间件，实现请求到 SQL 的全链路追踪
// HTTP Request -> Gin Handler -> Service -> GORM -> PostgreSQL
// 每一层都会生成 Span，自动串联

package observability

import (
    "gorm.io/gorm"
    "gorm.io/plugin/prometheus"
)

func SetupPrometheus(db *gorm.DB) error {
    return db.Use(prometheus.New(prometheus.Config{
        DBName:          "primary",       // 指标标签
        RefreshInterval: 15,              // 刷新间隔（秒）
        StartServer:     true,            // 启动独立 HTTP 端口暴露指标
        HTTPServerPort:  9090,            // Prometheus 抓取端口
        MetricsCollector: []prometheus.MetricsCollector{
            &prometheus.MySQL{
                VariableNames: []string{"Threads_running", "Threads_connected"},
                Prefix:        "gorm_mysql_",
                Interval:      100,
            },
        },
    }))
}

// 暴露的指标包括：
// gorm_dbstats_max_open_connections    -- 最大连接数
// gorm_dbstats_open_connections        -- 当前打开连接数
// gorm_dbstats_in_use                  -- 使用中连接数
// gorm_dbstats_idle                    -- 空闲连接数
// gorm_dbstats_wait_count              -- 等待连接的总次数
// gorm_dbstats_wait_duration           -- 等待连接的总时间
// gorm_dbstats_max_idle_closed         -- 因 MaxIdleConns 关闭的连接数
// gorm_dbstats_max_lifetime_closed     -- 因 ConnMaxLifetime 关闭的连接数
// gorm_dbstats_max_idletime_closed     -- 因 ConnMaxIdleTime 关闭的连接数

package observability

import (
    "log"
    "os"
    "time"

    "gorm.io/gorm"
    "gorm.io/gorm/logger"
)

func NewProductionLogger() logger.Interface {
    return logger.New(
        log.New(os.Stdout, "\r\n", log.LstdFlags),
        logger.Config{
            SlowThreshold:             200 * time.Millisecond, // 慢查询阈值
            LogLevel:                  logger.Warn,            // 生产环境只记录 Warn 及以上
            IgnoreRecordNotFoundError: true,                   // 忽略 ErrRecordNotFound
            Colorful:                  false,                  // JSON 日志不需要颜色
        },
    )
}

// 自定义结构化日志（集成 slog/zap）
type StructuredLogger struct {
    SlowThreshold time.Duration
}

func (l *StructuredLogger) LogMode(level logger.LogLevel) logger.Interface { return l }

func (l *StructuredLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
    elapsed := time.Since(begin)
    sql, rows := fc()

    fields := map[string]interface{}{
        "elapsed_ms": elapsed.Milliseconds(),
        "rows":       rows,
        "sql":        sql,
    }

    switch {
    case err != nil && !errors.Is(err, gorm.ErrRecordNotFound):
        slog.ErrorContext(ctx, "gorm query error", fields)
    case elapsed > l.SlowThreshold:
        slog.WarnContext(ctx, "gorm slow query", fields)
    default:
        slog.DebugContext(ctx, "gorm query", fields)
    }
}