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,34 @@ runs:
         sudo make install # Use sudo to install globally
         cd ${{ github.action_path }} # Return to action directory
 
+    - name: Start PostgreSQL
+      shell: bash
+      run: |
+        PORT=$(shuf -i 10000-65000 -n 1)
+        echo "PG_PORT=$PORT" >> $GITHUB_ENV
+        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
+        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

@@ -100,6 +100,10 @@ export class QueryOptimizer extends EventEmitter<EventMap> {
     return this._finish.promise;
   }
 
+  get statisticsMode(): StatisticsMode {
+    return this.target?.statistics.mode ?? QueryOptimizer.defaultStatistics;
+  }
+
   getDisabledIndexes(): PgIdentifier[] {
     return [...this.disabledIndexes];
   }

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> {

src/reporters/site-api.ts

@@ -1,7 +1,7 @@
 import * as github from "@actions/github";
 import type { IndexRecommendation, Nudge, SQLCommenterTag, TableReference } from "@query-doctor/core";
 import { DEFAULT_CONFIG, type AnalyzerConfig } from "../config.ts";
-import type { QueryProcessResult } from "../runner.ts";
+import type { OptimizedQuery } from "../sql/recent-query.ts";
 
 interface CiRunPayload {
   repo: string;
@@ -25,25 +25,25 @@ export interface CiQueryPayload {
 
 export type CiOptimization =
   | {
-      state: "improvements_available";
-      cost: number;
-      optimizedCost: number;
-      costReductionPercentage: number;
-      indexRecommendations: CiIndexRecommendation[];
-      indexesUsed: string[];
-      explainPlan?: object;
-      optimizedExplainPlan?: object;
-    }
+    state: "improvements_available";
+    cost: number;
+    optimizedCost: number;
+    costReductionPercentage: number;
+    indexRecommendations: CiIndexRecommendation[];
+    indexesUsed: string[];
+    explainPlan?: object;
+    optimizedExplainPlan?: object;
+  }
   | {
-      state: "no_improvement_found";
-      cost: number;
-      indexesUsed: string[];
-      explainPlan?: object;
-    }
+    state: "no_improvement_found";
+    cost: number;
+    indexesUsed: string[];
+    explainPlan?: object;
+  }
   | {
-      state: "error";
-      error: string;
-    };
+    state: "error";
+    error: string;
+  };
 
 interface CiIndexRecommendation {
   schema: string;
@@ -114,105 +114,25 @@ function mapIndexRecommendation(rec: IndexRecommendation): CiIndexRecommendation
   };
 }
 
-function mapResultToQuery(result: QueryProcessResult): CiQueryPayload | null {
-  switch (result.kind) {
-    case "recommendation":
-      return {
-        hash: result.recommendation.fingerprint,
-        query: result.rawQuery,
-        formattedQuery: result.recommendation.formattedQuery,
-        nudges: result.nudges,
-        tags: result.tags,
-        tableReferences: result.referencedTables ?? [],
-        optimization: {
-          state: "improvements_available",
-          cost: result.recommendation.baseCost,
-          optimizedCost: result.recommendation.optimizedCost,
-          costReductionPercentage:
-            result.recommendation.baseCost > 0
-              ? ((result.recommendation.baseCost - result.recommendation.optimizedCost) /
-                  result.recommendation.baseCost) *
-                100
-              : 0,
-          indexRecommendations: result.indexRecommendations.map(mapIndexRecommendation),
-          indexesUsed: result.recommendation.existingIndexes,
-          explainPlan: result.recommendation.baseExplainPlan,
-          optimizedExplainPlan: result.recommendation.explainPlan,
-        },
-      };
-
-    case "no_improvement":
-      return {
-        hash: result.fingerprint,
-        query: result.rawQuery,
-        formattedQuery: result.formattedQuery,
-        nudges: result.nudges,
-        tags: result.tags,
-        tableReferences: result.referencedTables ?? [],
-        optimization: {
-          state: "no_improvement_found",
-          cost: result.cost,
-          indexesUsed: result.existingIndexes,
-          explainPlan: result.explainPlan,
-        },
-      };
-
-    case "zero_cost_plan":
-      return {
-        hash: result.fingerprint,
-        query: result.rawQuery,
-        formattedQuery: result.formattedQuery,
-        nudges: result.nudges,
-        tags: result.tags,
-        tableReferences: result.referencedTables ?? [],
-        optimization: {
-          state: "no_improvement_found",
-          cost: 0,
-          indexesUsed: [],
-          explainPlan: result.explainPlan,
-        },
-      };
-
-    case "error":
-      return {
-        hash: result.fingerprint,
-        query: result.rawQuery,
-        formattedQuery: result.formattedQuery,
-        nudges: result.nudges,
-        tags: result.tags,
-        tableReferences: result.referencedTables ?? [],
-        optimization: {
-          state: "error",
-          error: result.error.message,
-        },
-      };
-
-    case "cost_past_threshold":
-      return {
-        hash: result.warning.fingerprint,
-        query: result.rawQuery,
-        formattedQuery: result.warning.formattedQuery,
-        nudges: result.nudges,
-        tags: result.tags,
-        tableReferences: result.referencedTables ?? [],
-        optimization: result.warning.optimization
-          ? {
-              state: "no_improvement_found",
-              cost: result.warning.baseCost,
-              indexesUsed: result.warning.optimization.existingIndexes,
-              explainPlan: result.warning.explainPlan,
-            }
-          : {
-              state: "no_improvement_found",
-              cost: result.warning.baseCost,
-              indexesUsed: [],
-              explainPlan: result.warning.explainPlan,
-            },
-      };
-
-    case "invalid":
-      return null;
+function mapResultToQuery(result: OptimizedQuery): CiQueryPayload | null {
+  const { optimization } = result;
+  if (
+    optimization.state === "waiting" ||
+    optimization.state === "optimizing" ||
+    optimization.state === "not_supported" ||
+    optimization.state === "timeout"
+  ) {
+    return null;
   }
+  return {
+    hash: result.hash,
+    query: result.query,
+    formattedQuery: result.formattedQuery,
+    nudges: result.nudges,
+    tags: result.tags,
+    tableReferences: result.tableReferences ?? [],
+    optimization,
+  };
 }
 
 function getQueryCost(q: CiQueryPayload): number | null {
@@ -228,7 +148,7 @@ function getQueryIndexes(q: CiQueryPayload): string[] {
 }
 
 export function buildQueries(
-  results: QueryProcessResult[],
+  results: OptimizedQuery[],
   config: AnalyzerConfig = DEFAULT_CONFIG,
 ): CiQueryPayload[] {
   const ignoredSet = new Set(config.ignoredQueryHashes);