feat: toggle indexes

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",

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"]);
+});

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<string>();
+
+  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<PgIdentifier> {
+    return this.disabledIndexNames.values().map((indexName) =>
+      PgIdentifier.fromString(indexName)
+    );
+  }
+}

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,