feat: use the query-doctor postgres fork for CI runs

.github/workflows/custom.yaml

@@ -43,8 +43,11 @@ jobs:
 
       - name: Run local GitHub Action
         uses: ./
+        with:
+          # TODO: support pg 16 explicitly
+          postgres-version: 17
         env:
           GITHUB_TOKEN: ${{ github.token }}
-          POSTGRES_URL: postgres://query_doctor@localhost:5432/testing
+          SOURCE_DATABASE_URL: postgres://query_doctor@localhost:5432/testing
           LOG_PATH: /var/log/postgresql/postgres.log
           SITE_API_ENDPOINT: ${{ vars.SITE_API_ENDPOINT }}

action.yaml

@@ -8,6 +8,12 @@ branding:
   icon: "database"
   color: "blue"
 
+inputs:
+  postgres-version:
+    description: "PostgreSQL version to use (14, 17, or 18)"
+    required: true
+    default: "18"
+
 runs:
   using: "composite"
   steps:
@@ -51,13 +57,46 @@ runs:
         sudo make install # Use sudo to install globally
         cd ${{ github.action_path }} # Return to action directory
 
+    - name: Cache postgresql-client
+      uses: actions/cache@v4
+      id: pg-client-cache
+      with:
+        path: /usr/lib/postgresql
+        key: pg-client-${{ runner.os }}-18
+
+    - name: Install postgresql-client
+      shell: bash
+      if: steps.pg-client-cache.outputs.cache-hit != 'true'
+      run: |
+        sudo apt-get install -y curl ca-certificates
+        curl -fsSL https://www.postgresql.org/media/keys/ACCC4CF8.asc | sudo gpg --dearmor -o /usr/share/keyrings/postgresql.gpg
+        echo "deb [signed-by=/usr/share/keyrings/postgresql.gpg] https://apt.postgresql.org/pub/repos/apt $(lsb_release -cs)-pgdg main" | sudo tee /etc/apt/sources.list.d/pgdg.list
+        sudo apt-get update
+        sudo apt-get install -y postgresql-client-18 || sudo apt-get install -y postgresql-client-17
+
+    - name: Start PostgreSQL
+      shell: bash
+      run: |
+        PORT=$(shuf -i 10000-65000 -n 1)
+        echo "PG_PORT=$PORT" >> $GITHUB_ENV
+        PG_CLIENT_VERSION=$(ls /usr/lib/postgresql | sort -rn | head -1)
+        echo "PG_DUMP_BINARY=/usr/lib/postgresql/${PG_CLIENT_VERSION}/bin/pg_dump" >> $GITHUB_ENV
+        echo "PG_RESTORE_BINARY=/usr/lib/postgresql/${PG_CLIENT_VERSION}/bin/pg_restore" >> $GITHUB_ENV
+        docker run -d \
+          --name query-doctor-postgres \
+          -p $PORT:5432 \
+          ghcr.io/query-doctor/postgres:pg-${{ inputs.postgres-version }}
+        until docker exec query-doctor-postgres pg_isready -U postgres; do
+          sleep 1
+        done
+
     # Run the application
     - name: Run Analyzer
       shell: bash
       working-directory: ${{ github.action_path }}
       run: npm run start
       env:
-        PG_DUMP_BINARY: /usr/bin/pg_dump
         CI: "true"
         GITHUB_TOKEN: ${{ env.GITHUB_TOKEN }}
         SITE_API_ENDPOINT: ${{ env.SITE_API_ENDPOINT }}
+        POSTGRES_URL: postgresql://postgres@localhost:${{ env.PG_PORT }}/postgres

src/main.ts

