init repo

internal/storage/dao/task.go

@@ -0,0 +1,94 @@
+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
+}