fix(budgets): address review findings — explain sig, concurrency, dead code, shouldBlock

- Fix DriverWithExplain to use single-object `explain(request)` matching the
  actual postgres driver signature (was silently failing via catch block)
- Replace shared `observedRows` closure variable with per-plan WeakMap tracking
  to prevent cross-request row count interference in concurrent executions
- Remove dead `void ctx.now()` call whose return value was discarded
- Align latency shouldBlock to use OR logic (like row budgets) so strict mode
  consistently blocks for all budget types

Author: wmadden

packages/2-sql/5-runtime/src/plugins/budgets.ts

@@ -17,10 +17,10 @@ export interface BudgetsOptions {
 }
 
 interface DriverWithExplain {
-  explain?(
-    sql: string,
-    params: unknown[],
-  ): Promise<{ rows: ReadonlyArray<Record<string, unknown>> }>;
+  explain?(request: {
+    sql: string;
+    params: readonly unknown[];
+  }): Promise<{ rows: ReadonlyArray<Record<string, unknown>> }>;
 }
 
 async function computeEstimatedRows(
@@ -32,7 +32,7 @@ async function computeEstimatedRows(
   }
 
   try {
-    const result = await driver.explain(plan.sql, [...plan.params]);
+    const result = await driver.explain({ sql: plan.sql, params: plan.params });
     return extractEstimatedRows(result.rows);
   } catch {
     return undefined;
@@ -196,14 +196,13 @@ export function budgets<TContract = unknown, TAdapter = unknown, TDriver = unkno
   const rowSeverity = options?.severities?.rowCount ?? 'error';
   const latencySeverity = options?.severities?.latency ?? 'warn';
 
-  let observedRows = 0;
+  const observedRowsByPlan = new WeakMap<ExecutionPlan, { count: number }>();
 
   return Object.freeze({
     name: 'budgets',
 
     async beforeExecute(plan: ExecutionPlan, ctx: PluginContext<TContract, TAdapter, TDriver>) {
-      observedRows = 0;
-      void ctx.now();
+      observedRowsByPlan.set(plan, { count: 0 });
 
       if (isQueryAst(plan.ast)) {
         if (plan.ast.kind === 'select') {
@@ -217,17 +216,16 @@ export function budgets<TContract = unknown, TAdapter = unknown, TDriver = unkno
 
     async onRow(
       _row: Record<string, unknown>,
-      _plan: ExecutionPlan,
+      plan: ExecutionPlan,
       _ctx: PluginContext<TContract, TAdapter, TDriver>,
     ) {
-      void _row;
-      void _plan;
-      void _ctx;
-      observedRows += 1;
-      if (observedRows > maxRows) {
+      const state = observedRowsByPlan.get(plan);
+      if (!state) return;
+      state.count += 1;
+      if (state.count > maxRows) {
         throw budgetError('BUDGET.ROWS_EXCEEDED', 'Observed row count exceeds budget', {
           source: 'observed',
-          observedRows,
+          observedRows: state.count,
           maxRows,
         });
       }
@@ -240,7 +238,7 @@ export function budgets<TContract = unknown, TAdapter = unknown, TDriver = unkno
     ) {
       const latencyMs = result.latencyMs;
       if (latencyMs > maxLatencyMs) {
-        const shouldBlock = latencySeverity === 'error' && ctx.mode === 'strict';
+        const shouldBlock = latencySeverity === 'error' || ctx.mode === 'strict';
         emitBudgetViolation(
           budgetError('BUDGET.TIME_EXCEEDED', 'Query latency exceeds budget', {
             latencyMs,

packages/2-sql/5-runtime/test/budgets.test.ts

@@ -182,6 +182,33 @@ describe('budgets plugin', () => {
       },
       timeouts.default,
     );
+
+    it(
+      'tracks row counts independently per execution plan',
+      async () => {
+        const plugin = budgets({ maxRows: 2 });
+        const planA = createPlan({ sql: 'INSERT INTO "user" (id) VALUES ($1)' });
+        const planB = createPlan({ sql: 'INSERT INTO "user" (id) VALUES ($2)' });
+        const ctxA = createPluginContext();
+        const ctxB = createPluginContext();
+
+        await plugin.beforeExecute?.(planA, ctxA);
+        await plugin.beforeExecute?.(planB, ctxB);
+
+        await plugin.onRow?.({}, planA, ctxA);
+        await plugin.onRow?.({}, planB, ctxB);
+        await plugin.onRow?.({}, planA, ctxA);
+        await plugin.onRow?.({}, planB, ctxB);
+
+        await expect(plugin.onRow?.({}, planA, ctxA)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+        });
+        await expect(plugin.onRow?.({}, planB, ctxB)).rejects.toMatchObject({
+          code: 'BUDGET.ROWS_EXCEEDED',
+        });
+      },
+      timeouts.default,
+    );
   });
 
   describe('latency budget (afterExecute)', () => {
@@ -217,6 +244,22 @@ describe('budgets plugin', () => {
       timeouts.default,
     );
 
+    it(
+      'throws when latency exceeds budget in strict mode even with warn severity',
+      async () => {
+        const plugin = budgets({ maxLatencyMs: 100, severities: { latency: 'warn' } });
+        const plan = createPlan({ sql: 'SELECT 1', meta: { annotations: { limit: 1 } } });
+        const ctx = createPluginContext({ mode: 'strict' });
+        const result: AfterExecuteResult = { rowCount: 1, latencyMs: 200, completed: true };
+
+        await expect(plugin.afterExecute?.(plan, result, ctx)).rejects.toMatchObject({
+          code: 'BUDGET.TIME_EXCEEDED',
+          category: 'BUDGET',
+        });
+      },
+      timeouts.default,
+    );
+
     it(
       'does not warn when latency is within budget',
       async () => {
@@ -243,6 +286,7 @@ describe('budgets plugin', () => {
         };
         const plan = createPlan({
           sql: 'SELECT id FROM "user" LIMIT 100',
+          params: ['a', 'b'],
           meta: { annotations: { limit: 100 } },
         });
         const plugin = budgets({ maxRows: 10_000, explain: { enabled: true } });
@@ -252,6 +296,10 @@ describe('budgets plugin', () => {
           code: 'BUDGET.ROWS_EXCEEDED',
           details: expect.objectContaining({ source: 'explain' }),
         });
+        expect(explainDriver.explain).toHaveBeenCalledWith({
+          sql: plan.sql,
+          params: plan.params,
+        });
       },
       timeouts.default,
     );