diff --git a/deno.json b/deno.json index 0483b8a..afc58f4 100644 --- a/deno.json +++ b/deno.json @@ -25,7 +25,7 @@ "@libpg-query/parser": "npm:@libpg-query/parser@^17.6.3", "@opentelemetry/api": "jsr:@opentelemetry/api@^1.9.0", "@pgsql/types": "npm:@pgsql/types@^17.6.1", - "@query-doctor/core": "npm:@query-doctor/core@^0.1.7", + "@query-doctor/core": "npm:@query-doctor/core@^0.1.9", "@rabbit-company/rate-limiter": "jsr:@rabbit-company/rate-limiter@^3.0.0", "@std/assert": "jsr:@std/assert@^1.0.14", "@std/collections": "jsr:@std/collections@^1.1.3", diff --git a/deno.lock b/deno.lock index 2ec4707..a65efc3 100644 --- a/deno.lock +++ b/deno.lock @@ -20,7 +20,7 @@ "npm:@actions/github@^6.0.1": "6.0.1_@octokit+core@5.2.2", "npm:@libpg-query/parser@^17.6.3": "17.6.3", "npm:@pgsql/types@^17.6.1": "17.6.1", - "npm:@query-doctor/core@~0.1.7": "0.1.7", + "npm:@query-doctor/core@~0.1.9": "0.1.10", "npm:@testcontainers/postgresql@^11.9.0": "11.9.0", "npm:@types/node@^24.9.1": "24.10.1", "npm:@types/nunjucks@^3.2.6": "3.2.6", @@ -350,8 +350,8 @@ "@protobufjs/utf8@1.1.0": { "integrity": "sha512-Vvn3zZrhQZkkBE8LSuW3em98c0FwgO4nxzv6OdSxPKJIEKY2bGbHn+mhGIPerzI4twdxaP8/0+06HBpwf345Lw==" }, - "@query-doctor/core@0.1.7": { - "integrity": "sha512-dIb3J5hGKevkIDyVgqrpjpGNT1C2NkZqfTVfnadYkbx1EMDaKoSRJO7I8aeJr3AkI5vSdC1TLpwpTMb5FQ5Mzg==", + "@query-doctor/core@0.1.10": { + "integrity": "sha512-nhzeNFUnVX0uBt4OwLcGNc3U6bDL0Ow9uOkUEzlJGWrDCbv+fAYND1+pD5agR5wRqXQRkAnFNX2YbfYW5tuJpQ==", "dependencies": [ "@pgsql/types", "colorette", @@ -1681,7 +1681,7 @@ "npm:@actions/github@^6.0.1", "npm:@libpg-query/parser@^17.6.3", "npm:@pgsql/types@^17.6.1", - "npm:@query-doctor/core@~0.1.7", + "npm:@query-doctor/core@~0.1.9", "npm:@testcontainers/postgresql@^11.9.0", "npm:@types/node@^24.9.1", "npm:@types/nunjucks@^3.2.6", diff --git a/src/remote/disabled-indexes.test.ts b/src/remote/disabled-indexes.test.ts new file mode 100644 index 0000000..0b1762f --- /dev/null +++ b/src/remote/disabled-indexes.test.ts @@ -0,0 +1,54 @@ +import { assertEquals } from "@std/assert/equals"; +import { DisabledIndexes } from "./disabled-indexes.ts"; +import { PgIdentifier } from "@query-doctor/core"; + +Deno.test("DisabledIndexes.add adds an index", () => { + const indexes = new DisabledIndexes(); + const indexName = PgIdentifier.fromString("my_index"); + indexes.add(indexName); + const result = [...indexes]; + assertEquals(result.length, 1); + assertEquals(result[0].toString(), "my_index"); +}); + +Deno.test("DisabledIndexes.remove removes an existing index", () => { + const indexes = new DisabledIndexes(); + const indexName = PgIdentifier.fromString("my_index"); + indexes.add(indexName); + const removed = indexes.remove(indexName); + assertEquals(removed, true); + assertEquals([...indexes].length, 0); +}); + +Deno.test("DisabledIndexes.remove returns false for non-existent index", () => { + const indexes = new DisabledIndexes(); + const indexName = PgIdentifier.fromString("my_index"); + const removed = indexes.remove(indexName); + assertEquals(removed, false); +}); + +Deno.test("DisabledIndexes.toggle disables an enabled index", () => { + const indexes = new DisabledIndexes(); + const indexName = PgIdentifier.fromString("my_index"); + const isDisabled = indexes.toggle(indexName); + assertEquals(isDisabled, true); + assertEquals([...indexes].length, 1); +}); + +Deno.test("DisabledIndexes.toggle enables a disabled index", () => { + const indexes = new DisabledIndexes(); + const indexName = PgIdentifier.fromString("my_index"); + indexes.add(indexName); + const isDisabled = indexes.toggle(indexName); + assertEquals(isDisabled, false); + assertEquals([...indexes].length, 0); +}); + +Deno.test("DisabledIndexes iterator returns PgIdentifier instances", () => { + const indexes = new DisabledIndexes(); + indexes.add(PgIdentifier.fromString("index_a")); + indexes.add(PgIdentifier.fromString("index_b")); + const result = [...indexes]; + assertEquals(result.length, 2); + assertEquals(result.map((i) => i.toString()).sort(), ["index_a", "index_b"]); +}); diff --git a/src/remote/disabled-indexes.ts b/src/remote/disabled-indexes.ts new file mode 100644 index 0000000..d234d72 --- /dev/null +++ b/src/remote/disabled-indexes.ts @@ -0,0 +1,40 @@ +import { PgIdentifier } from "@query-doctor/core"; + +/** + * A class representing disabled indexes for the + * sole purpose of exposing a {@link PgIdentifier} interface. + */ +export class DisabledIndexes { + private readonly disabledIndexNames = new Set(); + + add(indexName: PgIdentifier): void { + this.disabledIndexNames.add(indexName.toString()); + } + + remove(indexName: PgIdentifier): boolean { + return this.disabledIndexNames.delete(indexName.toString()); + } + + has(indexName: PgIdentifier): boolean { + return this.disabledIndexNames.has(indexName.toString()); + } + + /** + * Toggles the visibility of the index + * @returns is the boolean disabled? + */ + toggle(indexName: PgIdentifier): boolean { + const deleted = this.remove(indexName); + if (!deleted) { + this.add(indexName); + return true; + } + return false; + } + + [Symbol.iterator](): Iterator { + return this.disabledIndexNames.values().map((indexName) => + PgIdentifier.fromString(indexName) + ); + } +} diff --git a/src/remote/query-optimizer.test.ts b/src/remote/query-optimizer.test.ts index 59bbff1..630d761 100644 --- a/src/remote/query-optimizer.test.ts +++ b/src/remote/query-optimizer.test.ts @@ -182,6 +182,136 @@ Deno.test({ }, }); +Deno.test({ + name: "disabling an index removes it from indexesUsed and recommends it", + sanitizeOps: false, + sanitizeResources: false, + fn: async () => { + const pg = await new PostgreSqlContainer("postgres:17") + .withCopyContentToContainer([ + { + content: ` + create table users(id int, email text); + insert into users (id, email) select i, 'user' || i || '@example.com' from generate_series(1, 1000) i; + create index "users_email_idx" on users(email); + create extension pg_stat_statements; + select * from users where email = 'test@example.com'; + `, + target: "/docker-entrypoint-initdb.d/init.sql", + }, + ]) + .withCommand([ + "-c", + "shared_preload_libraries=pg_stat_statements", + "-c", + "autovacuum=off", + "-c", + "track_counts=off", + "-c", + "track_io_timing=off", + "-c", + "track_activities=off", + ]) + .start(); + + const manager = ConnectionManager.forLocalDatabase(); + + const conn = Connectable.fromString(pg.getConnectionUri()); + const optimizer = new QueryOptimizer(manager, conn); + const connector = manager.getConnectorFor(conn); + + const statsMode = { + kind: "fromStatisticsExport" as const, + source: { kind: "inline" as const }, + stats: [{ + tableName: "users", + schemaName: "public", + relpages: 10000, + reltuples: 10_000_000, + relallvisible: 1, + columns: [ + { columnName: "id", stats: null }, + { columnName: "email", stats: null }, + ], + indexes: [{ + indexName: "users_email_idx", + relpages: 100, + reltuples: 10_000_000, + relallvisible: 1, + }], + }], + }; + + try { + const recentQueries = await connector.getRecentQueries(); + const emailQuery = recentQueries.find((q) => + q.query.includes("email") && q.query.includes("users") + ); + assert(emailQuery, "Expected to find email query in recent queries"); + + await optimizer.start([emailQuery], statsMode); + await optimizer.finish; + + const queriesAfterFirstRun = optimizer.getQueries(); + const emailQueryResult = queriesAfterFirstRun.find((q) => + q.query.includes("email") + ); + assert(emailQueryResult, "Expected email query in results"); + assert( + emailQueryResult.optimization.state === "no_improvement_found", + `Expected no_improvement_found but got ${emailQueryResult.optimization.state}`, + ); + assertArrayIncludes( + emailQueryResult.optimization.indexesUsed, + ["users_email_idx"], + ); + const costWithIndex = emailQueryResult.optimization.cost; + + const { PgIdentifier } = await import("@query-doctor/core"); + optimizer.toggleIndex(PgIdentifier.fromString("users_email_idx")); + const disabledIndexes = optimizer.getDisabledIndexes(); + assert( + disabledIndexes.some((i) => i.toString() === "users_email_idx"), + `Expected users_email_idx to be disabled`, + ); + + await optimizer.addQueries([emailQuery]); + await optimizer.finish; + + const queriesAfterToggle = optimizer.getQueries(); + const emailQueryAfterToggle = queriesAfterToggle.find((q) => + q.query.includes("email") + ); + assert(emailQueryAfterToggle, "Expected email query after toggle"); + assert( + emailQueryAfterToggle.optimization.state === "improvements_available", + `Expected improvements_available after toggle but got ${emailQueryAfterToggle.optimization.state}`, + ); + assert( + !emailQueryAfterToggle.optimization.indexesUsed.includes( + "users_email_idx", + ), + "Expected users_email_idx to NOT be in indexesUsed after disabling", + ); + assertGreater( + emailQueryAfterToggle.optimization.cost, + costWithIndex, + "Expected cost without index to be higher than cost with index", + ); + const recommendations = + emailQueryAfterToggle.optimization.indexRecommendations; + assert( + recommendations.some((r) => + r.columns.some((c) => c.column === "email") + ), + "Expected recommendation for email column after disabling the index", + ); + } finally { + await pg.stop(); + } + }, +}); + Deno.test({ name: "hypertable optimization includes index recommendations", sanitizeOps: false, diff --git a/src/remote/query-optimizer.ts b/src/remote/query-optimizer.ts index 55dc86c..39e29db 100644 --- a/src/remote/query-optimizer.ts +++ b/src/remote/query-optimizer.ts @@ -5,17 +5,22 @@ import { ConnectionManager } from "../sync/connection-manager.ts"; import { Sema } from "async-sema"; import { Analyzer, + dropIndex, + IndexedTable, IndexOptimizer, IndexRecommendation, OptimizeResult, + PgIdentifier, PostgresExplainStage, PostgresQueryBuilder, + PostgresTransaction, PostgresVersion, Statistics, StatisticsMode, } from "@query-doctor/core"; import { Connectable } from "../sync/connectable.ts"; import { parse } from "@libpg-query/parser"; +import { DisabledIndexes } from "./disabled-indexes.ts"; const MINIMUM_COST_CHANGE_PERCENTAGE = 5; const QUERY_TIMEOUT_MS = 10000; @@ -43,6 +48,8 @@ export class QueryOptimizer extends EventEmitter { reltuples: 10_000, }; private readonly queries = new Map(); + private readonly disabledIndexes = new DisabledIndexes(); + private target?: Target; private semaphore = new Sema(QueryOptimizer.MAX_CONCURRENCY); private _finish = Promise.withResolvers(); @@ -83,6 +90,10 @@ export class QueryOptimizer extends EventEmitter { return this._finish.promise; } + getDisabledIndexes(): PgIdentifier[] { + return [...this.disabledIndexes]; + } + /** * Start optimizing a new set of queries * @returns Promise of array of queries that were considered for optimization. @@ -104,9 +115,8 @@ export class QueryOptimizer extends EventEmitter { statsMode, ); const existingIndexes = await statistics.getExistingIndexes(); - const optimizer = new IndexOptimizer(pg, statistics, existingIndexes, { - // we're not running on our pg fork (yet) - // so traces have to be disabled + const filteredIndexes = this.filterDisabledIndexes(existingIndexes); + const optimizer = new IndexOptimizer(pg, statistics, filteredIndexes, { trace: false, }); this.target = { optimizer, statistics }; @@ -126,13 +136,33 @@ export class QueryOptimizer extends EventEmitter { this._validQueriesProcessed = 0; } - restart() { + async restart({ clearQueries } = { clearQueries: false }) { this.semaphore = new Sema(QueryOptimizer.MAX_CONCURRENCY); - this.queries.clear(); + if (clearQueries) { + this.queries.clear(); + } else { + this.resetQueryOptimizationState(); + } this._finish = Promise.withResolvers(); - this._allQueries = 0; this._invalidQueries = 0; this._validQueriesProcessed = 0; + if (this.target) { + // update the indexes the optimizer knows about + // to exclude the disabled ones + this.target.optimizer.transformIndexes((indexes) => + this.filterDisabledIndexes(indexes) + ); + } + await this.work(); + } + + toggleIndex(identifier: PgIdentifier): boolean { + const disabled = this.disabledIndexes.toggle(identifier); + // TODO: Instead of blindly restarting the query optimizer + // we should introspect the index and only reset the queries + // that touch the table the index is defined on + this.restart(); + return disabled; } /** @@ -141,9 +171,29 @@ export class QueryOptimizer extends EventEmitter { */ async addQueries(queries: RecentQuery[]) { this.appendQueries(queries); - if (!this.running) { - await this.work(); + await this.work(); + } + + private resetQueryOptimizationState() { + for (const [hash, query] of this.queries) { + const status = this.checkQueryUnsupported(query); + let optimization: LiveQueryOptimization; + switch (status.type) { + case "ok": + optimization = { state: "waiting" }; + break; + case "not_supported": + optimization = this.onQueryUnsupported(status.reason); + break; + case "ignored": + continue; + } + this.queries.set( + hash, + query.withOptimization(optimization), + ); } + this._allQueries = this.queries.size; } private appendQueries(queries: RecentQuery[]): OptimizedQuery[] { @@ -181,6 +231,10 @@ export class QueryOptimizer extends EventEmitter { return; } + if (this.running) { + return; + } + this.running = true; try { while (true) { @@ -243,6 +297,13 @@ export class QueryOptimizer extends EventEmitter { } } + private filterDisabledIndexes(indexes: IndexedTable[]): IndexedTable[] { + return indexes.filter((idx) => { + const indexName = PgIdentifier.fromString(idx.index_name); + return !this.disabledIndexes.has(indexName); + }); + } + private checkQueryUnsupported( query: RecentQuery, ): { type: "ok" } | { type: "ignored" } | { @@ -300,7 +361,11 @@ export class QueryOptimizer extends EventEmitter { let result: OptimizeResult; try { result = await withTimeout( - target.optimizer.run(builder, indexes), + target.optimizer.run( + builder, + indexes, + (tx) => this.dropDisabledIndexes(tx), + ), timeoutMs, ); } catch (error) { @@ -320,6 +385,12 @@ export class QueryOptimizer extends EventEmitter { return this.onOptimizeReady(result, recent, explainPlan); } + private async dropDisabledIndexes(tx: PostgresTransaction): Promise { + for (const indexName of this.disabledIndexes) { + await dropIndex(tx, indexName); + } + } + private onOptimizeReady( result: OptimizeResult, recent: OptimizedQuery, diff --git a/src/remote/remote-controller.dto.ts b/src/remote/remote-controller.dto.ts new file mode 100644 index 0000000..3a06885 --- /dev/null +++ b/src/remote/remote-controller.dto.ts @@ -0,0 +1,10 @@ +import { z } from "zod"; +import { PgIdentifier } from "@query-doctor/core"; + +export const ToggleIndexDto = z.object({ + indexName: z.string().min(1).max(64).transform((value) => + PgIdentifier.fromString(value) + ), +}); + +export type ToggleIndexDto = z.infer; diff --git a/src/remote/remote-controller.test.ts b/src/remote/remote-controller.test.ts index eeeb30c..0f20eab 100644 --- a/src/remote/remote-controller.test.ts +++ b/src/remote/remote-controller.test.ts @@ -6,7 +6,7 @@ import { assertEquals } from "@std/assert/equals"; import { RemoteController } from "./remote-controller.ts"; import { ConnectionManager } from "../sync/connection-manager.ts"; import { RemoteSyncRequest } from "./remote.dto.ts"; -import { assertSpyCalls, spy } from "@std/testing/mock"; +import { spy } from "@std/testing/mock"; import { setTimeout } from "node:timers/promises"; import { assertGreaterOrEqual } from "@std/assert"; diff --git a/src/remote/remote-controller.ts b/src/remote/remote-controller.ts index b3bc96c..fc9de27 100644 --- a/src/remote/remote-controller.ts +++ b/src/remote/remote-controller.ts @@ -4,6 +4,9 @@ import { RemoteSyncRequest } from "./remote.dto.ts"; import { Remote } from "./remote.ts"; import * as errors from "../sync/errors.ts"; import type { OptimizedQuery } from "../sql/recent-query.ts"; +import { ToggleIndexDto } from "./remote-controller.dto.ts"; +import { ZodError } from "zod"; +import { PgIdentifier } from "@query-doctor/core"; const SyncStatus = { NOT_STARTED: "notStarted", @@ -43,10 +46,21 @@ export class RemoteController { } else if (request.method === "GET") { return this.getStatus(); } - } else if ( - url.pathname === "/postgres/reset" && request.method === "POST" - ) { - return await this.onReset(request); + return methodNotAllowed(); + } + + if (url.pathname === "/postgres/indexes/toggle") { + if (request.method === "POST") { + return await this.toggleIndex(request); + } + return methodNotAllowed(); + } + + if (url.pathname === "/postgres/reset") { + if (request.method === "POST") { + return await this.onReset(request); + } + return methodNotAllowed(); } } @@ -62,12 +76,39 @@ export class RemoteController { remote.on("restoreLog", this.makeLoggingHandler("pg_restore").bind(this)); } + private async toggleIndex(request: Request): Promise { + try { + const data = await request.json(); + const index = ToggleIndexDto.decode(data); + const isDisabled = this.remote.optimizer.toggleIndex(index.indexName); + return Response.json({ isDisabled }); + } catch (error) { + if (error instanceof ZodError) { + return Response.json({ + type: "error", + error: "invalid_body", + message: error.message, + }, { + status: 400, + }); + } + return Response.json({ + type: "error", + error: env.HOSTED ? "Internal Server Error" : error, + message: "Failed to sync database", + }, { + status: 500, + }); + } + } + private getStatus(): Response { if (!this.syncResponse || this.syncStatus !== SyncStatus.COMPLETED) { return Response.json({ status: this.syncStatus }); } const { schema, meta } = this.syncResponse; const queries = this.remote.optimizer.getQueries(); + const disabledIndexes = this.remote.optimizer.getDisabledIndexes(); this.remote.pollQueriesOnce().catch((error) => { log.error("Failed to poll queries", "remote-controller"); console.error(error); @@ -77,6 +118,7 @@ export class RemoteController { meta, schema, queries: { type: "ok", value: queries }, + disabledIndexes: { type: "ok", value: disabledIndexes }, }); } @@ -179,3 +221,7 @@ export class RemoteController { ); } } + +function methodNotAllowed(): Response { + return Response.json("Method not allowed", { status: 405 }); +} diff --git a/src/remote/remote.ts b/src/remote/remote.ts index ba860e3..aeb5327 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -274,7 +274,7 @@ export class Remote extends EventEmitter { async resetPgStatStatements(source: Connectable): Promise { const connector = this.sourceManager.getConnectorFor(source); await connector.resetPgStatStatements(); - this.optimizer.restart(); + this.optimizer.restart({ clearQueries: true }); } /**