@malloy-publisher/server 0.0.198 → 0.0.200

package/src/query_timeout.spec.ts

@@ -0,0 +1,224 @@
+import { afterEach, beforeEach, describe, expect, it } from "bun:test";
+
+import { QueryTimeoutError } from "./errors";
+import {
+   PUBLISHER_QUERY_TIMEOUT_REASON,
+   resetQueryTimeoutTelemetryForTesting,
+   runWithQueryTimeout,
+} from "./query_timeout";
+import {
+   startMetricsHarness,
+   type MetricsHarness,
+} from "./test_helpers/metrics_harness";
+
+describe("runWithQueryTimeout", () => {
+   it("returns the inner result when fn finishes before the timeout", async () => {
+      const result = await runWithQueryTimeout(async (signal) => {
+         expect(signal).toBeInstanceOf(AbortSignal);
+         expect(signal.aborted).toBe(false);
+         return "ok";
+      }, 1000);
+      expect(result).toBe("ok");
+   });
+
+   it("hands fn a signal even when the timeout is disabled (uniform contract)", async () => {
+      // timeoutMs=0 is the "operator opted out" path. We still want
+      // callers to be able to forward `signal` unconditionally; the
+      // signal must exist, be an AbortSignal, and never fire.
+      let observed: AbortSignal | undefined;
+      const result = await runWithQueryTimeout(async (signal) => {
+         observed = signal;
+         await new Promise((r) => setTimeout(r, 10));
+         return "still-here";
+      }, 0);
+      expect(result).toBe("still-here");
+      expect(observed).toBeInstanceOf(AbortSignal);
+      expect(observed?.aborted).toBe(false);
+   });
+
+   it("aborts the signal and throws QueryTimeoutError when fn exceeds the budget", async () => {
+      let observedSignal: AbortSignal | undefined;
+      await expect(
+         runWithQueryTimeout(async (signal) => {
+            observedSignal = signal;
+            // Mimic a slow driver: resolve only when the signal aborts
+            // (so the test doesn't hang). The thrown error from the
+            // driver is irrelevant — runWithQueryTimeout owns the
+            // verdict once the timer has fired.
+            await new Promise<void>((_resolve, reject) => {
+               signal.addEventListener("abort", () =>
+                  reject(new Error("driver: aborted by AbortSignal")),
+               );
+            });
+         }, 25),
+      ).rejects.toBeInstanceOf(QueryTimeoutError);
+      expect(observedSignal?.aborted).toBe(true);
+      // The reason sentinel lets composed helpers (e.g. streaming
+      // cap-abort) distinguish "publisher timeout" from "their own
+      // abort" without coupling to message strings.
+      expect(observedSignal?.reason).toBe(PUBLISHER_QUERY_TIMEOUT_REASON);
+   });
+
+   it("includes the configured timeout in the QueryTimeoutError message (operator can grep logs)", async () => {
+      let caught: unknown;
+      try {
+         await runWithQueryTimeout(async (signal) => {
+            await new Promise<void>((_resolve, reject) => {
+               signal.addEventListener("abort", () =>
+                  reject(new Error("driver aborted")),
+               );
+            });
+         }, 17);
+      } catch (err) {
+         caught = err;
+      }
+      expect(caught).toBeInstanceOf(QueryTimeoutError);
+      expect((caught as Error).message).toContain("17ms");
+      expect((caught as Error).message).toContain("PUBLISHER_QUERY_TIMEOUT_MS");
+   });
+
+   it("re-throws non-timeout errors verbatim (does not mask driver failures)", async () => {
+      // Driver fails *before* the timeout; we must surface its error
+      // unchanged so the controller's normal error mapping kicks in
+      // (502 ConnectionError, 400 BadRequestError, etc.) — wrapping
+      // every failure in QueryTimeoutError would lie to clients.
+      const driverError = new Error("upstream connection refused");
+      await expect(
+         runWithQueryTimeout(async () => {
+            throw driverError;
+         }, 1000),
+      ).rejects.toBe(driverError);
+   });
+
+   it("does not wrap an error that happens to mention 'abort' if the timer never fired", async () => {
+      // Edge case: a driver might surface its own AbortError for an
+      // unrelated reason (e.g. caller canceled, transport reset). If
+      // the publisher's timer never fired, the error is not ours and
+      // must not be re-cast as QueryTimeoutError.
+      const fakeAbort = Object.assign(new Error("aborted by something else"), {
+         name: "AbortError",
+      });
+      await expect(
+         runWithQueryTimeout(async () => {
+            throw fakeAbort;
+         }, 1000),
+      ).rejects.toBe(fakeAbort);
+   });
+
+   describe("telemetry", () => {
+      let harness: MetricsHarness;
+      beforeEach(async () => {
+         harness = await startMetricsHarness();
+         // Drop cached instruments so they re-init against the new
+         // provider; otherwise this test's writes go to a counter
+         // bound to the *previous* provider's reader.
+         resetQueryTimeoutTelemetryForTesting();
+      });
+      afterEach(async () => {
+         delete process.env.PUBLISHER_QUERY_TIMEOUT_MS;
+         await harness.shutdown();
+         resetQueryTimeoutTelemetryForTesting();
+      });
+
+      /**
+       * Install telemetry without firing a timeout. After the
+       * `ensureTimeoutTelemetry` fix any successful call is enough
+       * — both the counter and the gauge register on every entry,
+       * not just the timeout branch.
+       */
+      async function primeTelemetry(): Promise<void> {
+         await runWithQueryTimeout(async () => 0, 10_000);
+      }
+
+      it("publisher_query_timeout_total ticks each time the timer fires", async () => {
+         // Establish baseline (0) before the trigger so this test
+         // isn't sensitive to whatever else has happened earlier.
+         expect(
+            await harness.collectCounter("publisher_query_timeout_total"),
+         ).toBe(0);
+         await expect(
+            runWithQueryTimeout(async (signal) => {
+               await new Promise<void>((_resolve, reject) => {
+                  signal.addEventListener("abort", () =>
+                     reject(new Error("driver aborted")),
+                  );
+               });
+            }, 15),
+         ).rejects.toBeInstanceOf(QueryTimeoutError);
+         expect(
+            await harness.collectCounter("publisher_query_timeout_total"),
+         ).toBe(1);
+      });
+
+      it("does NOT tick the counter on non-timeout errors (driver failed before deadline)", async () => {
+         await expect(
+            runWithQueryTimeout(async () => {
+               throw new Error("upstream broken");
+            }, 1000),
+         ).rejects.toThrow("upstream broken");
+         // A driver failure is not a timeout — the counter must
+         // stay at zero or operators will chase phantom timeouts.
+         expect(
+            await harness.collectCounter("publisher_query_timeout_total"),
+         ).toBe(0);
+      });
+
+      it("publisher_query_timeout_ms gauge is registered after the FIRST call, not just after a timeout fires", async () => {
+         // Regression test for the lazy-init bug where the gauge
+         // installed only inside the timeout branch — leaving
+         // `publisher_query_timeout_ms` absent from `/metrics`
+         // until the first 504. Operators tuning the timeout
+         // BEFORE getting paged need this visible.
+         process.env.PUBLISHER_QUERY_TIMEOUT_MS = "30000";
+         await runWithQueryTimeout(async () => "ok", 60_000);
+         expect(await harness.collectGauge("publisher_query_timeout_ms")).toBe(
+            30000,
+         );
+      });
+
+      it("publisher_query_timeout_ms gauge reports the current config", async () => {
+         process.env.PUBLISHER_QUERY_TIMEOUT_MS = "42000";
+         await primeTelemetry();
+         expect(await harness.collectGauge("publisher_query_timeout_ms")).toBe(
+            42000,
+         );
+      });
+
+      it("publisher_query_timeout_ms gauge reports 0 when the timeout is opted out", async () => {
+         process.env.PUBLISHER_QUERY_TIMEOUT_MS = "0";
+         await primeTelemetry();
+         expect(await harness.collectGauge("publisher_query_timeout_ms")).toBe(
+            0,
+         );
+      });
+
+      it("publisher_query_timeout_ms gauge surfaces -1 on misconfig instead of crashing the scrape", async () => {
+         process.env.PUBLISHER_QUERY_TIMEOUT_MS = "garbage";
+         await primeTelemetry();
+         // Operators must be able to *see* misconfig in dashboards
+         // — silently dropping the data point would hide the
+         // problem until a query timed out.
+         expect(await harness.collectGauge("publisher_query_timeout_ms")).toBe(
+            -1,
+         );
+      });
+   });
+
+   it("clears the timer on success so the event loop can exit promptly", async () => {
+      // Hard to assert directly without leaking implementation
+      // internals. Proxy assertion: the call returns quickly and a
+      // subsequent timer doesn't race us. If the timer leaked we'd
+      // also see a QueryTimeoutError on the *next* call below — the
+      // signal would already be aborted.
+      const result = await runWithQueryTimeout(async () => "fast", 5000);
+      expect(result).toBe("fast");
+      // Run a second call: a leaked timer would have aborted the
+      // first signal, but the second call gets its own signal so
+      // this is really a smoke check.
+      const result2 = await runWithQueryTimeout(async (signal) => {
+         expect(signal.aborted).toBe(false);
+         return "also-fast";
+      }, 5000);
+      expect(result2).toBe("also-fast");
+   });
+});

