npm - @malloy-publisher/server - Versions diffs - 0.0.161 → 0.0.165 - Mend

@malloy-publisher/server 0.0.161 → 0.0.165

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (27) hide show

package/dist/app/api-doc.yaml +78 -0
package/dist/app/assets/HomePage-QekMXs8r.js +1 -0
package/dist/app/assets/{MainPage-CzAy2LSX.js → MainPage-DAyUfYba.js} +1 -1
package/dist/app/assets/{ModelPage-uZf0eDBu.js → ModelPage-CrMryV1s.js} +1 -1
package/dist/app/assets/PackagePage-DDaABD2A.js +1 -0
package/dist/app/assets/ProjectPage-FAYUFGhL.js +1 -0
package/dist/app/assets/RouteError-BKYctANX.js +1 -0
package/dist/app/assets/{WorkbookPage-CBJ0x84n.js → WorkbookPage-DZEVYGW3.js} +1 -1
package/dist/app/assets/{index-BoduKqUa.js → index-BvVmB5sv.js} +98 -98
package/dist/app/assets/{index-D_8e14NE.js → index-CsC07BYd.js} +1 -1
package/dist/app/assets/{index-Bl1fe3A7.js → index-DWhjtyBB.js} +1 -1
package/dist/app/assets/{index.umd-DRsjT6d4.js → index.umd-DvM-lTQa.js} +1 -1
package/dist/app/index.html +1 -1
package/dist/instrumentation.js +4 -2
package/dist/server.js +150 -18
package/k6-tests/load-test/load-test.ts +30 -0
package/package.json +5 -1
package/src/health.ts +181 -0
package/src/logger.ts +9 -1
package/src/server.ts +42 -1
package/src/service/model.spec.ts +5 -5
package/src/service/model.ts +56 -19
package/src/service/project_store.ts +5 -0
package/dist/app/assets/HomePage-CgSTNEQW.js +0 -1
package/dist/app/assets/PackagePage-BHvPFsmZ.js +0 -1
package/dist/app/assets/ProjectPage-ndnv5nGp.js +0 -1
package/dist/app/assets/RouteError-CZC9OOKh.js +0 -1

package/dist/app/index.html CHANGED Viewed

@@ -12,7 +12,7 @@
          href="https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;500;700&display=swap"
       />
       <title>Malloy Publisher</title>
-     <script type="module" crossorigin src="/assets/index-BoduKqUa.js"></script>
+     <script type="module" crossorigin src="/assets/index-BvVmB5sv.js"></script>
      <link rel="stylesheet" crossorigin href="/assets/index-CMlGQMcl.css">
    </head>
    <body>

package/dist/instrumentation.js CHANGED Viewed

