Files
Atlas/internal/store/tasks.go
T
renjue ad38beeb9e
CI / docker (push) Successful in 11m2s
实现 Atlas AI 平台一期能力并完成品牌化改造。
引入用户端对话/生图与管理端 Provider 配置,统一 invoke 编排、Gateway、请求日志与模型管理;前端按设计规范拆分布局,Go module 更名为 github.com/rose_cat707/Atlas;生图长任务进度改为 SSE 推送;重写 README;移除 Dockerfile syntax 指令以修复 CI 构建超时。

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-29 20:23:13 +08:00

94 lines
3.7 KiB
Go

package store
import "context"
type Task struct {
ID int64 `json:"id"`
SessionID string `json:"session_id"`
ServiceType string `json:"service_type"`
Protocol string `json:"protocol"`
ProviderID int64 `json:"provider_id"`
Status string `json:"status"`
Progress int `json:"progress"`
InputJSON string `json:"input_json"`
ResultJSON string `json:"result_json"`
Error string `json:"error"`
ExternalID string `json:"external_id"`
CreatedAt string `json:"created_at"`
UpdatedAt string `json:"updated_at"`
CompletedAt *string `json:"completed_at"`
}
func (s *Store) InsertTask(ctx context.Context, t *Task) (int64, error) {
res, err := s.db.ExecContext(ctx,
`INSERT INTO tasks (session_id, service_type, protocol, provider_id, status, progress, input_json, result_json, error, external_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
t.SessionID, t.ServiceType, t.Protocol, t.ProviderID, t.Status, t.Progress, t.InputJSON, t.ResultJSON, t.Error, t.ExternalID,
)
if err != nil {
return 0, err
}
return res.LastInsertId()
}
func (s *Store) UpdateTask(ctx context.Context, t *Task) error {
_, err := s.db.ExecContext(ctx,
`UPDATE tasks SET status=?, progress=?, result_json=?, error=?, external_id=?, updated_at=datetime('now'), completed_at=? WHERE id=? AND session_id=?`,
t.Status, t.Progress, t.ResultJSON, t.Error, t.ExternalID, t.CompletedAt, t.ID, t.SessionID,
)
return err
}
func (s *Store) GetTask(ctx context.Context, sessionID string, id int64) (*Task, error) {
var t Task
err := s.db.QueryRowContext(ctx,
`SELECT id, session_id, service_type, protocol, provider_id, status, progress, input_json, result_json, error, external_id, created_at, updated_at, completed_at FROM tasks WHERE id = ? AND session_id = ?`, id, sessionID,
).Scan(&t.ID, &t.SessionID, &t.ServiceType, &t.Protocol, &t.ProviderID, &t.Status, &t.Progress, &t.InputJSON, &t.ResultJSON, &t.Error, &t.ExternalID, &t.CreatedAt, &t.UpdatedAt, &t.CompletedAt)
if err != nil {
return nil, err
}
return &t, nil
}
func (s *Store) ListActiveTasks(ctx context.Context, sessionID, serviceType string) ([]Task, error) {
q := `SELECT id, session_id, service_type, protocol, provider_id, status, progress, input_json, result_json, error, external_id, created_at, updated_at, completed_at FROM tasks WHERE session_id = ? AND status IN ('pending','running')`
args := []any{sessionID}
if serviceType != "" {
q += ` AND service_type = ?`
args = append(args, serviceType)
}
q += ` ORDER BY id DESC`
rows, err := s.db.QueryContext(ctx, q, args...)
if err != nil {
return nil, err
}
defer rows.Close()
var items []Task
for rows.Next() {
var t Task
if err := rows.Scan(&t.ID, &t.SessionID, &t.ServiceType, &t.Protocol, &t.ProviderID, &t.Status, &t.Progress, &t.InputJSON, &t.ResultJSON, &t.Error, &t.ExternalID, &t.CreatedAt, &t.UpdatedAt, &t.CompletedAt); err != nil {
return nil, err
}
items = append(items, t)
}
return items, rows.Err()
}
func (s *Store) ListPendingTasks(ctx context.Context) ([]Task, error) {
rows, err := s.db.QueryContext(ctx,
`SELECT id, session_id, service_type, protocol, provider_id, status, progress, input_json, result_json, error, external_id, created_at, updated_at, completed_at FROM tasks WHERE status IN ('pending','running') ORDER BY id ASC`,
)
if err != nil {
return nil, err
}
defer rows.Close()
var items []Task
for rows.Next() {
var t Task
if err := rows.Scan(&t.ID, &t.SessionID, &t.ServiceType, &t.Protocol, &t.ProviderID, &t.Status, &t.Progress, &t.InputJSON, &t.ResultJSON, &t.Error, &t.ExternalID, &t.CreatedAt, &t.UpdatedAt, &t.CompletedAt); err != nil {
return nil, err
}
items = append(items, t)
}
return items, rows.Err()
}