142 lines
4.3 KiB
Go
142 lines
4.3 KiB
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"
|
||
)
|
||
|
||
// genTaskDao 文生图任务表 DAO：进度更新（done 计数）高频，单行 UPDATE 原子。
|
||
type genTaskDao struct{}
|
||
|
||
var GenTask = &genTaskDao{}
|
||
|
||
func init() {
|
||
ctx := context.Background()
|
||
_, err := g.DB().Exec(ctx, `CREATE TABLE IF NOT EXISTS gen_task (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
dataset_id BIGINT NOT NULL,
|
||
status VARCHAR(20) NOT NULL DEFAULT 'running',
|
||
total INTEGER NOT NULL DEFAULT 0,
|
||
done INTEGER NOT NULL DEFAULT 0,
|
||
rejected INTEGER NOT NULL DEFAULT 0,
|
||
error TEXT,
|
||
created_at TIMESTAMP NOT NULL,
|
||
finished_at TIMESTAMP
|
||
)`)
|
||
if err != nil {
|
||
panic(err)
|
||
}
|
||
// 存量库迁移：负样本生成空检剔除计数列（2026-09-07）
|
||
common.EnsureColumn(ctx, consts.TableGenTask, "rejected", "rejected INTEGER NOT NULL DEFAULT 0")
|
||
}
|
||
|
||
// Insert 创建生成任务，返回自增 id
|
||
func (d *genTaskDao) Insert(ctx context.Context, m *entity.GenTask) (int64, error) {
|
||
res, err := g.DB().Model(consts.TableGenTask).Ctx(ctx).Data(g.Map{
|
||
"dataset_id": m.DatasetId,
|
||
"status": m.Status,
|
||
"total": m.Total,
|
||
"done": m.Done,
|
||
"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 *genTaskDao) GetById(ctx context.Context, id int64) (*entity.GenTask, error) {
|
||
var e entity.GenTask
|
||
err := g.DB().Model(consts.TableGenTask).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 *genTaskDao) ListRunning(ctx context.Context) ([]*entity.GenTask, error) {
|
||
var list []*entity.GenTask
|
||
err := g.DB().Model(consts.TableGenTask).Ctx(ctx).
|
||
Where("status", consts.GenTaskRunning).OrderAsc("id").Scan(&list)
|
||
if err != nil {
|
||
if common.IsNoRows(err) {
|
||
return []*entity.GenTask{}, nil
|
||
}
|
||
return nil, err
|
||
}
|
||
return list, nil
|
||
}
|
||
|
||
// GetRunningByDataset 某数据集 running 任务（并发检查用）
|
||
func (d *genTaskDao) GetRunningByDataset(ctx context.Context, datasetId int64) (*entity.GenTask, error) {
|
||
var e entity.GenTask
|
||
err := g.DB().Model(consts.TableGenTask).Ctx(ctx).
|
||
Where("dataset_id", datasetId).Where("status", consts.GenTaskRunning).Scan(&e)
|
||
if err != nil {
|
||
if common.IsNoRows(err) {
|
||
return nil, nil
|
||
}
|
||
return nil, err
|
||
}
|
||
return &e, nil
|
||
}
|
||
|
||
// UpdateProgress 更新生成进度
|
||
func (d *genTaskDao) UpdateProgress(ctx context.Context, id int64, done int) error {
|
||
_, err := g.DB().Model(consts.TableGenTask).Ctx(ctx).Where("id", id).
|
||
Data(g.Map{"done": done}).Update()
|
||
return err
|
||
}
|
||
|
||
// UpdateRejected 更新空检剔除计数（负样本生成收尾时一次性覆盖写）
|
||
func (d *genTaskDao) UpdateRejected(ctx context.Context, id int64, rejected int) error {
|
||
_, err := g.DB().Model(consts.TableGenTask).Ctx(ctx).Where("id", id).
|
||
Data(g.Map{"rejected": rejected}).Update()
|
||
return err
|
||
}
|
||
|
||
// Finish 完成任务（done 状态 + 完成时间 + 失败原因；errMsg 非空时为 failed）
|
||
func (d *genTaskDao) Finish(ctx context.Context, id int64, errMsg string) error {
|
||
status := consts.GenTaskDone
|
||
if errMsg != "" {
|
||
status = consts.GenTaskFailed
|
||
}
|
||
_, err := g.DB().Model(consts.TableGenTask).Ctx(ctx).Where("id", id).
|
||
Data(g.Map{"status": status, "finished_at": gtime.Now(), "error": errMsg}).Update()
|
||
return err
|
||
}
|
||
|
||
// LatestByDataset 某数据集最近一次任务（无则 nil）
|
||
func (d *genTaskDao) LatestByDataset(ctx context.Context, datasetId int64) (*entity.GenTask, error) {
|
||
var e entity.GenTask
|
||
err := g.DB().Model(consts.TableGenTask).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 *genTaskDao) DeleteByDataset(ctx context.Context, datasetId int64) error {
|
||
_, err := g.DB().Model(consts.TableGenTask).Ctx(ctx).
|
||
Where("dataset_id", datasetId).Delete()
|
||
return err
|
||
}
|