package/src/query_timeout.ts

@@ -0,0 +1,178 @@
+import { metrics, type Counter } from "@opentelemetry/api";
+
+import { getQueryTimeoutMs } from "./config";
+import { QueryTimeoutError } from "./errors";
+
+/**
+ * Lazy-initialized telemetry. Instruments are created on first use
+ * rather than at module load so unit tests that install a real
+ * MeterProvider AFTER the module is imported still record their
+ * data — OTel JS's `ProxyMeter` binds counters created before
+ * `setGlobalMeterProvider` to NoOp instruments (see
+ * https://github.com/open-telemetry/opentelemetry-js/issues/3505).
+ * In production the first request triggers initialization, which
+ * is well after `server.ts` boot and any OTel SDK setup, so the
+ * lazy-init has no observable latency cost.
+ */
+let queryTimeoutCounter: Counter | null = null;
+let timeoutTelemetryInitialized = false;
+/**
+ * Idempotent installer for the timeout counter + config gauge.
+ * Called at the top of {@link runWithQueryTimeout} (every path,
+ * including the opt-out branch) so the gauge is registered with
+ * the OTel SDK as soon as the publisher serves its first query —
+ * not just on the first timeout firing. Returns the counter for
+ * convenience.
+ */
+function ensureTimeoutTelemetry(): Counter {
+   if (queryTimeoutCounter && timeoutTelemetryInitialized) {
+      return queryTimeoutCounter;
+   }
+   const meter = metrics.getMeter("publisher");
+   if (!queryTimeoutCounter) {
+      queryTimeoutCounter = meter.createCounter(
+         "publisher_query_timeout_total",
+         {
+            description:
+               "Queries aborted because PUBLISHER_QUERY_TIMEOUT_MS elapsed before the underlying SDK call completed",
+         },
+      );
+   }
+   if (!timeoutTelemetryInitialized) {
+      // Observable gauge so dashboards can render the *configured*
+      // timeout alongside actual query durations from
+      // `malloy_model_query_duration` /
+      // `http_server_request_duration_ms`. Read live on each
+      // scrape so an env-var change between scrapes is visible
+      // without a restart.
+      meter
+         .createObservableGauge("publisher_query_timeout_ms", {
+            description:
+               "Current effective PUBLISHER_QUERY_TIMEOUT_MS (0 = disabled)",
+            unit: "ms",
+         })
+         .addCallback((observation) => {
+            try {
+               observation.observe(getQueryTimeoutMs());
+            } catch {
+               // A misconfigured env var should fail the request
+               // that observes it, not the metric scrape. Surface
+               // as -1 so dashboards reveal the misconfig rather
+               // than silently dropping the sample.
+               observation.observe(-1);
+            }
+         });
+      timeoutTelemetryInitialized = true;
+   }
+   return queryTimeoutCounter;
+}
+
+/**
+ * Visible for tests so they can re-trigger lazy init against a
+ * freshly-installed MeterProvider between cases. Do NOT call from
+ * production code.
+ */
+export function resetQueryTimeoutTelemetryForTesting(): void {
+   queryTimeoutCounter = null;
+   timeoutTelemetryInitialized = false;
+}
+
+/**
+ * Per-query wall-clock guard. Hands an {@link AbortSignal} to `fn`
+ * and arms a `setTimeout` for `timeoutMs`. When the timer fires the
+ * signal is aborted with reason `Symbol.for("publisher.queryTimeout")`
+ * so downstream catch blocks can distinguish a publisher-issued
+ * timeout from a caller cancel or a driver-internal abort.
+ *
+ * Contract:
+ * - `fn` MUST forward `signal` into the underlying SDK call
+ *   (`runSQLOptions.abortSignal`, `runnable.run({ abortSignal })`,
+ *   etc.) so the abort actually cancels the work — not just unblocks
+ *   the awaiter. Failure to forward leaks the query for `timeoutMs`
+ *   beyond the 504 response.
+ * - `timeoutMs === 0` opts out (no timer is armed); the signal is
+ *   still passed for consistency. Use when an operator deliberately
+ *   sets `PUBLISHER_QUERY_TIMEOUT_MS=0`.
+ * - On timeout AND a subsequent rejection from `fn`, this throws
+ *   {@link QueryTimeoutError}. If `fn` happens to resolve cleanly
+ *   between "timer fired" and "we entered the catch" (a race that
+ *   any driver can win), the success value is returned to the
+ *   caller — a query that completed is more useful than a 504 with
+ *   an already-materialized result. The timeout counter ticks only
+ *   when 504 is actually emitted.
+ * - On non-timeout error, the underlying error is re-thrown
+ *   unmodified.
+ */
+export async function runWithQueryTimeout<T>(
+   fn: (signal: AbortSignal) => Promise<T>,
+   timeoutMs: number,
+): Promise<T> {
+   // Install telemetry on every call (idempotent) so the
+   // `publisher_query_timeout_ms` gauge shows up in `/metrics` as
+   // soon as the publisher serves its first query, even if no
+   // timeout ever fires. Without this, the gauge would be absent
+   // until the first 504 — useless for "tune the timeout BEFORE
+   // you get paged" workflows.
+   ensureTimeoutTelemetry();
+
+   if (timeoutMs <= 0) {
+      // Opt-out path: no timer, no abort. We still pass a never-aborts
+      // signal so `fn`'s signature is uniform and forwarding stays
+      // mechanical — no per-call branching for "did we get a timeout?".
+      const ac = new AbortController();
+      return fn(ac.signal);
+   }
+
+   const ac = new AbortController();
+   const reason = PUBLISHER_QUERY_TIMEOUT_REASON;
+   let timedOut = false;
+   const timer = setTimeout(() => {
+      timedOut = true;
+      // `abort(reason)` propagates the reason through `signal.reason`
+      // so a downstream catch can `signal.reason === reason` to tell
+      // "publisher timeout" from "client disconnect" from "driver
+      // internal error" without string-matching error messages.
+      ac.abort(reason);
+   }, timeoutMs);
+   // Match HTTP request lifecycle: don't keep the event loop alive
+   // just for the timer. If the process is shutting down and the
+   // query has already resolved, we don't want this hanging the
+   // graceful-shutdown.
+   timer.unref?.();
+
+   try {
+      return await fn(ac.signal);
+   } catch (error) {
+      if (timedOut) {
+         // Increment before throwing so the counter ticks even if
+         // the controller swallows the error (the MCP tool's catch
+         // surfaces failures as content payloads, for instance).
+         // Carry the configured timeout as a label so dashboards
+         // can pivot a flapping pod between "we tuned the env var
+         // down" and "queries got slower".
+         ensureTimeoutTelemetry().add(1, { timeout_ms: timeoutMs });
+         throw new QueryTimeoutError(
+            `Query exceeded PUBLISHER_QUERY_TIMEOUT_MS (${timeoutMs}ms) and was aborted. Refine the query (add a more selective WHERE, lower LIMIT, or simplify joins) or raise the timeout.`,
+         );
+      }
+      throw error;
+   } finally {
+      clearTimeout(timer);
+   }
+}
+
+/**
+ * Sentinel attached to `AbortSignal.reason` when a publisher-issued
+ * query timeout fires. Currently consumed by `runWithQueryTimeout`
+ * itself (and by tests verifying the wiring); exported so future
+ * call sites composing this signal with their own (e.g. a custom
+ * driver wrapper) can write `if (signal.reason === PUBLISHER_QUERY_TIMEOUT_REASON)`
+ * to detect "this was the publisher's timeout, not the cap" without
+ * coupling to error-message strings. `runWithQueryTimeout`'s own
+ * timeout-vs-other-error distinction uses the local `timedOut`
+ * flag rather than this symbol, so consumers can rely on the
+ * symbol being attached even if the implementation changes.
+ */
+export const PUBLISHER_QUERY_TIMEOUT_REASON = Symbol.for(
+   "publisher.queryTimeout",
+);

