Perform the REINDEX and cleanup failed previous attempts, add unit tests (AI commit)

StephenHinck · zinic · commit 628620baa8f5 · 2026-05-13T07:54:56.000-07:00
diff --git a/README.md b/README.md
@@ -64,7 +64,7 @@ The PostgreSQL driver's schema bootstrap (`drivers/pg/query/sql/schema_up.sql`)
 |---------------|----------|----------------------------------------------------------------------------------------------------------|
 | `pg_trgm`     | yes      | GIN trigram indexes for `contains` / `starts with` / `ends with` lookups on graph entity properties.     |
 | `intarray`    | yes      | Extended integer array operations used when maintaining node kind arrays.                                |
-| `pgstattuple` | no       | Measures btree leaf density and fragmentation for the driver-managed index optimization assessment.      |
+| `pgstattuple` | no       | Measures btree leaf density and fragmentation for the driver-managed index optimization.                 |
 
 `pgstattuple` is treated as best-effort. Installing it typically requires a superuser role and on some managed
 Postgres deployments the contrib package is not exposed at all. Bootstrap wraps the install in an exception
@@ -75,3 +75,24 @@ optimization in such environments, have an administrator install the extension o
 ```sql
 create extension if not exists pgstattuple;
 ```
+
+### Index Optimization
+
+When `pgstattuple` is available the driver's `Optimize` method (exposed through the optional `graph.Optimizer`
+interface) performs the following on each invocation against btree indexes on partitions of the `node` and `edge`
+tables:
+
+1. Drops any `INVALID` indexes whose name matches the `_ccnew[N]` pattern, which Postgres leaves behind when a
+   prior `REINDEX CONCURRENTLY` was aborted. Cleanup failures are logged at `WARN` and never fatal; an orphan
+   wastes disk but does not block productive rebuilds.
+2. Measures every candidate index with `pgstatindex` and flags those whose average leaf density falls below
+   `60%` or whose leaf-page fragmentation reaches `40%`. Thresholds are calibrated against production samples
+   and a freshly rebuilt baseline (~73.8% density).
+3. Rebuilds each flagged index with `REINDEX INDEX CONCURRENTLY`, smallest first so that an early cancellation
+   still produces the maximum number of completed rebuilds. Per-index failures are logged at `WARN` and the
+   loop continues with the next candidate. Context cancellation aborts before the next candidate; an in-flight
+   `REINDEX CONCURRENTLY` runs to completion (interrupting it would leave a `_ccnew` artifact for the next pass
+   to reap).
+
+The pass is not bounded by wall-clock time or candidate count. Callers should serialize `Optimize` against
+their own scheduling loop.
diff --git a/drivers/pg/optimize.go b/drivers/pg/optimize.go
@@ -4,6 +4,10 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"sort"
+	"time"
+
+	"github.com/jackc/pgx/v5"
 )
 
 const (
@@ -43,6 +47,28 @@ const (
 		order by c.relname, i.relname`
 
 	sqlSelectIndexBloatMetrics = `select avg_leaf_density, leaf_fragmentation from pgstatindex($1::regclass)`
+
+	// sqlSelectOrphanedReindexArtifacts identifies INVALID btree indexes left
+	// behind by aborted REINDEX CONCURRENTLY runs on partitions of the node
+	// and edge tables. Postgres names the in-progress copy <original>_ccnew or
+	// <original>_ccnew<N>; on failure or cancellation those copies remain in
+	// pg_index with indisvalid = false until explicitly dropped.
+	sqlSelectOrphanedReindexArtifacts = `
+		select n.nspname as schema_name,
+		       i.relname as index_name
+		from pg_inherits inh
+		         join pg_class p on p.oid = inh.inhparent
+		         join pg_class c on c.oid = inh.inhrelid
+		         join pg_namespace n on n.oid = c.relnamespace
+		         join pg_index x on x.indrelid = c.oid
+		         join pg_class i on i.oid = x.indexrelid
+		         join pg_am a on a.oid = i.relam
+		where p.relname in ('node', 'edge')
+		  and p.relnamespace = n.oid
+		  and a.amname = 'btree'
+		  and x.indisvalid = false
+		  and i.relname ~ '_ccnew[0-9]*$'
+		order by n.nspname, i.relname`
 )
 
 // indexRow is a candidate index discovered by the listing query, prior to
@@ -64,10 +90,12 @@ type indexCandidate struct {
 	reason        string
 }
 
-// Optimize satisfies the graph.Optimizer interface. The current phase performs
-// a read-only assessment: it identifies btree indexes on partitions of the
-// node and edge tables whose leaf density or fragmentation cross the rebuild
-// thresholds and logs the candidates. No DDL is executed in this phase.
+// Optimize satisfies the graph.Optimizer interface. It performs a pre-flight
+// sweep of orphaned REINDEX CONCURRENTLY artifacts, identifies btree indexes
+// on partitions of the node and edge tables whose leaf density or
+// fragmentation cross the rebuild thresholds, and rebuilds each flagged index
+// with REINDEX INDEX CONCURRENTLY. Per-candidate failures are logged and do
+// not abort the pass; the loop honors ctx cancellation between candidates.
 func (s *Driver) Optimize(ctx context.Context) error {
 	if installed, err := s.pgstattupleInstalled(ctx); err != nil {
 		return fmt.Errorf("checking pgstattuple extension: %w", err)
@@ -76,6 +104,8 @@ func (s *Driver) Optimize(ctx context.Context) error {
 		return nil
 	}
 
+	s.cleanupOrphanedReindexArtifacts(ctx)
+
 	indexes, err := s.listGraphPartitionBtreeIndexes(ctx)
 	if err != nil {
 		return fmt.Errorf("listing graph partition btree indexes: %w", err)
@@ -111,6 +141,14 @@ func (s *Driver) Optimize(ctx context.Context) error {
 		"Index optimization assessment complete: %d candidate(s) totaling %d bytes",
 		len(candidates), totalBytes,
 	))
+
+	// Process smallest candidates first so that a mid-pass ctx cancellation
+	// still results in the maximum number of completed rebuilds.
+	sort.SliceStable(candidates, func(i, j int) bool {
+		return candidates[i].sizeBytes < candidates[j].sizeBytes
+	})
+
+	s.reindexCandidates(ctx, candidates)
 	return nil
 }
 
@@ -160,3 +198,99 @@ func (s *Driver) measureIndexBloat(ctx context.Context, indexOID uint32) (float6
 	}
 	return density, fragmentation, nil
 }
+
+// orphanedReindexArtifact is an INVALID btree index left behind by an aborted
+// REINDEX CONCURRENTLY run on a node or edge partition.
+type orphanedReindexArtifact struct {
+	schema string
+	name   string
+}
+
+// cleanupOrphanedReindexArtifacts scans for and drops INVALID _ccnew indexes
+// left behind by previously aborted REINDEX CONCURRENTLY runs. Failures are
+// logged at WARN and never fatal: an orphan wastes disk but does not block
+// productive rebuilds, so a stuck cleanup must not gate the rest of the pass.
+func (s *Driver) cleanupOrphanedReindexArtifacts(ctx context.Context) {
+	orphans, err := s.listOrphanedReindexArtifacts(ctx)
+	if err != nil {
+		slog.WarnContext(ctx, fmt.Sprintf("Index optimization cleanup: failed to scan for orphaned reindex artifacts; continuing: %v", err))
+		return
+	}
+	if len(orphans) == 0 {
+		return
+	}
+
+	slog.InfoContext(ctx, fmt.Sprintf("Index optimization cleanup: dropping %d orphaned reindex artifact(s)", len(orphans)))
+	for _, o := range orphans {
+		if _, err := s.pool.Exec(ctx, buildDropInvalidIndexSQL(o.schema, o.name)); err != nil {
+			slog.WarnContext(ctx, fmt.Sprintf("Index optimization cleanup: failed to drop orphaned reindex artifact %s.%s; continuing: %v", o.schema, o.name, err))
+			continue
+		}
+		slog.InfoContext(ctx, fmt.Sprintf("Index optimization cleanup: dropped orphaned reindex artifact %s.%s", o.schema, o.name))
+	}
+}
+
+func (s *Driver) listOrphanedReindexArtifacts(ctx context.Context) ([]orphanedReindexArtifact, error) {
+	rows, err := s.pool.Query(ctx, sqlSelectOrphanedReindexArtifacts)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var out []orphanedReindexArtifact
+	for rows.Next() {
+		var o orphanedReindexArtifact
+		if err := rows.Scan(&o.schema, &o.name); err != nil {
+			return nil, err
+		}
+		out = append(out, o)
+	}
+	return out, rows.Err()
+}
+
+// reindexCandidates rebuilds each flagged index with REINDEX INDEX
+// CONCURRENTLY, in the order supplied (Optimize sorts ascending by size).
+// Per-candidate failures are logged at WARN and the loop continues; ctx
+// cancellation aborts further candidates but in-flight REINDEX statements
+// must run to completion in Postgres to avoid leaving _ccnew artifacts.
+func (s *Driver) reindexCandidates(ctx context.Context, candidates []indexCandidate) {
+	for _, c := range candidates {
+		if err := ctx.Err(); err != nil {
+			slog.WarnContext(ctx, fmt.Sprintf("Index optimization rebuild cancelled before processing %s.%s: %v", c.schema, c.index, err))
+			return
+		}
+
+		slog.InfoContext(ctx, fmt.Sprintf(
+			"Index optimization rebuild starting: %s.%s on %s (size=%d bytes, leaf_density=%.1f%%, fragmentation=%.1f%%)",
+			c.schema, c.index, c.table, c.sizeBytes, c.leafDensity, c.fragmentation,
+		))
+
+		started := time.Now()
+		if _, err := s.pool.Exec(ctx, buildReindexSQL(c.schema, c.index)); err != nil {
+			slog.WarnContext(ctx, fmt.Sprintf(
+				"Index optimization rebuild failed for %s.%s after %s; continuing with next candidate: %v",
+				c.schema, c.index, time.Since(started), err,
+			))
+			continue
+		}
+
+		slog.InfoContext(ctx, fmt.Sprintf(
+			"Index optimization rebuild complete: %s.%s in %s",
+			c.schema, c.index, time.Since(started),
+		))
+	}
+}
+
+// buildReindexSQL composes a REINDEX INDEX CONCURRENTLY statement with a
+// safely quoted schema-qualified identifier. REINDEX does not accept query
+// parameters; quoting is the only defense against malformed identifiers.
+func buildReindexSQL(schema, name string) string {
+	return "reindex index concurrently " + pgx.Identifier{schema, name}.Sanitize()
+}
+
+// buildDropInvalidIndexSQL composes a DROP INDEX CONCURRENTLY IF EXISTS
+// statement for cleaning up an orphaned _ccnew artifact. IF EXISTS prevents
+// a race where another session has already dropped the same orphan.
+func buildDropInvalidIndexSQL(schema, name string) string {
+	return "drop index concurrently if exists " + pgx.Identifier{schema, name}.Sanitize()
+}
diff --git a/drivers/pg/optimize_test.go b/drivers/pg/optimize_test.go
@@ -14,7 +14,8 @@ var _ graph.Optimizer = (*Driver)(nil)
 // TestNeedsReindex exercises the threshold logic that decides whether a
 // measured index is flagged as a rebuild candidate. The function is pure;
 // integration coverage of the surrounding pg_extension / pg_inherits /
-// pgstatindex queries lands with the Phase 5 REINDEX work.
+// pgstatindex queries and the REINDEX execution itself is exercised under
+// make test_integration against a real Postgres backend.
 func TestNeedsReindex(t *testing.T) {
 	cases := []struct {
 		name          string
@@ -111,3 +112,66 @@ func TestThresholdsAreOrdered(t *testing.T) {
 	assert.Greater(t, highIndexFragmentationThreshold, 0.0, "fragmentation threshold must be positive")
 	assert.Less(t, highIndexFragmentationThreshold, 100.0, "fragmentation threshold must be a percentage")
 }
+
+// TestBuildReindexSQL verifies the REINDEX statement is produced with both
+// schema and index name quoted via pgx.Identifier.Sanitize, including the
+// edge case where an identifier itself contains a double quote.
+func TestBuildReindexSQL(t *testing.T) {
+	cases := []struct {
+		name   string
+		schema string
+		index  string
+		want   string
+	}{
+		{
+			name:   "ordinary identifiers are double-quoted",
+			schema: "graph",
+			index:  "edge_1_pkey",
+			want:   `reindex index concurrently "graph"."edge_1_pkey"`,
+		},
+		{
+			name:   "mixed-case identifiers preserve case under quoting",
+			schema: "Graph",
+			index:  "Edge_1_PKey",
+			want:   `reindex index concurrently "Graph"."Edge_1_PKey"`,
+		},
+		{
+			name:   "embedded double quote is escaped by doubling",
+			schema: `evil"schema`,
+			index:  "edge_1_pkey",
+			want:   `reindex index concurrently "evil""schema"."edge_1_pkey"`,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.want, buildReindexSQL(tc.schema, tc.index))
+		})
+	}
+}
+
+// TestBuildDropInvalidIndexSQL verifies the orphan cleanup statement is
+// produced with IF EXISTS and CONCURRENTLY guarding identifier handling.
+func TestBuildDropInvalidIndexSQL(t *testing.T) {
+	got := buildDropInvalidIndexSQL("graph", "edge_1_pkey_ccnew")
+	assert.Equal(t, `drop index concurrently if exists "graph"."edge_1_pkey_ccnew"`, got)
+
+	// Identifier sanitization must apply to the orphan name as well, since
+	// _ccnew suffixes are read directly from pg_class.relname.
+	got = buildDropInvalidIndexSQL("graph", `weird"_ccnew`)
+	assert.Equal(t, `drop index concurrently if exists "graph"."weird""_ccnew"`, got)
+}
+
+// TestOrphanedReindexArtifactQuery_FiltersByValidityAndNameAndAm guards the
+// SQL string that scans for cleanup candidates against accidental loosening
+// of its filters, since this query controls the blast radius of DROP INDEX.
+func TestOrphanedReindexArtifactQuery_FiltersByValidityAndNameAndAm(t *testing.T) {
+	for _, fragment := range []string{
+		"x.indisvalid = false",
+		"a.amname = 'btree'",
+		"i.relname ~ '_ccnew[0-9]*$'",
+		"p.relname in ('node', 'edge')",
+	} {
+		assert.Contains(t, sqlSelectOrphanedReindexArtifacts, fragment,
+			"orphan-cleanup SQL is missing required filter %q; relaxing this filter risks dropping unrelated indexes", fragment)
+	}
+}
