Files
doc_ai_backed/internal/storage/dao/task.go
liuyuanchuang 48e63894eb init repo
2025-12-10 18:33:37 +08:00

95 lines
3.4 KiB
Go

package dao
import (
"time"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
type TaskStatus int
type TaskType string
const (
TaskStatusPending TaskStatus = 0
TaskStatusProcessing TaskStatus = 1
TaskStatusCompleted TaskStatus = 2
TaskStatusFailed TaskStatus = 3
TaskTypeFormula TaskType = "FORMULA"
TaskTypeText TaskType = "TEXT"
TaskTypeTable TaskType = "TABLE"
TaskTypeLayout TaskType = "LAYOUT"
)
func (t TaskType) String() string {
return string(t)
}
func (t TaskStatus) String() string {
return []string{"PENDING", "PROCESSING", "COMPLETED", "FAILED"}[t]
}
type RecognitionTask struct {
BaseModel
UserID int64 `gorm:"column:user_id;not null;default:0;comment:用户ID" json:"user_id"`
TaskUUID string `gorm:"column:task_uuid;varchar(64);not null;default:'';comment:任务唯一标识" json:"task_uuid"`
FileName string `gorm:"column:file_name;varchar(256);not null;default:'';comment:文件名" json:"file_name"`
FileHash string `gorm:"column:file_hash;varchar(64);not null;default:'';comment:文件hash" json:"file_hash"`
FileURL string `gorm:"column:file_url;varchar(128);not null;comment:oss文件地址;default:''" json:"file_url"`
TaskType TaskType `gorm:"column:task_type;varchar(16);not null;comment:任务类型;default:''" json:"task_type"`
Status TaskStatus `gorm:"column:status;tinyint(2);not null;comment:任务状态;default:0" json:"status"`
CompletedAt time.Time `gorm:"column:completed_at;not null;default:current_timestamp;comment:完成时间" json:"completed_at"`
Remark string `gorm:"column:remark;varchar(64);comment:备注;not null;default:''" json:"remark"`
IP string `gorm:"column:ip;varchar(16);comment:IP地址;not null;default:''" json:"ip"`
}
func (t *RecognitionTask) TableName() string {
return "recognition_tasks"
}
type RecognitionTaskDao struct{}
func NewRecognitionTaskDao() *RecognitionTaskDao {
return &RecognitionTaskDao{}
}
// 模型方法
func (dao *RecognitionTaskDao) Create(tx *gorm.DB, data *RecognitionTask) error {
return tx.Create(data).Error
}
func (dao *RecognitionTaskDao) Update(tx *gorm.DB, filter map[string]interface{}, data map[string]interface{}) error {
return tx.Model(RecognitionTask{}).Where(filter).Updates(data).Error
}
func (dao *RecognitionTaskDao) GetByTaskNo(tx *gorm.DB, taskUUID string) (task *RecognitionTask, err error) {
task = &RecognitionTask{}
err = tx.Model(RecognitionTask{}).Where("task_uuid = ?", taskUUID).First(task).Error
return
}
func (dao *RecognitionTaskDao) GetTaskByFileURL(tx *gorm.DB, userID int64, fileHash string) (task *RecognitionTask, err error) {
task = &RecognitionTask{}
err = tx.Model(RecognitionTask{}).Where("user_id = ? AND file_hash = ?", userID, fileHash).First(task).Error
return
}
func (dao *RecognitionTaskDao) GetTaskByID(tx *gorm.DB, id int64) (task *RecognitionTask, err error) {
task = &RecognitionTask{}
err = tx.Model(RecognitionTask{}).Where("id = ?", id).First(task).Error
if err != nil {
if err == gorm.ErrRecordNotFound {
return nil, nil
}
return nil, err
}
return task, nil
}
func (dao *RecognitionTaskDao) GetTaskList(tx *gorm.DB, taskType TaskType, page int, pageSize int) (tasks []*RecognitionTask, err error) {
offset := (page - 1) * pageSize
err = tx.Model(RecognitionTask{}).Where("task_type = ?", taskType).Offset(offset).Limit(pageSize).Order(clause.OrderByColumn{Column: clause.Column{Name: "id"}, Desc: true}).Find(&tasks).Error
return
}