iotlcm/internal/repository/batch/repository.go

package batch

import (
	"context"
	"strings"

	"github.com/gogf/gf/v2/database/gdb"
	"github.com/gogf/gf/v2/errors/gerror"
	"github.com/gogf/gf/v2/os/gtime"

	batchdomain "git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/batch"
	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/dao"
	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model/do"
	"git.5iot.com/iotlcm-public/iotlcm-worker-svr/internal/model/entity"
)

const (
	AsyncStatusSuccess          = 3
	RunResultSuccess            = 2
	batchDetailErrorReasonLimit = 512
)

type Repository struct {
	db gdb.DB
}

func New() *Repository {
	return &Repository{db: dao.BatchMain.DB()}
}

func (r *Repository) FindMain(ctx context.Context, batchNo string) (*entity.BatchMain, error) {
	var main *entity.BatchMain
	if err := r.db.Model(dao.BatchMain.Table()).Ctx(ctx).
		Where(dao.BatchMain.Columns().BatchNo, batchNo).
		Scan(&main); err != nil {
		return nil, gerror.Wrap(err, "query batch main failed")
	}
	return main, nil
}

func (r *Repository) MarkProcessing(ctx context.Context, batchNo string, batchType int) (claimed bool, err error) {
	err = r.db.Transaction(ctx, func(ctx context.Context, tx gdb.TX) error {
		mainColumns := dao.BatchMain.Columns()
		result, updateErr := tx.Ctx(ctx).Model(dao.BatchMain.Table()).
			Where(mainColumns.BatchNo, batchNo).
			Where(mainColumns.BatchType, batchType).
			Where(mainColumns.BatchAsyncStatus, batchdomain.AsyncStatusPending).
			Data(do.BatchMain{
				BatchAsyncStatus: batchdomain.AsyncStatusProcessing,
				RunResult:        batchdomain.RunResultRunning,
				StartTime:        gtime.Now(),
			}).
			Update()
		if updateErr != nil {
			return gerror.Wrap(updateErr, "CAS batch main failed")
		}
		affected, affectedErr := result.RowsAffected()
		if affectedErr != nil {
			return gerror.Wrap(affectedErr, "read batch main CAS result failed")
		}
		if affected == 0 {
			return nil
		}

		detailColumns := dao.BatchDetail.Columns()
		if _, updateErr = tx.Ctx(ctx).Model(dao.BatchDetail.Table()).
			Where(detailColumns.BatchNo, batchNo).
			Where(detailColumns.BatchType, batchType).
			Where(detailColumns.DetailStatus, batchdomain.DetailStatusPending).
			Data(do.BatchDetail{DetailStatus: batchdomain.DetailStatusRunning}).
			Update(); updateErr != nil {
			return gerror.Wrap(updateErr, "mark pending batch details running failed")
		}
		claimed = true
		return nil
	})
	return claimed, gerror.Wrap(err, "claim batch transaction failed")
}

func (r *Repository) ReclaimProcessing(ctx context.Context, batchNo string, batchType int, startedAt *gtime.Time) (claimed bool, err error) {
	columns := dao.BatchMain.Columns()
	model := r.db.Model(dao.BatchMain.Table()).Ctx(ctx).
		Where(columns.BatchNo, batchNo).
		Where(columns.BatchType, batchType).
		Where(columns.BatchAsyncStatus, batchdomain.AsyncStatusProcessing)
	if startedAt == nil || startedAt.IsZero() {
		model = model.Where(columns.StartTime, nil)
	} else {
		model = model.Where(columns.StartTime, startedAt)
	}

	result, err := model.Data(do.BatchMain{
		StartTime: gtime.Now(),
		RunResult: batchdomain.RunResultRunning,
	}).Update()
	if err != nil {
		return false, gerror.Wrap(err, "reclaim stale batch failed")
	}
	affected, err := result.RowsAffected()
	if err != nil {
		return false, gerror.Wrap(err, "read stale batch reclaim result failed")
	}
	return affected == 1, nil
}

