go-policy-service/models/task.go

package models

import (
	"context"
	"go-policy-service/utils"

	"gorm.io/gorm"
)

type HgFlightSearchTask struct {
	ID         uint   `gorm:"primaryKey;column:id"`
	Dep        string `gorm:"type:varchar(12);column:dep"`
	Arr        string `gorm:"type:varchar(12);column:arr"`
	DateRange  string `gorm:"type:varchar(12);column:date_range"`
	TripType   string `gorm:"type:char(2);column:trip_type"`
	FlightNo   string `gorm:"type:varchar(12);column:flight_no"`
	Flight     string `gorm:"type:char(2);column:flight"`
	CreatedAt  uint   `gorm:"column:created_at"` // 如果你想用 time.Time，可以改成 time.Time 并加解析
	UpdatedAt  uint   `gorm:"column:updated_at"`
	MerchantID uint   `gorm:"column:merchant_id;default:3"`
	ServiceTag string `gorm:"type:varchar(45);column:service_tag"`
}

// TableName 设置表名
func (HgFlightSearchTask) TableName() string {
	return "hg_flight_search_task"
}

type ProcessedData struct {
	gorm.Model
	TaskID uint
	Data   string `gorm:"type:text"`
}

func GetPendingTasks(ctx context.Context, limit int) ([]HgFlightSearchTask, error) {
	var hgShFlNotasks []HgFlightSearchTask
	result := utils.Db.WithContext(ctx).Model(&HgFlightSearchTask{}).Where("status = ?", 1)
	if limit > 0 {
		result = result.Limit(limit).Find(&hgShFlNotasks)
	} else {
		result = result.Find(&hgShFlNotasks)
	}
	return hgShFlNotasks, result.Error
}

func SaveProcessedData(ctx context.Context, data *ProcessedData) error {
	return utils.Db.WithContext(ctx).Model(&HgFlightSearchTask{}).Create(data).Error
}

func UpdateTaskStatus(ctx context.Context, taskID uint, status string, attempts int) error {
	return utils.Db.WithContext(ctx).Model(&HgFlightSearchTask{}).
		Where("id = ?", taskID).
		Updates(map[string]interface{}{
			"status":   status,
			"attempts": attempts,
		}).Error
}