@malloy-publisher/server 0.0.199 → 0.0.200

package/src/controller/connection.controller.ts

@@ -1,8 +1,19 @@
 import { Connection, RunSQLOptions, TableSourceDef } from "@malloydata/malloy";
 import { PersistSQLResults } from "@malloydata/malloy/connection";
 import { components } from "../api";
-import { BadRequestError, ConnectionError } from "../errors";
+import {
+   getMaxQueryRows,
+   getMaxResponseBytes,
+   getQueryTimeoutMs,
+} from "../config";
+import {
+   BadRequestError,
+   ConnectionError,
+   PayloadTooLargeError,
+} from "../errors";
+import { recordQueryCapExceeded } from "../query_cap_metrics";
 import { logger } from "../logger";
+import { runWithQueryTimeout } from "../query_timeout";
 import { testConnectionConfig } from "../service/connection";
 import { validateDuckdbApiSurface } from "../service/connection_config";
 import { ConnectionService } from "../service/connection_service";
@@ -12,6 +23,7 @@ import {
 } from "../service/db_utils";
 import type { Environment } from "../service/environment";
 import { EnvironmentStore } from "../service/environment_store";
+import { isStreamingConnection, streamSqlWithBudget } from "../stream_helpers";
 
 type ApiConnection = components["schemas"]["Connection"];
 type ApiConnectionStatus = components["schemas"]["ConnectionStatus"];
