diff --git a/deno.json b/deno.json index 0483b8a..123ea45 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.8", "@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..3b5f669 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.8": "0.1.8", "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.8": { + "integrity": "sha512-NvW1yjXqA8TIZSspa9jsrafadNMs2TzyHPQVMvRTHl1/Deb+CRuQszIoqytW+o7dyDNOiGtU9D8Y19+mWG6oEA==", "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.8", "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..156c7ef --- /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 wasDisabled = indexes.toggle(indexName); + assertEquals(wasDisabled, false); + 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 wasDisabled = indexes.toggle(indexName); + assertEquals(wasDisabled, true); + 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..135a788 --- /dev/null +++ b/src/remote/disabled-indexes.ts @@ -0,0 +1,36 @@ +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()); + } + + /** + * Toggles the visibility of the index + * @returns did the index get disabled? + */ + toggle(indexName: PgIdentifier): boolean { + const deleted = this.remove(indexName); + if (!deleted) { + this.add(indexName); + return false; + } + return true; + } + + [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 e67c41e..6a19912 100644 --- a/src/remote/query-optimizer.test.ts +++ b/src/remote/query-optimizer.test.ts @@ -182,6 +182,134 @@ 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 optimizer = new QueryOptimizer(manager); + + const conn = Connectable.fromString(pg.getConnectionUri()); + 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(conn, [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 f1bf79b..26bbff3 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; @@ -32,6 +37,7 @@ type Target = { connectable: Connectable; optimizer: IndexOptimizer; statistics: Statistics; + existingIndexes: IndexedTable[]; }; export class QueryOptimizer extends EventEmitter { @@ -42,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(); @@ -78,6 +86,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. @@ -100,12 +112,11 @@ 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 = { connectable: conn, optimizer, statistics }; + this.target = { connectable: conn, optimizer, statistics, existingIndexes }; this._allQueries = this.queries.size; await this.work(); @@ -122,13 +133,23 @@ export class QueryOptimizer extends EventEmitter { this._validQueriesProcessed = 0; } - restart() { + async restart() { this.semaphore = new Sema(QueryOptimizer.MAX_CONCURRENCY); this.queries.clear(); this._finish = Promise.withResolvers(); this._allQueries = 0; this._invalidQueries = 0; this._validQueriesProcessed = 0; + if (this.target) { + const { statistics, connectable, existingIndexes } = this.target; + const pg = this.manager.getOrCreateConnection(connectable); + const filteredIndexes = this.filterDisabledIndexes(existingIndexes); + const optimizer = new IndexOptimizer(pg, statistics, filteredIndexes, { + trace: false, + }); + this.target = { connectable, optimizer, statistics, existingIndexes }; + } + await this.work(); } /** @@ -137,9 +158,7 @@ export class QueryOptimizer extends EventEmitter { */ async addQueries(queries: RecentQuery[]) { this.appendQueries(queries); - if (!this.running) { - await this.work(); - } + await this.work(); } private appendQueries(queries: RecentQuery[]): OptimizedQuery[] { @@ -177,6 +196,10 @@ export class QueryOptimizer extends EventEmitter { return; } + if (this.running) { + return; + } + this.running = true; try { while (true) { @@ -224,6 +247,28 @@ export class QueryOptimizer extends EventEmitter { return Array.from(this.queries.values()); } + toggleIndex(indexName: PgIdentifier): boolean { + const status = this.toggleIndexStatus(indexName); + // optimizer has to be restarted to allow re-evaluating + // queries based on the newly toggled index + this.restart(); + return status; + } + + /** + * @returns whether the index is enabled + */ + private toggleIndexStatus(indexName: PgIdentifier): boolean { + return this.disabledIndexes.toggle(indexName); + } + + private filterDisabledIndexes(indexes: IndexedTable[]): IndexedTable[] { + const disabledNames = new Set( + [...this.disabledIndexes].map((i) => i.toString()), + ); + return indexes.filter((idx) => !disabledNames.has(idx.index_name)); + } + private checkQueryUnsupported( query: RecentQuery, ): { type: "ok" } | { type: "ignored" } | { @@ -281,7 +326,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) { @@ -301,6 +350,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.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 }); +}