diff --git a/package.json b/package.json index 89fb477..e9959a3 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 812a7c5..cd84c54 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 f6347fb..a4b583e 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 0000000..1af49fa --- /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 f567171..ae5303d 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";