feat: accept stats imports in POST /postgres/stats

src/remote/query-optimizer.ts

@@ -115,25 +115,25 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
   ): Promise<OptimizedQuery[]> {
     this.stop();
     const validQueries = this.appendQueries(allRecentQueries);
+    await this.setStatistics(statsMode);
+    this._allQueries = this.queries.size;
+    await this.work();
+    return validQueries;
+  }
+
+  async setStatistics(
+    statsMode: StatisticsMode = QueryOptimizer.defaultStatistics,
+  ): Promise<void> {
     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,
     });
     this.target = { optimizer, statistics };
-
-    this._allQueries = this.queries.size;
-    await this.work();
-    return validQueries;
   }
 
   stop() {

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<HandlerResult> {
+    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.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<HandlerResult> {
     const body = RemoteSyncRequest.safeDecode(rawBody);
     if (!body.success) {

src/remote/remote.ts

@@ -178,7 +178,7 @@ export class Remote extends EventEmitter<RemoteEvents> {
           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);
@@ -319,6 +319,11 @@ export class Remote extends EventEmitter<RemoteEvents> {
     return connector.getDatabaseInfo();
   }
 
+  async applyStatistics(statsMode: StatisticsMode): Promise<void> {
+    await this.optimizer.setStatistics(statsMode);
+    await this.optimizer.restart();
+  }
+
   async resetPgStatStatements(source: Connectable): Promise<void> {
     const connector = this.sourceManager.getConnectorFor(source);
     await connector.resetPgStatStatements();

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 });