Unify storage ID queries under one implementation

This commit is contained in:
Ian Gulliver
2026-10-21 11:58:42 -07:00
parent 069474e6aa
commit abe4d7c9d5
6 changed files with 37 additions and 117 deletions
-19
View File
@@ -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()
}
+2 -34
View File
@@ -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 := ""
+30
View File
@@ -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()
}
+1 -16
View File
@@ -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 {
+3 -32
View File
@@ -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 {
+1 -16
View File
@@ -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 {