feat: run and optimize queries

devenv.lock

@@ -16,62 +16,6 @@
         "type": "github"
       }
     },
-    "flake-compat": {
-      "flake": false,
-      "locked": {
-        "lastModified": 1765121682,
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "rev": "65f23138d8d09a92e30f1e5c87611b23ef451bf3",
-        "type": "github"
-      },
-      "original": {
-        "owner": "edolstra",
-        "repo": "flake-compat",
-        "type": "github"
-      }
-    },
-    "git-hooks": {
-      "inputs": {
-        "flake-compat": "flake-compat",
-        "gitignore": "gitignore",
-        "nixpkgs": [
-          "nixpkgs"
-        ]
-      },
-      "locked": {
-        "lastModified": 1765016596,
-        "owner": "cachix",
-        "repo": "git-hooks.nix",
-        "rev": "548fc44fca28a5e81c5d6b846e555e6b9c2a5a3c",
-        "type": "github"
-      },
-      "original": {
-        "owner": "cachix",
-        "repo": "git-hooks.nix",
-        "type": "github"
-      }
-    },
-    "gitignore": {
-      "inputs": {
-        "nixpkgs": [
-          "git-hooks",
-          "nixpkgs"
-        ]
-      },
-      "locked": {
-        "lastModified": 1762808025,
-        "owner": "hercules-ci",
-        "repo": "gitignore.nix",
-        "rev": "cb5e3fdca1de58ccbc3ef53de65bd372b48f567c",
-        "type": "github"
-      },
-      "original": {
-        "owner": "hercules-ci",
-        "repo": "gitignore.nix",
-        "type": "github"
-      }
-    },
     "nixpkgs": {
       "locked": {
         "lastModified": 1764580874,
@@ -90,14 +34,10 @@
     "root": {
       "inputs": {
         "devenv": "devenv",
-        "git-hooks": "git-hooks",
-        "nixpkgs": "nixpkgs",
-        "pre-commit-hooks": [
-          "git-hooks"
-        ]
+        "nixpkgs": "nixpkgs"
       }
     }
   },
   "root": "root",
   "version": 7
-}
+}

package.json

@@ -10,7 +10,7 @@
     "dev:docker": "docker compose -f docker-compose.dev.yml up --build",
     "test": "vitest",
     "typecheck": "tsc --noEmit",
-    "build": "esbuild src/main.ts --bundle --platform=node --format=esm --outfile=dist/main.mjs --packages=external && cp src/reporters/github/success.md.j2 src/sync/schema_dump.sql dist/"
+    "build": "npm run typecheck && esbuild src/main.ts --bundle --platform=node --format=esm --outfile=dist/main.mjs --packages=external && cp src/reporters/github/success.md.j2 src/sync/schema_dump.sql dist/"
   },
   "dependencies": {
     "@actions/core": "^3.0.0",

src/remote/api-client.ts

@@ -1,6 +1,6 @@
 import { newWebSocketRpcSession, RpcTarget } from "capnweb";
 import type { RpcStub } from "capnweb";
-import type { ConnectionMode, UnauthenticatedServerApi, ClientApi, IndexDefinition, ServerApi } from "@query-doctor/core";
+import type { ConnectionMode, UnauthenticatedServerApi, ClientApi, IndexDefinition, LiveQueryOptimization, ServerApi } from "@query-doctor/core";
 import type { ExportedStats } from "@query-doctor/core";
 import { PgIdentifier, Statistics } from "@query-doctor/core";
 import { log } from "../log.ts";
@@ -106,7 +106,7 @@
   static async connect(endpoint: string, token: string, mode: ConnectionMode, remote: Remote, onBroken: (err: unknown) => void): Promise<ApiConnection> {
     const wsEndpoint = `${endpoint}/relay`.replace(/^http/, "ws");
     const unauthenticated = newWebSocketRpcSession<UnauthenticatedServerApi>(wsEndpoint);
     const api = await unauthenticated.authenticate(token, new this(remote), mode) as unknown as RpcStub<ServerApi>;
     let broken = false;
     const triggerBroken = (err: unknown) => {
       if (broken) return;
@@ -195,7 +195,7 @@
     );
   }
 
-  async runQuery(_query: string): Promise<void> {
-    log.warn("runQuery is not implemented", ApiClient.name);
+  async runQuery(query: string): Promise<LiveQueryOptimization> {
+    return this.remote.runAdHocQuery(query);
   }
 }

src/remote/query-optimizer.test.ts

@@ -635,3 +635,92 @@ test("optimizer does not treat ASC index as duplicate of DESC candidate", async
       await pg.stop();
     }
 });