@@ -442,6 +454,28 @@ export class ConnectionController {
       options: string,
       packageName?: string,
    ): Promise<ApiQueryData> {
+      // Express parses repeated query parameters (?sqlStatement=a&sqlStatement=b)
+      // and array-shaped JSON bodies as `string[]`, not `string`. The route
+      // handlers up-cast for TypeScript's benefit, so re-validate here at the
+      // controller boundary before forwarding to the connector.
+      if (typeof sqlStatement !== "string") {
+         throw new BadRequestError("sqlStatement must be a string");
+      }
+      if (
+         options !== undefined &&
+         options !== null &&
+         typeof options !== "string"
+      ) {
+         throw new BadRequestError("options must be a string");
+      }
+
+      // Shed load before any disk / DB work; same rationale as
+      // QueryController. The env-fetch is cached so this is effectively
+      // a single O(1) boolean read on the hot path.
+      const environmentForAdmission =
+         await this.environmentStore.getEnvironment(environmentName, false);
+      environmentForAdmission.assertCanAdmitQuery();
+
       const malloyConnection = await this.getMalloyConnection(
          environmentName,
          connectionName,
@@ -450,7 +484,27 @@ export class ConnectionController {
 
       let runSQLOptions: RunSQLOptions = {};
       if (options) {
-         runSQLOptions = JSON.parse(options) as RunSQLOptions;
+         // JSON.parse happily produces `null`, `42`, `"foo"`, `[1,2,3]`, etc.
+         // Anything that isn't a plain object would either crash later when
+         // we touch `.abortSignal` / `.rowLimit` (null), or mutate a caller
+         // string / array and pass it to the connector. Reject at the
+         // boundary alongside the other CodeQL type guards above.
+         let parsed: unknown;
+         try {
+            parsed = JSON.parse(options);
+         } catch {
+            throw new BadRequestError("options must be valid JSON");
+         }
+         if (
+            parsed === null ||
+            typeof parsed !== "object" ||
+            Array.isArray(parsed)
+         ) {
+            throw new BadRequestError(
+               'options must be a JSON object (e.g. {"rowLimit": 100})',
+            );
+         }
+         runSQLOptions = parsed as RunSQLOptions;
       }
       if (runSQLOptions.abortSignal) {
          // Add support for abortSignal in the future
@@ -458,15 +512,96 @@ export class ConnectionController {
          runSQLOptions.abortSignal = undefined;
       }
 
-      try {
-         return {
-            data: JSON.stringify(
-               await malloyConnection.runSQL(sqlStatement, runSQLOptions),
-            ),
+      // Bound the response with two layered caps:
+      //
+      //   - Row cap (PUBLISHER_MAX_QUERY_ROWS) — pushed to the driver as
+      //     RunSQLOptions.rowLimit = min(callerLimit, cap+1). The +1 is a
+      //     sentinel: a response of cap+1 rows means the original would have
+      //     overflowed, so we fail the request with HTTP 413 rather than
+      //     serializing it. A caller-supplied rowLimit lower than cap+1 is
+      //     kept; we only enforce the ceiling. Cap of 0 disables the bound.
+      //
+      //   - Byte cap (PUBLISHER_MAX_RESPONSE_BYTES) — the actual memory
+      //     bound, since a single 10 MB JSON column can blow past the
+      //     row cap's safe envelope. Only enforceable on streaming
+      //     connections; row-buffered `runSQL` has already done the
+      //     damage by the time we'd count bytes.
+      //
+      // We trust connectors to honor RunSQLOptions.rowLimit; connectors
+      // that don't are upstream bugs.
+      const maxRows = getMaxQueryRows();
+      const maxBytes = getMaxResponseBytes();
+      if (maxRows > 0) {
+         runSQLOptions.rowLimit = Math.min(
+            runSQLOptions.rowLimit ?? maxRows + 1,
+            maxRows + 1,
+         );
+      }
+
+      // Streaming-capable connections (Postgres, DuckDB, ...) go through
+      // streamSqlWithBudget so the byte cap can fire mid-stream. Other
+      // connections fall back to the buffered path; client-side byte
+      // counting there would be security theatre.
+      //
+      // The whole driver call is wrapped in runWithQueryTimeout so a
+      // hung query is surfaced as HTTP 504 instead of holding the
+      // HTTP socket open until the load balancer cuts the connection.
+      // The timeout signal is plumbed into RunSQLOptions.abortSignal
+      // so the abort actually cancels the underlying network call —
+      // not just unblocks the awaiter.
+      if (isStreamingConnection(malloyConnection)) {
+         const streamed = await runWithQueryTimeout(async (signal) => {
+            const optionsWithSignal: RunSQLOptions = {
+               ...runSQLOptions,
+               abortSignal: signal,
+            };
+            try {
+               return await streamSqlWithBudget(
+                  malloyConnection,
+                  sqlStatement,
+                  optionsWithSignal,
+                  { maxRows, maxBytes },
+               );
+            } catch (error) {
+               if (error instanceof PayloadTooLargeError) throw error;
+               // If runWithQueryTimeout is about to wrap this in a
+               // QueryTimeoutError (because the timer fired), the
+               // ConnectionError we'd throw here is discarded — the
+               // timeout verdict wins. So this branch only matters
+               // for genuine driver failures.
+               throw new ConnectionError((error as Error).message);
+            }
+         }, getQueryTimeoutMs());
+         return { data: JSON.stringify(streamed) };
+      }
+
+      const result = await runWithQueryTimeout(async (signal) => {
+         const optionsWithSignal: RunSQLOptions = {
+            ...runSQLOptions,
+            abortSignal: signal,
          };
-      } catch (error) {
-         throw new ConnectionError((error as Error).message);
+         try {
+            return await malloyConnection.runSQL(
+               sqlStatement,
+               optionsWithSignal,
+            );
+         } catch (error) {
+            throw new ConnectionError((error as Error).message);
+         }
+      }, getQueryTimeoutMs());
+
+      if (maxRows > 0 && result.rows.length > maxRows) {
+         // Tick the cap-exceeded counter on the buffered path too
+         // so `publisher_query_cap_exceeded_total{cap_type="rows",
+         // source="connection_sql"}` captures rejections from
+         // *all* connectors, not just streaming-capable ones.
+         recordQueryCapExceeded("rows", "connection_sql");
+         throw new PayloadTooLargeError(
+            `Query returned more than ${maxRows} rows. Refine the query (add a LIMIT or more selective WHERE) or raise PUBLISHER_MAX_QUERY_ROWS.`,
+         );
       }
+
+      return { data: JSON.stringify(result) };
    }
 
    public async getConnectionTemporaryTable(
@@ -475,23 +610,75 @@ export class ConnectionController {
       sqlStatement: string,
       packageName?: string,
    ): Promise<ApiTemporaryTable> {
+      // Express parses repeated query parameters / array-shaped JSON
+      // bodies as `string[]`. The route handlers up-cast for
+      // TypeScript's benefit; re-validate here at the controller
+      // boundary before forwarding to the connector — same pattern
+      // as `getConnectionQueryData`.
+      if (typeof sqlStatement !== "string") {
+         throw new BadRequestError("sqlStatement must be a string");
+      }
+
+      // Shed load before any disk / DB work; same rationale as
+      // `getConnectionQueryData`. `manifestTemporaryTable` issues a
+      // real `CREATE TEMPORARY TABLE AS (<sql>)` against the
+      // connector, so a back-pressured publisher must reject these
+      // with HTTP 503 just like ad-hoc SELECTs.
+      const environmentForAdmission =
+         await this.environmentStore.getEnvironment(environmentName, false);
+      environmentForAdmission.assertCanAdmitQuery();
+
       const malloyConnection = await this.getMalloyConnection(
          environmentName,
          connectionName,
          packageName,
       );
 
-      try {
-         return {
-            table: JSON.stringify(
-               await (
-                  malloyConnection as PersistSQLResults
-               ).manifestTemporaryTable(sqlStatement),
-            ),
-         };
-      } catch (error) {
-         throw new ConnectionError((error as Error).message);
-      }
+      // `manifestTemporaryTable(sqlCommand: string): Promise<string>`
+      // does not accept an `abortSignal`, so we cannot push
+      // cancellation down to the driver. Instead we race the call
+      // against the publisher's timeout signal so the HTTP response
+      // unblocks at the wall-clock deadline regardless of whether
+      // the underlying DDL responds. The DDL continues running in
+      // the database — the publisher just stops waiting and
+      // returns HTTP 504. This is strictly better than holding the
+      // socket and concurrency slot open until the load balancer
+      // kills the connection, but it does mean a runaway DDL is
+      // not actually cancelled (driver-side enhancement needed in
+      // `@malloydata/db-*` to plumb abort through).
+      //
+      // The `runWithQueryTimeout` wrapper's catch handler checks its
+      // `timedOut` flag at error time, so any rejection that
+      // propagates after the timer has fired is re-emitted as
+      // `QueryTimeoutError` regardless of its shape — that's why
+      // we re-throw the abort reason here unwrapped.
+      return await runWithQueryTimeout(async (signal) => {
+         try {
+            const tableName = await Promise.race([
+               (malloyConnection as PersistSQLResults).manifestTemporaryTable(
+                  sqlStatement,
+               ),
+               new Promise<never>((_resolve, reject) => {
+                  if (signal.aborted) {
+                     reject(signal.reason);
+                     return;
+                  }
+                  signal.addEventListener(
+                     "abort",
+                     () => reject(signal.reason),
+                     { once: true },
+                  );
+               }),
+            ]);
+            return { table: JSON.stringify(tableName) };
+         } catch (error) {
+            // If the abort fired first, the timeout wrapper above
+            // will convert this to QueryTimeoutError on its own
+            // — don't bury the reason in ConnectionError.
+            if (signal.aborted) throw error;
+            throw new ConnectionError((error as Error).message);
+         }
+      }, getQueryTimeoutMs());
    }
 
    public async testConnectionConfiguration(

package/src/controller/model.controller.ts

@@ -1,5 +1,7 @@
 import { components } from "../api";
+import { getQueryTimeoutMs } from "../config";
 import { ModelNotFoundError } from "../errors";
+import { runWithQueryTimeout } from "../query_timeout";
 import { EnvironmentStore } from "../service/environment_store";
 import type { FilterParams } from "../service/filter";
 import type { GivenValue } from "@malloydata/malloy";
@@ -110,6 +112,10 @@ export class ModelController {
          environmentName,
          false,
       );
+      // Shed load before any disk / DB work; same rationale as
+      // QueryController.getQuery — already-loaded packages bypass the
+      // package-load admission gate.
+      environment.assertCanAdmitQuery();
       const p = await environment.getPackage(packageName, false);
       const model = p.getModel(notebookPath);
       if (!model) {
@@ -119,11 +125,16 @@ export class ModelController {
          throw new ModelNotFoundError(`${notebookPath} is a model`);
       }
 
-      return model.executeNotebookCell(
-         cellIndex,
-         filterParams,
-         bypassFilters,
-         givens,
+      return runWithQueryTimeout(
+         (abortSignal) =>
+            model.executeNotebookCell(
+               cellIndex,
+               filterParams,
+               bypassFilters,
+               givens,
+               abortSignal,
+            ),
+         getQueryTimeoutMs(),
       );
    }
 }

package/src/controller/query.controller.ts

@@ -1,7 +1,9 @@
 import { validateRenderTags } from "@malloydata/render-validator";
 import { components } from "../api";
+import { getQueryTimeoutMs } from "../config";
 import { API_PREFIX } from "../constants";
 import { ModelNotFoundError } from "../errors";
+import { runWithQueryTimeout } from "../query_timeout";
 import { EnvironmentStore } from "../service/environment_store";
 import type { FilterParams } from "../service/filter";
 import type { GivenValue } from "@malloydata/malloy";
@@ -39,19 +41,30 @@ export class QueryController {
          environmentName,
          false,
       );
+      // Shed load before any disk / DB work so the publisher returns 503
+      // immediately under memory pressure instead of starting a query and
+      // crashing partway through. Already-loaded packages bypass the
+      // package-load admission gate, so this is the only thing protecting
+      // query traffic on a hot pod.
+      environment.assertCanAdmitQuery();
       const p = await environment.getPackage(packageName, false);
       const model = p.getModel(modelPath);
 
       if (!model) {
          throw new ModelNotFoundError(`${modelPath} does not exist`);
       } else {
-         const { result, compactResult } = await model.getQueryResults(
-            sourceName,
-            queryName,
-            query,
-            filterParams,
-            bypassFilters,
-            givens,
+         const { result, compactResult } = await runWithQueryTimeout(
+            (abortSignal) =>
+               model.getQueryResults(
+                  sourceName,
+                  queryName,
+                  query,
+                  filterParams,
+                  bypassFilters,
+                  givens,
+                  abortSignal,
+               ),
+            getQueryTimeoutMs(),
          );
          const renderLogs = validateRenderTags(result);
          return {

package/src/controller/watch-mode.controller.ts

@@ -1,8 +1,9 @@
 import chokidar, { FSWatcher } from "chokidar";
 import { RequestHandler } from "express";
-import path from "path";
 import { components } from "../api";
+import { internalErrorToHttpError } from "../errors";
 import { logger } from "../logger";
+import { assertSafePackageName, safeJoinUnderRoot } from "../path_safety";
 import { EnvironmentStore } from "../service/environment_store";
 
 type StartWatchReq = components["schemas"]["StartWatchRequest"];
@@ -32,6 +33,14 @@ export class WatchModeController {
       res,
    ) => {
       const watchName = req.body.environmentName ?? "";
+      try {
+         assertSafePackageName(watchName);
+      } catch (error) {
+         logger.error(error);
+         const { status } = internalErrorToHttpError(error as Error);
+         res.status(status).json({ error: (error as Error).message });
+         return;
+      }
       const environmentManifest =
          await EnvironmentStore.reloadEnvironmentManifest(
             this.environmentStore.serverRootPath,
@@ -53,7 +62,7 @@ export class WatchModeController {
          return;
       }
 
-      this.watchingPath = path.join(
+      this.watchingPath = safeJoinUnderRoot(
          this.environmentStore.serverRootPath,
          watchName,
       );

package/src/errors.spec.ts

@@ -4,6 +4,9 @@ import {
    ConnectionAuthError,
    ConnectionError,
    internalErrorToHttpError,
+   PayloadTooLargeError,
+   QueryTimeoutError,
+   ServiceUnavailableError,
 } from "./errors";
 
 describe("internalErrorToHttpError", () => {
@@ -39,4 +42,45 @@ describe("internalErrorToHttpError", () => {
       expect(status).toBe(500);
       expect(json.message).toBe("boom");
    });
+
+   it("maps PayloadTooLargeError to 413", () => {
+      const { status, json } = internalErrorToHttpError(
+         new PayloadTooLargeError(
+            "Query returned more than 100000 rows; refine the query or raise PUBLISHER_MAX_QUERY_ROWS.",
+         ),
+      );
+      expect(status).toBe(413);
+      expect(json).toEqual({
+         code: 413,
+         message:
+            "Query returned more than 100000 rows; refine the query or raise PUBLISHER_MAX_QUERY_ROWS.",
+      });
+   });
+
+   it("maps ServiceUnavailableError to 503 (load shedding / back-pressure)", () => {
+      const { status, json } = internalErrorToHttpError(
+         new ServiceUnavailableError(
+            "Pod at max concurrent queries (32); retry later.",
+         ),
+      );
+      expect(status).toBe(503);
+      expect(json).toEqual({
+         code: 503,
+         message: "Pod at max concurrent queries (32); retry later.",
+      });
+   });
+
+   it("maps QueryTimeoutError to 504 (gateway timeout, distinct from 503 back-pressure)", () => {
+      const { status, json } = internalErrorToHttpError(
+         new QueryTimeoutError(
+            "Query exceeded PUBLISHER_QUERY_TIMEOUT_MS (300000ms) and was aborted.",
+         ),
+      );
+      expect(status).toBe(504);
+      expect(json).toEqual({
+         code: 504,
+         message:
+            "Query exceeded PUBLISHER_QUERY_TIMEOUT_MS (300000ms) and was aborted.",
+      });
+   });
 });

package/src/errors.ts

@@ -30,6 +30,10 @@ export function internalErrorToHttpError(error: Error) {
       return httpError(409, error.message);
    } else if (error instanceof ServiceUnavailableError) {
       return httpError(503, error.message);
+   } else if (error instanceof PayloadTooLargeError) {
+      return httpError(413, error.message);
+   } else if (error instanceof QueryTimeoutError) {
+      return httpError(504, error.message);
    } else {
       return httpError(500, error.message);
    }
@@ -135,3 +139,33 @@ export class ServiceUnavailableError extends Error {
       super(message);
    }
 }
+
+/**
+ * Thrown when a response would exceed a server-side size cap (e.g. an
+ * ad-hoc connection SQL query that returned more than
+ * `PUBLISHER_MAX_QUERY_ROWS` rows). Mapped to HTTP 413 so callers know
+ * the request was well-formed but the result is too large for the
+ * publisher to materialize; the remediation is "refine the query" or
+ * "raise the cap", not "retry".
+ */
+export class PayloadTooLargeError extends Error {
+   constructor(message: string) {
+      super(message);
+   }
+}
+
+/**
+ * Thrown when a query exceeded the configured wall-clock budget
+ * (`PUBLISHER_QUERY_TIMEOUT_MS`) and the publisher aborted it
+ * mid-execution. Mapped to HTTP 504 (`Gateway Timeout`) because the
+ * publisher acts as a gateway to the underlying database — the
+ * upstream caller did nothing wrong, but the downstream query took
+ * too long. Distinct from {@link ServiceUnavailableError} so clients
+ * can distinguish "back off, the pod is loaded" (503, retryable)
+ * from "this specific query is too expensive" (504, refine it).
+ */
+export class QueryTimeoutError extends Error {
+   constructor(message: string) {
+      super(message);
+   }
+}

package/src/heap_check.spec.ts

@@ -0,0 +1,144 @@
+import { afterEach, beforeEach, describe, expect, it, mock } from "bun:test";
+
+import {
+   checkHeapConfiguration,
+   resetHeapTelemetryForTesting,
+} from "./heap_check";
+import {
+   startMetricsHarness,
+   type MetricsHarness,
+} from "./test_helpers/metrics_harness";
+
+function makeLogStub(): {
+   warn: ReturnType<typeof mock>;
+   info: ReturnType<typeof mock>;
+} {
+   return {
+      warn: mock(() => undefined),
+      info: mock(() => undefined),
+   };
+}
+
+describe("checkHeapConfiguration", () => {
+   it("warns + reports warned=true when heap limit is below the 2 GiB threshold", () => {
+      const log = makeLogStub();
+      // 1.5 GiB — below the recommended floor.
+      const { warned } = checkHeapConfiguration({
+         getHeapStatistics: () => ({
+            heap_size_limit: 1.5 * 1024 * 1024 * 1024,
+         }),
+         log,
+      });
+      expect(warned).toBe(true);
+      expect(log.warn).toHaveBeenCalledTimes(1);
+      expect(log.info).not.toHaveBeenCalled();
+      const args = log.warn.mock.calls[0] as unknown as [
+         string,
+         Record<string, unknown>,
+      ];
+      // Operator must be able to grep for the offending env-var
+      // hint without guessing the surrounding sentence.
+      expect(args[0]).toContain("--max-old-space-size");
+      expect(args[0]).toContain("MiB");
+      expect(args[1].heapSizeLimitBytes).toBe(1.5 * 1024 * 1024 * 1024);
+   });
+
+   it("logs at info + reports warned=false when heap limit meets the recommendation", () => {
+      const log = makeLogStub();
+      const { warned, heapSizeLimitBytes } = checkHeapConfiguration({
+         getHeapStatistics: () => ({
+            heap_size_limit: 4 * 1024 * 1024 * 1024,
+         }),
+         log,
+      });
+      expect(warned).toBe(false);
+      expect(heapSizeLimitBytes).toBe(4 * 1024 * 1024 * 1024);
+      expect(log.warn).not.toHaveBeenCalled();
+      expect(log.info).toHaveBeenCalledTimes(1);
+   });
+
+   it("treats exactly-the-recommended-value as 'meets threshold' (no warn)", () => {
+      // Boundary: the comparison is `<`, so equality should pass.
+      // Lock the boundary so a future tightening to `<=` requires
+      // an explicit test change rather than silently regressing.
+      const log = makeLogStub();
+      const { warned } = checkHeapConfiguration({
+         getHeapStatistics: () => ({
+            heap_size_limit: 2 * 1024 * 1024 * 1024,
+         }),
+         log,
+      });
+      expect(warned).toBe(false);
+      expect(log.warn).not.toHaveBeenCalled();
+   });
+
+   it("does not throw or warn under tiny heaps; the cap helpers still work", () => {
+      // Pathological "this pod has 128 MiB" case: we want a noisy
+      // warning, not a process crash, so the publisher still boots
+      // and the operator sees the message in pod logs.
+      const log = makeLogStub();
+      expect(() =>
+         checkHeapConfiguration({
+            getHeapStatistics: () => ({ heap_size_limit: 128 * 1024 * 1024 }),
+            log,
+         }),
+      ).not.toThrow();
+      expect(log.warn).toHaveBeenCalledTimes(1);
+   });
+
+   it("uses live v8.getHeapStatistics when no override is provided (smoke check)", () => {
+      // No assertion on warn/info — the result depends on how Node
+      // was started — just that the call resolves and returns a
+      // sensible structure. Locks the production code path against
+      // accidental coupling to the injected getter.
+      const result = checkHeapConfiguration();
+      expect(result.heapSizeLimitBytes).toBeGreaterThan(0);
+      expect(typeof result.warned).toBe("boolean");
+   });
+
+   describe("telemetry", () => {
+      let harness: MetricsHarness;
+
+      beforeEach(async () => {
+         harness = await startMetricsHarness();
+         resetHeapTelemetryForTesting();
+      });
+
+      afterEach(async () => {
+         resetHeapTelemetryForTesting();
+         await harness.shutdown();
+      });
+
+      it("publisher_heap_size_limit_bytes reports the live V8 heap_size_limit", async () => {
+         const log = makeLogStub();
+         checkHeapConfiguration({ log });
+         const value = await harness.collectGauge(
+            "publisher_heap_size_limit_bytes",
+         );
+         // Live value — we just assert it's a sensible positive
+         // number so the test isn't sensitive to how this Bun
+         // process was launched.
+         expect(typeof value).toBe("number");
+         expect(value).toBeGreaterThan(0);
+      });
+
+      it("publisher_heap_used_bytes reports the live V8 used_heap_size", async () => {
+         const log = makeLogStub();
+         checkHeapConfiguration({ log });
+         const value = await harness.collectGauge("publisher_heap_used_bytes");
+         expect(typeof value).toBe("number");
+         expect(value).toBeGreaterThan(0);
+      });
+
+      it("does not register the gauges before checkHeapConfiguration is called (lazy install)", async () => {
+         // The gauges are wired up via `installHeapGauges`, which
+         // is intentionally called from `checkHeapConfiguration`
+         // so the OTel SDK is fully up before instruments are
+         // registered. Without that, instruments would bind to
+         // NoOp during module load and never emit data.
+         expect(
+            await harness.collectGauge("publisher_heap_size_limit_bytes"),
+         ).toBeUndefined();
+      });
+   });
+});