diff --git a/src/sync/pg-connector.test.ts b/src/sync/pg-connector.test.ts new file mode 100644 index 0000000..4ec9f23 --- /dev/null +++ b/src/sync/pg-connector.test.ts @@ -0,0 +1,96 @@ +import { test, expect } from "vitest"; +import { PostgreSqlContainer } from "@testcontainers/postgresql"; +import { ConnectionManager } from "./connection-manager.ts"; +import { Connectable } from "./connectable.ts"; + +test("getRecentQueries resolves pg_stat_statements in a non-default schema", async () => { + const pg = await new PostgreSqlContainer("postgres:17") + .withCopyContentToContainer([ + { + content: ` + CREATE SCHEMA monitoring; + CREATE EXTENSION pg_stat_statements SCHEMA monitoring; + + CREATE TABLE users(id int, name text); + INSERT INTO users (id, name) VALUES (1, 'alice'); + SELECT * FROM users WHERE id = 1; + `, + target: "/docker-entrypoint-initdb.d/init.sql", + }, + ]) + .withCommand([ + "-c", + "shared_preload_libraries=pg_stat_statements", + "-c", + "autovacuum=off", + "-c", + "track_counts=off", + "-c", + "track_io_timing=off", + "-c", + "track_activities=off", + ]) + .start(); + + const manager = ConnectionManager.forLocalDatabase(); + const conn = Connectable.fromString(pg.getConnectionUri()); + const connector = manager.getConnectorFor(conn); + + try { + const recentQueries = await connector.getRecentQueries(); + const userQuery = recentQueries.find((q) => + q.query.includes("users") + ); + expect(userQuery, "Expected to find a query involving 'users' table").toBeTruthy(); + } finally { + await manager.closeAll(); + await pg.stop(); + } +}); + +test("resetPgStatStatements works with a non-default schema", async () => { + const pg = await new PostgreSqlContainer("postgres:17") + .withCopyContentToContainer([ + { + content: ` + CREATE SCHEMA monitoring; + CREATE EXTENSION pg_stat_statements SCHEMA monitoring; + + CREATE TABLE users(id int, name text); + INSERT INTO users (id, name) VALUES (1, 'alice'); + SELECT * FROM users WHERE id = 1; + `, + target: "/docker-entrypoint-initdb.d/init.sql", + }, + ]) + .withCommand([ + "-c", + "shared_preload_libraries=pg_stat_statements", + "-c", + "autovacuum=off", + "-c", + "track_counts=off", + "-c", + "track_io_timing=off", + "-c", + "track_activities=off", + ]) + .start(); + + const manager = ConnectionManager.forLocalDatabase(); + const conn = Connectable.fromString(pg.getConnectionUri()); + const connector = manager.getConnectorFor(conn); + + try { + const before = await connector.getRecentQueries(); + expect(before.length, "Expected queries before reset").toBeGreaterThan(0); + + await connector.resetPgStatStatements(); + + const after = await connector.getRecentQueries(); + expect(after.length, "Expected 0 queries after reset").toEqual(0); + } finally { + await manager.closeAll(); + await pg.stop(); + } +}); diff --git a/src/sync/pg-connector.ts b/src/sync/pg-connector.ts index 8ebbb14..441d894 100644 --- a/src/sync/pg-connector.ts +++ b/src/sync/pg-connector.ts @@ -90,6 +90,7 @@ export type ResetPgStatStatementsResult = export class PostgresConnector implements DatabaseConnector { private static readonly QUERY_DOCTOR_USER = "query_doctor_db_link"; private readonly tupleEstimates = new Map(); + private pssSchema: PgIdentifier | null = null; /** * The minimum size for a table to be considered for sampling. * Otherwise we use the `order by random()` instead. @@ -471,6 +472,21 @@ ORDER BY }; } + private async getPssSchema(): Promise { + if (this.pssSchema) return this.pssSchema; + const result = await this.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' + `); + if (result.length === 0) { + throw new ExtensionNotInstalledError("pg_stat_statements"); + } + this.pssSchema = PgIdentifier.fromString(result[0].schema); + return this.pssSchema; + } + /** * Get the latest queries using pg_stat_statements * @throws {ExtensionNotInstalledError} - pg_stat_statements is not installed @@ -478,6 +494,7 @@ ORDER BY */ public async getRecentQueries(): Promise { try { + const pssSchema = await this.getPssSchema(); const results = await this.db.exec(` SELECT 'unknown_user' as "username", @@ -486,7 +503,7 @@ ORDER BY calls, rows, toplevel as "topLevel" - FROM pg_stat_statements + FROM ${pssSchema}.pg_stat_statements WHERE query not like '%pg_stat_statements%' -- and dbid = (select oid from pg_database where datname = current_database()) and query not like '%@qd_introspection%' @@ -498,9 +515,7 @@ ORDER BY results, ); } catch (err) { - if (err instanceof ExtensionNotInstalledError) { - throw err; - } + if (err instanceof ExtensionNotInstalledError) throw err; if ( err instanceof Error && err.message.includes('relation "pg_stat_statements" does not exist') @@ -518,12 +533,14 @@ ORDER BY */ public async resetPgStatStatements(): Promise { try { + const pssSchema = await this.getPssSchema(); await this.db.exec(` - SELECT pg_stat_statements_reset(); -- @qd_introspection + SELECT ${pssSchema}.pg_stat_statements_reset(); -- @qd_introspection `); this.segmentedQueryCache.reset(this.db); } catch (err) { + if (err instanceof ExtensionNotInstalledError) throw err; if ( err instanceof Error && err.message.includes(