@malloy-publisher/server 0.0.198 → 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,7 +1,10 @@
 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";
 
 type ApiNotebook = components["schemas"]["Notebook"];
 type ApiModel = components["schemas"]["Model"];
@@ -97,6 +100,7 @@ export class ModelController {
       cellIndex: number,
       filterParams?: FilterParams,
       bypassFilters?: boolean,
+      givens?: Record<string, GivenValue>,
    ): Promise<{
       type: "code" | "markdown";
       text: string;
@@ -108,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) {
@@ -117,6 +125,16 @@ export class ModelController {
          throw new ModelNotFoundError(`${notebookPath} is a model`);
       }
 
-      return model.executeNotebookCell(cellIndex, filterParams, bypassFilters);
+      return runWithQueryTimeout(
+         (abortSignal) =>
+            model.executeNotebookCell(
+               cellIndex,
+               filterParams,
+               bypassFilters,
+               givens,
+               abortSignal,
+            ),
+         getQueryTimeoutMs(),
+      );
    }
 }

package/src/controller/query.controller.ts

@@ -1,9 +1,12 @@
 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";
 
 type ApiQuery = components["schemas"]["QueryResult"];
 
@@ -32,23 +35,36 @@ export class QueryController {
       compactJson: boolean = false,
       filterParams?: FilterParams,
       bypassFilters?: boolean,
+      givens?: Record<string, GivenValue>,
    ): Promise<ApiQuery> {
       const environment = await this.environmentStore.getEnvironment(
          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,
+         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/health.spec.ts

@@ -0,0 +1,90 @@
+import { afterEach, beforeEach, describe, expect, it, spyOn } from "bun:test";
+import { Server } from "http";
+import { performGracefulShutdownAfterDrain } from "./health";
+import { logger } from "./logger";
+
+// Regression test for the graceful-shutdown ordering bug that caused
+//   [winston] Attempt to write logs with no transports: {"message":"Waiting 50 seconds..."}
+// to appear in production logs. logger.close() must run after every
+// logger.* call, including the "Waiting ... seconds after server close
+// before exit..." message.
+//
+// Tests call performGracefulShutdownAfterDrain directly rather than
+// emitting SIGTERM, so module-level operationalState is not mutated
+// and the spec stays isolated from sibling tests in the same process.
+describe("performGracefulShutdownAfterDrain: shutdown ordering", () => {
+   const originalExit = process.exit;
+   let callOrder: string[];
+
+   beforeEach(() => {
+      callOrder = [];
+
+      spyOn(logger, "info").mockImplementation(((msg: string) => {
+         callOrder.push(`info:${msg}`);
+         return logger;
+      }) as never);
+      spyOn(logger, "close").mockImplementation((() => {
+         callOrder.push("close");
+         return logger;
+      }) as never);
+      // Silence warn/error calls so spec output stays clean. They are
+      // not load-bearing for these assertions.
+      spyOn(logger, "warn").mockImplementation((() => logger) as never);
+      spyOn(logger, "error").mockImplementation((() => logger) as never);
+
+      process.exit = ((_code?: number) => {
+         callOrder.push("exit");
+      }) as never;
+   });
+
+   afterEach(() => {
+      process.exit = originalExit;
+   });
+
+   const fakeServer = (): Server => ({ listening: false }) as unknown as Server;
+
+   it("logs the 'Waiting ...' message before closing the logger", async () => {
+      await performGracefulShutdownAfterDrain(fakeServer(), fakeServer(), 0.05);
+
+      const waitingIdx = callOrder.findIndex((entry) =>
+         entry.startsWith("info:Waiting"),
+      );
+      const closeIdx = callOrder.indexOf("close");
+      const exitIdx = callOrder.indexOf("exit");
+
+      expect(waitingIdx).toBeGreaterThanOrEqual(0);
+      expect(closeIdx).toBeGreaterThanOrEqual(0);
+      expect(exitIdx).toBeGreaterThanOrEqual(0);
+      expect(waitingIdx).toBeLessThan(closeIdx);
+      expect(closeIdx).toBeLessThan(exitIdx);
+   });
+
+   it("emits no logger.info calls after logger.close", async () => {
+      await performGracefulShutdownAfterDrain(fakeServer(), fakeServer(), 0.05);
+
+      const closeIdx = callOrder.indexOf("close");
+      const lateInfoIdx = callOrder.findIndex(
+         (entry, idx) => idx > closeIdx && entry.startsWith("info:"),
+      );
+      expect(closeIdx).toBeGreaterThanOrEqual(0);
+      expect(lateInfoIdx).toBe(-1);
+   });
+
+   it("closes the logger exactly once", async () => {
+      await performGracefulShutdownAfterDrain(fakeServer(), fakeServer(), 0.05);
+
+      const closes = callOrder.filter((entry) => entry === "close").length;
+      expect(closes).toBe(1);
+   });
+
+   it("skips the 'Waiting ...' message when gracefulCloseTimeoutSeconds is 0", async () => {
+      await performGracefulShutdownAfterDrain(fakeServer(), fakeServer(), 0);
+
+      const waitingCalls = callOrder.filter((entry) =>
+         entry.startsWith("info:Waiting"),
+      );
+      expect(waitingCalls.length).toBe(0);
+      expect(callOrder.indexOf("close")).toBeGreaterThanOrEqual(0);
+      expect(callOrder.indexOf("exit")).toBeGreaterThanOrEqual(0);
+   });
+});