3588AdminBackend/internal/storage/tasks_repo.go

150 lines
3.3 KiB
Go

package storage
import (
"database/sql"
"encoding/json"
"time"
"3588AdminBackend/internal/models"
)
type TasksRepo struct {
db *sql.DB
}
func NewTasksRepo(db *sql.DB) *TasksRepo {
return &TasksRepo{db: db}
}
func (r *TasksRepo) Save(task *models.Task) error {
if r == nil || r.db == nil || task == nil {
return nil
}
task.Mu.RLock()
payload, err := json.Marshal(task.Payload)
if err != nil {
task.Mu.RUnlock()
return err
}
status := task.Status
devices := make([]models.DeviceTaskStatus, 0, len(task.Devices))
for _, ds := range task.Devices {
if ds == nil {
continue
}
devices = append(devices, models.DeviceTaskStatus{
DeviceID: ds.DeviceID,
Status: ds.Status,
Progress: ds.Progress,
Error: ds.Error,
})
}
task.Mu.RUnlock()
now := time.Now().Format(time.RFC3339)
finishedAt := ""
if status == models.TaskSuccess || status == models.TaskFailed {
finishedAt = now
}
tx, err := r.db.Begin()
if err != nil {
return err
}
defer tx.Rollback()
_, err = tx.Exec(`
INSERT INTO tasks(task_id, type, payload_json, status, created_at, finished_at)
VALUES(?, ?, ?, ?, COALESCE((SELECT created_at FROM tasks WHERE task_id = ?), ?), ?)
ON CONFLICT(task_id) DO UPDATE SET
type=excluded.type,
payload_json=excluded.payload_json,
status=excluded.status,
finished_at=excluded.finished_at
`, task.ID, task.Type, string(payload), string(status), task.ID, now, finishedAt)
if err != nil {
return err
}
if _, err := tx.Exec(`DELETE FROM task_devices WHERE task_id = ?`, task.ID); err != nil {
return err
}
for _, ds := range devices {
if _, err := tx.Exec(`
INSERT INTO task_devices(task_id, device_id, status, progress, error_text)
VALUES(?, ?, ?, ?, ?)
`, task.ID, ds.DeviceID, string(ds.Status), ds.Progress, ds.Error); err != nil {
return err
}
}
return tx.Commit()
}
func (r *TasksRepo) List() ([]models.Task, error) {
if r == nil || r.db == nil {
return nil, nil
}
rows, err := r.db.Query(`
SELECT task_id, type, payload_json, status
FROM tasks
ORDER BY created_at DESC, task_id DESC
`)
if err != nil {
return nil, err
}
defer rows.Close()
var out []models.Task
for rows.Next() {
var (
id, tType, payloadJSON, status string
)
if err := rows.Scan(&id, &tType, &payloadJSON, &status); err != nil {
return nil, err
}
var payload any
if payloadJSON != "" {
if err := json.Unmarshal([]byte(payloadJSON), &payload); err != nil {
return nil, err
}
}
task := models.Task{
ID: id,
Type: tType,
Payload: payload,
Status: models.TaskStatus(status),
Devices: map[string]*models.DeviceTaskStatus{},
}
deviceRows, err := r.db.Query(`
SELECT device_id, status, progress, error_text
FROM task_devices
WHERE task_id = ?
ORDER BY rowid ASC
`, id)
if err != nil {
return nil, err
}
for deviceRows.Next() {
var did, dsStatus, errText string
var progress float64
if err := deviceRows.Scan(&did, &dsStatus, &progress, &errText); err != nil {
deviceRows.Close()
return nil, err
}
task.DeviceIDs = append(task.DeviceIDs, did)
task.Devices[did] = &models.DeviceTaskStatus{
DeviceID: did,
Status: models.TaskStatus(dsStatus),
Progress: progress,
Error: errText,
}
}
deviceRows.Close()
out = append(out, task)
}
return out, rows.Err()
}