@malloy-publisher/server 0.0.198 → 0.0.200

package/src/service/model_limits.spec.ts

@@ -0,0 +1,181 @@
+import { describe, expect, it } from "bun:test";
+
+import { PayloadTooLargeError } from "../errors";
+import {
+   assertWithinModelResponseLimits,
+   resolveModelQueryRowLimit,
+} from "./model_limits";
+
+describe("resolveModelQueryRowLimit", () => {
+   it("uses the user's LIMIT when set and below the maxRows ceiling", () => {
+      expect(
+         resolveModelQueryRowLimit(500, {
+            defaultLimit: 1000,
+            maxRows: 100_000,
+         }),
+      ).toBe(500);
+   });
+
+   it("falls back to defaultLimit when the user's LIMIT is undefined", () => {
+      expect(
+         resolveModelQueryRowLimit(undefined, {
+            defaultLimit: 1000,
+            maxRows: 100_000,
+         }),
+      ).toBe(1000);
+   });
+
+   it("falls back to defaultLimit when the user's LIMIT is 0 (Malloy returns 0 for 'no limit')", () => {
+      // Malloy's PreparedResult returns 0 from `resultExplore.limit` when the
+      // query has no LIMIT clause; treat that as "no user limit".
+      expect(
+         resolveModelQueryRowLimit(0, {
+            defaultLimit: 1000,
+            maxRows: 100_000,
+         }),
+      ).toBe(1000);
+   });
+
+   it("clamps a too-high user LIMIT to the maxRows + 1 sentinel", () => {
+      expect(
+         resolveModelQueryRowLimit(1_000_000, {
+            defaultLimit: 1000,
+            maxRows: 100_000,
+         }),
+      ).toBe(100_001);
+   });
+
+   it("clamps a too-high defaultLimit to the maxRows + 1 sentinel", () => {
+      // Operator misconfigured DEFAULT > MAX; the hard cap still wins.
+      expect(
+         resolveModelQueryRowLimit(undefined, {
+            defaultLimit: 1_000_000,
+            maxRows: 50,
+         }),
+      ).toBe(51);
+   });
+
+   it("returns the requested limit unchanged when maxRows is 0 (cap disabled)", () => {
+      expect(
+         resolveModelQueryRowLimit(1_000_000, {
+            defaultLimit: 1000,
+            maxRows: 0,
+         }),
+      ).toBe(1_000_000);
+   });
+
+   it("returns the default unchanged when maxRows is 0 and no user limit", () => {
+      expect(
+         resolveModelQueryRowLimit(undefined, {
+            defaultLimit: 1000,
+            maxRows: 0,
+         }),
+      ).toBe(1000);
+   });
+
+   it("rejects negative user limits by treating them as 'no limit'", () => {
+      // Defensive — shouldn't happen in practice, but a -1 from a malformed
+      // PreparedResult shouldn't propagate as a negative rowLimit to the driver.
+      expect(
+         resolveModelQueryRowLimit(-1, {
+            defaultLimit: 1000,
+            maxRows: 100_000,
+         }),
+      ).toBe(1000);
+   });
+});
+
+describe("assertWithinModelResponseLimits", () => {
+   it("does not throw when both counts are below their caps", () => {
+      expect(() =>
+         assertWithinModelResponseLimits(
+            500,
+            1_000,
+            { maxRows: 1000, maxBytes: 10_000 },
+            "model_query",
+         ),
+      ).not.toThrow();
+   });
+
+   it("does not throw when row count equals the cap exactly (sentinel hasn't fired)", () => {
+      expect(() =>
+         assertWithinModelResponseLimits(
+            1000,
+            1_000,
+            { maxRows: 1000, maxBytes: 10_000 },
+            "model_query",
+         ),
+      ).not.toThrow();
+   });
+
+   it("throws PayloadTooLargeError with the row-cap message on row overflow", () => {
+      expect(() =>
+         assertWithinModelResponseLimits(
+            1001,
+            1_000,
+            { maxRows: 1000, maxBytes: 10_000 },
+            "model_query",
+         ),
+      ).toThrow(PayloadTooLargeError);
+      expect(() =>
+         assertWithinModelResponseLimits(
+            1001,
+            1_000,
+            { maxRows: 1000, maxBytes: 10_000 },
+            "model_query",
+         ),
+      ).toThrow("more than 1000 rows");
+   });
+
+   it("throws PayloadTooLargeError with the byte-cap message on byte overflow", () => {
+      expect(() =>
+         assertWithinModelResponseLimits(
+            10,
+            50_000,
+            { maxRows: 1000, maxBytes: 10_000 },
+            "model_query",
+         ),
+      ).toThrow(PayloadTooLargeError);
+      expect(() =>
+         assertWithinModelResponseLimits(
+            10,
+            50_000,
+            { maxRows: 1000, maxBytes: 10_000 },
+            "model_query",
+         ),
+      ).toThrow("exceeded 10000 bytes");
+   });
+
+   it("prefers the row-cap message when both caps would have fired (row check runs first)", () => {
+      expect(() =>
+         assertWithinModelResponseLimits(
+            2000,
+            50_000,
+            { maxRows: 1000, maxBytes: 10_000 },
+            "model_query",
+         ),
+      ).toThrow("more than 1000 rows");
+   });
+
+   it("disables row cap when maxRows is 0", () => {
+      expect(() =>
+         assertWithinModelResponseLimits(
+            1_000_000,
+            1_000,
+            { maxRows: 0, maxBytes: 10_000 },
+            "model_query",
+         ),
+      ).not.toThrow();
+   });
+
+   it("disables byte cap when maxBytes is 0", () => {
+      expect(() =>
+         assertWithinModelResponseLimits(
+            10,
+            1_000_000_000,
+            { maxRows: 1000, maxBytes: 0 },
+            "model_query",
+         ),
+      ).not.toThrow();
+   });
+});

