// Package database define database operation functions package database import ( "context" "time" "modelRT/orm" "github.com/gofrs/uuid" "gorm.io/gorm" ) // UpdateTaskStarted updates task start time and status to running func UpdateTaskStarted(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, startedAt int64) error { cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Model(&orm.AsyncTask{}). Where("task_id = ?", taskID). Updates(map[string]any{ "status": orm.AsyncTaskStatusRunning, "started_at": startedAt, }) return result.Error } // UpdateTaskRetryInfo updates task retry information func UpdateTaskRetryInfo(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, retryCount int, nextRetryTime int64) error { cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() updateData := map[string]any{ "retry_count": retryCount, } if nextRetryTime <= 0 { updateData["next_retry_time"] = nil } else { updateData["next_retry_time"] = nextRetryTime } result := tx.WithContext(cancelCtx). Model(&orm.AsyncTask{}). Where("task_id = ?", taskID). Updates(updateData) return result.Error } // UpdateTaskErrorInfo updates task error information func UpdateTaskErrorInfo(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, errorMsg, stackTrace string) error { cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Model(&orm.AsyncTask{}). Where("task_id = ?", taskID). Updates(map[string]any{ "failure_reason": errorMsg, "stack_trace": stackTrace, }) return result.Error } // UpdateTaskExecutionTime updates task execution time func UpdateTaskExecutionTime(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, executionTime int64) error { cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Model(&orm.AsyncTask{}). Where("task_id = ?", taskID). Update("execution_time", executionTime) return result.Error } // UpdateTaskWorkerID updates the worker ID that is processing the task func UpdateTaskWorkerID(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, workerID string) error { cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Model(&orm.AsyncTask{}). Where("task_id = ?", taskID). Update("worker_id", workerID) return result.Error } // UpdateTaskPriority updates task priority func UpdateTaskPriority(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, priority int) error { cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Model(&orm.AsyncTask{}). Where("task_id = ?", taskID). Update("priority", priority) return result.Error } // UpdateTaskQueueName updates task queue name func UpdateTaskQueueName(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, queueName string) error { cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Model(&orm.AsyncTask{}). Where("task_id = ?", taskID). Update("queue_name", queueName) return result.Error } // UpdateTaskCreatedBy updates task creator information func UpdateTaskCreatedBy(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, createdBy string) error { cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Model(&orm.AsyncTask{}). Where("task_id = ?", taskID). Update("created_by", createdBy) return result.Error } // UpdateTaskResultWithMetrics updates task result with execution metrics func UpdateTaskResultWithMetrics(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, executionTime int64, memoryUsage *int64, cpuUsage *float64, retryCount int, completedAt int64) error { cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Model(&orm.AsyncTaskResult{}). Where("task_id = ?", taskID). Updates(map[string]any{ "execution_time": executionTime, "memory_usage": memoryUsage, "cpu_usage": cpuUsage, "retry_count": retryCount, "completed_at": completedAt, }) return result.Error } // GetTasksForRetry retrieves tasks that are due for retry func GetTasksForRetry(ctx context.Context, tx *gorm.DB, limit int) ([]orm.AsyncTask, error) { var tasks []orm.AsyncTask cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() now := time.Now().Unix() result := tx.WithContext(cancelCtx). Where("status = ? AND next_retry_time IS NOT NULL AND next_retry_time <= ?", orm.AsyncTaskStatusFailed, now). Order("next_retry_time ASC"). Limit(limit). Find(&tasks) if result.Error != nil { return nil, result.Error } return tasks, nil } // GetTasksByPriority retrieves tasks by priority order func GetTasksByPriority(ctx context.Context, tx *gorm.DB, status orm.AsyncTaskStatus, limit int) ([]orm.AsyncTask, error) { var tasks []orm.AsyncTask cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Where("status = ?", status). Order("priority DESC, created_at ASC"). Limit(limit). Find(&tasks) if result.Error != nil { return nil, result.Error } return tasks, nil } // GetTasksByWorkerID retrieves tasks being processed by a specific worker func GetTasksByWorkerID(ctx context.Context, tx *gorm.DB, workerID string) ([]orm.AsyncTask, error) { var tasks []orm.AsyncTask cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) defer cancel() result := tx.WithContext(cancelCtx). Where("worker_id = ? AND status = ?", workerID, orm.AsyncTaskStatusRunning). Find(&tasks) if result.Error != nil { return nil, result.Error } return tasks, nil } // CleanupStaleTasks marks tasks as failed if they have been running for too long func CleanupStaleTasks(ctx context.Context, tx *gorm.DB, timeoutSeconds int64) (int64, error) { cancelCtx, cancel := context.WithTimeout(ctx, 30*time.Second) defer cancel() threshold := time.Now().Unix() - timeoutSeconds result := tx.WithContext(cancelCtx). Model(&orm.AsyncTask{}). Where("status = ? AND started_at IS NOT NULL AND started_at < ?", orm.AsyncTaskStatusRunning, threshold). Updates(map[string]any{ "status": orm.AsyncTaskStatusFailed, "failure_reason": "task timeout", "finished_at": time.Now().Unix(), }) return result.RowsAffected, result.Error }