observer/server/biz/dao/label_task.go

package dao

import (
	"context"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/gtime"

	"observer-server/biz/consts"
	"observer-server/biz/model/entity"
	"observer-server/common"
)

// LabelTask 标注任务表 DAO：进度更新（done 计数）高频，单行 UPDATE 原子。
type labelTaskDao struct{}

var LabelTask = &labelTaskDao{}

func init() {
	ctx := context.Background()
	_, err := g.DB().Exec(ctx, `CREATE TABLE IF NOT EXISTS label_task (
  id          INTEGER PRIMARY KEY AUTOINCREMENT,
  dataset_id  INTEGER NOT NULL,
  status      TEXT NOT NULL DEFAULT 'running',
  total       INTEGER NOT NULL DEFAULT 0,
  done        INTEGER NOT NULL DEFAULT 0,
  boxes_file  TEXT,
  filenames   TEXT,
  error       TEXT,
  created_at  TEXT NOT NULL,
  finished_at TEXT
)`)
	if err != nil {
		panic(err)
	}
}

// Insert 创建标注任务，返回自增 id
func (d *labelTaskDao) Insert(ctx context.Context, m *entity.LabelTask) (int64, error) {
	res, err := g.DB().Model(consts.TableLabelTask).Ctx(ctx).Data(g.Map{
		"dataset_id":  m.DatasetId,
		"status":      m.Status,
		"total":       m.Total,
		"done":        m.Done,
		"boxes_file":  m.BoxesFile,
		"filenames":   m.Filenames,
		"error":       m.Error,
		"created_at":  m.CreatedAt,
		"finished_at": m.FinishedAt,
	}).Insert()
	if err != nil {
		return 0, err
	}
	return res.LastInsertId()
}

// GetById 按主键查询，不存在返回 nil
func (d *labelTaskDao) GetById(ctx context.Context, id int64) (*entity.LabelTask, error) {
	var e entity.LabelTask
	err := g.DB().Model(consts.TableLabelTask).Ctx(ctx).Where("id", id).Scan(&e)
	if err != nil {
		if common.IsNoRows(err) {
			return nil, nil
		}
		return nil, err
	}
	return &e, nil
}

// ListRunning 全部 running 任务（服务重启恢复用）
func (d *labelTaskDao) ListRunning(ctx context.Context) ([]*entity.LabelTask, error) {
	var list []*entity.LabelTask
	err := g.DB().Model(consts.TableLabelTask).Ctx(ctx).
		Where("status", consts.LabelTaskRunning).OrderAsc("id").Scan(&list)
	if err != nil {
		if common.IsNoRows(err) {
			return []*entity.LabelTask{}, nil
		}
		return nil, err
	}
	return list, nil
}

// GetRunningByDataset 某数据集 running 任务（并发检查用）
func (d *labelTaskDao) GetRunningByDataset(ctx context.Context, datasetId int64) (*entity.LabelTask, error) {
	var e entity.LabelTask
	err := g.DB().Model(consts.TableLabelTask).Ctx(ctx).
		Where("dataset_id", datasetId).Where("status", consts.LabelTaskRunning).Scan(&e)
	if err != nil {
		if common.IsNoRows(err) {
			return nil, nil
		}
		return nil, err
	}
	return &e, nil
}

// UpdateProgress 更新处理进度
func (d *labelTaskDao) UpdateProgress(ctx context.Context, id int64, done int) error {
	_, err := g.DB().Model(consts.TableLabelTask).Ctx(ctx).Where("id", id).
		Data(g.Map{"done": done}).Update()
	return err
}

// Finish 完成任务（done 状态 + 完成时间 + 候选框文件路径 + 失败原因）
func (d *labelTaskDao) Finish(ctx context.Context, id int64, boxesFile, errMsg string) error {
	_, err := g.DB().Model(consts.TableLabelTask).Ctx(ctx).Where("id", id).
		Data(g.Map{"status": consts.LabelTaskDone, "finished_at": gtime.Now(), "boxes_file": boxesFile, "error": errMsg}).Update()
	return err
}

// LatestByDataset 某数据集最近一次任务（无则 nil）
func (d *labelTaskDao) LatestByDataset(ctx context.Context, datasetId int64) (*entity.LabelTask, error) {
	var e entity.LabelTask
	err := g.DB().Model(consts.TableLabelTask).Ctx(ctx).
		Where("dataset_id", datasetId).OrderDesc("id").Limit(1).Scan(&e)
	if err != nil {
		if common.IsNoRows(err) {
			return nil, nil
		}
		return nil, err
	}
	return &e, nil
}

// DeleteByDataset 删除数据集关联的全部标注任务（数据集删除时级联清理孤儿记录）
func (d *labelTaskDao) DeleteByDataset(ctx context.Context, datasetId int64) error {
	_, err := g.DB().Model(consts.TableLabelTask).Ctx(ctx).
		Where("dataset_id", datasetId).Delete()
	return err
}

// Page 标注任务分页：创建时间倒序
func (d *labelTaskDao) Page(ctx context.Context, page, size int) ([]*entity.LabelTask, int64, error) {
	base := g.DB().Model(consts.TableLabelTask).Ctx(ctx)
	total, err := base.Count()
	if err != nil {
		return nil, 0, err
	}
	var list []*entity.LabelTask
	err = base.OrderDesc("id").Limit((page-1)*size, size).Scan(&list)
	if err != nil {
		if common.IsNoRows(err) {
			return []*entity.LabelTask{}, int64(total), nil
		}
		return nil, 0, err
	}
	return list, int64(total), nil
}