pig/pig-farm-controller
1
0
Fork 0
Code Issues 11 Pull Requests Actions Packages Projects Releases Wiki Activity

修改infra.repository包

Browse Source
This commit is contained in:
huang
2025-11-05 23:00:07 +08:00
parent 97aea66f7c
commit 10b123ab93
25 changed files with 877 additions and 608 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
internal/infra/repository/pending_task_repository.go
View File

@@ -1,63 +1,69 @@
package repository
import (
"context"
"errors"
"fmt"
"time"
"git.huangwc.com/pig/pig-farm-controller/internal/infra/logs"
"git.huangwc.com/pig/pig-farm-controller/internal/infra/models"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
// PendingTaskRepository 定义了与待执行任务队列交互的接口。
type PendingTaskRepository interface {
FindAllPendingTasks() ([]models.PendingTask, error)
FindPendingTriggerByPlanID(planID uint) (*models.PendingTask, error)
DeletePendingTasksByIDs(ids []uint) error
CreatePendingTask(task *models.PendingTask) error
CreatePendingTasksInBatch(tasks []*models.PendingTask) error
FindAllPendingTasks(ctx context.Context) ([]models.PendingTask, error)
FindPendingTriggerByPlanID(ctx context.Context, planID uint) (*models.PendingTask, error)
DeletePendingTasksByIDs(ctx context.Context, ids []uint) error
CreatePendingTask(ctx context.Context, task *models.PendingTask) error
CreatePendingTasksInBatch(ctx context.Context, tasks []*models.PendingTask) error
// UpdatePendingTaskExecuteAt 更新指定待执行任务的执行时间
UpdatePendingTaskExecuteAt(id uint, executeAt time.Time) error
UpdatePendingTaskExecuteAt(ctx context.Context, id uint, executeAt time.Time) error
// ClearAllPendingTasks 清空所有待执行任务
ClearAllPendingTasks() error
ClearAllPendingTasks(ctx context.Context) error
// ClaimNextAvailableTask 原子地认领下一个可用的任务。
// 它会同时返回被认领任务对应的日志对象，以及被删除的待办任务对象的内存副本。
ClaimNextAvailableTask(excludePlanIDs []uint) (*models.TaskExecutionLog, *models.PendingTask, error)
ClaimNextAvailableTask(ctx context.Context, excludePlanIDs []uint) (*models.TaskExecutionLog, *models.PendingTask, error)
// RequeueTask 安全地将一个任务重新放回队列。
RequeueTask(originalPendingTask *models.PendingTask) error
RequeueTask(ctx context.Context, originalPendingTask *models.PendingTask) error
// FindPendingTasksByTaskLogIDs 根据 TaskExecutionLogID 列表查找对应的待执行任务
FindPendingTasksByTaskLogIDs(taskLogIDs []uint) ([]models.PendingTask, error)
FindPendingTasksByTaskLogIDs(ctx context.Context, taskLogIDs []uint) ([]models.PendingTask, error)
// DeletePendingTasksByPlanLogID 删除与指定计划执行日志ID相关的所有待执行任务
DeletePendingTasksByPlanLogID(planLogID uint) error
DeletePendingTasksByPlanLogID(ctx context.Context, planLogID uint) error
}
// gormPendingTaskRepository 是使用 GORM 的具体实现。
type gormPendingTaskRepository struct {
db *gorm.DB
ctx context.Context
db *gorm.DB
}
// NewGormPendingTaskRepository 创建一个新的待执行任务队列仓库。
func NewGormPendingTaskRepository(db *gorm.DB) PendingTaskRepository {
return &gormPendingTaskRepository{db: db}
func NewGormPendingTaskRepository(ctx context.Context, db *gorm.DB) PendingTaskRepository {
return &gormPendingTaskRepository{ctx: ctx, db: db}
}
func (r *gormPendingTaskRepository) FindAllPendingTasks() ([]models.PendingTask, error) {
func (r *gormPendingTaskRepository) FindAllPendingTasks(ctx context.Context) ([]models.PendingTask, error) {
repoCtx := logs.AddFuncName(ctx, r.ctx, "FindAllPendingTasks")
var tasks []models.PendingTask
// 预加载 Task 以便后续访问 Task.PlanID
// 预加载 TaskExecutionLog 以便后续访问 PlanExecutionLogID
err := r.db.Preload("Task").Preload("TaskExecutionLog").Find(&tasks).Error
err := r.db.WithContext(repoCtx).Preload("Task").Preload("TaskExecutionLog").Find(&tasks).Error
return tasks, err
}
func (r *gormPendingTaskRepository) FindPendingTriggerByPlanID(planID uint) (*models.PendingTask, error) {
func (r *gormPendingTaskRepository) FindPendingTriggerByPlanID(ctx context.Context, planID uint) (*models.PendingTask, error) {
repoCtx := logs.AddFuncName(ctx, r.ctx, "FindPendingTriggerByPlanID")
var pendingTask models.PendingTask
// 关键修改：通过 JOIN tasks 表并查询 parameters JSON 字段来查找触发器，而不是依赖 task.plan_id
err := r.db.
err := r.db.WithContext(repoCtx).
Joins("JOIN tasks ON tasks.id = pending_tasks.task_id").
Where("tasks.type = ? AND tasks.parameters->>'plan_id' = ?", models.TaskPlanAnalysis, fmt.Sprintf("%d", planID)).
First(&pendingTask).Error
@@ -67,42 +73,48 @@ func (r *gormPendingTaskRepository) FindPendingTriggerByPlanID(planID uint) (*mo
return &pendingTask, err
}
func (r *gormPendingTaskRepository) DeletePendingTasksByIDs(ids []uint) error {
func (r *gormPendingTaskRepository) DeletePendingTasksByIDs(ctx context.Context, ids []uint) error {
repoCtx := logs.AddFuncName(ctx, r.ctx, "DeletePendingTasksByIDs")
if len(ids) == 0 {
return nil
}
return r.db.Where("id IN ?", ids).Delete(&models.PendingTask{}).Error
return r.db.WithContext(repoCtx).Where("id IN ?", ids).Delete(&models.PendingTask{}).Error
}
func (r *gormPendingTaskRepository) CreatePendingTask(task *models.PendingTask) error {
return r.db.Create(task).Error
func (r *gormPendingTaskRepository) CreatePendingTask(ctx context.Context, task *models.PendingTask) error {
repoCtx := logs.AddFuncName(ctx, r.ctx, "CreatePendingTask")
return r.db.WithContext(repoCtx).Create(task).Error
}
// CreatePendingTasksInBatch 在一次数据库调用中创建多个待执行任务条目。
func (r *gormPendingTaskRepository) CreatePendingTasksInBatch(tasks []*models.PendingTask) error {
func (r *gormPendingTaskRepository) CreatePendingTasksInBatch(ctx context.Context, tasks []*models.PendingTask) error {
repoCtx := logs.AddFuncName(ctx, r.ctx, "CreatePendingTasksInBatch")
if len(tasks) == 0 {
return nil
}
return r.db.Create(&tasks).Error
return r.db.WithContext(repoCtx).Create(&tasks).Error
}
// UpdatePendingTaskExecuteAt 更新指定待执行任务的执行时间
func (r *gormPendingTaskRepository) UpdatePendingTaskExecuteAt(id uint, executeAt time.Time) error {
return r.db.Model(&models.PendingTask{}).Where("id = ?", id).Update("execute_at", executeAt).Error
func (r *gormPendingTaskRepository) UpdatePendingTaskExecuteAt(ctx context.Context, id uint, executeAt time.Time) error {
repoCtx := logs.AddFuncName(ctx, r.ctx, "UpdatePendingTaskExecuteAt")
return r.db.WithContext(repoCtx).Model(&models.PendingTask{}).Where("id = ?", id).Update("execute_at", executeAt).Error
}
// ClearAllPendingTasks 清空所有待执行任务
func (r *gormPendingTaskRepository) ClearAllPendingTasks() error {
return r.db.Where("1 = 1").Delete(&models.PendingTask{}).Error
func (r *gormPendingTaskRepository) ClearAllPendingTasks(ctx context.Context) error {
repoCtx := logs.AddFuncName(ctx, r.ctx, "ClearAllPendingTasks")
return r.db.WithContext(repoCtx).Where("1 = 1").Delete(&models.PendingTask{}).Error
}
// ClaimNextAvailableTask 以原子方式认领下一个可用的任务。
func (r *gormPendingTaskRepository) ClaimNextAvailableTask(excludePlanIDs []uint) (*models.TaskExecutionLog, *models.PendingTask, error) {
func (r *gormPendingTaskRepository) ClaimNextAvailableTask(ctx context.Context, excludePlanIDs []uint) (*models.TaskExecutionLog, *models.PendingTask, error) {
repoCtx := logs.AddFuncName(ctx, r.ctx, "ClaimNextAvailableTask")
var log models.TaskExecutionLog
var pendingTask models.PendingTask
err := r.db.Transaction(func(tx *gorm.DB) error {
query := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
err := r.db.WithContext(repoCtx).Transaction(func(tx *gorm.DB) error {
query := tx.WithContext(repoCtx).Clauses(clause.Locking{Strength: "UPDATE"}).
Where("execute_at <= ?", time.Now()).
Order("execute_at ASC")
@@ -114,7 +126,7 @@ func (r *gormPendingTaskRepository) ClaimNextAvailableTask(excludePlanIDs []uint
return err
}
if err := tx.Unscoped().Delete(&pendingTask).Error; err != nil {
if err := tx.WithContext(repoCtx).Unscoped().Delete(&pendingTask).Error; err != nil {
return err
}
@@ -122,12 +134,12 @@ func (r *gormPendingTaskRepository) ClaimNextAvailableTask(excludePlanIDs []uint
"status": models.ExecutionStatusStarted,
"started_at": time.Now(),
}
if err := tx.Model(&models.TaskExecutionLog{}).Where("id = ?", pendingTask.TaskExecutionLogID).Updates(updates).Error; err != nil {
if err := tx.WithContext(repoCtx).Model(&models.TaskExecutionLog{}).Where("id = ?", pendingTask.TaskExecutionLogID).Updates(updates).Error; err != nil {
return err
}
// 关键修改：在 Preload("Task") 时，使用 Unscoped() 来忽略 Task 的软删除状态
if err := tx.Preload("Task", func(db *gorm.DB) *gorm.DB {
if err := tx.WithContext(repoCtx).Preload("Task", func(db *gorm.DB) *gorm.DB {
return db.Unscoped()
}).First(&log, pendingTask.TaskExecutionLogID).Error; err != nil {
return err
@@ -145,10 +157,11 @@ func (r *gormPendingTaskRepository) ClaimNextAvailableTask(excludePlanIDs []uint
// RequeueTask 安全地将一个任务重新放回队列。
// 它通过将原始 PendingTask 的 ID 重置为 0，并重新创建它来实现。
func (r *gormPendingTaskRepository) RequeueTask(originalPendingTask *models.PendingTask) error {
return r.db.Transaction(func(tx *gorm.DB) error {
func (r *gormPendingTaskRepository) RequeueTask(ctx context.Context, originalPendingTask *models.PendingTask) error {
repoCtx := logs.AddFuncName(ctx, r.ctx, "RequeueTask")
return r.db.WithContext(repoCtx).Transaction(func(tx *gorm.DB) error {
// 1. 将日志状态恢复为 waiting
if err := tx.Model(&models.TaskExecutionLog{}).Where("id = ?", originalPendingTask.TaskExecutionLogID).Update("status", models.ExecutionStatusWaiting).Error; err != nil {
if err := tx.WithContext(repoCtx).Model(&models.TaskExecutionLog{}).Where("id = ?", originalPendingTask.TaskExecutionLogID).Update("status", models.ExecutionStatusWaiting).Error; err != nil {
return err
}
@@ -157,25 +170,27 @@ func (r *gormPendingTaskRepository) RequeueTask(originalPendingTask *models.Pend
originalPendingTask.ID = 0
// 3. 重新创建待办任务。GORM 会忽略掉已被重置的 ID，并让数据库生成一个新的主键。
return tx.Create(originalPendingTask).Error
return tx.WithContext(repoCtx).Create(originalPendingTask).Error
})
}
// FindPendingTasksByTaskLogIDs 根据 TaskExecutionLogID 列表查找对应的待执行任务
func (r *gormPendingTaskRepository) FindPendingTasksByTaskLogIDs(taskLogIDs []uint) ([]models.PendingTask, error) {
func (r *gormPendingTaskRepository) FindPendingTasksByTaskLogIDs(ctx context.Context, taskLogIDs []uint) ([]models.PendingTask, error) {
repoCtx := logs.AddFuncName(ctx, r.ctx, "FindPendingTasksByTaskLogIDs")
if len(taskLogIDs) == 0 {
return []models.PendingTask{}, nil
}
var pendingTasks []models.PendingTask
err := r.db.Where("task_execution_log_id IN ?", taskLogIDs).Find(&pendingTasks).Error
err := r.db.WithContext(repoCtx).Where("task_execution_log_id IN ?", taskLogIDs).Find(&pendingTasks).Error
return pendingTasks, err
}
// DeletePendingTasksByPlanLogID 删除与指定计划执行日志ID相关的所有待执行任务
func (r *gormPendingTaskRepository) DeletePendingTasksByPlanLogID(planLogID uint) error {
func (r *gormPendingTaskRepository) DeletePendingTasksByPlanLogID(ctx context.Context, planLogID uint) error {
repoCtx := logs.AddFuncName(ctx, r.ctx, "DeletePendingTasksByPlanLogID")
// 使用子查询找到所有与 planLogID 相关的 task_execution_log_id
subQuery := r.db.Model(&models.TaskExecutionLog{}).Select("id").Where("plan_execution_log_id = ?", planLogID)
subQuery := r.db.WithContext(repoCtx).Model(&models.TaskExecutionLog{}).Select("id").Where("plan_execution_log_id = ?", planLogID)
// 使用子查询的结果来删除待执行任务
return r.db.Where("task_execution_log_id IN (?)", subQuery).Delete(&models.PendingTask{}).Error
return r.db.WithContext(repoCtx).Where("task_execution_log_id IN (?)", subQuery).Delete(&models.PendingTask{}).Error
}