293 lines
12 KiB
Go
293 lines
12 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"git.gobha.me/xcaliber/chat-switchboard/models"
|
|
)
|
|
|
|
type TaskStore struct{}
|
|
|
|
func NewTaskStore() *TaskStore { return &TaskStore{} }
|
|
|
|
// ── Full column list (single source of truth) ──────────────────
|
|
// Used by GetByID, GetByTriggerToken, and list(). Every query must
|
|
// SELECT and Scan the same columns in the same order to prevent drift.
|
|
//
|
|
// scanTask is the ONLY function that maps columns → struct fields.
|
|
// If a column is added or reordered, update taskColumns + scanTask.
|
|
|
|
// taskColumns is the canonical SELECT list for tasks.
|
|
// COALESCE wraps nullable columns that scan into non-pointer Go types
|
|
// (string, json.RawMessage). database/sql returns "unsupported Scan,
|
|
// storing driver.Value type <nil>" for these without COALESCE.
|
|
const taskColumns = `id, owner_id, team_id, name, description, scope,
|
|
task_type, COALESCE(system_function, '') AS system_function,
|
|
persona_id, model_id, system_prompt, user_prompt,
|
|
workflow_id, COALESCE(tool_grants, '[]'::jsonb) AS tool_grants,
|
|
schedule, timezone, is_active,
|
|
COALESCE(trigger_token, '') AS trigger_token,
|
|
max_tokens, max_tool_calls, max_wall_clock, output_mode,
|
|
output_channel_id,
|
|
COALESCE(webhook_url, '') AS webhook_url,
|
|
COALESCE(webhook_secret, '') AS webhook_secret,
|
|
provider_config_id,
|
|
notify_on_complete, notify_on_failure,
|
|
last_run_at, next_run_at, run_count, created_at, updated_at`
|
|
|
|
// scanTask maps one row from taskColumns into a Task struct.
|
|
// Single source of truth — used by GetByID, GetByTriggerToken, and list().
|
|
func scanTask(scanner interface{ Scan(...interface{}) error }, t *models.Task) error {
|
|
return scanner.Scan(
|
|
&t.ID, &t.OwnerID, &t.TeamID, &t.Name, &t.Description, &t.Scope,
|
|
&t.TaskType, &t.SystemFunction,
|
|
&t.PersonaID, &t.ModelID, &t.SystemPrompt, &t.UserPrompt,
|
|
&t.WorkflowID, &t.ToolGrants, &t.Schedule, &t.Timezone, &t.IsActive,
|
|
&t.TriggerToken,
|
|
&t.MaxTokens, &t.MaxToolCalls, &t.MaxWallClock, &t.OutputMode,
|
|
&t.OutputChannelID, &t.WebhookURL, &t.WebhookSecret, &t.ProviderConfigID,
|
|
&t.NotifyOnComplete, &t.NotifyOnFailure,
|
|
&t.LastRunAt, &t.NextRunAt, &t.RunCount, &t.CreatedAt, &t.UpdatedAt,
|
|
)
|
|
}
|
|
|
|
// ── CRUD ───────────────────────────────────────
|
|
|
|
func (s *TaskStore) Create(ctx context.Context, t *models.Task) error {
|
|
toolGrants := jsonOrNull(t.ToolGrants)
|
|
return DB.QueryRowContext(ctx, `
|
|
INSERT INTO tasks (owner_id, team_id, name, description, scope,
|
|
task_type, system_function,
|
|
persona_id, model_id, system_prompt, user_prompt,
|
|
workflow_id, tool_grants, schedule, timezone, is_active,
|
|
trigger_token,
|
|
max_tokens, max_tool_calls, max_wall_clock, output_mode,
|
|
output_channel_id, webhook_url, webhook_secret, provider_config_id,
|
|
notify_on_complete, notify_on_failure, next_run_at)
|
|
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23,$24,$25,$26,$27,$28)
|
|
RETURNING id, created_at, updated_at`,
|
|
t.OwnerID, t.TeamID, t.Name, t.Description, t.Scope,
|
|
t.TaskType, t.SystemFunction,
|
|
t.PersonaID, t.ModelID, t.SystemPrompt, t.UserPrompt,
|
|
t.WorkflowID, toolGrants, t.Schedule, t.Timezone, t.IsActive,
|
|
nilIfEmpty(t.TriggerToken),
|
|
t.MaxTokens, t.MaxToolCalls, t.MaxWallClock, t.OutputMode,
|
|
t.OutputChannelID, t.WebhookURL, t.WebhookSecret, t.ProviderConfigID,
|
|
t.NotifyOnComplete, t.NotifyOnFailure, t.NextRunAt,
|
|
).Scan(&t.ID, &t.CreatedAt, &t.UpdatedAt)
|
|
}
|
|
|
|
func (s *TaskStore) GetByID(ctx context.Context, id string) (*models.Task, error) {
|
|
t := &models.Task{}
|
|
row := DB.QueryRowContext(ctx,
|
|
`SELECT `+taskColumns+` FROM tasks WHERE id = $1`, id,
|
|
)
|
|
if err := scanTask(row, t); err != nil {
|
|
return nil, err
|
|
}
|
|
return t, nil
|
|
}
|
|
|
|
func (s *TaskStore) GetByTriggerToken(ctx context.Context, token string) (*models.Task, error) {
|
|
t := &models.Task{}
|
|
row := DB.QueryRowContext(ctx,
|
|
`SELECT `+taskColumns+` FROM tasks WHERE trigger_token = $1`, token,
|
|
)
|
|
if err := scanTask(row, t); err != nil {
|
|
return nil, err
|
|
}
|
|
return t, nil
|
|
}
|
|
|
|
func (s *TaskStore) Update(ctx context.Context, id string, p models.TaskPatch) error {
|
|
// Dynamic PATCH: only update non-nil fields
|
|
q := "UPDATE tasks SET updated_at = NOW()"
|
|
args := []interface{}{}
|
|
i := 1
|
|
if p.Name != nil { q += comma(i, "name"); args = append(args, *p.Name); i++ }
|
|
if p.Description != nil { q += comma(i, "description"); args = append(args, *p.Description); i++ }
|
|
if p.PersonaID != nil { q += comma(i, "persona_id"); args = append(args, *p.PersonaID); i++ }
|
|
if p.ModelID != nil { q += comma(i, "model_id"); args = append(args, *p.ModelID); i++ }
|
|
if p.SystemPrompt != nil { q += comma(i, "system_prompt"); args = append(args, *p.SystemPrompt); i++ }
|
|
if p.UserPrompt != nil { q += comma(i, "user_prompt"); args = append(args, *p.UserPrompt); i++ }
|
|
if p.WorkflowID != nil { q += comma(i, "workflow_id"); args = append(args, *p.WorkflowID); i++ }
|
|
if p.ToolGrants != nil { q += comma(i, "tool_grants"); args = append(args, jsonOrNull(*p.ToolGrants)); i++ }
|
|
if p.Schedule != nil { q += comma(i, "schedule"); args = append(args, *p.Schedule); i++ }
|
|
if p.Timezone != nil { q += comma(i, "timezone"); args = append(args, *p.Timezone); i++ }
|
|
if p.IsActive != nil { q += comma(i, "is_active"); args = append(args, *p.IsActive); i++ }
|
|
if p.MaxTokens != nil { q += comma(i, "max_tokens"); args = append(args, *p.MaxTokens); i++ }
|
|
if p.MaxToolCalls != nil { q += comma(i, "max_tool_calls"); args = append(args, *p.MaxToolCalls); i++ }
|
|
if p.MaxWallClock != nil { q += comma(i, "max_wall_clock"); args = append(args, *p.MaxWallClock); i++ }
|
|
if p.OutputMode != nil { q += comma(i, "output_mode"); args = append(args, *p.OutputMode); i++ }
|
|
if p.OutputChannelID != nil { q += comma(i, "output_channel_id"); args = append(args, *p.OutputChannelID); i++ }
|
|
if p.WebhookURL != nil { q += comma(i, "webhook_url"); args = append(args, *p.WebhookURL); i++ }
|
|
if p.ProviderConfigID != nil { q += comma(i, "provider_config_id"); args = append(args, *p.ProviderConfigID); i++ }
|
|
if p.NotifyOnComplete != nil { q += comma(i, "notify_on_complete"); args = append(args, *p.NotifyOnComplete); i++ }
|
|
if p.NotifyOnFailure != nil { q += comma(i, "notify_on_failure"); args = append(args, *p.NotifyOnFailure); i++ }
|
|
|
|
q += pgWhere(i, "id")
|
|
args = append(args, id)
|
|
_, err := DB.ExecContext(ctx, q, args...)
|
|
return err
|
|
}
|
|
|
|
func (s *TaskStore) Delete(ctx context.Context, id string) error {
|
|
_, err := DB.ExecContext(ctx, `DELETE FROM tasks WHERE id = $1`, id)
|
|
return err
|
|
}
|
|
|
|
func (s *TaskStore) ListByOwner(ctx context.Context, ownerID string) ([]models.Task, error) {
|
|
return s.list(ctx, `WHERE owner_id = $1 ORDER BY created_at DESC`, ownerID)
|
|
}
|
|
|
|
func (s *TaskStore) ListByTeam(ctx context.Context, teamID string) ([]models.Task, error) {
|
|
return s.list(ctx, `WHERE team_id = $1 ORDER BY created_at DESC`, teamID)
|
|
}
|
|
|
|
func (s *TaskStore) ListAll(ctx context.Context) ([]models.Task, error) {
|
|
return s.list(ctx, `ORDER BY created_at DESC`)
|
|
}
|
|
|
|
func (s *TaskStore) ListDue(ctx context.Context, limit int) ([]models.Task, error) {
|
|
return s.list(ctx, `WHERE is_active = true AND next_run_at <= NOW() ORDER BY next_run_at ASC LIMIT $1`, limit)
|
|
}
|
|
|
|
// ── Scheduler bookkeeping ──────────────────────
|
|
|
|
func (s *TaskStore) SetNextRun(ctx context.Context, id string, nextRun interface{}) error {
|
|
_, err := DB.ExecContext(ctx, `UPDATE tasks SET next_run_at = $1, updated_at = NOW() WHERE id = $2`, nextRun, id)
|
|
return err
|
|
}
|
|
|
|
func (s *TaskStore) SetLastRun(ctx context.Context, id string) error {
|
|
_, err := DB.ExecContext(ctx, `UPDATE tasks SET last_run_at = NOW(), updated_at = NOW() WHERE id = $1`, id)
|
|
return err
|
|
}
|
|
|
|
func (s *TaskStore) IncrementRunCount(ctx context.Context, id string) error {
|
|
_, err := DB.ExecContext(ctx, `UPDATE tasks SET run_count = run_count + 1, updated_at = NOW() WHERE id = $1`, id)
|
|
return err
|
|
}
|
|
|
|
// ── Run History ─────────────────────────────────
|
|
|
|
func (s *TaskStore) CreateRun(ctx context.Context, r *models.TaskRun) error {
|
|
triggerPayload := nilIfEmpty(r.TriggerPayload)
|
|
return DB.QueryRowContext(ctx, `
|
|
INSERT INTO task_runs (task_id, channel_id, status, trigger_payload)
|
|
VALUES ($1, $2, $3, $4) RETURNING id, started_at`,
|
|
r.TaskID, r.ChannelID, r.Status, triggerPayload,
|
|
).Scan(&r.ID, &r.StartedAt)
|
|
}
|
|
|
|
func (s *TaskStore) UpdateRun(ctx context.Context, id, status string, tokensUsed, toolCalls, wallClock int, errMsg string) error {
|
|
_, err := DB.ExecContext(ctx, `
|
|
UPDATE task_runs SET status = $1, tokens_used = $2, tool_calls = $3,
|
|
wall_clock = $4, error = $5, completed_at = NOW()
|
|
WHERE id = $6`, status, tokensUsed, toolCalls, wallClock, errMsg, id)
|
|
return err
|
|
}
|
|
|
|
func (s *TaskStore) TransitionRunStatus(ctx context.Context, id string, status string) error {
|
|
_, err := DB.ExecContext(ctx, `UPDATE task_runs SET status = $1 WHERE id = $2`, status, id)
|
|
return err
|
|
}
|
|
|
|
func (s *TaskStore) GetActiveRun(ctx context.Context, taskID string) (*models.TaskRun, error) {
|
|
r := &models.TaskRun{}
|
|
err := DB.QueryRowContext(ctx, `
|
|
SELECT id, task_id, channel_id, status, COALESCE(trigger_payload, ''), started_at
|
|
FROM task_runs WHERE task_id = $1 AND status = 'running'
|
|
LIMIT 1`, taskID,
|
|
).Scan(&r.ID, &r.TaskID, &r.ChannelID, &r.Status, &r.TriggerPayload, &r.StartedAt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return r, nil
|
|
}
|
|
|
|
func (s *TaskStore) GetQueuedRun(ctx context.Context, taskID string) (*models.TaskRun, error) {
|
|
r := &models.TaskRun{}
|
|
err := DB.QueryRowContext(ctx, `
|
|
SELECT id, task_id, channel_id, status, COALESCE(trigger_payload, ''), started_at
|
|
FROM task_runs WHERE task_id = $1 AND status = 'queued'
|
|
ORDER BY started_at ASC LIMIT 1`, taskID,
|
|
).Scan(&r.ID, &r.TaskID, &r.ChannelID, &r.Status, &r.TriggerPayload, &r.StartedAt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return r, nil
|
|
}
|
|
|
|
func (s *TaskStore) ListRuns(ctx context.Context, taskID string, limit int) ([]models.TaskRun, error) {
|
|
rows, err := DB.QueryContext(ctx, `
|
|
SELECT id, task_id, channel_id, status, COALESCE(trigger_payload, ''),
|
|
started_at, completed_at, tokens_used, tool_calls, wall_clock,
|
|
COALESCE(error, '')
|
|
FROM task_runs WHERE task_id = $1
|
|
ORDER BY started_at DESC LIMIT $2`, taskID, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var runs []models.TaskRun
|
|
for rows.Next() {
|
|
var r models.TaskRun
|
|
if err := rows.Scan(&r.ID, &r.TaskID, &r.ChannelID, &r.Status, &r.TriggerPayload,
|
|
&r.StartedAt, &r.CompletedAt, &r.TokensUsed, &r.ToolCalls, &r.WallClock,
|
|
&r.Error); err != nil {
|
|
continue
|
|
}
|
|
runs = append(runs, r)
|
|
}
|
|
return runs, nil
|
|
}
|
|
|
|
// ── list helper (single source of truth for SELECT + Scan) ─────
|
|
|
|
func (s *TaskStore) list(ctx context.Context, where string, args ...interface{}) ([]models.Task, error) {
|
|
q := `SELECT ` + taskColumns + ` FROM tasks ` + where
|
|
|
|
rows, err := DB.QueryContext(ctx, q, args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var tasks []models.Task
|
|
for rows.Next() {
|
|
var t models.Task
|
|
if err := scanTask(rows, &t); err != nil {
|
|
continue
|
|
}
|
|
tasks = append(tasks, t)
|
|
}
|
|
return tasks, nil
|
|
}
|
|
|
|
// ── Helpers ─────────────────────────────────────
|
|
|
|
// comma builds ", column = $N"
|
|
func comma(i int, col string) string {
|
|
return ", " + col + " = " + pgArg(i)
|
|
}
|
|
|
|
// pgArg returns "$N"
|
|
func pgArg(i int) string {
|
|
return fmt.Sprintf("$%d", i)
|
|
}
|
|
|
|
// pgWhere returns " WHERE col = $N"
|
|
func pgWhere(i int, col string) string {
|
|
return " WHERE " + col + " = " + pgArg(i)
|
|
}
|
|
|
|
// nilIfEmpty returns nil for empty strings (prevents empty-string vs NULL issues).
|
|
func nilIfEmpty(s string) interface{} {
|
|
if s == "" {
|
|
return nil
|
|
}
|
|
return s
|
|
}
|