@@ -15,9 +15,9 @@ import { formatCost, queryPreview } from "./reporters/github/github.ts";
 import { DEFAULT_CONFIG, fetchAnalyzerConfig } from "./config.ts";
 
 async function runInCI(
-  postgresUrl: Connectable,
+  targetPostgresUrl: Connectable,
+  sourcePostgresUrl: Connectable,
   logPath: string,
-  statisticsPath?: string,
   maxCost?: number,
 ) {
   const siteApiEndpoint = env.SITE_API_ENDPOINT;
@@ -31,8 +31,8 @@ async function runInCI(
       : DEFAULT_CONFIG;
 
   const runner = await Runner.build({
-    postgresUrl,
-    statisticsPath,
+    targetPostgresUrl,
+    sourcePostgresUrl,
     logPath,
     maxCost,
     ignoredQueryHashes: config.ignoredQueryHashes,
@@ -85,8 +85,8 @@ async function runInCI(
       log.info(
         "main",
         `No baseline found on branch "${comparisonBranch}". Comparison will be skipped. ` +
-          `To establish a baseline, run the analyzer on pushes to "${comparisonBranch}" ` +
-          `(add "push: branches: [${comparisonBranch}]" to your workflow trigger).`,
+        `To establish a baseline, run the analyzer on pushes to "${comparisonBranch}" ` +
+        `(add "push: branches: [${comparisonBranch}]" to your workflow trigger).`,
       );
     }
   }
@@ -100,6 +100,7 @@ async function runInCI(
     );
   }
 
+  console.log("Creating report...")
   // Generate PR comment with comparison data
   await runner.report(reportContext);
 
@@ -153,19 +154,28 @@ async function main() {
       core.setFailed("POSTGRES_URL environment variable is not set");
       process.exit(1);
     }
+    if (!env.SOURCE_DATABASE_URL) {
+      core.setFailed("SOURCE_DATABASE_URL environment variable is not set");
+      process.exit(1);
+    }
     if (!env.LOG_PATH) {
       core.setFailed("LOG_PATH environment variable is not set");
       process.exit(1);
     }
     await runInCI(
       Connectable.fromString(env.POSTGRES_URL),
+      Connectable.fromString(env.SOURCE_DATABASE_URL),
       env.LOG_PATH,
-      env.STATISTICS_PATH,
       typeof env.MAX_COST === "number" ? env.MAX_COST : undefined,
     );
   } else {
     await runOutsideCI();
   }
 }
 
-await main();
+try {
+  await main();
+} catch (error) {
+  console.error(error);
+  process.exit(1);
+}

src/remote/query-optimizer.ts

@@ -49,6 +49,7 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
   private readonly queries = new Map<QueryHash, OptimizedQuery>();
   private readonly disabledIndexes = new DisabledIndexes();
 
+  private existingIndexes: IndexedTable[] = [];
   private target?: Target;
   private semaphore = new Sema(QueryOptimizer.MAX_CONCURRENCY);
   private _finish = Promise.withResolvers<void>();
@@ -100,6 +101,14 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
     return this._finish.promise;
   }
 
+  get statisticsMode(): StatisticsMode {
+    return this.target?.statistics.mode ?? QueryOptimizer.defaultStatistics;
+  }
+
+  getExistingIndexes(): IndexedTable[] {
+    return this.existingIndexes;
+  }
+
   getDisabledIndexes(): PgIdentifier[] {
     return [...this.disabledIndexes];
   }
@@ -128,8 +137,8 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
     const pg = this.manager.getOrCreateConnection(this.connectable);
     const ownStats = await Statistics.dumpStats(pg, version, "full");
     const statistics = new Statistics(pg, version, ownStats, statsMode);
-    const existingIndexes = await statistics.getExistingIndexes();
-    const filteredIndexes = this.filterDisabledIndexes(existingIndexes);
+    this.existingIndexes = await statistics.getExistingIndexes();
+    const filteredIndexes = this.filterDisabledIndexes(this.existingIndexes);
     const optimizer = new IndexOptimizer(pg, statistics, filteredIndexes, {
       trace: false,
     });

src/remote/remote.ts

@@ -71,6 +71,7 @@ export class Remote extends EventEmitter<RemoteEvents> {
     /** The manager for ONLY the source db connections */
     private readonly sourceManager: ConnectionManager = ConnectionManager
       .forRemoteDatabase(),
+    private readonly options: { disableQueryLoader: boolean } = { disableQueryLoader: false }
   ) {
     super();
     this.baseDbURL = targetURL.withDatabaseName(Remote.baseDbName);
@@ -108,6 +109,10 @@ export class Remote extends EventEmitter<RemoteEvents> {
         this.getDatabaseInfo(source),
       ]);
 
+    if (restoreResult.status === "rejected") {
+      throw new Error(`Schema sync failed: ${restoreResult.reason}`);
+    }
+
     if (fullSchema.status === "fulfilled") {
       this.schemaLoader?.update(fullSchema.value);
     }
@@ -200,7 +205,7 @@ export class Remote extends EventEmitter<RemoteEvents> {
    * there isn't already an in-flight request
    */
   private async pollQueriesOnce() {
-    if (this.queryLoader && !this.isPolling) {
+    if (this.queryLoader && !this.isPolling && !this.options.disableQueryLoader) {
       try {
         this.isPolling = true;
         await this.queryLoader.poll();
@@ -373,7 +378,9 @@ export class Remote extends EventEmitter<RemoteEvents> {
       log.error("Query loader exited", "remote");
       this.queryLoader = undefined;
     });
-    this.queryLoader.start();
+    if (!this.options.disableQueryLoader) {
+      this.queryLoader.start();
+    }
   }
 
   async cleanup(): Promise<void> {