pig-farm-controller/internal/infra/repository/unit_of_work.go

package repository

import (
	"context"
	"fmt"

	"git.huangwc.com/pig/pig-farm-controller/internal/infra/logs"

	"gorm.io/gorm"
)

// UnitOfWork 定义了工作单元接口，用于抽象事务管理
type UnitOfWork interface {
	// ExecuteInTransaction 在一个数据库事务中执行给定的函数。
	// 如果函数返回错误，事务将被回滚；否则，事务将被提交。
	// tx 参数是当前事务的 GORM DB 实例，应传递给所有仓库方法。
	ExecuteInTransaction(ctx context.Context, fn func(tx *gorm.DB) error) error
}

// gormUnitOfWork 是 UnitOfWork 接口的 GORM 实现
type gormUnitOfWork struct {
	ctx context.Context
	db  *gorm.DB
}

// NewGormUnitOfWork 创建一个新的 gormUnitOfWork 实例
func NewGormUnitOfWork(ctx context.Context, db *gorm.DB) UnitOfWork {
	return &gormUnitOfWork{
		ctx: ctx,
		db:  db,
	}
}

// ExecuteInTransaction 实现了 UnitOfWork 接口的事务执行逻辑
func (u *gormUnitOfWork) ExecuteInTransaction(ctx context.Context, fn func(tx *gorm.DB) error) error {
	uowCtx, logger := logs.Trace(ctx, u.ctx, "ExecuteInTransaction")

	tx := u.db.WithContext(uowCtx).Begin()
	if tx.Error != nil {
		logger.Errorf("开启数据库事务失败: %v", tx.Error) // 记录错误日志
		return fmt.Errorf("开启事务失败: %w", tx.Error)
	}

	defer func() {
		if r := recover(); r != nil {
			tx.Rollback()
			logger.Errorf("事务中发生 panic，已回滚: %v", r) // 记录 panic 日志
		} else if tx.Error != nil { // 如果函数执行过程中返回错误，或者事务本身有错误，则回滚
			tx.Rollback()
			logger.Errorf("事务执行失败，已回滚: %v", tx.Error) // 记录错误日志
		}
	}()

	// 执行业务逻辑函数
	if err := fn(tx); err != nil {
		tx.Rollback()
		return err // 返回业务逻辑函数中的错误
	}

	// 提交事务
	if err := tx.Commit().Error; err != nil {
		logger.Errorf("提交数据库事务失败: %v", err) // 记录错误日志
		return fmt.Errorf("提交事务失败: %w", err)
	}

	return nil
}