feat: retry failed queries with increasing timeout

Dockerfile

@@ -131,6 +131,8 @@ RUN sed -i 's|nobody:/|nobody:/home|' /etc/passwd && chown nobody:nobody /home
 
 ENV POSTGRES_URL=postgresql://postgres@localhost/postgres?host=/tmp
 
+EXPOSE 5432
+
 # Development command - starts both PostgreSQL and the analyzer
 CMD ["/bin/bash", "-c", "\
     su-exec postgres initdb -D $PGDATA || true && \

deno.json

@@ -45,7 +45,7 @@
     "prettier": "npm:prettier@^3",
     "prettier-plugin-sql": "npm:prettier-plugin-sql@^0.19",
     "sql-highlight": "npm:sql-highlight@^6.1.0",
-    "postgresjs": "https://deno.land/x/postgresjs@v3.4.7/mod.js",
+    "postgresjs": "https://deno.land/x/postgresjs@v3.4.3/mod.js",
     "zod": "npm:zod@^4.1.12"
   }
 }

pg14.Dockerfile

@@ -145,6 +145,8 @@ RUN su-exec postgres initdb -D $PGDATA || true && \
 
 USER postgres
 
+EXPOSE 2345
+
 CMD ["/bin/bash", "-c", "\
     pg_ctl -D $PGDATA -l $PGDATA/logfile start || (cat $PGDATA/logfile && exit 1) && \
     until pg_isready -h /tmp; do sleep 0.5; done && \

src/remote/optimization.ts

@@ -19,7 +19,10 @@ export const LiveQueryOptimization = z.discriminatedUnion("state", [
   z.object({
     state: z.literal("waiting"),
   }),
-  z.object({ state: z.literal("optimizing") }),
+  z.object({
+    state: z.literal("optimizing"),
+    retries: z.number().nonnegative(),
+  }),
   z.object({ state: z.literal("not_supported"), reason: z.string() }),
   z.object({
     state: z.literal("improvements_available"),
@@ -37,7 +40,11 @@ export const LiveQueryOptimization = z.discriminatedUnion("state", [
     indexesUsed: z.array(z.string()),
     explainPlan: z.custom<PostgresExplainStage>(),
   }),
-  z.object({ state: z.literal("timeout") }),
+  z.object({
+    state: z.literal("timeout"),
+    waitedMs: z.number(),
+    retries: z.number().nonnegative(),
+  }),
   z.object({
     state: z.literal("error"),
     error: z.string(),

src/remote/query-optimizer.test.ts

@@ -4,7 +4,7 @@ import { ConnectionManager } from "../sync/connection-manager.ts";
 import { Connectable } from "../sync/connectable.ts";
 import { setTimeout } from "node:timers/promises";
 import { assertArrayIncludes } from "@std/assert/array-includes";
-import { assert, assertGreater } from "@std/assert";
+import { assert, assertEquals, assertGreater } from "@std/assert";
 import { type OptimizedQuery, RecentQuery } from "../sql/recent-query.ts";
 
 Deno.test({
@@ -446,3 +446,101 @@ Deno.test({
     }
   },
 });
+
+Deno.test({
+  name:
+    "timed out queries are retried with exponential backoff up to maxRetries",
+  sanitizeOps: false,
+  sanitizeResources: false,
+  fn: async () => {
+    const pg = await new PostgreSqlContainer("postgres:17")
+      .withCopyContentToContainer([
+        {
+          content: `
+            create table slow_table(id int, data text);
+            insert into slow_table (id, data) select i, repeat('x', 1000) from generate_series(1, 100) i;
+            create extension pg_stat_statements;
+            select * from slow_table 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 maxRetries = 2;
+    const queryTimeoutMs = 1;
+
+    const manager = ConnectionManager.forLocalDatabase();
+    const conn = Connectable.fromString(pg.getConnectionUri());
+    const optimizer = new QueryOptimizer(manager, conn, {
+      maxRetries,
+      queryTimeoutMs,
+      queryTimeoutMaxMs: 100,
+    });
+
+    const timeoutEvents: { query: OptimizedQuery; waitedMs: number }[] = [];
+    optimizer.addListener("timeout", (query, waitedMs) => {
+      timeoutEvents.push({ query, waitedMs });
+    });
+
+    const connector = manager.getConnectorFor(conn);
+    try {
+      const recentQueries = await connector.getRecentQueries();
+      const slowQuery = recentQueries.find((q) =>
+        q.query.includes("slow_table") && q.query.startsWith("select")
+      );
+      assert(slowQuery, "Expected to find slow_table query");
+
+      await optimizer.start([slowQuery], {
+        kind: "fromStatisticsExport",
+        source: { kind: "inline" },
+        stats: [{
+          tableName: "slow_table",
+          schemaName: "public",
+          relpages: 1000,
+          reltuples: 1_000_000,
+          relallvisible: 1,
+          columns: [
+            { columnName: "id", stats: null },
+            { columnName: "data", stats: null },
+          ],
+          indexes: [],
+        }],
+      });
+
+      await optimizer.finish;
+
+      const queries = optimizer.getQueries();
+      const resultQuery = queries.find((q) => q.query.includes("slow_table"));
+      assert(resultQuery, "Expected slow_table query in results");
+
+      assertEquals(
+        resultQuery.optimization.state,
+        "timeout",
+        "Expected query to be in timeout state",
+      );
+
+      if (resultQuery.optimization.state === "timeout") {
+        assertEquals(
+          resultQuery.optimization.retries,
+          maxRetries,
+          `Expected ${maxRetries} retries`,
+        );
+      }
+    } finally {
+      await pg.stop();
+    }
+  },
+});