From 7ff6087326f64917e628bbdd0c88737403066d45 Mon Sep 17 00:00:00 2001 From: Xetera Date: Wed, 15 Jul 2026 16:57:16 +0300 Subject: [PATCH] feat: run and optimize queries --- devenv.lock | 64 +------------------ package.json | 2 +- src/remote/api-client.ts | 6 +- src/remote/query-optimizer.test.ts | 89 +++++++++++++++++++++++++++ src/remote/query-optimizer.ts | 99 +++++++++++++++++++++++++----- src/remote/remote.ts | 5 ++ 6 files changed, 183 insertions(+), 82 deletions(-) diff --git a/devenv.lock b/devenv.lock index 2d377d47..89a60eda 100644 --- a/devenv.lock +++ b/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 -} +} \ No newline at end of file diff --git a/package.json b/package.json index b15d8fb6..0991c543 100644 --- a/package.json +++ b/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", diff --git a/src/remote/api-client.ts b/src/remote/api-client.ts index 9f048ea9..696a72b9 100644 --- a/src/remote/api-client.ts +++ b/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"; @@ -195,7 +195,7 @@ export class ApiClient extends RpcTarget implements ClientApi { ); } - async runQuery(_query: string): Promise { - log.warn("runQuery is not implemented", ApiClient.name); + async runQuery(query: string): Promise { + return this.remote.runAdHocQuery(query); } } diff --git a/src/remote/query-optimizer.test.ts b/src/remote/query-optimizer.test.ts index 0fe1cd63..da3d7ca4 100644 --- a/src/remote/query-optimizer.test.ts +++ b/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(); + } +}); diff --git a/src/remote/query-optimizer.ts b/src/remote/query-optimizer.ts index 7739eef6..648ce152 100644 --- a/src/remote/query-optimizer.ts +++ b/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 { 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 { + 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 { private async optimizeQuery( recent: OptimizedQuery, target: Target, - options: { timeoutMs: number }, + options: { timeoutMs: number; sideEffects?: boolean }, ): Promise { + 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 { 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 { @@ -470,6 +523,7 @@ export class QueryOptimizer extends EventEmitter { private onOptimizeReady( result: OptimizeResult, recent: OptimizedQuery, + sideEffects = true, ): LiveQueryOptimization { switch (result.kind) { case "ok": { @@ -477,12 +531,14 @@ export class QueryOptimizer extends EventEmitter { 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, @@ -490,7 +546,9 @@ export class QueryOptimizer extends EventEmitter { 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 { } // 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 { 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 { 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 }; } } diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 8d070a98..607139d2 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -17,6 +17,7 @@ import { ExtensionNotInstalledError } from "../sync/errors.ts"; import { type OptimizedQuery, type RecentQuery } from "../sql/recent-query.ts"; import { type Op } from "jsondiffpatch/formatters/jsonpatch"; import { type RemoteSyncFullSchemaResponse } from "./remote.dto.ts"; +import { type LiveQueryOptimization } from "./optimization.ts"; import { QueryOptimizer } from "./query-optimizer.ts"; import { EventEmitter } from "node:events"; import { log } from "../log.ts"; @@ -387,6 +388,10 @@ export class Remote extends EventEmitter { return connector.getDatabaseInfo(); } + async runAdHocQuery(query: string): Promise { + return this.optimizer.optimizeAdHoc(query); + } + async applyStatistics(statsMode: StatisticsMode): Promise { await this.optimizer.setStatistics(statsMode); const stats = this.optimizer.ownMetadata;