package/src/server-old.ts

@@ -41,7 +41,8 @@ import {
    NotImplementedError,
 } from "./errors";
 import { logger } from "./logger";
-import { normalizeQueryArray } from "./server";
+import { queryConcurrency } from "./query_concurrency";
+import { normalizeQueryArray } from "./query_param_utils";
 import { EnvironmentStore } from "./service/environment_store";
 
 const LEGACY_API_PREFIX = "/api/v0";
@@ -459,8 +460,18 @@ export function registerLegacyRoutes(
 
    // queryData (deprecated GET) + sqlQuery (supported POST), per-project +
    // per-package
+   // Legacy `/projects/...` query routes keep the GET `queryData`
+   // endpoints (unlike the modern `/environments/...` surface, which
+   // removed them) so existing SDK clients are not broken. The
+   // missing protection is concurrency: without `queryConcurrency()`
+   // a flood of legacy clients can saturate the pod even while the
+   // modern routes are properly gated. Apply the same per-pod cap
+   // here so the legacy surface respects PUBLISHER_MAX_CONCURRENT_QUERIES.
+   // Admission, timeout, and row/byte caps are already enforced by
+   // the shared controllers downstream.
    app.get(
       `${LEGACY_API_PREFIX}/projects/:projectName/connections/:connectionName/queryData`,
+      queryConcurrency(),
       async (req, res) => {
          try {
             res.status(200).json(
@@ -481,6 +492,7 @@ export function registerLegacyRoutes(
 
    app.get(
       `${LEGACY_API_PREFIX}/projects/:projectName/packages/:packageName/connections/:connectionName/queryData`,
+      queryConcurrency(),
       async (req, res) => {
          try {
             res.status(200).json(
@@ -502,6 +514,7 @@ export function registerLegacyRoutes(
 
    app.post(
       `${LEGACY_API_PREFIX}/projects/:projectName/connections/:connectionName/sqlQuery`,
+      queryConcurrency(),
       async (req, res) => {
          try {
             let options: string | ParsedQs | (string | ParsedQs)[] | undefined;
@@ -528,6 +541,7 @@ export function registerLegacyRoutes(
 
    app.post(
       `${LEGACY_API_PREFIX}/projects/:projectName/packages/:packageName/connections/:connectionName/sqlQuery`,
+      queryConcurrency(),
       async (req, res) => {
          try {
             let options: string | ParsedQs | (string | ParsedQs)[] | undefined;
@@ -556,6 +570,7 @@ export function registerLegacyRoutes(
    // temporaryTable (deprecated GET) + sqlTemporaryTable (supported POST)
    app.get(
       `${LEGACY_API_PREFIX}/projects/:projectName/connections/:connectionName/temporaryTable`,
+      queryConcurrency(),
       async (req, res) => {
          try {
             res.status(200).json(
@@ -575,6 +590,7 @@ export function registerLegacyRoutes(
 
    app.get(
       `${LEGACY_API_PREFIX}/projects/:projectName/packages/:packageName/connections/:connectionName/temporaryTable`,
+      queryConcurrency(),
       async (req, res) => {
          try {
             res.status(200).json(
@@ -595,6 +611,7 @@ export function registerLegacyRoutes(
 
    app.post(
       `${LEGACY_API_PREFIX}/projects/:projectName/connections/:connectionName/sqlTemporaryTable`,
+      queryConcurrency(),
       async (req, res) => {
          try {
             res.status(200).json(
@@ -614,6 +631,7 @@ export function registerLegacyRoutes(
 
    app.post(
       `${LEGACY_API_PREFIX}/projects/:projectName/packages/:packageName/connections/:connectionName/sqlTemporaryTable`,
+      queryConcurrency(),
       async (req, res) => {
          try {
             res.status(200).json(
@@ -780,6 +798,7 @@ export function registerLegacyRoutes(
 
    app.post(
       `${LEGACY_API_PREFIX}/projects/:projectName/packages/:packageName/models/*?/query`,
+      queryConcurrency(),
       async (req, res) => {
          if (req.body.versionId) {
             setVersionIdError(res);
@@ -856,6 +875,7 @@ export function registerLegacyRoutes(
    // Cell execution route comes BEFORE the general getNotebook wildcard
    app.get(
       `${LEGACY_API_PREFIX}/projects/:projectName/packages/:packageName/notebooks/*/cells/:cellIndex`,
+      queryConcurrency(),
       async (req, res) => {
          if (req.query.versionId) {
             setVersionIdError(res);