+
+test("optimizeAdHoc returns an error before any sync and optimizes an unobserved query after", async () => {
+    const pg = await new PostgreSqlContainer("postgres:17")
+      .withCopyContentToContainer([
+        {
+          content: `
+            create table orders(id int, customer_id int, status text);
+            insert into orders (id, customer_id, status)
+              select i, i % 500, 'open' from generate_series(1, 1000) i;
+            create extension pg_stat_statements;
+          `,
+          target: "/docker-entrypoint-initdb.d/init.sql",
+        },
+      ])
+      .withCommand([
+        "-c",
+        "shared_preload_libraries=pg_stat_statements",
+        "-c",
+        "autovacuum=off",
+      ])
+      .start();
+
+    const manager = ConnectionManager.forLocalDatabase();
+    const conn = Connectable.fromString(pg.getConnectionUri());
+    const optimizer = new QueryOptimizer(manager, conn);
+
+    const statsMode = {
+      kind: "fromStatisticsExport" as const,
+      source: { kind: "inline" as const },
+      stats: [{
+        tableName: "orders",
+        schemaName: "public",
+        relpages: 10000,
+        reltuples: 10_000_000,
+        relallvisible: 1,
+        columns: [
+          { columnName: "id", stats: null, attlen: null },
+          { columnName: "customer_id", stats: null, attlen: null },
+          { columnName: "status", stats: null, attlen: null },
+        ],
+        indexes: [],
+      }],
+    };
+
+    try {
+      // Never observed in pg_stat_statements, and no CREATE INDEX exists on
+      // customer_id, so the optimizer should recommend one once it can run.
+      const adHoc = "select * from orders where customer_id = 7;";
+
+      const beforeSync = await optimizer.optimizeAdHoc(adHoc);
+      assert(
+        beforeSync.state === "error",
+        `Expected error before sync but got ${beforeSync.state}`,
+      );
+
+      await optimizer.start([], statsMode);
+      await optimizer.finish;
+
+      const result = await optimizer.optimizeAdHoc(adHoc);
+      assert(
+        result.state === "improvements_available",
+        `Expected improvements_available but got ${result.state}`,
+      );
+      expect(
+        result.indexRecommendations.some((r) =>
+          r.columns.some((c) => c.column === "customer_id")
+        ),
+        "Expected a recommendation on customer_id",
+      ).toBeTruthy();
+
+      // The ad-hoc query must not join the tracked query set.
+      expect(
+        optimizer.getQueries().some((q) => q.query.includes("customer_id")),
+        "Ad-hoc query should not be added to the tracked set",
+      ).toBeFalsy();
+
+      const unsupported = await optimizer.optimizeAdHoc(
+        "update orders set status = 'closed' where id = 1;",
+      );
+      assert(
+        unsupported.state === "not_supported",
+        `Expected not_supported for a non-SELECT but got ${unsupported.state}`,
+      );
+    } finally {
+      optimizer.stop();
+      await manager.closeAll();
+      await pg.stop();
+    }
+});

src/remote/query-optimizer.ts

@@ -1,5 +1,6 @@
 import EventEmitter from "node:events";
 import { OptimizedQuery, QueryHash, RecentQuery } from "../sql/recent-query.ts";
+import { syncQueries } from "../sync/query-sync.ts";
 import type { LiveQueryOptimization } from "./optimization.ts";
 import { ConnectionManager } from "../sync/connection-manager.ts";
 import { Sema } from "async-sema";
@@ -219,6 +220,57 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
     this.restart();
   }
 