package/src/service/model_limits.ts

@@ -0,0 +1,110 @@
+/**
+ * Memory guards for the Malloy model-query path (the `runnable.run`
+ * flow used by `getQueryResults` and notebook cell execution).
+ *
+ * Two layered defenses:
+ *
+ *   1. {@link resolveModelQueryRowLimit} — compute the effective
+ *      `rowLimit` to push down to `runnable.run`. The user's Malloy
+ *      `LIMIT` clause wins when present; otherwise the operator-
+ *      tunable default ({@link getDefaultQueryRowLimit}) fills in.
+ *      Either way the result is clamped to `maxRows + 1` so the
+ *      database itself stops producing rows when a user-supplied
+ *      `LIMIT 1_000_000` would otherwise blow up the process.
+ *
+ *   2. {@link assertWithinModelResponseLimits} — post-run overflow
+ *      detection. If the connector returned `maxRows + 1` rows
+ *      (the sentinel) or the JSON-serialized response exceeds the
+ *      byte cap, throw `PayloadTooLargeError` so the caller sees a
+ *      clean HTTP 413.
+ *
+ * Caveat on the byte cap: this path runs `runnable.run` (buffered),
+ * not `runStream`, so by the time we measure bytes the result has
+ * already been materialized in memory. The byte cap here is loud-
+ * failure detection — it surfaces oversize responses with a 413
+ * instead of letting the client receive a half-transmitted payload
+ * — not OOM prevention. True prevention requires streaming +
+ * `Result` reconstruction from `DataRecord`s, which is out of scope
+ * for this step (the model-query streaming path entangles with
+ * Malloy's `Result` schema metadata in non-trivial ways).
+ *
+ * Both helpers are pure so they can be unit-tested without spinning
+ * up a model runtime; the caller injects the env-derived limits.
+ */
+
+import { PayloadTooLargeError } from "../errors";
+import {
+   recordQueryCapExceeded,
+   type QueryCapSource,
+} from "../query_cap_metrics";
+
+export interface ResolveRowLimitConfig {
+   /**
+    * Result of {@link getDefaultQueryRowLimit}. Applied when the
+    * user's Malloy query doesn't carry a `LIMIT` clause.
+    */
+   defaultLimit: number;
+   /**
+    * Result of {@link getMaxQueryRows}. The effective row limit is
+    * clamped to `maxRows + 1` so a sentinel-count overflow check can
+    * distinguish "ran right up to the cap" from "would have
+    * overflowed". A value of `0` disables the cap.
+    */
+   maxRows: number;
+}
+
+/**
+ * Compute the `rowLimit` to pass to `runnable.run`. The +1 sentinel
+ * mirrors the Step 1 / Step 2 patterns on the connection-query path
+ * so behavior is uniform across all query surfaces.
+ */
+export function resolveModelQueryRowLimit(
+   userLimit: number | undefined,
+   { defaultLimit, maxRows }: ResolveRowLimitConfig,
+): number {
+   const requested = userLimit && userLimit > 0 ? userLimit : defaultLimit;
+   if (maxRows <= 0) return requested;
+   return Math.min(requested, maxRows + 1);
+}
+
+export interface ModelResponseLimitsConfig {
+   /** Result of {@link getMaxQueryRows}. `0` disables the row cap. */
+   maxRows: number;
+   /** Result of {@link getMaxResponseBytes}. `0` disables the byte cap. */
+   maxBytes: number;
+}
+
+/**
+ * Throw {@link PayloadTooLargeError} (HTTP 413) when a model-query
+ * response exceeds either configured cap. `rowCount` should be the
+ * raw row count Malloy actually fetched (typically
+ * `result._queryResult.data.rawData.length`); `serializedBytes`
+ * should be the byte length of the JSON-stringified response that
+ * would otherwise be returned to the client.
+ *
+ * Row check uses the `> maxRows` sentinel (not `>= maxRows`), since
+ * {@link resolveModelQueryRowLimit} asked the connector for
+ * `maxRows + 1` and we want to fail only when that sentinel fires.
+ */
+export function assertWithinModelResponseLimits(
+   rowCount: number,
+   serializedBytes: number,
+   { maxRows, maxBytes }: ModelResponseLimitsConfig,
+   source: QueryCapSource,
+): void {
+   if (maxRows > 0 && rowCount > maxRows) {
+      // Tick the counter *before* throwing so it reflects the
+      // event even if a downstream `catch` swallows the error
+      // (notebook handlers and MCP tools both do this in places).
+      recordQueryCapExceeded("rows", source);
+      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.`,
+      );
+   }
+   if (maxBytes > 0 && serializedBytes > maxBytes) {
+      recordQueryCapExceeded("bytes", source);
+      throw new PayloadTooLargeError(
+         `Query response exceeded ${maxBytes} bytes (was ${serializedBytes}). Project fewer columns, add a LIMIT, or raise PUBLISHER_MAX_RESPONSE_BYTES.`,
+      );
+   }
+}

package/src/service/package.spec.ts

@@ -1,3 +1,5 @@
+import { DuckDBConnection } from "@malloydata/db-duckdb";
+import "@malloydata/db-duckdb/native";
 import { afterEach, beforeEach, describe, expect, it } from "bun:test";
 import { Stats } from "fs";
 import fs from "fs/promises";
@@ -11,7 +13,7 @@ import { Package } from "./package";
 type PartialModel = Pick<Model, "getPath">;
 
 describe("service/package", () => {
-   const testPackageDirectory = "testPackage";
+   const testPackageDirectory = resolve("testPackage");
 
    beforeEach(async () => {
       await fs.mkdir(testPackageDirectory, { recursive: true });
@@ -98,11 +100,7 @@ describe("service/package", () => {
                   testPackageDirectory,
                   new Map(),
                ),
-            ).rejects.toThrowError(
-               new PackageNotFoundError(
-                  "Package manifest for testPackage does not exist.",
-               ),
-            );
+            ).rejects.toBeInstanceOf(PackageNotFoundError);
          });
          it(
             "should return a Package object if the package exists",
@@ -340,10 +338,18 @@ describe("service/package", () => {
          it("should return the size of the database file", async () => {
             sinon.stub(fs, "stat").resolves({ size: 13 } as Stats);
 
+            // `getDatabaseInfo` now requires the caller to pass in the
+            // shared DuckDB connection (resolved once by `readDatabases`
+            // off the package's MalloyConfig). For this isolated unit
+            // test we mint a fresh ephemeral one — production paths
+            // reuse a single connection per package via `Package.create`.
+            const conn = new DuckDBConnection("duckdb");
+
             // @ts-expect-error Accessing private static method for testing
             const info = await Package.getDatabaseInfo(
                testPackageDirectory,
                "database.csv",
+               conn,
             );
 
             expect(info).toEqual({