Fix deduplicate plan timeout: index + batch RFC822 duplicate-group lookup

internal/dedup/dedup.go

@@ -324,6 +324,17 @@ func (e *Engine) Scan(ctx context.Context) (*Report, error) {
 		return nil, fmt.Errorf("find duplicates: %w", err)
 	}
 
+	rfc822IDs := make([]string, len(storeGroups))
+	for i, sg := range storeGroups {
+		rfc822IDs[i] = sg.RFC822MessageID
+	}
+	msgsByGroup, err := e.store.GetDuplicateGroupMessagesBatchContext(
+		ctx, rfc822IDs, e.config.AccountSourceIDs...,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("get duplicate group messages: %w", err)
+	}
+
 	report := &Report{
 		TotalMessages:   totalMessages,
 		BySourcePair:    make(map[string]int),
@@ -334,15 +345,7 @@ func (e *Engine) Scan(ctx context.Context) (*Report, error) {
 		if ctx.Err() != nil {
 			return nil, ctx.Err()
 		}
-		msgs, err := e.store.GetDuplicateGroupMessages(
-			sg.RFC822MessageID, e.config.AccountSourceIDs...,
-		)
-		if err != nil {
-			return nil, fmt.Errorf(
-				"get group messages for %s: %w",
-				sg.RFC822MessageID, err,
-			)
-		}
+		msgs := msgsByGroup[sg.RFC822MessageID]
 		if len(msgs) < 2 {
 			continue
 		}
@@ -380,6 +383,11 @@ func (e *Engine) Scan(ctx context.Context) (*Report, error) {
 		report.Groups = append(report.Groups, group)
 		report.BySourcePair[sourcePairKey(group.Messages)]++
 	}
+	// Release the batch result map now that every group has been copied
+	// into report.Groups: on a large archive it can hold tens of
+	// thousands of rows, and the content-hash pass below can run long.
+	//nolint:ineffassign,wastedassign // deliberately drops the reference so the GC can reclaim it before the content-hash pass, not a leftover assignment
+	msgsByGroup = nil
 
 	if e.config.ContentHashFallback {
 		// Exclude only losers (messages already selected for pruning) from

internal/store/dedup.go

@@ -107,11 +107,14 @@ func (s *Store) FindDuplicatesByRFC822ID(sourceIDs ...int64) ([]DuplicateGroupKe
 	return groups, rows.Err()
 }
 
-func (s *Store) GetDuplicateGroupMessages(
-	rfc822ID string, sourceIDs ...int64,
-) ([]DuplicateMessageRow, error) {
-	query := `
-		SELECT m.id, m.source_id, s.source_type, s.identifier,
+// duplicateGroupMessageColumns is the SELECT column list shared by
+// GetDuplicateGroupMessages and GetDuplicateGroupMessagesBatch: the message
+// metadata, the two correlated subqueries (label count, from address), and
+// the EXISTS clause used to detect the Gmail SENT label. Both methods build
+// their SELECT clause from this constant so the two queries can't drift
+// apart; GetDuplicateGroupMessagesBatch prepends m.rfc822_message_id (needed
+// to key its result map) since it has no per-call rfc822ID to bind.
+const duplicateGroupMessageColumns = `m.id, m.source_id, s.source_type, s.identifier,
 		       m.source_message_id,
 		       COALESCE(m.subject, ''), m.sent_at, m.archived_at,
 		       (CASE WHEN mr.message_id IS NOT NULL THEN 1 ELSE 0 END) AS has_raw,
@@ -132,7 +135,19 @@ func (s *Store) GetDuplicateGroupMessages(
 		           WHERE mr_from.message_id = m.id
 		             AND mr_from.recipient_type = 'from'
 		           LIMIT 1
-		       ), '') AS from_email
+		       ), '') AS from_email`
+
+// GetDuplicateGroupMessages fetches every message row for a single RFC822
+// duplicate group in one query. It is retained as the reference
+// implementation that the equivalence test in dedup_test.go checks
+// GetDuplicateGroupMessagesBatch against, and it is still exercised by its
+// own pre-existing direct tests, even though Engine.Scan now calls
+// GetDuplicateGroupMessagesBatch instead.
+func (s *Store) GetDuplicateGroupMessages(
+	rfc822ID string, sourceIDs ...int64,
+) ([]DuplicateMessageRow, error) {
+	query := `
+		SELECT ` + duplicateGroupMessageColumns + `
 		FROM messages m
 		JOIN sources s ON s.id = m.source_id
 		LEFT JOIN message_raw mr ON mr.message_id = m.id
@@ -181,6 +196,89 @@ func (s *Store) GetDuplicateGroupMessages(
 	return msgs, rows.Err()
 }
 
+// GetDuplicateGroupMessagesBatch is the batched form of
+// GetDuplicateGroupMessages: it fetches every message row for many RFC822
+// duplicate groups in a handful of chunked queries instead of one query per
+// group (see kenn-io/msgvault#510 — 22,025 groups meant 22,025 unindexed
+// queries). Returns a map keyed by RFC822 message ID; each value preserves
+// the same per-group id-ascending order as GetDuplicateGroupMessages.
+func (s *Store) GetDuplicateGroupMessagesBatch(
+	rfc822IDs []string, sourceIDs ...int64,
+) (map[string][]DuplicateMessageRow, error) {
+	return s.GetDuplicateGroupMessagesBatchContext(
+		context.Background(), rfc822IDs, sourceIDs...,
+	)
+}
+
+// GetDuplicateGroupMessagesBatchContext is the request-aware form of
+// GetDuplicateGroupMessagesBatch.
+func (s *Store) GetDuplicateGroupMessagesBatchContext(
+	ctx context.Context, rfc822IDs []string, sourceIDs ...int64,
+) (map[string][]DuplicateMessageRow, error) {
+	result := make(map[string][]DuplicateMessageRow)
+	if len(rfc822IDs) == 0 {
+		return result, nil
+	}
+
+	const selectCols = "\n\t\tm.rfc822_message_id, " + duplicateGroupMessageColumns
+
+	// queryInChunks binds prefixArgs BEFORE the chunked %s placeholder, so
+	// the source_id filter (prefixArgs) must appear textually before the
+	// rfc822_message_id IN (%s) clause below — the reverse of
+	// GetDuplicateGroupMessages's clause order, which binds rfc822ID first
+	// via a plain "=" and has no ordering constraint to satisfy.
+	var prefixArgs []any
+	sourceFilter := ""
+	if len(sourceIDs) > 0 {
+		placeholders := make([]string, len(sourceIDs))
+		for i, id := range sourceIDs {
+			placeholders[i] = "?"
+			prefixArgs = append(prefixArgs, id)
+		}
+		sourceFilter = "m.source_id IN (" + strings.Join(placeholders, ",") + ") AND "
+	}
+
+	queryTemplate := `
+		SELECT` + selectCols + `
+		FROM messages m
+		JOIN sources s ON s.id = m.source_id
+		LEFT JOIN message_raw mr ON mr.message_id = m.id
+		WHERE ` + sourceFilter + `m.rfc822_message_id IN (%s)
+		  AND ` + LiveMessagesWhere("m", true) + `
+		ORDER BY m.rfc822_message_id, m.id`
+
+	err := queryInChunksContext(ctx, s.db, rfc822IDs, prefixArgs, queryTemplate,
+		func(rows *loggedRows) error {
+			var dm DuplicateMessageRow
+			var rfc822ID string
+			var sentAt, archivedAt sql.NullTime
+			var hasRaw, isFromMe, hasSent int
+			if err := rows.Scan(
+				&rfc822ID, &dm.ID, &dm.SourceID, &dm.SourceType, &dm.SourceIdentifier,
+				&dm.SourceMessageID, &dm.Subject, &sentAt, &archivedAt,
+				&hasRaw, &dm.LabelCount, &isFromMe, &hasSent,
+				&dm.FromEmail,
+			); err != nil {
+				return err
+			}
+			if sentAt.Valid {
+				dm.SentAt = sentAt.Time
+			}
+			if archivedAt.Valid {
+				dm.ArchivedAt = archivedAt.Time
+			}
+			dm.HasRawMIME = hasRaw == 1
+			dm.IsFromMe = isFromMe == 1
+			dm.HasSentLabel = hasSent == 1
+			result[rfc822ID] = append(result[rfc822ID], dm)
+			return nil
+		})
+	if err != nil {
+		return nil, fmt.Errorf("get duplicate group messages batch: %w", err)
+	}
+	return result, nil
+}
+
 func (s *Store) MergeDuplicates(
 	survivorID int64, duplicateIDs []int64, batchID string,
 ) (*MergeResult, error) {

internal/store/dedup_index_test.go

@@ -0,0 +1,46 @@
+package store
+
+import (
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// TestGetDuplicateGroupMessages_UsesRFC822Index verifies InitSchema creates
+// idx_messages_rfc822_message_id and that the dedup per-group lookup query
+// plans as an index search, not a full table scan. Locks in the fix for
+// kenn-io/msgvault#510: 22,025 unindexed lookups (~190ms each) burned the
+// CLI's 30-minute plan-request timeout before content-hash comparison
+// started.
+func TestGetDuplicateGroupMessages_UsesRFC822Index(t *testing.T) {
+	require := require.New(t)
+	assert := assert.New(t)
+
+	dir := t.TempDir()
+	s, err := OpenForTest(filepath.Join(dir, "dedup_index.db"))
+	require.NoError(err)
+	defer func() { _ = s.Close() }()
+	require.NoError(s.InitSchema())
+
+	var idxCount int
+	require.NoError(s.db.QueryRow(
+		`SELECT COUNT(*) FROM sqlite_master WHERE type='index' AND name='idx_messages_rfc822_message_id'`,
+	).Scan(&idxCount))
+	assert.Equal(1, idxCount, "idx_messages_rfc822_message_id should be created by InitSchema")
+
+	plan := explainPlan(t, s,
+		`SELECT id FROM messages WHERE rfc822_message_id = ?`, "some-id")
+	assert.Contains(plan, "idx_messages_rfc822_message_id",
+		"lookup by rfc822_message_id should use the index, not a full scan:\n%s", plan)
+	assert.NotContains(plan, "SCAN messages",
+		"lookup should not do a full table scan:\n%s", plan)
+
+	batchPlan := explainPlan(t, s,
+		`SELECT id FROM messages WHERE rfc822_message_id IN (?,?)`, "a", "b")
+	assert.Contains(batchPlan, "idx_messages_rfc822_message_id",
+		"batched IN lookup should also use the index:\n%s", batchPlan)
+	assert.NotContains(batchPlan, "SCAN messages",
+		"batched lookup should not do a full table scan:\n%s", batchPlan)
+}

internal/store/dedup_test.go

@@ -1,6 +1,7 @@
 package store_test
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"testing"
@@ -332,3 +333,72 @@ func TestStore_GetAllRawMIMECandidates_PreservesFromCase(t *testing.T) {
 	require.NotNil(got, "test message %d not in candidates: %+v", id, cands)
 	assert.Equal(t, mxid, got.FromEmail, "FromEmail (case must be preserved)")
 }
+
+func TestStore_GetDuplicateGroupMessagesBatch_MatchesPerGroupQuery(t *testing.T) {
+	require := require.New(t)
+	assert := assert.New(t)
+	f := storetest.New(t)
+
+	// 600 groups exceeds queryInChunks's chunk size of 500, so this
+	// exercises at least two chunk rounds and would catch a bug where the
+	// result map is reset per chunk instead of accumulated across chunks.
+	const numGroups = 600
+	rfc822IDs := make([]string, numGroups)
+	for i := range numGroups {
+		rfc822ID := fmt.Sprintf("rfc822-batch-%d", i)
+		rfc822IDs[i] = rfc822ID
+		newRFC822Message(t, f, fmt.Sprintf("src-%d-a", i), rfc822ID)
+		newRFC822Message(t, f, fmt.Sprintf("src-%d-b", i), rfc822ID)
+	}
+
+	batched, err := f.Store.GetDuplicateGroupMessagesBatch(rfc822IDs, f.Source.ID)
+	require.NoError(err, "GetDuplicateGroupMessagesBatch")
+	require.Len(batched, numGroups, "batched group count")
+
+	for _, rfc822ID := range rfc822IDs {
+		want, err := f.Store.GetDuplicateGroupMessages(rfc822ID, f.Source.ID)
+		require.NoError(err, "GetDuplicateGroupMessages reference for %s", rfc822ID)
+		assert.Equal(want, batched[rfc822ID], "mismatch for group %s", rfc822ID)
+	}
+}
+
+func TestStore_GetDuplicateGroupMessagesBatch_EmptyInput(t *testing.T) {
+	require := require.New(t)
+	assert := assert.New(t)
+	f := storetest.New(t)
+
+	batched, err := f.Store.GetDuplicateGroupMessagesBatch(nil)
+	require.NoError(err, "GetDuplicateGroupMessagesBatch with nil input")
+	assert.Empty(batched, "no groups requested, no groups returned")
+}
+
+func TestStore_GetDuplicateGroupMessagesBatchContext_Canceled(t *testing.T) {
+	require := require.New(t)
+	f := storetest.New(t)
+	newRFC822Message(t, f, "src-a", "rfc822-canceled")
+	newRFC822Message(t, f, "src-b", "rfc822-canceled")
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	_, err := f.Store.GetDuplicateGroupMessagesBatchContext(
+		ctx, []string{"rfc822-canceled"}, f.Source.ID,
+	)
+	require.ErrorIs(err, context.Canceled)
+}
+
+func TestStore_GetDuplicateGroupMessagesBatch_FiltersBySourceID(t *testing.T) {
+	require := require.New(t)
+	assert := assert.New(t)
+	f := storetest.New(t)
+	idA := newRFC822Message(t, f, "src-a", "rfc822-scoped")
+	idB := newRFC822Message(t, f, "src-b", "rfc822-scoped")
+
+	scoped, err := f.Store.GetDuplicateGroupMessagesBatch(
+		[]string{"rfc822-scoped"}, f.Source.ID,
+	)
+	require.NoError(err, "GetDuplicateGroupMessagesBatch scoped to source")
+	require.Len(scoped["rfc822-scoped"], 2)
+	assert.Equal(idA, scoped["rfc822-scoped"][0].ID)
+	assert.Equal(idB, scoped["rfc822-scoped"][1].ID)
+}

internal/store/store.go

@@ -603,12 +603,21 @@ func (s *Store) runMaintenance(ctx context.Context, fn func(ctx context.Context,
 // so streaming-query timing reflects scan-close, not just prepare.
 type chunkQuerier interface {
 	Query(query string, args ...any) (*loggedRows, error)
+	QueryContext(ctx context.Context, query string, args ...any) (*loggedRows, error)
 	Exec(query string, args ...any) (sql.Result, error)
 }
 
 func queryInChunks[T any](db chunkQuerier, ids []T, prefixArgs []any, queryTemplate string, fn func(*loggedRows) error) error {
+	return queryInChunksContext(context.Background(), db, ids, prefixArgs, queryTemplate, fn)
+}
+
+func queryInChunksContext[T any](ctx context.Context, db chunkQuerier, ids []T, prefixArgs []any, queryTemplate string, fn func(*loggedRows) error) error {
 	const chunkSize = 500
 	for i := 0; i < len(ids); i += chunkSize {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+
 		end := min(i+chunkSize, len(ids))
 		chunk := ids[i:end]
 
@@ -620,7 +629,7 @@ func queryInChunks[T any](db chunkQuerier, ids []T, prefixArgs []any, queryTempl
 		}
 
 		query := fmt.Sprintf(queryTemplate, strings.Join(placeholders, ","))
-		rows, err := db.Query(query, args...)
+		rows, err := db.QueryContext(ctx, query, args...)
 		if err != nil {
 			return err
 		}
@@ -940,6 +949,30 @@ func (s *Store) InitSchema() error {
 		}
 	}
 
+	// Index over rfc822_message_id serves dedup's per-group message lookup
+	// (GetDuplicateGroupMessages / GetDuplicateGroupMessagesBatch). Without
+	// it, each lookup was a full scan of the messages table — measured at
+	// ~190ms/lookup, with one lookup per duplicate group, so a scan with
+	// 22k groups burned the entire 30-minute CLI plan-request timeout
+	// before content-hash comparison even started (kenn-io/msgvault#510).
+	// Plain (non-partial) index: a partial WHERE rfc822_message_id IS NOT
+	// NULL AND != '' form is not usable by the planner for this table's
+	// bound `= ?` / `IN (...)` lookups — SQLite can't prove col = ? implies
+	// col != '' since ? could bind to '' — so it would silently fall back
+	// to SCAN (verified via EXPLAIN QUERY PLAN before writing this).
+	// Identical DDL on both backends; runMaintenance already handles the
+	// PostgreSQL statement_timeout exemption internally (finding S1). IF
+	// NOT EXISTS is idempotent per start.
+	if err := s.runMaintenance(context.Background(), func(ctx context.Context, tx *loggedTx) error {
+		_, err := tx.ExecContext(ctx, `
+			CREATE INDEX IF NOT EXISTS idx_messages_rfc822_message_id
+			    ON messages(rfc822_message_id)
+		`)
+		return err
+	}); err != nil {
+		return fmt.Errorf("create rfc822 message id index: %w", err)
+	}
+
 	// Backfill last_modified for rows that predate the column. SQLite cannot
 	// ADD COLUMN with a non-constant default, so the legacy ADD COLUMN above
 	// leaves existing rows NULL; this one-shot UPDATE sets them to