+  /**
+   * Optimize a one-off SQL string that was never observed in pg_stat_statements —
+   * an agent handing us a query to check on demand. It is NOT added to the tracked
+   * query set; the result is computed and returned, leaving the background
+   * optimization loop untouched. Requires a prior sync (a `target` must exist);
+   * returns an `error` state if none has happened yet. Runs under the same
+   * semaphore as the loop so it never shares the optimizer connection concurrently.
+   */
+  async optimizeAdHoc(query: string): Promise<LiveQueryOptimization> {
+    const target = this.target;
+    if (!target) {
+      return {
+        state: "error",
+        error:
+          "No database has been synced yet. Connect a database before running ad-hoc queries.",
+      };
+    }
+    const [recent] = await syncQueries([{
+      username: "",
+      query,
+      formattedQuery: query,
+      meanTime: 0,
+      calls: "0",
+      rows: "0",
+      topLevel: true,
+    }]);
+    if (!recent) {
+      return { state: "error", error: "Failed to analyze query" };
+    }
+    const status = this.checkQueryUnsupported(recent);
+    if (status.type === "ignored") {
+      return {
+        state: "not_supported",
+        reason: "Query is an introspection or system query",
+      };
+    }
+    if (status.type === "not_supported") {
+      return { state: "not_supported", reason: status.reason };
+    }
+    const optimized = recent.withOptimization({ state: "waiting" });
+    const token = await this.semaphore.acquire();
+    try {
+      return await this.optimizeQuery(optimized, target, {
+        timeoutMs: this.queryTimeoutMs,
+        sideEffects: false,
+      });
+    } finally {
+      this.semaphore.release(token);
+    }
+  }
+
   /**
    * Insert new queries to be processed. The {@link start} method must
    * have been called previously for this to take effect
@@ -429,8 +481,9 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
   private async optimizeQuery(
     recent: OptimizedQuery,
     target: Target,
-    options: { timeoutMs: number },
+    options: { timeoutMs: number; sideEffects?: boolean },
   ): Promise<LiveQueryOptimization> {
+    const sideEffects = options.sideEffects ?? true;
     const builder = new PostgresQueryBuilder(recent.query);
     const indexes = this.getPotentialIndexCandidates(
       target.statistics,
@@ -450,15 +503,15 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
       console.error("Error with optimization", error);
       console.error("[query-optimizer] failed query:", recent.query);
       if (error instanceof TimeoutError) {
-        return this.onTimeout(recent, options.timeoutMs);
+        return this.onTimeout(recent, options.timeoutMs, sideEffects);
       } else if (error instanceof Error) {
-        return this.onError(recent, error.message);
+        return this.onError(recent, error.message, sideEffects);
       } else {
-        return this.onError(recent, "Internal error");
+        return this.onError(recent, "Internal error", sideEffects);
       }
     }
 
-    return this.onOptimizeReady(result, recent);
+    return this.onOptimizeReady(result, recent, sideEffects);
   }
 
   private async dropDisabledIndexes(tx: PostgresTransaction): Promise<void> {
@@ -470,27 +523,32 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
   private onOptimizeReady(
     result: OptimizeResult,
     recent: OptimizedQuery,
+    sideEffects = true,
   ): 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,
-          );
+          if (sideEffects) {
+            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);
+          if (sideEffects) {
+            this.onImprovementsAvailable(recent, result, result.baseExplainPlan);
+          }
           return {
             state: "improvements_available",
             cost: result.baseCost,
@@ -505,7 +563,7 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
       }
       // unlikely to hit if we've already checked the base plan for zero cost
       case "zero_cost_plan":
-        return this.onZeroCostPlan(recent, result.explainPlan);
+        return this.onZeroCostPlan(recent, result.explainPlan, sideEffects);
     }
   }
 
@@ -585,8 +643,11 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
   private onZeroCostPlan(
     recent: OptimizedQuery,
     explainPlan?: PostgresExplainStage,
+    sideEffects = true,
   ): LiveQueryOptimization {
-    this.emit("zeroCostPlan", recent);
+    if (sideEffects) {
+      this.emit("zeroCostPlan", recent);
+    }
     return {
       state: "error",
       error:
@@ -598,21 +659,27 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
   private onError(
     recent: OptimizedQuery,
     errorMessage: string,
+    sideEffects = true,
   ): LiveQueryOptimization {
     const error = new Error(errorMessage);
-    this.emit("error", error, recent);
+    if (sideEffects) {
+      this.emit("error", error, recent);
+    }
     return { state: "error", error: error.message };
   }
 
   private onTimeout(
     recent: OptimizedQuery,
     waitedMs: number,
+    sideEffects = true,
   ): LiveQueryOptimization {
     let retries = 0;
     if ("retries" in recent.optimization) {
       retries = recent.optimization.retries + 1;
     }
-    this.emit("timeout", recent, waitedMs);
+    if (sideEffects) {
+      this.emit("timeout", recent, waitedMs);
+    }
     return { state: "timeout", waitedMs, retries };
   }
 }