From abe4d7c9d50be84ce40503a268664acf661a0de0 Mon Sep 17 00:00:00 2001 From: Ian Gulliver Date: Wed, 21 Oct 2026 11:58:42 -0700 Subject: [PATCH] Unify storage ID queries under one implementation --- internal/store/coordinate.go | 19 ------------------- internal/store/inbox.go | 36 ++---------------------------------- internal/store/query.go | 30 ++++++++++++++++++++++++++++++ internal/store/review.go | 17 +---------------- internal/store/store.go | 35 +++-------------------------------- internal/store/workflow.go | 17 +---------------- 6 files changed, 37 insertions(+), 117 deletions(-) create mode 100644 internal/store/query.go diff --git a/internal/store/coordinate.go b/internal/store/coordinate.go index 90dc6ad..ecdb361 100644 --- a/internal/store/coordinate.go +++ b/internal/store/coordinate.go @@ -48,22 +48,3 @@ func (s *Store) Coordinate(ctx context.Context) (Work, error) { }) return work, err } - -func queryIDs(ctx context.Context, tx *sql.Tx, query string) ([]string, error) { - rows, err := tx.QueryContext(ctx, query) - if err != nil { - return nil, err - } - defer rows.Close() - - ids := []string{} - for rows.Next() { - id := "" - err = rows.Scan(&id) - if err != nil { - return nil, err - } - ids = append(ids, id) - } - return ids, rows.Err() -} diff --git a/internal/store/inbox.go b/internal/store/inbox.go index 4c7912a..891b188 100644 --- a/internal/store/inbox.go +++ b/internal/store/inbox.go @@ -95,23 +95,7 @@ func (s *Store) ListInbox(ctx context.Context, f InboxFilter) ([]api.InboxItem, args = append(args, f.SessionID) } - rows, err := s.db.QueryContext(ctx, `SELECT id FROM inbox_items WHERE `+strings.Join(where, " AND ")+` ORDER BY created_at`, args...) - if err != nil { - return nil, err - } - - ids := []string{} - for rows.Next() { - id := "" - err = rows.Scan(&id) - if err != nil { - rows.Close() - return nil, err - } - ids = append(ids, id) - } - rows.Close() - err = rows.Err() + ids, err := queryIDs(ctx, s.db, `SELECT id FROM inbox_items WHERE `+strings.Join(where, " AND ")+` ORDER BY created_at`, args...) if err != nil { return nil, err } @@ -230,21 +214,10 @@ func (s *Store) DismissItem(ctx context.Context, id, actor string) error { } func dismissTaskInbox(ctx context.Context, tx *sql.Tx, emit emitFunc, taskID, actor string) error { - ids := []string{} - rows, err := tx.QueryContext(ctx, `SELECT id FROM inbox_items WHERE task_id = ? AND state = ?`, taskID, string(api.Open)) + ids, err := queryIDs(ctx, tx, `SELECT id FROM inbox_items WHERE task_id = ? AND state = ?`, taskID, string(api.Open)) if err != nil { return err } - for rows.Next() { - id := "" - err = rows.Scan(&id) - if err != nil { - rows.Close() - return err - } - ids = append(ids, id) - } - rows.Close() for _, id := range ids { _, err = tx.ExecContext(ctx, `UPDATE inbox_items SET state = ?, resolved_by = ?, resolved_at = ? WHERE id = ?`, @@ -346,11 +319,6 @@ func (s *Store) MarkUndelivered(ctx context.Context, responseIDs []string) error }) } -type querier interface { - QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row - QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) -} - func getItem(ctx context.Context, q querier, id string) (api.InboxItem, error) { item := api.InboxItem{} payload := "" diff --git a/internal/store/query.go b/internal/store/query.go new file mode 100644 index 0000000..c28a8d6 --- /dev/null +++ b/internal/store/query.go @@ -0,0 +1,30 @@ +package store + +import ( + "context" + "database/sql" +) + +type querier interface { + QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row + QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) +} + +func queryIDs(ctx context.Context, q querier, query string, args ...any) ([]string, error) { + rows, err := q.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + ids := []string{} + for rows.Next() { + id := "" + err = rows.Scan(&id) + if err != nil { + return nil, err + } + ids = append(ids, id) + } + return ids, rows.Err() +} diff --git a/internal/store/review.go b/internal/store/review.go index b4a868c..28afb92 100644 --- a/internal/store/review.go +++ b/internal/store/review.go @@ -78,26 +78,11 @@ func (s *Store) RunningReviews(ctx context.Context) ([]api.Review, error) { } func (s *Store) LatestDoneStages(ctx context.Context) ([]string, error) { - rows, err := s.db.QueryContext(ctx, ` + return queryIDs(ctx, s.db, ` SELECT r.id FROM reviews r JOIN tasks t ON t.id = r.task_id WHERE t.status = ? AND r.state = ? AND r.definition_id IS NOT NULL AND r.snapshot_n = (SELECT max(snapshot_n) FROM reviews WHERE task_id = r.task_id) ORDER BY r.created_at`, string(api.TaskStatusActive), string(api.ReviewStateDone)) - if err != nil { - return nil, err - } - defer rows.Close() - - ids := []string{} - for rows.Next() { - id := "" - err = rows.Scan(&id) - if err != nil { - return nil, err - } - ids = append(ids, id) - } - return ids, rows.Err() } func (s *Store) FinishReview(ctx context.Context, id string, state api.ReviewState, result map[string]any, errText string) error { diff --git a/internal/store/store.go b/internal/store/store.go index 441334d..6e6e7d4 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -706,24 +706,9 @@ func (s *Store) CancelMergeAfterReview(ctx context.Context, id, actor string) er } func (s *Store) FinishedMergeAfterReviews(ctx context.Context) ([]string, error) { - rows, err := s.db.QueryContext(ctx, ` + return queryIDs(ctx, s.db, ` SELECT tasks.merge_after_review_id FROM tasks JOIN reviews ON reviews.id = tasks.merge_after_review_id WHERE tasks.merge_after_review_by != '' AND reviews.state != ?`, string(api.ReviewStateRunning)) - if err != nil { - return nil, err - } - defer rows.Close() - - ids := []string{} - for rows.Next() { - id := "" - err = rows.Scan(&id) - if err != nil { - return nil, err - } - ids = append(ids, id) - } - return ids, rows.Err() } func (s *Store) TakeMergeAfterReview(ctx context.Context, id string, reviewIDs []string) (MergeAfterReview, error) { @@ -1679,25 +1664,11 @@ func (s *Store) EndReviewRequest(ctx context.Context, taskID string) error { func (s *Store) ClearReviewRequests(ctx context.Context) error { return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error { - rows, err := tx.QueryContext(ctx, `UPDATE tasks SET review_requests = 0 WHERE review_requests > 0 RETURNING id`) - if err != nil { - return err - } - ids := []string{} - for rows.Next() { - id := "" - err = rows.Scan(&id) - if err != nil { - rows.Close() - return err - } - ids = append(ids, id) - } - rows.Close() - err = rows.Err() + ids, err := queryIDs(ctx, tx, `UPDATE tasks SET review_requests = 0 WHERE review_requests > 0 RETURNING id`) if err != nil { return err } + for _, id := range ids { err = settleTaskState(ctx, tx, emit, id, "byre") if err != nil { diff --git a/internal/store/workflow.go b/internal/store/workflow.go index d236605..79a6145 100644 --- a/internal/store/workflow.go +++ b/internal/store/workflow.go @@ -257,22 +257,7 @@ func (s *Store) StageOutput(ctx context.Context, sessionID string) (string, erro } func (s *Store) FinishedStages(ctx context.Context) ([]string, error) { - rows, err := s.db.QueryContext(ctx, `SELECT id FROM agent_sessions WHERE stage_finished = 1 ORDER BY created_at`) - if err != nil { - return nil, err - } - defer rows.Close() - - ids := []string{} - for rows.Next() { - id := "" - err = rows.Scan(&id) - if err != nil { - return nil, err - } - ids = append(ids, id) - } - return ids, rows.Err() + return queryIDs(ctx, s.db, `SELECT id FROM agent_sessions WHERE stage_finished = 1 ORDER BY created_at`) } func (s *Store) ClearStage(ctx context.Context, sessionID string) error {