From 03343c079e81e59c45ebb69d854a4285f152227e Mon Sep 17 00:00:00 2001 From: Xetera Date: Fri, 13 Mar 2026 18:18:34 +0300 Subject: [PATCH 1/2] feat: accept stats imports in `POST /postgres/stats` --- src/remote/query-optimizer.ts | 22 +++++++++++--------- src/remote/remote-controller.ts | 36 +++++++++++++++++++++++++++++++++ src/remote/remote.ts | 2 +- src/server/http.ts | 11 +++++++--- 4 files changed, 57 insertions(+), 14 deletions(-) diff --git a/src/remote/query-optimizer.ts b/src/remote/query-optimizer.ts index 8bce5c8..1ed3ca8 100644 --- a/src/remote/query-optimizer.ts +++ b/src/remote/query-optimizer.ts @@ -115,25 +115,27 @@ export class QueryOptimizer extends EventEmitter { ): Promise { this.stop(); const validQueries = this.appendQueries(allRecentQueries); + await this.applyStatistics(statsMode); + this._allQueries = this.queries.size; + await this.work(); + return validQueries; + } + + async applyStatistics( + statsMode: StatisticsMode = QueryOptimizer.defaultStatistics, + ): Promise { const version = PostgresVersion.parse("17"); const pg = this.manager.getOrCreateConnection(this.connectable); const ownStats = await Statistics.dumpStats(pg, version, "full"); - const statistics = new Statistics( - pg, - version, - ownStats, - statsMode, - ); + const statistics = new Statistics(pg, version, ownStats, statsMode); const existingIndexes = await statistics.getExistingIndexes(); const filteredIndexes = this.filterDisabledIndexes(existingIndexes); const optimizer = new IndexOptimizer(pg, statistics, filteredIndexes, { trace: false, }); + // ordering is important here this.target = { optimizer, statistics }; - - this._allQueries = this.queries.size; - await this.work(); - return validQueries; + await this.restart(); } stop() { diff --git a/src/remote/remote-controller.ts b/src/remote/remote-controller.ts index f9f9f71..12f0a23 100644 --- a/src/remote/remote-controller.ts +++ b/src/remote/remote-controller.ts @@ -10,6 +10,7 @@ import { ToggleIndexDto, } from "./remote-controller.dto.ts"; import { ZodError } from "zod"; +import { ExportedStats, Statistics } from "@query-doctor/core"; const SyncStatus = { NOT_STARTED: "notStarted", @@ -157,6 +158,41 @@ export class RemoteController { } } + async onImportStats(body: unknown): Promise { + let stats: ExportedStats[]; + try { + stats = ExportedStats.array().parse(body); + } catch (error) { + if (error instanceof ZodError) { + return { + status: 400, + body: { type: "error", error: "invalid_body", message: error.message }, + }; + } + return { + status: 400, + body: { type: "error", error: "invalid_body", message: "body must be an array of ExportedStats" }, + }; + } + + try { + await this.remote.optimizer.applyStatistics( + Statistics.statsModeFromExport(stats), + ); + return { status: 200, body: { success: true } }; + } catch (error) { + console.error(error); + return { + status: 500, + body: { + type: "error", + error: env.HOSTED ? "Internal Server Error" : error, + message: "Failed to import stats", + }, + }; + } + } + async onReset(rawBody: string): Promise { const body = RemoteSyncRequest.safeDecode(rawBody); if (!body.success) { diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 351c675..2f01c4b 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -178,7 +178,7 @@ export class Remote extends EventEmitter { throw error; }, ) ?? - [] as Op[], /* no panic in case schemaLoader has not loaded in yet */ + [] as Op[], /* no panic in case schemaLoader has not loaded in yet */ this.pollQueriesOnce().catch((error) => { log.error("Failed to poll queries", "remote"); console.error(error); diff --git a/src/server/http.ts b/src/server/http.ts index 367b8b2..6808db6 100644 --- a/src/server/http.ts +++ b/src/server/http.ts @@ -147,8 +147,8 @@ export async function createServer( const remoteController = targetDb ? new RemoteController( - new Remote(targetDb, optimizingDbConnectionManager), - ) + new Remote(targetDb, optimizingDbConnectionManager), + ) : undefined; fastify.get("/", async (_request, reply) => { @@ -188,7 +188,7 @@ export async function createServer( return reply.send(result); }); - fastify.register(async function (app) { + fastify.register(async function(app) { app.get( "/postgres/ws", { websocket: true }, @@ -218,6 +218,11 @@ export async function createServer( return reply.status(result.status).send(result.body); }); + fastify.post("/postgres/stats", async (request, reply) => { + log.info(`[POST] /postgres/stats`, "http"); + const result = await remoteController.onImportStats(request.body); + return reply.status(result.status).send(result.body); + }); } await fastify.listen({ host: hostname, port }); From 0cbd0c24afa461e78612ad649a4574267e0d8f43 Mon Sep 17 00:00:00 2001 From: Xetera Date: Fri, 13 Mar 2026 19:02:56 +0300 Subject: [PATCH 2/2] fix: move restart logic out of query-optimizer --- src/remote/query-optimizer.ts | 6 ++---- src/remote/remote-controller.ts | 2 +- src/remote/remote.ts | 5 +++++ 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/remote/query-optimizer.ts b/src/remote/query-optimizer.ts index 1ed3ca8..daabacf 100644 --- a/src/remote/query-optimizer.ts +++ b/src/remote/query-optimizer.ts @@ -115,13 +115,13 @@ export class QueryOptimizer extends EventEmitter { ): Promise { this.stop(); const validQueries = this.appendQueries(allRecentQueries); - await this.applyStatistics(statsMode); + await this.setStatistics(statsMode); this._allQueries = this.queries.size; await this.work(); return validQueries; } - async applyStatistics( + async setStatistics( statsMode: StatisticsMode = QueryOptimizer.defaultStatistics, ): Promise { const version = PostgresVersion.parse("17"); @@ -133,9 +133,7 @@ export class QueryOptimizer extends EventEmitter { const optimizer = new IndexOptimizer(pg, statistics, filteredIndexes, { trace: false, }); - // ordering is important here this.target = { optimizer, statistics }; - await this.restart(); } stop() { diff --git a/src/remote/remote-controller.ts b/src/remote/remote-controller.ts index 12f0a23..6317827 100644 --- a/src/remote/remote-controller.ts +++ b/src/remote/remote-controller.ts @@ -176,7 +176,7 @@ export class RemoteController { } try { - await this.remote.optimizer.applyStatistics( + await this.remote.applyStatistics( Statistics.statsModeFromExport(stats), ); return { status: 200, body: { success: true } }; diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 2f01c4b..69892ba 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -319,6 +319,11 @@ export class Remote extends EventEmitter { return connector.getDatabaseInfo(); } + async applyStatistics(statsMode: StatisticsMode): Promise { + await this.optimizer.setStatistics(statsMode); + await this.optimizer.restart(); + } + async resetPgStatStatements(source: Connectable): Promise { const connector = this.sourceManager.getConnectorFor(source); await connector.resetPgStatStatements();