@malloy-publisher/server 0.0.198 → 0.0.200

package/src/server.ts

@@ -1,4 +1,5 @@
 // Pre-load the instrumentation module; the instrumentation module must be loaded before the other imports.
+import type { GivenValue } from "@malloydata/malloy";
 import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 import bodyParser from "body-parser";
 import cors from "cors";
@@ -34,6 +35,8 @@ import {
 import { logger, loggerMiddleware } from "./logger";
 
 import { getMemoryGovernorConfig } from "./config";
+import { checkHeapConfiguration } from "./heap_check";
+import { queryConcurrency } from "./query_concurrency";
 import { ManifestController } from "./controller/manifest.controller";
 import { MaterializationController } from "./controller/materialization.controller";
 import { initializeMcpServer } from "./mcp/server";
@@ -41,14 +44,10 @@ import { registerLegacyRoutes } from "./server-old";
 import { EnvironmentStore } from "./service/environment_store";
 import { ManifestService } from "./service/manifest_service";
 import { MaterializationService } from "./service/materialization_service";
+import { normalizeQueryArray } from "./query_param_utils";
 import { PackageMemoryGovernor } from "./service/package_memory_governor";
 
-/** Normalize an Express query param into a string[] or undefined. */
-export function normalizeQueryArray(value: unknown): string[] | undefined {
-   if (value === undefined || value === null) return undefined;
-   if (Array.isArray(value)) return value.map(String);
-   return [String(value)];
-}
+export { normalizeQueryArray } from "./query_param_utils";
 
 // Parse command line arguments
 function parseArgs() {
@@ -122,10 +121,12 @@ function parseArgs() {
    // Zero-config invocation (`npx @malloy-publisher/server`) opts in to
    // the bundled DuckDB-only sample config so the Quick Start works
    // without any flags. Any explicit --server_root or --config disables
-   // this — the user told us where to look. Skip in NODE_ENV=test so
-   // specs that import this module for utility helpers (e.g.
-   // db_utils.spec.ts -> normalizeQueryArray) don't get the bundled
-   // default leaked into their EnvironmentStore construction.
+   // this — the user told us where to look. Skip in NODE_ENV=test as a
+   // belt-and-suspenders so any spec that ends up evaluating this
+   // module doesn't accidentally pin the EnvironmentStore to the
+   // bundled malloy-samples config; query-param helpers have been
+   // moved to `./query_param_utils` precisely so unit specs no longer
+   // need to import this module at all.
    if (!sawServerRoot && !sawConfig && process.env.NODE_ENV !== "test") {
       process.env.PUBLISHER_USE_BUNDLED_DEFAULT = "true";
    }
@@ -155,6 +156,12 @@ const isDevelopment = process.env["NODE_ENV"] === "development";
 export const app = express();
 app.use(loggerMiddleware);
 app.use(httpMetricsMiddleware);
+// Probe the V8 heap ceiling once at startup and warn if it's below
+// the recommended floor. The row/byte caps from Steps 1–3 still
+// bound per-request memory; this is a "your --max-old-space-size
+// looks low for the default caps" advisory so operators don't
+// chase OOMKills before checking the obvious config.
+checkHeapConfiguration();
 const environmentStore = new EnvironmentStore(SERVER_ROOT);
 const manifestService = new ManifestService(environmentStore);
 const watchModeController = new WatchModeController(environmentStore);
@@ -714,55 +721,18 @@ app.post(
    },
 );
 
-/**
- * @deprecated Use /environments/:environmentName/connections/:connectionName/sqlQuery POST method instead
- */
-app.get(
-   `${API_PREFIX}/environments/:environmentName/connections/:connectionName/queryData`,
-   async (req, res) => {
-      try {
-         res.status(200).json(
-            await connectionController.getConnectionQueryData(
-               req.params.environmentName,
-               req.params.connectionName,
-               req.query.sqlStatement as string,
-               req.query.options as string,
-            ),
-         );
-      } catch (error) {
-         logger.error(error);
-         const { json, status } = internalErrorToHttpError(error as Error);
-         res.status(status).json(json);
-      }
-   },
-);
-
-/**
- * @deprecated Use /environments/:environmentName/packages/:packageName/connections/:connectionName/sqlQuery
- */
-app.get(
-   `${API_PREFIX}/environments/:environmentName/packages/:packageName/connections/:connectionName/queryData`,
-   async (req, res) => {
-      try {
-         res.status(200).json(
-            await connectionController.getConnectionQueryData(
-               req.params.environmentName,
-               req.params.connectionName,
-               req.query.sqlStatement as string,
-               req.query.options as string,
-               req.params.packageName,
-            ),
-         );
-      } catch (error) {
-         logger.error(error);
-         const { json, status } = internalErrorToHttpError(error as Error);
-         res.status(status).json(json);
-      }
-   },
-);
-
+// NOTE: The deprecated `GET …/connections/:connectionName/queryData`
+// and `GET …/packages/:packageName/connections/:connectionName/queryData`
+// routes were removed in the operational-guards changeset.
+// They had been marked `@deprecated` for several releases; clients
+// must now use the POST `…/sqlQuery` endpoints below, which take the
+// SQL in the request body so the row/byte caps and query-timeout
+// signals introduced in the OOM-mitigation work apply uniformly.
+// The legacy `GET /projects/…/queryData` twins under `server-old.ts`
+// remain in place for now.
 app.post(
    `${API_PREFIX}/environments/:environmentName/connections/:connectionName/sqlQuery`,
+   queryConcurrency(),
    async (req, res) => {
       try {
          let options: string | ParsedQs | (string | ParsedQs)[] | undefined;
@@ -792,6 +762,7 @@ app.post(
 
 app.post(
    `${API_PREFIX}/environments/:environmentName/packages/:packageName/connections/:connectionName/sqlQuery`,
+   queryConcurrency(),
    async (req, res) => {
       try {
          let options: string | ParsedQs | (string | ParsedQs)[] | undefined;
@@ -822,6 +793,7 @@ app.post(
  */
 app.get(
    `${API_PREFIX}/environments/:environmentName/connections/:connectionName/temporaryTable`,
+   queryConcurrency(),
    async (req, res) => {
       try {
          res.status(200).json(
@@ -844,6 +816,7 @@ app.get(
  */
 app.get(
    `${API_PREFIX}/environments/:environmentName/packages/:packageName/connections/:connectionName/temporaryTable`,
+   queryConcurrency(),
    async (req, res) => {
       try {
          res.status(200).json(
@@ -864,6 +837,7 @@ app.get(
 
 app.post(
    `${API_PREFIX}/environments/:environmentName/connections/:connectionName/sqlTemporaryTable`,
+   queryConcurrency(),
    async (req, res) => {
       try {
          res.status(200).json(
@@ -883,6 +857,7 @@ app.post(
 
 app.post(
    `${API_PREFIX}/environments/:environmentName/packages/:packageName/connections/:connectionName/sqlTemporaryTable`,
+   queryConcurrency(),
    async (req, res) => {
       try {
          res.status(200).json(
@@ -1077,6 +1052,7 @@ app.get(
 // to avoid the wildcard matching incorrectly
 app.get(
    `${API_PREFIX}/environments/:environmentName/packages/:packageName/notebooks/*/cells/:cellIndex`,
+   queryConcurrency(),
    async (req, res) => {
       if (req.query.versionId) {
          setVersionIdError(res);
@@ -1110,6 +1086,18 @@ app.get(
          const bypassFilters =
             req.query.bypass_filters === "true" ? true : undefined;
 
+         let givens: Record<string, GivenValue> | undefined;
+         if (typeof req.query.givens === "string") {
+            try {
+               givens = JSON.parse(req.query.givens);
+            } catch {
+               res.status(400).json({
+                  error: "Invalid givens: must be valid JSON",
+               });
+               return;
+            }
+         }
+
          res.status(200).json(
             await modelController.executeNotebookCell(
                req.params.environmentName,
@@ -1118,6 +1106,7 @@ app.get(
                cellIndex,
                filterParams,
                bypassFilters,
+               givens,
             ),
          );
       } catch (error) {
@@ -1156,6 +1145,7 @@ app.get(
 
 app.post(
    `${API_PREFIX}/environments/:environmentName/packages/:packageName/models/*?/query`,
+   queryConcurrency(),
    async (req, res) => {
       if (req.body.versionId) {
          setVersionIdError(res);
@@ -1178,6 +1168,7 @@ app.post(
                   | Record<string, string | string[]>
                   | undefined,
                req.body.bypassFilters === true ? true : undefined,
+               req.body.givens as Record<string, GivenValue> | undefined,
             ),
          );
       } catch (error) {
@@ -1221,6 +1212,7 @@ app.post(
             req.params.modelName,
             req.body.source,
             req.body.includeSql === true,
+            req.body.givens as Record<string, GivenValue> | undefined,
          );
          res.status(200).json(result);
       } catch (error) {
@@ -1431,6 +1423,18 @@ app.use(
    },
 );
 
+// Eagerly construct the package-load worker pool so we fail fast at
+// boot if PACKAGE_LOAD_WORKERS is misconfigured (e.g. set to 0, the
+// removed in-process fallback). Surfacing the bad config here is much
+// friendlier than surfacing it on the first package load, which could
+// be hours after start.
+{
+   const { getPackageLoadPool } = await import(
+      "./package_load/package_load_pool"
+   );
+   getPackageLoadPool();
+}
+
 const mainServer = http.createServer({ maxHeaderSize: 262144 }, app);
 
 mainServer.timeout = 600000;

package/src/service/connection.ts

@@ -22,10 +22,14 @@ import {
 import type { LookupConnection } from "@malloydata/malloy/connection";
 import { AxiosError } from "axios";
 import fs from "fs/promises";
-import path from "path";
 import { components } from "../api";
 import { logAxiosError, logger } from "../logger";
 import { redactPgSecrets } from "../pg_helpers";
+import {
+   assertSafeEnvironmentPath,
+   assertSafePackageName,
+   safeJoinUnderRoot,
+} from "../path_safety";
 import {
    assembleEnvironmentConnections,
    CoreConnectionEntry,
@@ -814,7 +818,9 @@ export async function deleteDuckLakeConnectionFile(
    connectionName: string,
    environmentPath: string,
 ): Promise<void> {
-   const ducklakePath = path.join(
+   assertSafePackageName(connectionName);
+   assertSafeEnvironmentPath(environmentPath);
+   const ducklakePath = safeJoinUnderRoot(
       environmentPath,
       `${connectionName}_ducklake.duckdb`,
    );

package/src/service/db_utils.spec.ts

@@ -12,7 +12,7 @@ mock.module("@google-cloud/bigquery", () => ({
 }));
 
 import { Connection } from "@malloydata/malloy";
-import { normalizeQueryArray } from "../server";
+import { normalizeQueryArray } from "../query_param_utils";
 import {
    extractErrorDataFromError,
    getSchemasForConnection,

package/src/service/environment.ts

@@ -1,5 +1,6 @@
-import type { LogMessage } from "@malloydata/malloy";
+import type { GivenValue, LogMessage } from "@malloydata/malloy";
 import { MalloyError, Runtime } from "@malloydata/malloy";
+import { metrics } from "@opentelemetry/api";
 import { Mutex } from "async-mutex";
 import crypto from "crypto";
 import * as fs from "fs";
@@ -69,6 +70,49 @@ type RetiredConnectionGeneration = {
 
 const RETIRED_CONNECTION_DRAIN_MS = 30_000;
 
+/**
+ * Module-scoped admission-rejection counters. Lazy-initialized so
+ * the OTel JS `ProxyMeter` cannot strand them on a NoOp instrument
+ * created before the SDK MeterProvider was registered (a real risk
+ * in unit tests; see comment in `query_timeout.ts`). Environment
+ * name is attached as a label so dashboards can identify hot
+ * environments without grepping logs.
+ */
+import { type Counter } from "@opentelemetry/api";
+let queryAdmissionRejectionsCounter: Counter | null = null;
+let packageAdmissionRejectionsCounter: Counter | null = null;
+function getQueryAdmissionRejectionsCounter(): Counter {
+   if (queryAdmissionRejectionsCounter) return queryAdmissionRejectionsCounter;
+   queryAdmissionRejectionsCounter = metrics
+      .getMeter("publisher")
+      .createCounter("publisher_query_admission_rejections_total", {
+         description:
+            "Queries rejected with 503 because Environment.assertCanAdmitQuery() observed memory back-pressure",
+      });
+   return queryAdmissionRejectionsCounter;
+}
+function getPackageAdmissionRejectionsCounter(): Counter {
+   if (packageAdmissionRejectionsCounter) {
+      return packageAdmissionRejectionsCounter;
+   }
+   packageAdmissionRejectionsCounter = metrics
+      .getMeter("publisher")
+      .createCounter("publisher_package_admission_rejections_total", {
+         description:
+            "Package loads rejected with 503 because Environment.assertCanAdmitNewPackage() observed memory back-pressure",
+      });
+   return packageAdmissionRejectionsCounter;
+}
+
+/**
+ * Visible for tests; production code never calls this. Resets the
+ * lazy caches so a fresh MeterProvider can capture future writes.
+ */
+export function resetAdmissionTelemetryForTesting(): void {
+   queryAdmissionRejectionsCounter = null;
+   packageAdmissionRejectionsCounter = null;
+}
+
 export class Environment {
    private packages: Map<string, Package> = new Map();
    // Lock ordering: connectionMutex (environment) MUST be acquired before any
@@ -176,6 +220,7 @@ export class Environment {
       environmentPath: string,
       connections: ApiConnection[],
    ): Promise<Environment> {
+      assertSafeEnvironmentPath(environmentPath);
       if (!(await fs.promises.stat(environmentPath))?.isDirectory()) {
          throw new EnvironmentNotFoundError(
             `Environment path ${environmentPath} not found`,
@@ -218,7 +263,7 @@ export class Environment {
       try {
          readme = (
             await fs.promises.readFile(
-               path.join(this.environmentPath, README_NAME),
+               safeJoinUnderRoot(this.environmentPath, README_NAME),
             )
          ).toString();
       } catch {
@@ -238,6 +283,7 @@ export class Environment {
       modelName: string,
       source: string,
       includeSql: boolean = false,
+      givens?: Record<string, GivenValue>,
    ): Promise<{ problems: LogMessage[]; sql?: string }> {
       assertSafePackageName(packageName);
       assertSafeRelativeModelPath(modelName);
@@ -308,7 +354,7 @@ export class Environment {
             if (includeSql) {
                try {
                   const queryMaterializer = modelMaterializer.loadFinalQuery();
-                  sql = await queryMaterializer.getSQL();
+                  sql = await queryMaterializer.getSQL({ givens });
                } catch {
                   // Source may not contain a runnable query (e.g. only source definitions),
                   // in which case we simply omit the sql field.
@@ -578,8 +624,43 @@ export class Environment {
    ): void {
       if (allowAdmission) return;
       if (!this.memoryGovernor?.isBackpressured()) return;
+      // Increment *before* throwing so the metric ticks even on
+      // the not-uncommon "caught and swallowed" path. The label
+      // shape mirrors `assertCanAdmitQuery` so a dashboard panel
+      // can sum both rejection kinds by environment.
+      getPackageAdmissionRejectionsCounter().add(1, {
+         environment: this.environmentName,
+         reason,
+      });
+      throw new ServiceUnavailableError(
+         `Publisher is under memory pressure and cannot ${reason} (package "${packageName}", environment "${this.environmentName}"). Retry after the server's memory usage drops below the low-water mark (PUBLISHER_MEMORY_LOW_WATER_FRACTION of PUBLISHER_MAX_MEMORY_BYTES), or raise PUBLISHER_MAX_MEMORY_BYTES if you have headroom.`,
+      );
+   }
+
+   /**
+    * Reject incoming queries with HTTP 503 when the memory governor
+    * has tripped its high-water mark. Used by every query controller
+    * (connection SQL, model query, notebook cell, MCP `execute_query`)
+    * to shed load before the query runs — complementing
+    * {@link assertCanAdmitNewPackage}, which only fires on cache-miss
+    * package loads and so leaves already-loaded packages fully
+    * queryable under pressure. With this in place, "back-pressured"
+    * means "no new work of any kind" until the governor's low-water
+    * mark is crossed.
+    *
+    * Cheap O(1) boolean read; no allocation when happy.
+    */
+   public assertCanAdmitQuery(): void {
+      if (!this.memoryGovernor?.isBackpressured()) return;
+      // Tick first so the counter reflects every rejection even
+      // when the controller's catch block swallows the error (e.g.
+      // an MCP tool surfaces it as a content payload rather than
+      // letting it bubble to the HTTP error mapper).
+      getQueryAdmissionRejectionsCounter().add(1, {
+         environment: this.environmentName,
+      });
       throw new ServiceUnavailableError(
-         `Publisher is under memory pressure and cannot ${reason} (package "${packageName}", environment "${this.environmentName}"). Retry after the server's memory usage drops below the configured low-water mark.`,
+         `Publisher is under memory pressure and cannot accept new queries (environment "${this.environmentName}"). Retry after the server's memory usage drops below the low-water mark (PUBLISHER_MEMORY_LOW_WATER_FRACTION of PUBLISHER_MAX_MEMORY_BYTES), or raise PUBLISHER_MAX_MEMORY_BYTES if you have headroom.`,
       );
    }
 

package/src/service/environment_admission.spec.ts

@@ -4,8 +4,12 @@ import * as os from "os";
 import * as path from "path";
 
 import { ServiceUnavailableError } from "../errors";
+import {
+   startMetricsHarness,
+   type MetricsHarness,
+} from "../test_helpers/metrics_harness";
 import { buildEnvironmentMalloyConfig } from "./connection";
-import { Environment } from "./environment";
+import { Environment, resetAdmissionTelemetryForTesting } from "./environment";
 import type { PackageMemoryGovernor } from "./package_memory_governor";
 
 /**
@@ -178,3 +182,163 @@ describe("Environment admission gate (memory governor choke point)", () => {
       expect(caught).not.toBeInstanceOf(ServiceUnavailableError);
    });
 });
+
+describe("Environment.assertCanAdmitQuery (query-path back-pressure)", () => {
+   let envDir: string;
+
+   beforeEach(() => {
+      envDir = fs.mkdtempSync(
+         path.join(os.tmpdir(), "publisher-env-query-admission-"),
+      );
+   });
+
+   afterEach(() => {
+      fs.rmSync(envDir, { recursive: true, force: true });
+   });
+
+   it("is a no-op when no governor is attached", () => {
+      const env = makeEnvironment(envDir);
+      // No governor — must not throw. Equivalent of an OSS / non-Docker
+      // deployment that never opted into PUBLISHER_MAX_MEMORY_BYTES.
+      expect(() => env.assertCanAdmitQuery()).not.toThrow();
+   });
+
+   it("is a no-op when the governor is happy (under high-water mark)", () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = false;
+
+      expect(() => env.assertCanAdmitQuery()).not.toThrow();
+   });
+
+   it("throws ServiceUnavailableError (→503) when back-pressured", () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      expect(() => env.assertCanAdmitQuery()).toThrow(ServiceUnavailableError);
+   });
+
+   it("error message names the environment so operators can pinpoint the hot pod's load", () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      let caught: unknown;
+      try {
+         env.assertCanAdmitQuery();
+      } catch (err) {
+         caught = err;
+      }
+      expect(caught).toBeInstanceOf(ServiceUnavailableError);
+      expect((caught as Error).message).toContain('environment "test-env"');
+      expect((caught as Error).message).toContain("memory pressure");
+   });
+
+   it("clearing back-pressure immediately re-admits queries (matches governor hysteresis)", () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+
+      governor.backpressured = true;
+      expect(() => env.assertCanAdmitQuery()).toThrow(ServiceUnavailableError);
+
+      governor.backpressured = false;
+      expect(() => env.assertCanAdmitQuery()).not.toThrow();
+   });
+
+   it("detaching the governor reverts to legacy admit-everything", () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      env.setMemoryGovernor(null);
+      // Even with the stub still claiming back-pressure, a detached
+      // governor leaves nothing to consult. Mirrors the package-admission
+      // detach behavior.
+      expect(() => env.assertCanAdmitQuery()).not.toThrow();
+   });
+});
+
+describe("Environment admission telemetry", () => {
+   let envDir: string;
+   let harness: MetricsHarness;
+
+   beforeEach(async () => {
+      envDir = fs.mkdtempSync(
+         path.join(os.tmpdir(), "publisher-env-admission-telemetry-"),
+      );
+      harness = await startMetricsHarness();
+      resetAdmissionTelemetryForTesting();
+   });
+
+   afterEach(async () => {
+      fs.rmSync(envDir, { recursive: true, force: true });
+      resetAdmissionTelemetryForTesting();
+      await harness.shutdown();
+   });
+
+   it("publisher_query_admission_rejections_total ticks per query rejection, labeled by environment", async () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      // Drive three rejections so the counter is unambiguous (a
+      // single-tick assertion can pass against a leaked counter
+      // from a different test; three is harder to fake).
+      for (let i = 0; i < 3; i++) {
+         expect(() => env.assertCanAdmitQuery()).toThrow(
+            ServiceUnavailableError,
+         );
+      }
+      expect(
+         await harness.collectCounter(
+            "publisher_query_admission_rejections_total",
+            { environment: "test-env" },
+         ),
+      ).toBe(3);
+   });
+
+   it("counter stays at zero when the governor is happy (no spurious ticks)", async () => {
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      // Not back-pressured.
+      env.assertCanAdmitQuery();
+      env.assertCanAdmitQuery();
+      // No rejection should have been recorded; verifies that the
+      // counter doesn't tick on the happy-path admission either.
+      expect(
+         await harness.collectCounter(
+            "publisher_query_admission_rejections_total",
+         ),
+      ).toBe(0);
+   });
+
+   it("publisher_package_admission_rejections_total ticks per package-load rejection", async () => {
+      // Ensure the package directory exists so the 404 doesn't
+      // short-circuit ahead of the back-pressure gate.
+      const pkgName = "real-pkg";
+      fs.mkdirSync(path.join(envDir, pkgName));
+
+      const env = makeEnvironment(envDir);
+      const governor = new StubGovernor();
+      env.setMemoryGovernor(governor as unknown as PackageMemoryGovernor);
+      governor.backpressured = true;
+
+      await expect(env.addPackage(pkgName)).rejects.toBeInstanceOf(
+         ServiceUnavailableError,
+      );
+      expect(
+         await harness.collectCounter(
+            "publisher_package_admission_rejections_total",
+            { environment: "test-env", reason: "add a new package" },
+         ),
+      ).toBe(1);
+   });
+});