fix(sync): install pg_stat_statements into a named schema and verify through the resolver

src/remote/remote-controller.ts

@@ -1,7 +1,10 @@
 import { WebSocket } from "ws";
 import { env } from "../env.ts";
 import { log } from "../log.ts";
-import { RemoteSyncRequest } from "./remote.dto.ts";
+import {
+  InstallPgStatStatementsRequest,
+  RemoteSyncRequest,
+} from "./remote.dto.ts";
 import { Remote } from "./remote.ts";
 import * as errors from "../sync/errors.ts";
 import type { OptimizedQuery } from "../sql/recent-query.ts";
@@ -251,14 +254,24 @@ export class RemoteController {
   }
 
   async onInstallPgStatStatements(rawBody: string): Promise<HandlerResult> {
-    const body = RemoteSyncRequest.safeDecode(rawBody);
+    const body = InstallPgStatStatementsRequest.safeDecode(rawBody);
     if (!body.success) {
       return { status: 400, body: body.error };
     }
 
     try {
-      const result = await this.remote.installPgStatStatements(body.data.db);
-      return { status: 200, body: { success: true, preloadUpdated: result.preloadUpdated } };
+      const result = await this.remote.installPgStatStatements(
+        body.data.db,
+        body.data.schema,
+      );
+      return {
+        status: 200,
+        body: {
+          success: true,
+          preloadUpdated: result.preloadUpdated,
+          schema: result.schema,
+        },
+      };
     } catch (error) {
       console.error(error);
       if (error instanceof errors.PostgresError) {

src/remote/remote.dto.ts

@@ -17,6 +17,29 @@ export const RemoteSyncRequest = z.codec(
   },
 );
 
+/**
+ * `schema` is where the extension gets created. Omitting it takes the
+ * connector's default, which is deliberately not `public`.
+ */
+export const InstallPgStatStatementsRequest = z.codec(
+  z.string(),
+  z.object({
+    db: z.custom<Connectable>(),
+    schema: z.string().min(1).optional(),
+  }),
+  {
+    encode: (value) =>
+      JSON.stringify({ db: value.db.toString(), schema: value.schema }),
+    decode: (value) => {
+      const parsed = JSON.parse(value);
+      return {
+        db: Connectable.fromString(parsed.db),
+        schema: typeof parsed.schema === "string" ? parsed.schema : undefined,
+      };
+    },
+  },
+);
+
 export const RemoteSyncFullSchemaResponse = z.discriminatedUnion("type", [
   z.object({ type: z.literal("ok"), value: z.custom<FullSchema>() }),
   z.object({ type: z.literal("error"), error: z.string() }),

src/remote/remote.ts

@@ -638,9 +638,12 @@ export class Remote extends EventEmitter<RemoteEvents> {
     this.optimizer.restart({ clearQueries: true });
   }
 
-  async installPgStatStatements(source: Connectable): Promise<{ preloadUpdated: boolean }> {
+  async installPgStatStatements(
+    source: Connectable,
+    schema?: string,
+  ): Promise<{ preloadUpdated: boolean; schema: string }> {
     const connector = this.sourceManager.getConnectorFor(source);
-    return connector.installPgStatStatements();
+    return connector.installPgStatStatements({ schema });
   }
 
   /**

src/server/http.ts

@@ -12,7 +12,10 @@ import { SyncResult } from "../sync/syncer.ts";
 import * as errors from "../sync/errors.ts";
 import { RemoteController } from "../remote/remote-controller.ts";
 import { Connectable } from "../sync/connectable.ts";
-import { RemoteSyncRequest } from "../remote/remote.dto.ts";
+import {
+  InstallPgStatStatementsRequest,
+  RemoteSyncRequest,
+} from "../remote/remote.dto.ts";
 import { ConnectionManager } from "../sync/connection-manager.ts";
 import { Remote } from "../remote/remote.ts";
 
@@ -170,14 +173,18 @@ export async function createServer(
 
   fastify.post("/postgres/extensions/pg_stat_statements", async (request, reply) => {
     log.info(`[POST] /postgres/extensions/pg_stat_statements`, "http");
-    const body = RemoteSyncRequest.safeDecode(JSON.stringify(request.body));
+    const body = InstallPgStatStatementsRequest.safeDecode(
+      JSON.stringify(request.body),
+    );
     if (!body.success) {
       return reply.status(400).send(body.error);
     }
     try {
       const db = await body.data.db.resolveDockerHost();
       const connector = sourceConnectionManager.getConnectorFor(db);
-      const result = await connector.installPgStatStatements();
+      const result = await connector.installPgStatStatements({
+        schema: body.data.schema,
+      });
       return reply.status(200).send({ success: true, ...result });
     } catch (error) {
       return reply.status(500).send(makeUnexpectedErrorResult(error).body);

src/sync/pg-connector.test.ts

@@ -1,8 +1,9 @@
-import { test, expect } from "vitest";
+import { test, expect, vi } from "vitest";
 import { PostgreSqlContainer } from "@testcontainers/postgresql";
 import { ConnectionManager } from "./connection-manager.ts";
 import { Connectable } from "./connectable.ts";
 import { PgIdentifier } from "@query-doctor/core";
+import { isMissingExtensionObject } from "./pg-connector.ts";
 
 test("getRecentQueries resolves pg_stat_statements in a non-default schema", async () => {
     const pg = await new PostgreSqlContainer("postgres:17")
@@ -312,3 +313,225 @@ test("getTotalRowCount does not double-count when a table appears twice in the i
     await pg.stop();
   }
 });
+
+// Installing into `public` breaks any project whose migration tool reconciles
+// that schema: the extension owns `pg_stat_statements_info`, and a reconciler
+// that tries to drop it aborts half-applied with SQLSTATE 2BP01. The one path
+// where we choose the placement must not choose `public`.
+test("installPgStatStatements creates the extension outside public", async () => {
+  const pg = await new PostgreSqlContainer("postgres:17")
+    .withCommand(["-c", "shared_preload_libraries=pg_stat_statements"])
+    .start();
+
+  const manager = ConnectionManager.forLocalDatabase();
+  const conn = Connectable.fromString(pg.getConnectionUri());
+  const db = manager.getOrCreateConnection(conn);
+  const connector = manager.getConnectorFor(db);
+
+  try {
+    const result = await connector.installPgStatStatements();
+
+    expect(result.schema).toBe("query_doctor");
+
+    const [placement] = await db.exec<{ schema: string }>(`
+      SELECT n.nspname AS schema
+      FROM pg_extension e
+      JOIN pg_namespace n ON n.oid = e.extnamespace
+      WHERE e.extname = 'pg_stat_statements'
+    `);
+    expect(placement?.schema).toBe("query_doctor");
+
+    const [leftInPublic] = await db.exec<{ count: string }>(`
+      SELECT count(*) AS count
+      FROM pg_class c
+      JOIN pg_namespace n ON n.oid = c.relnamespace
+      WHERE n.nspname = 'public' AND c.relname LIKE 'pg_stat_statements%'
+    `);
+    expect(leftInPublic?.count).toBe("0");
+  } finally {
+    await manager.closeAll();
+    await pg.stop();
+  }
+});
+
+test("installPgStatStatements installs into a caller-supplied schema", async () => {
+  const pg = await new PostgreSqlContainer("postgres:17")
+    .withCommand(["-c", "shared_preload_libraries=pg_stat_statements"])
+    .start();
+
+  const manager = ConnectionManager.forLocalDatabase();
+  const conn = Connectable.fromString(pg.getConnectionUri());
+  const db = manager.getOrCreateConnection(conn);
+  const connector = manager.getConnectorFor(db);
+
+  try {
+    const result = await connector.installPgStatStatements({ schema: "ext" });
+
+    expect(result.schema).toBe("ext");
+
+    const [placement] = await db.exec<{ schema: string }>(`
+      SELECT n.nspname AS schema
+      FROM pg_extension e
+      JOIN pg_namespace n ON n.oid = e.extnamespace
+      WHERE e.extname = 'pg_stat_statements'
+    `);
+    expect(placement?.schema).toBe("ext");
+  } finally {
+    await manager.closeAll();
+    await pg.stop();
+  }
+});
+
+// The verify step used to probe an unqualified `pg_stat_statements`, which
+// resolves through the search_path and so returns 42P01 for exactly the
+// placement the install step now produces: a working install reported as a
+// failure.
+test("installPgStatStatements verifies an extension that sits off the search_path", async () => {
+  const pg = await new PostgreSqlContainer("postgres:17")
+    .withCopyContentToContainer([
+      {
+        content: `
+          CREATE SCHEMA monitoring;
+          CREATE EXTENSION pg_stat_statements SCHEMA monitoring;
+        `,
+        target: "/docker-entrypoint-initdb.d/init.sql",
+      },
+    ])
+    .withCommand(["-c", "shared_preload_libraries=pg_stat_statements"])
+    .start();
+
+  const manager = ConnectionManager.forLocalDatabase();
+  const conn = Connectable.fromString(pg.getConnectionUri());
+  const connector = manager.getConnectorFor(conn);
+
+  try {
+    const result = await connector.installPgStatStatements();
+
+    expect(result.schema).toBe("monitoring");
+  } finally {
+    await manager.closeAll();
+    await pg.stop();
+  }
+});
+
+// Today the user finds out about a `public` install through unrelated migration
+// failures. The resolver already knows, so say it where we read the extension.
+test("getRecentQueries warns when the extension sits in public", async () => {
+  const pg = await new PostgreSqlContainer("postgres:17")
+    .withCopyContentToContainer([
+      {
+        content: `
+          CREATE EXTENSION pg_stat_statements;
+          CREATE TABLE users(id int, name text);
+          SELECT * FROM users WHERE id = 1;
+        `,
+        target: "/docker-entrypoint-initdb.d/init.sql",
+      },
+    ])
+    .withCommand(["-c", "shared_preload_libraries=pg_stat_statements"])
+    .start();
+
+  const manager = ConnectionManager.forLocalDatabase();
+  const conn = Connectable.fromString(pg.getConnectionUri());
+  const connector = manager.getConnectorFor(conn);
+  const logged = vi.spyOn(console, "log").mockImplementation(() => {});
+
+  try {
+    await connector.getRecentQueries();
+    // A fresh connector over the same connection, which is what every poll
+    // does. The warning dedupes per connection, so this must stay silent.
+    await manager.getConnectorFor(conn).getRecentQueries();
+
+    const messages = logged.mock.calls.map((call) => String(call[0]));
+    const warnings = messages.filter((message) =>
+      message.includes("public schema")
+    );
+    expect(warnings).toHaveLength(1);
+    expect(
+      messages.some((message) =>
+        message.includes("pg_stat_statements") &&
+        message.includes("public") &&
+        message.includes("SET SCHEMA")
+      ),
+      `Expected a warning naming the fix. Logged:\n${messages.join("\n")}`,
+    ).toBe(true);
+  } finally {
+    logged.mockRestore();
+    await manager.closeAll();
+    await pg.stop();
+  }
+});
+
+// The catch blocks in getRecentQueries/resetPgStatStatements used to match the
+// literal `relation "pg_stat_statements" does not exist`. Every query that
+// reaches them is schema-qualified, so Postgres names the schema too and the
+// match never fired: a missing extension surfaced as a generic PostgresError
+// rather than ExtensionNotInstalledError, and the UI showed raw SQL text
+// instead of its install panel.
+test("a qualified read of a missing extension reports 42P01, not the bare relation name", async () => {
+  const pg = await new PostgreSqlContainer("postgres:17")
+    .withCopyContentToContainer([
+      {
+        content: `CREATE SCHEMA monitoring;`,
+        target: "/docker-entrypoint-initdb.d/init.sql",
+      },
+    ])
+    .start();
+
+  const manager = ConnectionManager.forLocalDatabase();
+  const conn = Connectable.fromString(pg.getConnectionUri());
+  const db = manager.getOrCreateConnection(conn);
+
+  try {
+    const err = await db
+      .exec("SELECT 1 FROM monitoring.pg_stat_statements LIMIT 1")
+      .then(() => null, (e: unknown) => e as { code?: string; message: string });
+
+    expect(err).toBeTruthy();
+    expect(err!.code).toBe("42P01");
+    expect(err!.message).not.toContain(
+      'relation "pg_stat_statements" does not exist',
+    );
+
+    const fnErr = await db
+      .exec("SELECT monitoring.pg_stat_statements_reset()")
+      .then(() => null, (e: unknown) => e as { code?: string; message: string });
+
+    expect(fnErr).toBeTruthy();
+    expect(fnErr!.code).toBe("42883");
+    expect(fnErr!.message).not.toContain(
+      "function pg_stat_statements_reset() does not exist",
+    );
+
+    // A function call through a schema that is gone entirely reports the
+    // schema rather than the function, so that code has to be recognised too.
+    const schemaErr = await db
+      .exec("SELECT absent_schema.pg_stat_statements_reset()")
+      .then(() => null, (e: unknown) => e as { code?: string });
+
+    expect(schemaErr?.code).toBe("3F000");
+  } finally {
+    await manager.closeAll();
+    await pg.stop();
+  }
+});
+
+test("isMissingExtensionObject classifies the codes Postgres uses, and nothing else", () => {
+  expect(isMissingExtensionObject({ code: "42P01" })).toBe(true);
+  expect(isMissingExtensionObject({ code: "42883" })).toBe(true);
+  expect(isMissingExtensionObject({ code: "3F000" })).toBe(true);
+
+  // Insufficient privilege and a syntax error are real failures the caller
+  // must see, not "the extension isn't installed".
+  expect(isMissingExtensionObject({ code: "42501" })).toBe(false);
+  expect(isMissingExtensionObject({ code: "42601" })).toBe(false);
+
+  // The message text the old check matched carries no code of its own.
+  expect(
+    isMissingExtensionObject(
+      new Error('relation "pg_stat_statements" does not exist'),
+    ),
+  ).toBe(false);
+  expect(isMissingExtensionObject(null)).toBe(false);
+  expect(isMissingExtensionObject(undefined)).toBe(false);
+});