@@ -106335,7 +106335,7 @@ var require_styles = __commonJS((exports2, module2) => {
 // ../../node_modules/@colors/colors/lib/system/has-flag.js
 var require_has_flag2 = __commonJS((exports2, module2) => {
   module2.exports = function(flag, argv) {
-    argv = argv || process.argv || [];
+    argv = argv || process.argv;
     var terminatorPos = argv.indexOf("--");
     var prefix = /^-{1,2}/.test(flag) ? "" : "--";
     var pos = argv.indexOf(prefix + flag);
@@ -116569,7 +116569,9 @@ var loggerMiddleware = (req, res, next) => {
     if (traceId) {
       logMetadata.traceId = traceId;
     }
-    logger.info(`${req.method} ${req.url}`, logMetadata);
+    if (req.url !== "/metrics" && req.url !== "/health" && req.url !== "/health/liveness" && req.url !== "/health/readiness") {
+      logger.info(`${req.method} ${req.url}`, logMetadata);
+    }
   });
   next();
 };

package/dist/server.js CHANGED Viewed

@@ -33301,7 +33301,7 @@ var require_styles = __commonJS((exports2, module2) => {
 // ../../node_modules/@colors/colors/lib/system/has-flag.js
 var require_has_flag2 = __commonJS((exports2, module2) => {
   module2.exports = function(flag, argv) {
-    argv = argv || process.argv || [];
+    argv = argv || process.argv;
     var terminatorPos = argv.indexOf("--");
     var prefix = /^-{1,2}/.test(flag) ? "" : "--";
     var pos = argv.indexOf(prefix + flag);
@@ -124960,7 +124960,9 @@ var loggerMiddleware = (req, res, next) => {
     if (traceId) {
       logMetadata.traceId = traceId;
     }
-    logger.info(`${req.method} ${req.url}`, logMetadata);
+    if (req.url !== "/metrics" && req.url !== "/health" && req.url !== "/health/liveness" && req.url !== "/health/readiness") {
+      logger.info(`${req.method} ${req.url}`, logMetadata);
+    }
   });
   next();
 };
@@ -135906,6 +135908,96 @@ var getProcessedPublisherConfig = (serverRoot) => {
   };
 };
+// src/health.ts
+var operationalState = "initializing";
+var ready = false;
+var preGracefulShutdownCompleted = false;
+function getOperationalState() {
+  return operationalState;
+}
+function markReady() {
+  if (operationalState !== "draining") {
+    operationalState = "serving";
+    ready = true;
+    logger.info("Service marked as ready");
+  } else {
+    logger.error("Service is already draining - cannot mark as ready");
+  }
+}
+function markNotReady() {
+  ready = false;
+  logger.info("Service marked as not ready - readiness probe will fail");
+}
+function registerSignalHandlers(server, mcpServer, shutdownDrainDurationSeconds = 0, shutdownGracefulCloseTimeoutSeconds = 0) {
+  process.once("SIGTERM", async () => {
+    logger.info("========== SIGTERM RECEIVED ==========");
+    markNotReady();
+    operationalState = "draining";
+    logger.info(`Service entering draining state for ${shutdownDrainDurationSeconds} seconds before closing servers...`);
+    await new Promise((resolve3) => setTimeout(() => {
+      preGracefulShutdownCompleted = true;
+      resolve3(true);
+    }, shutdownDrainDurationSeconds * 1000));
+    const closeServer = (server2, name) => new Promise((resolve3) => {
+      if (server2 && server2.listening) {
+        server2.close((err) => {
+          if (err) {
+            logger.error(`${name} close error:`, err);
+          } else {
+            logger.info(`${name} closed`);
+          }
+          resolve3();
+        });
+      } else {
+        resolve3();
+      }
+    });
+    await Promise.all([
+      closeServer(server, "Main server"),
+      closeServer(mcpServer, "MCP server")
+    ]);
+    try {
+      logger.close();
+    } catch (_error) {}
+    if (shutdownGracefulCloseTimeoutSeconds > 0) {
+      logger.info(`Waiting ${shutdownGracefulCloseTimeoutSeconds} seconds after server close before exit...`);
+      await new Promise((resolve3) => setTimeout(resolve3, shutdownGracefulCloseTimeoutSeconds * 1000));
+    }
+    process.exit(0);
+  });
+}
+function drainingGuard(req, res, next) {
+  if (operationalState === "draining" && preGracefulShutdownCompleted && !req.path.startsWith("/health") && !req.path.startsWith("/metrics")) {
+    res.status(503).json({
+      message: "Service is draining",
+      details: "The service is shutting down and hence is not available."
+    });
+    return;
+  }
+  next();
+}
+function registerHealthEndpoints(app) {
+  app.get("/health", (_req, res) => {
+    res.status(200).json({
+      status: "UP",
+      components: {
+        livenessState: { status: "UP" },
+        readinessState: { status: ready ? "UP" : "DOWN" }
+      },
+      groups: ["livenessState", "readinessState"]
+    });
+  });
+  app.get("/health/liveness", (_req, res) => {
+    res.status(200).json({ status: "UP" });
+  });
+  app.get("/health/readiness", (_req, res) => {
+    const isReady = ready;
+    res.status(isReady ? 200 : 503).json({
+      status: isReady ? "UP" : "DOWN"
+    });
+  });
+}
 // src/storage/duckdb/DuckDBConnection.ts
 var duckdb = __toESM(require("duckdb"));
 var path4 = __toESM(require("path"));
@@ -136747,28 +136839,51 @@ class Model {
   async getQueryResults(sourceName, queryName, query) {
     const startTime = performance.now();
     if (this.compilationError) {
-      throw this.compilationError;
+      if (this.compilationError instanceof import_malloy2.MalloyError || this.compilationError instanceof ModelCompilationError) {
+        throw this.compilationError;
+      }
+      throw new BadRequestError(`Model compilation failed: ${this.compilationError.message}`);
     }
     let runnable;
     if (!this.modelMaterializer || !this.modelDef || !this.modelInfo)
       throw new BadRequestError("Model has no queryable entities.");
-    if (!sourceName && !queryName && query) {
-      runnable = this.modelMaterializer.loadQuery(`
+    try {
+      if (!sourceName && !queryName && query) {
+        runnable = this.modelMaterializer.loadQuery(`
 ` + query);
-    } else if (queryName && !query) {
-      runnable = this.modelMaterializer.loadQuery(`
+      } else if (queryName && !query) {
+        runnable = this.modelMaterializer.loadQuery(`
 run: ${sourceName ? sourceName + "->" : ""}${queryName}`);
-    } else {
-      const endTime2 = performance.now();
-      const executionTime2 = endTime2 - startTime;
-      this.queryExecutionHistogram.record(executionTime2, {
-        "malloy.model.path": this.modelPath,
-        "malloy.model.query.name": queryName,
-        "malloy.model.query.source": sourceName,
-        "malloy.model.query.query": query,
-        "malloy.model.query.status": "error"
+      } else {
+        const endTime2 = performance.now();
+        const executionTime2 = endTime2 - startTime;
+        this.queryExecutionHistogram.record(executionTime2, {
+          "malloy.model.path": this.modelPath,
+          "malloy.model.query.name": queryName,
+          "malloy.model.query.source": sourceName,
+          "malloy.model.query.query": query,
+          "malloy.model.query.status": "error"
+        });
+        throw new BadRequestError("Invalid query request. (Query AND !sourceName) OR (queryName AND sourceName) must be defined.");
+      }
+    } catch (error) {
+      if (error instanceof BadRequestError) {
+        throw error;
+      }
+      if (error instanceof import_malloy2.MalloyError) {
+        throw error;
+      }
+      const errorMessage = error instanceof Error ? error.message : String(error);
+      logger.error("Query parsing error", {
+        error,
+        errorMessage,
+        projectName: this.packageName,
+        modelPath: this.modelPath,
+        query,
+        queryName,
+        sourceName
       });
-      throw new BadRequestError("Invalid query request. (Query AND !sourceName) OR (queryName AND sourceName) must be defined.");
+      throw new BadRequestError(`Invalid query: ${errorMessage}`);
     }
     const rowLimit = (await runnable.getPreparedResult()).resultExplore.limit || ROW_LIMIT;
     const endTime = performance.now();
@@ -137749,9 +137864,11 @@ class ProjectStore {
         }
       }
       this.isInitialized = true;
+      markReady();
       const initializationDuration = performance.now() - initialTime;
       logger.info(`Project store successfully initialized in ${formatDuration(initializationDuration)}`);
     } catch (error) {
+      markNotReady();
       logger.error("Error initializing project store", { error });
       process.exit(1);
     }
@@ -137984,7 +138101,8 @@ class ProjectStore {
       timestamp: Date.now(),
       projects: [],
       initialized: this.isInitialized,
-      frozenConfig: isPublisherConfigFrozen(this.serverRootPath)
+      frozenConfig: isPublisherConfigFrozen(this.serverRootPath),
+      operationalState: getOperationalState()
     };
     const projects = await this.listProjects(true);
     await Promise.all(projects.map(async (project) => {
@@ -142457,6 +142575,12 @@ function parseArgs() {
     } else if (arg === "--mcp_port" && args[i + 1]) {
       process.env.MCP_PORT = args[i + 1];
       i++;
+    } else if (arg === "--shutdown_drain_duration_seconds" && args[i + 1]) {
+      process.env.SHUTDOWN_DRAIN_DURATION_SECONDS = args[i + 1];
+      i++;
+    } else if (arg === "--shutdown_graceful_close_timeout_seconds" && args[i + 1]) {
+      process.env.SHUTDOWN_GRACEFUL_CLOSE_TIMEOUT_SECONDS = args[i + 1];
+      i++;
     } else if (arg === "--init") {
       process.env.INITIALIZE_STORAGE = "true";
     } else if (arg === "--help" || arg === "-h") {
@@ -142469,6 +142593,9 @@ function parseArgs() {
       console.log("  --host <string>        Host to bind the server to (default: localhost)");
       console.log("  --server_root <path>   Root directory to serve files from (default: .)");
       console.log("  --mcp_port <number>    Port for MCP server (default: 4040)");
+      console.log("  --shutdown_drain_duration_seconds <number>  Time in seconds to keep service in draining state before closing servers (default: 0)");
+      console.log("  --shutdown_graceful_close_timeout_seconds <number>  Time in seconds to wait after closing servers before exit (default: 0)");
+      console.log("  --init                 Initialize the storage (default: false)");
       console.log("  --help, -h             Show this help message");
       process.exit(0);
     }
@@ -142479,6 +142606,8 @@ var PUBLISHER_PORT = Number(process.env.PUBLISHER_PORT || 4000);
 var PUBLISHER_HOST = process.env.PUBLISHER_HOST || "0.0.0.0";
 var MCP_PORT = Number(process.env.MCP_PORT || 4040);
 var MCP_ENDPOINT = "/mcp";
+var SHUTDOWN_DRAIN_DURATION_SECONDS = Number(process.env.SHUTDOWN_DRAIN_DURATION_SECONDS || 0);
+var SHUTDOWN_GRACEFUL_CLOSE_TIMEOUT_SECONDS = Number(process.env.SHUTDOWN_GRACEFUL_CLOSE_TIMEOUT_SECONDS || 0);
 var ROOT;
 if (require.main) {
   ROOT = path10.join(path10.dirname(require.main.filename), "app");
@@ -142578,6 +142707,8 @@ app.use(import_cors.default({
   credentials: true
 }));
 app.use(bodyParser.json());
+registerHealthEndpoints(app);
+app.use(drainingGuard);
 app.get(`${API_PREFIX2}/status`, async (_req, res) => {
   try {
     const status = await projectStore.getStatus();
@@ -142974,3 +143105,4 @@ var mcpServer = mcpApp.listen(MCP_PORT, PUBLISHER_HOST, () => {
 mcpServer.timeout = 600000;
 mcpServer.keepAliveTimeout = 600000;
 mcpServer.headersTimeout = 600000;
+registerSignalHandlers(mainServer, mcpServer, SHUTDOWN_DRAIN_DURATION_SECONDS, SHUTDOWN_GRACEFUL_CLOSE_TIMEOUT_SECONDS);

package/k6-tests/load-test/load-test.ts CHANGED Viewed

@@ -440,6 +440,11 @@ export const loadTest: TestPreset = {
                check(response.response, {
                   "get connection request successful": (r) => r.status === 200,
                });
+               if (response.response.status !== 200) {
+                  logger.error(
+                     `get connection failed: status=${response.response.status}, body=${response.response.body}`,
+                  );
+               }
             });
          }
       });
@@ -590,6 +595,11 @@ export const loadTest: TestPreset = {
                      "post query data request successful": (r) =>
                         r.status === 200,
                   });
+                  if (response.response.status !== 200) {
+                     logger.error(
+                        `post query data failed: status=${response.response.status}, body=${response.response.body}`,
+                     );
+                  }
                } catch (error) {
                   logger.warn(`Failed to post query data: ${error}`);
                }
@@ -619,6 +629,11 @@ export const loadTest: TestPreset = {
                   check(schemasResponse.response, {
                      "list schemas request successful": (r) => r.status === 200,
                   });
+                  if (schemasResponse.response.status !== 200) {
+                     logger.error(
+                        `list schemas failed: status=${schemasResponse.response.status}, body=${schemasResponse.response.body}`,
+                     );
+                  }
                   sleep(0.1);
@@ -646,6 +661,11 @@ export const loadTest: TestPreset = {
                                  "list tables request successful": (r) =>
                                     r.status === 200,
                               });
+                              if (tablesResponse.response.status !== 200) {
+                                 logger.error(
+                                    `list tables failed: status=${tablesResponse.response.status}, body=${tablesResponse.response.body}`,
+                                 );
+                              }
                               sleep(0.1);
@@ -717,6 +737,11 @@ export const loadTest: TestPreset = {
                      "post sql source request successful": (r) =>
                         r.status === 200,
                   });
+                  if (response.response.status !== 200) {
+                     logger.error(
+                        `post sql source failed: status=${response.response.status}, body=${response.response.body}`,
+                     );
+                  }
                } catch (error) {
                   logger.warn(`Failed to post sql source: ${error}`);
                }
@@ -739,6 +764,11 @@ export const loadTest: TestPreset = {
                      "post temporary table request successful": (r) =>
                         r.status === 200,
                   });
+                  if (response.response.status !== 200) {
+                     logger.error(
+                        `post temporary table failed: status=${response.response.status}, body=${response.response.body}`,
+                     );
+                  }
                } catch (error) {
                   logger.warn(`Failed to post temporary table: ${error}`);
                }

package/package.json CHANGED Viewed

@@ -1,11 +1,15 @@
 {
    "name": "@malloy-publisher/server",
    "description": "Malloy Publisher Server",
-   "version": "0.0.161",
+   "version": "0.0.165",
    "main": "dist/server.js",
    "bin": {
       "malloy-publisher": "dist/server.js"
    },
+   "repository": {
+      "type": "git",
+      "url": "https://github.com/malloydata/publisher.git"
+   },
    "publishConfig": {
       "access": "public"
    },

package/src/health.ts ADDED Viewed

@@ -0,0 +1,181 @@
+/**
+ * Health check endpoints and graceful shutdown handling.
+ *
+ * Provides Kubernetes-compatible health probes:
+ * - /health - Overall health status
+ * - /health/liveness - Always returns 200 (process is alive)
+ * - /health/readiness - Returns 200 when ready, 503 when not ready
+ *
+ * On SIGTERM, the service enters draining state to allow in-flight requests
+ * to complete before shutdown.
+ */
+import { Express, NextFunction, Request, Response } from "express";
+import { Server } from "http";
+import { components } from "./api";
+import { logger } from "./logger";
+export type OperationalState =
+   components["schemas"]["ServerStatus"]["operationalState"];
+let operationalState: OperationalState = "initializing" as OperationalState;
+let ready: boolean = false;
+let preGracefulShutdownCompleted: boolean = false;
+/**
+ * Returns the current operational state of the service.
+ */
+export function getOperationalState(): OperationalState {
+   return operationalState;
+}
+/**
+ * Marks the service as ready to serve traffic.
+ */
+export function markReady(): void {
+   if (operationalState !== "draining") {
+      operationalState = "serving";
+      ready = true;
+      logger.info("Service marked as ready");
+   } else {
+      logger.error("Service is already draining - cannot mark as ready");
+   }
+}
+/**
+ * Marks the service as not ready (readiness probe will return 503).
+ */
+export function markNotReady(): void {
+   ready = false;
+   logger.info("Service marked as not ready - readiness probe will fail");
+}
+/**
+ * Registers SIGTERM handler for graceful shutdown.
+ *
+ * Shutdown sequence:
+ * 1. Marks service as not ready (readiness probe returns 503) and enters draining state
+ * 2. Waits shutdownDrainDurationSeconds to allow in-flight requests to complete
+ * 3. Sets preGracefulShutdownCompleted flag (enables drainingGuard middleware to reject new requests)
+ * 4. Closes main server and MCP server (stops accepting new connections)
+ * 5. Closes logger
+ * 6. Waits shutdownGracefulCloseTimeoutSeconds (if > 0) for final cleanup
+ * 7. Exits process
+ *
+ * Note: drainingGuard only rejects requests after step 3 completes. During step 2,
+ * the service is draining but still accepts requests (readiness probe returns 503).
+ *
+ * @param server - Main HTTP server instance
+ * @param mcpServer - MCP server instance
+ * @param shutdownDrainDurationSeconds - Duration in seconds to wait before closing servers
+ * @param shutdownGracefulCloseTimeoutSeconds - Duration in seconds to wait after closing servers before exit
+ */
+export function registerSignalHandlers(
+   server: Server,
+   mcpServer: Server,
+   shutdownDrainDurationSeconds: number = 0,
+   shutdownGracefulCloseTimeoutSeconds: number = 0,
+): void {
+   // Keep the process alive on SIGTERM — do not close the server.
+   // K8s will SIGKILL after terminationGracePeriodSeconds (which cannot be caught).
+   process.once("SIGTERM", async () => {
+      logger.info("========== SIGTERM RECEIVED ==========");
+      markNotReady();
+      operationalState = "draining" as OperationalState;
+      logger.info(
+         `Service entering draining state for ${shutdownDrainDurationSeconds} seconds before closing servers...`,
+      );
+      await new Promise((resolve) =>
+         setTimeout(() => {
+            preGracefulShutdownCompleted = true;
+            resolve(true);
+         }, shutdownDrainDurationSeconds * 1000),
+      );
+      const closeServer = (server: Server, name: string) =>
+         new Promise<void>((resolve) => {
+            if (server && server.listening) {
+               server.close((err) => {
+                  if (err) {
+                     logger.error(`${name} close error:`, err);
+                  } else {
+                     logger.info(`${name} closed`);
+                  }
+                  resolve();
+               });
+            } else {
+               resolve();
+            }
+         });
+      await Promise.all([
+         closeServer(server, "Main server"),
+         closeServer(mcpServer, "MCP server"),
+      ]);
+      try {
+         logger.close();
+      } catch (_error) {
+         /* do nothing */
+      }
+      if (shutdownGracefulCloseTimeoutSeconds > 0) {
+         logger.info(
+            `Waiting ${shutdownGracefulCloseTimeoutSeconds} seconds after server close before exit...`,
+         );
+         await new Promise((resolve) =>
+            setTimeout(resolve, shutdownGracefulCloseTimeoutSeconds * 1000),
+         );
+      }
+      process.exit(0);
+   });
+}
+/**
+ * Middleware that returns 503 for non-health and metrics requests when service is draining.
+ * Must be registered before application routes.
+ */
+export function drainingGuard(
+   req: Request,
+   res: Response,
+   next: NextFunction,
+): void {
+   if (
+      operationalState === "draining" &&
+      preGracefulShutdownCompleted &&
+      !req.path.startsWith("/health") &&
+      !req.path.startsWith("/metrics")
+   ) {
+      res.status(503).json({
+         message: "Service is draining",
+         details: "The service is shutting down and hence is not available.",
+      });
+      return;
+   }
+   next();
+}
+/**
+ * Registers health check endpoints
+ */
+export function registerHealthEndpoints(app: Express): void {
+   app.get("/health", (_req: Request, res: Response) => {
+      res.status(200).json({
+         status: "UP",
+         components: {
+            livenessState: { status: "UP" },
+            readinessState: { status: ready ? "UP" : "DOWN" },
+         },
+         groups: ["livenessState", "readinessState"],
+      });
+   });
+   app.get("/health/liveness", (_req: Request, res: Response) => {
+      res.status(200).json({ status: "UP" });
+   });
+   app.get("/health/readiness", (_req: Request, res: Response) => {
+      const isReady = ready;
+      res.status(isReady ? 200 : 503).json({
+         status: isReady ? "UP" : "DOWN",
+      });
+   });
+}

package/src/logger.ts CHANGED Viewed

@@ -119,7 +119,15 @@ export const loggerMiddleware: RequestHandler = (req, res, next) => {
          logMetadata.traceId = traceId;
       }
-      logger.info(`${req.method} ${req.url}`, logMetadata);
+      // Skip logging for metrics and health endpoints to reduce log noise
+      if (
+         req.url !== "/metrics" &&
+         req.url !== "/health" &&
+         req.url !== "/health/liveness" &&
+         req.url !== "/health/readiness"
+      ) {
+         logger.info(`${req.method} ${req.url}`, logMetadata);
+      }
    });
    next();
 };

package/src/server.ts CHANGED Viewed

@@ -13,10 +13,14 @@ import { PackageController } from "./controller/package.controller";
 import { QueryController } from "./controller/query.controller";
 import { WatchModeController } from "./controller/watch-mode.controller";
 import { internalErrorToHttpError, NotImplementedError } from "./errors";
+import {
+   drainingGuard,
+   registerHealthEndpoints,
+   registerSignalHandlers,
+} from "./health";
 import { logger, loggerMiddleware } from "./logger";
 import { initializeMcpServer } from "./mcp/server";
 import { ProjectStore } from "./service/project_store";
 // Parse command line arguments
 function parseArgs() {
    const args = process.argv.slice(2);
@@ -34,6 +38,15 @@ function parseArgs() {
       } else if (arg === "--mcp_port" && args[i + 1]) {
          process.env.MCP_PORT = args[i + 1];
          i++;
+      } else if (arg === "--shutdown_drain_duration_seconds" && args[i + 1]) {
+         process.env.SHUTDOWN_DRAIN_DURATION_SECONDS = args[i + 1];
+         i++;
+      } else if (
+         arg === "--shutdown_graceful_close_timeout_seconds" &&
+         args[i + 1]
+      ) {
+         process.env.SHUTDOWN_GRACEFUL_CLOSE_TIMEOUT_SECONDS = args[i + 1];
+         i++;
       } else if (arg === "--init") {
          process.env.INITIALIZE_STORAGE = "true";
       } else if (arg === "--help" || arg === "-h") {
@@ -54,6 +67,15 @@ function parseArgs() {
          console.log(
             "  --mcp_port <number>    Port for MCP server (default: 4040)",
          );
+         console.log(
+            "  --shutdown_drain_duration_seconds <number>  Time in seconds to keep service in draining state before closing servers (default: 0)",
+         );
+         console.log(
+            "  --shutdown_graceful_close_timeout_seconds <number>  Time in seconds to wait after closing servers before exit (default: 0)",
+         );
+         console.log(
+            "  --init                 Initialize the storage (default: false)",
+         );
          console.log("  --help, -h             Show this help message");
          process.exit(0);
       }
@@ -67,6 +89,12 @@ const PUBLISHER_PORT = Number(process.env.PUBLISHER_PORT || 4000);
 const PUBLISHER_HOST = process.env.PUBLISHER_HOST || "0.0.0.0";
 const MCP_PORT = Number(process.env.MCP_PORT || 4040);
 const MCP_ENDPOINT = "/mcp";
+const SHUTDOWN_DRAIN_DURATION_SECONDS = Number(
+   process.env.SHUTDOWN_DRAIN_DURATION_SECONDS || 0,
+);
+const SHUTDOWN_GRACEFUL_CLOSE_TIMEOUT_SECONDS = Number(
+   process.env.SHUTDOWN_GRACEFUL_CLOSE_TIMEOUT_SECONDS || 0,
+);
 // Find the app directory - handle NPX vs local execution
 let ROOT: string;
 if (require.main) {
@@ -212,6 +240,12 @@ app.use(
 );
 app.use(bodyParser.json());
+// Register health check endpoints
+registerHealthEndpoints(app);
+// Register draining guard middleware - must be after health endpoints but before other routes
+app.use(drainingGuard);
 app.get(`${API_PREFIX}/status`, async (_req, res) => {
    try {
       const status = await projectStore.getStatus();
@@ -905,3 +939,10 @@ const mcpServer = mcpApp.listen(MCP_PORT, PUBLISHER_HOST, () => {
 mcpServer.timeout = 600000;
 mcpServer.keepAliveTimeout = 600000;
 mcpServer.headersTimeout = 600000;
+registerSignalHandlers(
+   mainServer,
+   mcpServer,
+   SHUTDOWN_DRAIN_DURATION_SECONDS,
+   SHUTDOWN_GRACEFUL_CLOSE_TIMEOUT_SECONDS,
+);

package/src/service/model.spec.ts CHANGED Viewed

@@ -1,10 +1,10 @@
-import { expect, it, describe } from "bun:test";
 import { MalloyError, Runtime } from "@malloydata/malloy";
-import sinon from "sinon";
+import { describe, expect, it } from "bun:test";
 import fs from "fs/promises";
+import sinon from "sinon";
-import { Model, ModelType } from "./model";
 import { BadRequestError, ModelNotFoundError } from "../errors";
+import { Model, ModelType } from "./model";
 describe("service/model", () => {
    const packageName = "test-package";
@@ -190,7 +190,7 @@ describe("service/model", () => {
       });
       describe("getQueryResults", () => {
-         it("should throw ModelCompilationError if a compilation error exists", async () => {
+         it("should throw BadRequestError if a non-MalloyError compilation error exists", async () => {
             const error = new Error("Compilation error");
             const model = new Model(
                packageName,
@@ -208,7 +208,7 @@ describe("service/model", () => {
             await expect(async () => {
                await model.getQueryResults();
-            }).toThrowError(Error);
+            }).toThrowError(BadRequestError);
             sinon.restore();
          });