func (r *Repository) MarkFailed(ctx context.Context, batchNo string, batchType int, reason string) (marked bool, err error) {
	detailReason := truncateBatchDetailReason(reason)
	err = r.db.Transaction(ctx, func(ctx context.Context, tx gdb.TX) error {
		mainColumns := dao.BatchMain.Columns()
		result, updateErr := tx.Ctx(ctx).Model(dao.BatchMain.Table()).
			Where(mainColumns.BatchNo, batchNo).
			Where(mainColumns.BatchType, batchType).
			Where(mainColumns.BatchAsyncStatus, batchdomain.AsyncStatusProcessing).
			Data(do.BatchMain{
				BatchAsyncStatus: batchdomain.AsyncStatusFailed,
				RunResult:        batchdomain.RunResultFailed,
				FailReason:       reason,
				EndTime:          gtime.Now(),
			}).
			Update()
		if updateErr != nil {
			return gerror.Wrap(updateErr, "mark batch main failed failed")
		}
		affected, affectedErr := result.RowsAffected()
		if affectedErr != nil {
			return gerror.Wrap(affectedErr, "read batch failure status update result failed")
		}
		if affected == 0 {
			return nil
		}

		detailColumns := dao.BatchDetail.Columns()
		if _, updateErr = tx.Ctx(ctx).Model(dao.BatchDetail.Table()).
			Where(detailColumns.BatchNo, batchNo).
			Where(detailColumns.BatchType, batchType).
			Where(detailColumns.DetailStatus, batchdomain.DetailStatusRunning).
			Data(do.BatchDetail{
				DetailStatus: batchdomain.DetailStatusFailed,
				ErrorReason:  detailReason,
			}).
			Update(); updateErr != nil {
			return gerror.Wrap(updateErr, "mark running batch details failed failed")
		}
		marked = true
		return nil
	})
	return marked, gerror.Wrap(err, "mark batch failure transaction failed")
}

func truncateBatchDetailReason(reason string) string {
	runes := []rune(strings.TrimSpace(reason))
	if len(runes) <= batchDetailErrorReasonLimit {
		return string(runes)
	}
	return string(runes[:batchDetailErrorReasonLimit])
}

// MarkSuccess 将指定批次从 Processing 标记为 Success，并结算明细状态。
func (r *Repository) MarkSuccess(ctx context.Context, batchNo string, batchType int, success, fail int, reason string) (marked bool, err error) {
	err = r.db.Transaction(ctx, func(ctx context.Context, tx gdb.TX) error {
		mainColumns := dao.BatchMain.Columns()
		result, updateErr := tx.Ctx(ctx).Model(dao.BatchMain.Table()).
			Where(mainColumns.BatchNo, batchNo).
			Where(mainColumns.BatchType, batchType).
			Where(mainColumns.BatchAsyncStatus, batchdomain.AsyncStatusProcessing).
			Data(do.BatchMain{
				BatchAsyncStatus: AsyncStatusSuccess,
				RunResult:        RunResultSuccess,
				SuccessCount:     success,
				FailCount:        fail,
				EndTime:          gtime.Now(),
			}).
			Update()
		if updateErr != nil {
			return gerror.Wrap(updateErr, "mark batch main success failed")
		}
		affected, affectedErr := result.RowsAffected()
		if affectedErr != nil {
			return gerror.Wrap(affectedErr, "read batch success status update result failed")
		}
		if affected == 0 {
			return nil
		}
		detailColumns := dao.BatchDetail.Columns()
		if _, updateErr = tx.Ctx(ctx).Model(dao.BatchDetail.Table()).
			Where(detailColumns.BatchNo, batchNo).
			Where(detailColumns.BatchType, batchType).
			Where(detailColumns.DetailStatus, batchdomain.DetailStatusRunning).
			Data(do.BatchDetail{
				DetailStatus: batchdomain.DetailStatusSuccess,
				ErrorReason:  reason,
			}).
			Update(); updateErr != nil {
			return gerror.Wrap(updateErr, "mark running batch details success failed")
		}
		marked = true
		return nil
	})
	return marked, gerror.Wrap(err, "mark batch success transaction failed")
}