From 3411a308a50ad7e8f0c902d17c5938412b2269db Mon Sep 17 00:00:00 2001 From: Jean-Philippe Sirois Date: Fri, 14 Aug 2026 07:19:19 -0300 Subject: [PATCH] feat(remote): offer the rewrite when no index helps The analyzer derived index candidates and nothing else, so a query whose only real fix is a rewrite reported no_improvement_found and stopped there. That state means the index search came back empty, which is narrower than it sounds. deriveImprovements now derives the rewrites the query's shape allows, plans each one under the same statistics the original was planned under, and ranks both kinds on the one number a reader acts on. The improvements ride beside the state rather than changing it, so the CI gate still fails on what it failed on before. Rewrites are derived from the query the optimizer costed, not from the earlier analysis, which ran before the LIMIT substitution and the pg_stat_statements rewrite. The costing pass gets its own flat budget rather than the query's. Reusing options.timeoutMs would double how long one query can hold the single worker, and spending what the index search left over would starve the expensive queries a rewrite is the only fix for. onOptimizeReady now builds the optimization once. withOptimization mutates the query in place, so an event listener reads the object handed to the emit, and the second literal would have left the live relay without the improvements CI receives. Needs @query-doctor/core 0.29.0, which exports costRewrites and indexCandidates. Co-Authored-By: Claude --- package.json | 2 +- src/remote/optimization.ts | 4 +- src/remote/query-optimizer.ts | 202 +++++++++++++++--------- src/remote/rewrite-improvements.test.ts | 130 +++++++++++++++ src/reporters/site-api.ts | 4 +- 5 files changed, 260 insertions(+), 82 deletions(-) create mode 100644 src/remote/rewrite-improvements.test.ts diff --git a/package.json b/package.json index 89fb477e..e9959a3b 100644 --- a/package.json +++ b/package.json @@ -21,7 +21,7 @@ "@libpg-query/parser": "^17.6.3", "@opentelemetry/api": "^1.9.0", "@pgsql/types": "^17.6.2", - "@query-doctor/core": "^0.28.0", + "@query-doctor/core": "^0.29.0", "async-sema": "^3.1.1", "capnweb": "^0.7.0", "dedent": "^1.7.1", diff --git a/src/remote/optimization.ts b/src/remote/optimization.ts index 812a7c50..cd84c544 100644 --- a/src/remote/optimization.ts +++ b/src/remote/optimization.ts @@ -1,4 +1,4 @@ -import type { PostgresExplainStage } from "@query-doctor/core"; +import type { Improvement, PostgresExplainStage } from "@query-doctor/core"; import z from "zod"; const IndexRecommendation = z.object({ @@ -33,12 +33,14 @@ export const LiveQueryOptimization = z.discriminatedUnion("state", [ indexesUsed: z.array(z.string()), explainPlan: z.custom(), optimizedExplainPlan: z.custom(), + improvements: z.array(z.custom()).optional(), }), z.object({ state: z.literal("no_improvement_found"), cost: z.number(), indexesUsed: z.array(z.string()), explainPlan: z.custom(), + improvements: z.array(z.custom()).optional(), }), z.object({ state: z.literal("timeout"), diff --git a/src/remote/query-optimizer.ts b/src/remote/query-optimizer.ts index f6347fbc..a4b583e4 100644 --- a/src/remote/query-optimizer.ts +++ b/src/remote/query-optimizer.ts @@ -5,9 +5,14 @@ import { ConnectionManager } from "../sync/connection-manager.ts"; import { Sema } from "async-sema"; import { Analyzer, + costRewrites, + deriveRewrites, dropIndex, FullSchema, FullSchemaIndex, + type Improvement, + type ImprovementCandidate, + indexCandidates, IndexOptimizer, IndexRecommendation, OptimizeResult, @@ -17,6 +22,8 @@ import { PostgresQueryBuilder, PostgresTransaction, PostgresVersion, + rankImprovements, + type RewriteCandidate, Statistics, StatisticsMode, type ExportedStats, @@ -487,7 +494,29 @@ export class QueryOptimizer extends EventEmitter { } } - return this.onOptimizeReady(result, recent); + if (result.kind === "zero_cost_plan") { + return this.onZeroCostPlan(recent, result.explainPlan); + } + // A flat budget, not `options.timeoutMs`, which the retry ladder grows to + // 80s. Reusing it would let one query hold the single worker for twice + // that; spending only what the index search left over would starve the + // expensive queries a rewrite is the only fix for. Flat costs a first + // attempt 5s it did not spend before, and caps the worst case at the + // attempt's own budget plus 5s rather than twice the attempt's. + const indexRecommendations = mapIndexRecommandations(result); + const improvements = await deriveImprovements( + recent.query, + target.optimizer, + result, + indexRecommendations, + this.queryTimeoutMs, + ); + return this.onOptimizeReady( + result, + recent, + indexRecommendations, + improvements, + ); } private async dropDisabledIndexes(tx: PostgresTransaction): Promise { @@ -497,62 +526,47 @@ export class QueryOptimizer extends EventEmitter { } private onOptimizeReady( - result: OptimizeResult, + result: Extract, recent: OptimizedQuery, + indexRecommendations: IndexRecommendation[], + improvements: Improvement[], ): LiveQueryOptimization { - switch (result.kind) { - case "ok": { - const indexRecommendations = mapIndexRecommandations(result); - const indexesUsed = Array.from(result.existingIndexes); - const reduction = costReductionPercentage(result.baseCost, result.finalCost); - if (reduction < MINIMUM_COST_CHANGE_PERCENTAGE) { - this.onNoImprovements( - recent, - result.baseCost, - indexesUsed, - result.baseExplainPlan, - ); - return { - state: "no_improvement_found", - cost: result.baseCost, - indexesUsed, - explainPlan: result.baseExplainPlan, - }; - } else { - this.onImprovementsAvailable(recent, result, result.baseExplainPlan); - return { - state: "improvements_available", - cost: result.baseCost, - optimizedCost: result.finalCost, - costReductionPercentage: reduction, - indexRecommendations, - indexesUsed, - explainPlan: result.baseExplainPlan, - optimizedExplainPlan: result.explainPlan, - }; - } - } - // unlikely to hit if we've already checked the base plan for zero cost - case "zero_cost_plan": - return this.onZeroCostPlan(recent, result.explainPlan); + const indexesUsed = Array.from(result.existingIndexes); + const reduction = costReductionPercentage(result.baseCost, result.finalCost); + // Each branch builds its optimization once and both emits and returns it. + // `withOptimization` mutates the query in place, so a listener reads the + // object handed to the emit, and a second literal is a second answer. + if (reduction < MINIMUM_COST_CHANGE_PERCENTAGE) { + const optimization = { + state: "no_improvement_found", + cost: result.baseCost, + indexesUsed, + explainPlan: result.baseExplainPlan, + improvements, + } as const; + this.onNoImprovements(recent, optimization); + return optimization; } + const optimization = { + state: "improvements_available", + cost: result.baseCost, + optimizedCost: result.finalCost, + costReductionPercentage: reduction, + indexRecommendations, + indexesUsed, + explainPlan: result.baseExplainPlan, + optimizedExplainPlan: result.explainPlan, + improvements, + } as const; + this.onImprovementsAvailable(recent, optimization); + return optimization; } private onNoImprovements( recent: OptimizedQuery, - cost: number, - indexesUsed: string[], - explainPlan: PostgresExplainStage, + optimization: Extract, ) { - this.emit( - "noImprovements", - recent.withOptimization({ - state: "no_improvement_found", - cost, - indexesUsed, - explainPlan, - }), - ); + this.emit("noImprovements", recent.withOptimization(optimization)); } private getPotentialIndexCandidates( @@ -577,38 +591,14 @@ export class QueryOptimizer extends EventEmitter { private onImprovementsAvailable( recent: OptimizedQuery, - result: Extract, - explainPlan: PostgresExplainStage, + optimization: Extract< + LiveQueryOptimization, + { state: "improvements_available" } + >, ) { - const optimized = recent.withOptimization( - this.resultToImprovementsAvailable(result, explainPlan), - ); + const optimized = recent.withOptimization(optimization); this.emit("improvementsAvailable", optimized); - this.queries.set( - optimized.hash, - optimized, - ); - } - - private resultToImprovementsAvailable( - result: Extract, - explainPlan: PostgresExplainStage, - ): LiveQueryOptimization { - const indexesUsed = Array.from(result.existingIndexes); - const indexRecommendations = Array.from(result.newIndexes) - .map((n) => result.triedIndexes.get(n)) - .filter((n) => n !== undefined); - const reduction = costReductionPercentage(result.baseCost, result.finalCost); - return { - state: "improvements_available", - cost: result.baseCost, - optimizedCost: result.finalCost, - costReductionPercentage: reduction, - indexRecommendations, - indexesUsed, - explainPlan, - optimizedExplainPlan: result.explainPlan, - }; + this.queries.set(optimized.hash, optimized); } private onZeroCostPlan( @@ -646,6 +636,60 @@ export class QueryOptimizer extends EventEmitter { } } +/** + * Every proven way to make this query cheaper, ranked together: the index set + * the optimizer already costed, and each rewrite the query's shape allows, + * planned under the same statistics. Without the rewrites a query whose only + * real fix is one reports `no_improvement_found`, a state that has only ever + * meant the index search came back empty. + * + * Rewrites are derived from the query the optimizer costed, not from the + * analysis, which ran before the LIMIT substitution and the pg_stat_statements + * rewrite. Costing a rewrite of a different string compares two queries rather + * than two shapes. + */ +async function deriveImprovements( + query: string, + optimizer: IndexOptimizer, + result: Extract, + indexRecommendations: IndexRecommendation[], + timeoutMs: number, +): Promise { + const indexes = indexCandidates(indexRecommendations, result.finalCost); + const rewrites = await costCandidates(query, optimizer, timeoutMs); + return rankImprovements(result.baseCost, [...indexes, ...rewrites]); +} + +/** + * Most queries match no rule, so the deadline is armed only once there is + * something to plan. Deriving first also keeps a parse failure from reading as + * a costing failure. + */ +async function costCandidates( + query: string, + optimizer: IndexOptimizer, + timeoutMs: number, +): Promise { + try { + const candidates = await rewritesFor(query); + if (candidates.length === 0) return []; + return await withTimeout(costRewrites(candidates, optimizer), timeoutMs); + } catch (error) { + console.error("[query-optimizer] could not cost rewrites", error); + return []; + } +} + +/** + * The rewrites this query's own shape allows. A parse that yields no first + * statement yields nothing to rewrite. + */ +async function rewritesFor(query: string): Promise { + const ast = await parse(query); + const stmt = ast.stmts?.[0]?.stmt; + return stmt ? deriveRewrites(stmt) : []; +} + export class TimeoutError extends Error { constructor() { super("Timeout"); diff --git a/src/remote/rewrite-improvements.test.ts b/src/remote/rewrite-improvements.test.ts new file mode 100644 index 00000000..1af49fa7 --- /dev/null +++ b/src/remote/rewrite-improvements.test.ts @@ -0,0 +1,130 @@ +import { expect, test } from "vitest"; +import { assert, assertDefined } from "./test-utils.ts"; +import { PostgreSqlContainer } from "@testcontainers/postgresql"; +import { QueryOptimizer } from "./query-optimizer.ts"; +import { ConnectionManager } from "../sync/connection-manager.ts"; +import { Connectable } from "../sync/connectable.ts"; +import { + type OptimizedQuery, + QueryHash, + RecentQuery, +} from "../sql/recent-query.ts"; + +/** + * A correlated `EXISTS` that aggregates per outer row. `item_modifiers` already + * carries the only index the access path can use, so the index search comes back + * empty and the query settles on `no_improvement_found` — the state that used to + * end the analyzer's answer. + */ +const AGGREGATE_EXISTS_QUERY = + "select count(*) from items where exists (select 1 from item_modifiers im " + + "where im.item_id = items.id and im.stat = $1 having sum(im.value) >= $2);"; + +/** + * Real rows, because stock Postgres sizes a relation from its own page count and + * would price an empty fixture at nothing whatever statistics say. + */ +const SCHEMA = ` + create table items (id int primary key); + create table item_modifiers (item_id int not null, stat text not null, value int not null); + create index item_modifiers_item_id_stat_idx on item_modifiers (item_id, stat); + + insert into items (id) select g from generate_series(1, 3000) g; + insert into item_modifiers (item_id, stat, value) + select 1 + (g % 3000), (array['str','dex','vit'])[1 + (g % 3)], g % 40 + from generate_series(1, 9000) g; + vacuum analyze; +`; + +/** Through `analyze`, so the query costed is the one the pipeline would cost. */ +function recentQuery(query: string): Promise { + return RecentQuery.analyze( + { + calls: "1", + formattedQuery: query, + meanTime: 100, + query, + rows: "1", + topLevel: true, + username: "test", + }, + QueryHash.parse("aggregate-exists"), + QueryHash.parse("aggregate-exists"), + ); +} + +test("a query no index can fix still carries the rewrite that can", async () => { + const pg = await new PostgreSqlContainer("postgres:17") + .withCopyContentToContainer([ + { content: SCHEMA, target: "/docker-entrypoint-initdb.d/init.sql" }, + ]) + .withCommand(["-c", "autovacuum=off"]) + .start(); + + const manager = ConnectionManager.forLocalDatabase(); + const conn = Connectable.fromString(pg.getConnectionUri()); + const optimizer = new QueryOptimizer(manager, conn); + + const settled: OptimizedQuery[] = []; + optimizer.addListener("noImprovements", (query) => settled.push(query)); + optimizer.addListener("improvementsAvailable", (query) => settled.push(query)); + optimizer.addListener("error", (query, error) => { + console.error("error when running query", query); + throw error; + }); + + try { + await optimizer.start([], { + kind: "fromStatisticsExport", + source: { kind: "inline" }, + stats: [ + { + tableName: "items", + schemaName: "public", + relpages: 14, + reltuples: 3_000, + relallvisible: 14, + columns: [{ columnName: "id", stats: null, attlen: 4 }], + indexes: [], + }, + { + tableName: "item_modifiers", + schemaName: "public", + relpages: 49, + reltuples: 9_000, + relallvisible: 49, + columns: [ + { columnName: "item_id", stats: null, attlen: 4 }, + { columnName: "stat", stats: null, attlen: -1 }, + { columnName: "value", stats: null, attlen: 4 }, + ], + indexes: [{ + indexName: "item_modifiers_item_id_stat_idx", + relpages: 27, + reltuples: 9_000, + relallvisible: 27, + amname: "btree", + columns: [{ attlen: 4 }, { attlen: -1 }], + fillfactor: 0.9, + }], + }, + ], + }); + await optimizer.addQueries([await recentQuery(AGGREGATE_EXISTS_QUERY)]); + + expect(settled).toHaveLength(1); + const { optimization } = settled[0]; + assert(optimization.state === "no_improvement_found"); + + const [improvement] = optimization.improvements ?? []; + assertDefined(improvement); + assert(improvement.action.kind === "rewrite"); + expect(improvement.action.rule).toBe("HOIST_CORRELATED_EXISTS"); + // The saving is the reason the state is worth contradicting. Core owns what + // the rewrite says; this owns that a settled query carries it at all. + expect(improvement.costReductionPercentage).toBeGreaterThan(90); + } finally { + optimizer.stop(); + await pg.stop(); + } +}, 180_000); diff --git a/src/reporters/site-api.ts b/src/reporters/site-api.ts index f567171e..ae5303d7 100644 --- a/src/reporters/site-api.ts +++ b/src/reporters/site-api.ts @@ -3,7 +3,7 @@ import { gzip } from "node:zlib"; import { promisify } from "node:util"; import * as github from "@actions/github"; import { isTestOriginQuery } from "@query-doctor/core"; -import type { ComputedStats, FullSchema, IndexRecommendation, Nudge, SQLCommenterTag, StatisticsMode, TableReference } from "@query-doctor/core"; +import type { ComputedStats, FullSchema, Improvement, IndexRecommendation, Nudge, SQLCommenterTag, StatisticsMode, TableReference } from "@query-doctor/core"; import { DEFAULT_CONFIG, type AnalyzerConfig } from "../config.ts"; import { originsCompatible, shapeKey } from "./query-shape.ts"; import type { OptimizedQuery } from "../sql/recent-query.ts"; @@ -50,12 +50,14 @@ export type CiOptimization = indexesUsed: string[]; explainPlan?: object; optimizedExplainPlan?: object; + improvements?: Improvement[]; } | { state: "no_improvement_found"; cost: number; indexesUsed: string[]; explainPlan?: object; + improvements?: Improvement[]; } | { state: "error";