Unify storage ID queries under one implementation

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()
-}

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...)
+	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
-	}
-
-	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()
 	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{}
+	ids, err := queryIDs(ctx, tx, `SELECT id FROM inbox_items WHERE task_id = ? AND state = ?`, taskID, string(api.Open))
-	rows, err := tx.QueryContext(ctx, `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 := ""

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()
+}

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 {

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`)
+		ids, err := queryIDs(ctx, tx, `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()
 		if err != nil {
 			return err
 		}
+
 		for _, id := range ids {
 			err = settleTaskState(ctx, tx, emit, id, "byre")
 			if err != nil {

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`)
+	return queryIDs(ctx, s.db, `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()
 }
 
 func (s *Store) ClearStage(ctx context.Context, sessionID string) error {