@malloy-publisher/server 0.0.198-dev4 → 0.0.198-dev6

package/src/service/package.ts

@@ -1,7 +1,6 @@
 import * as fs from "fs/promises";
 import * as path from "path";
 
-import { DuckDBConnection } from "@malloydata/db-duckdb";
 import "@malloydata/db-duckdb/native";
 import {
    Connection,
@@ -10,19 +9,24 @@ import {
    EmptyURLReader,
    FixedConnectionMap,
    MalloyConfig,
+   MalloyError,
    SourceDef,
 } from "@malloydata/malloy";
 import { metrics } from "@opentelemetry/api";
 import recursive from "recursive-readdir";
 import { components } from "../api";
-import { getCompilePool } from "../compile/compile_pool";
+import { getPackageLoadPool } from "../package_load/package_load_pool";
 import {
    API_PREFIX,
    MODEL_FILE_SUFFIX,
    NOTEBOOK_FILE_SUFFIX,
    PACKAGE_MANIFEST_NAME,
 } from "../constants";
-import { PackageNotFoundError } from "../errors";
+import {
+   ModelCompilationError,
+   PackageNotFoundError,
+   ServiceUnavailableError,
+} from "../errors";
 import { formatDuration, logger } from "../logger";
 import { BuildManifest } from "../storage/DatabaseInterface";
 import { ignoreDotfiles } from "../utils";
@@ -95,23 +99,12 @@ export class Package {
       });
 
       try {
-         const packageConfig = await Package.readPackageConfig(packagePath);
-         const packageConfigTime = performance.now();
-         logger.info("Package config read completed", {
-            packageName,
-            duration: formatDuration(
-               packageConfigTime - manifestValidationTime,
-            ),
-         });
-         packageConfig.resource = `${API_PREFIX}/environments/${environmentName}/packages/${packageName}`;
-
-         const databases = await Package.readDatabases(packagePath);
-         const databasesTime = performance.now();
-         logger.info("Databases read completed", {
-            packageName,
-            databaseCount: databases.length,
-            duration: formatDuration(databasesTime - packageConfigTime),
-         });
+         // The MalloyConfig is always built on the main thread — it
+         // owns the live native connection handles the package needs
+         // to *serve queries* after load (workers can't share native
+         // handles across the V8 isolate boundary). The worker proxies
+         // non-duckdb connection lookups back through this MalloyConfig
+         // during compile.
          const malloyConfig = Package.buildPackageMalloyConfig(
             packagePath,
             typeof environmentMalloyConfig === "function"
@@ -119,68 +112,29 @@ export class Package {
                : () => Package.toMalloyConfig(environmentMalloyConfig),
          );
 
-         const models = await Package.loadModels(
-            packageName,
-            packagePath,
-            malloyConfig,
-         );
-         const modelsTime = performance.now();
-         logger.info("Models loaded", {
-            packageName,
-            modelCount: models.size,
-            duration: formatDuration(modelsTime - databasesTime),
-         });
-         for (const [modelPath, model] of models.entries()) {
-            const maybeModel = model as unknown as {
-               compilationError?: unknown;
-            };
-            if (maybeModel.compilationError) {
-               const err = maybeModel.compilationError;
-               const message =
-                  err instanceof Error
-                     ? err.message
-                     : `Unknown compilation error in ${modelPath}`;
-
-               logger.error("Model compilation failed", {
-                  packageName,
-                  modelPath,
-                  error: message,
-               });
-
-               this.packageLoadHistogram.record(performance.now() - startTime, {
-                  malloy_package_name: packageName,
-                  status: "compilation_error",
-               });
-               throw err;
-            }
-         }
-         const endTime = performance.now();
-         const executionTime = endTime - startTime;
-         this.packageLoadHistogram.record(executionTime, {
-            malloy_package_name: packageName,
-            status: "success",
-         });
-         logger.info(`Successfully loaded package ${packageName}`, {
-            packageName,
-            duration: formatDuration(executionTime),
-         });
-         return new Package(
+         return await Package.loadViaWorker(
             environmentName,
             packageName,
             packagePath,
-            packageConfig,
-            databases,
-            models,
             malloyConfig,
+            startTime,
+            manifestValidationTime,
          );
       } catch (error) {
          logger.error(`Error loading package ${packageName}`, { error });
          console.error(error);
          const endTime = performance.now();
          const executionTime = endTime - startTime;
+         const status =
+            error instanceof ModelCompilationError ||
+            error instanceof MalloyError
+               ? "compilation_error"
+               : error instanceof ServiceUnavailableError
+                 ? "pool_unavailable"
+                 : "error";
          this.packageLoadHistogram.record(executionTime, {
             malloy_package_name: packageName,
-            status: "error",
+            status,
          });
          // Clean up package directory on failure
          try {
@@ -198,6 +152,141 @@ export class Package {
       }
    }
 
+   /**
+    * Load the package via the package-load worker pool. The worker
+    * performs the CPU-bound bulk of the load off-thread (manifest
+    * read, every `.malloy` / `.malloynb` compile) and ships back a
+    * structured-clonable `LoadPackageOutcome`. Database probes
+    * (`.parquet` / `.csv`) run on the main thread, in parallel with
+    * the worker compile, against the package's existing DuckDB
+    * connection — they're async-IO-bound and don't compete with the
+    * worker for CPU.
+    *
+    * Pool-infrastructure failures (worker crash, RPC timeout, pool
+    * shutting down) are rewrapped as `ServiceUnavailableError` so
+    * the HTTP layer responds 503 (transient, retryable). Real compile
+    * errors (`MalloyError` / `ModelCompilationError`) propagate
+    * unchanged so they keep their 4xx mapping.
+    */
+   private static async loadViaWorker(
+      environmentName: string,
+      packageName: string,
+      packagePath: string,
+      malloyConfig: MalloyConfig,
+      startTime: number,
+      manifestValidationTime: number,
+   ): Promise<Package> {
+      const pool = getPackageLoadPool();
+      const dispatchTime = performance.now();
+      // Submit the worker job and run database probing on the main
+      // thread in parallel. We isolate the worker-job promise inside
+      // a wrapper so we can map pool-infrastructure failures (worker
+      // crash, RPC timeout, pool shutting down) to a 503 without
+      // accidentally re-mapping `readDatabases`'s own errors.
+      const workerOutcome = pool
+         .loadPackage({
+            packagePath,
+            packageName,
+            malloyConfig,
+            defaultConnectionName: "duckdb",
+         })
+         .catch((err: unknown) => {
+            // Compile errors surface in-band via
+            // `LoadPackageOutcome.models[i].compilationError`; if the
+            // pool itself rejects, it's an infra-side failure
+            // (shutting down, worker spawn failed, worker crashed,
+            // RPC timeout) and the client should retry. Real Malloy
+            // compile errors deserialised by the pool still carry
+            // their MalloyError / ModelCompilationError identity —
+            // let those bubble untouched so they keep their 4xx
+            // mapping in `errors.ts`.
+            const realError =
+               err instanceof Error
+                  ? err
+                  : new Error(
+                       `Package-load worker pool failure: ${String(err)}`,
+                    );
+            if (
+               realError instanceof MalloyError ||
+               realError instanceof ModelCompilationError
+            ) {
+               throw realError;
+            }
+            throw new ServiceUnavailableError(
+               `Package-load worker pool unavailable: ${realError.message}`,
+            );
+         });
+      const [outcome, databases] = await Promise.all([
+         workerOutcome,
+         Package.readDatabases(packagePath, malloyConfig),
+      ]);
+      const workerDoneTime = performance.now();
+      logger.info("Package load via worker pool completed", {
+         packageName,
+         manifestValidationMs: dispatchTime - manifestValidationTime,
+         workerDurationMs: outcome.loadDurationMs,
+         dispatchOverheadMs:
+            workerDoneTime - dispatchTime - outcome.loadDurationMs,
+         modelCount: outcome.models.length,
+         databaseCount: databases.length,
+      });
+
+      // Override the manifest-derived resource URI — the worker only
+      // returns name/description from publisher.json, but the rest of
+      // the API surface expects a `resource` field too.
+      const packageConfig: ApiPackage = {
+         name: outcome.packageMetadata.name,
+         description: outcome.packageMetadata.description,
+         resource: `${API_PREFIX}/environments/${environmentName}/packages/${packageName}`,
+      };
+
+      // Build live `Model`s from worker output. Any per-model compile
+      // failure aborts the load — matches the historical behaviour of
+      // `Package.create` failing the whole package on the first model
+      // error. (`Package.reloadAllModels` keeps the failed-model
+      // placeholders instead; that branch goes through a different
+      // hydration path.)
+      const models = new Map<string, Model>();
+      for (const sm of outcome.models) {
+         if (sm.compilationError) {
+            const err = Model.deserializeCompilationError(sm.compilationError);
+            logger.error("Model compilation failed", {
+               packageName,
+               modelPath: sm.modelPath,
+               error: err.message,
+            });
+            // The outer catch in Package.create records the metric +
+            // cleans the package directory.
+            throw err;
+         }
+         models.set(
+            sm.modelPath,
+            Model.fromSerialized(packageName, packagePath, malloyConfig, sm),
+         );
+      }
+
+      const endTime = performance.now();
+      const executionTime = endTime - startTime;
+      this.packageLoadHistogram.record(executionTime, {
+         malloy_package_name: packageName,
+         status: "success",
+      });
+      logger.info(`Successfully loaded package ${packageName}`, {
+         packageName,
+         duration: formatDuration(executionTime),
+      });
+
+      return new Package(
+         environmentName,
+         packageName,
+         packagePath,
+         packageConfig,
+         databases,
+         models,
+         malloyConfig,
+      );
+   }
+
    public getPackageName(): string {
       return this.packageName;
    }
@@ -232,6 +321,21 @@ export class Package {
       return Array.from(this.models.keys());
    }
 
+   /**
+    * Re-compile every model in the package against a new build
+    * manifest (called after a materialization build commits new
+    * physicalised tables). Runs through the package-load worker pool
+    * — same off-main-thread compile path as initial `Package.create`
+    * — so a reload of a large package can't block the K8s liveness
+    * probe.
+    *
+    * Unlike `Package.create`, a per-model compile failure here does
+    * NOT abort the reload: we keep the failed model as a placeholder
+    * (`Model.fromCompilationError`) in `this.models`, matching the
+    * historical reload semantics. Whole-pool failures (worker crash,
+    * timeout, pool shutting down) propagate as `ServiceUnavailableError`
+    * — the caller (manifest service) decides how to retry.
+    */
    public async reloadAllModels(
       buildManifest: BuildManifest["entries"],
    ): Promise<void> {
@@ -241,20 +345,62 @@ export class Package {
          modelCount: modelPaths.length,
          manifestEntryCount: Object.keys(buildManifest).length,
       });
-      const reloaded = await Promise.all(
-         modelPaths.map((modelPath) =>
-            Model.create(
-               this.packageName,
-               this.packagePath,
-               modelPath,
-               this.malloyConfig,
-               { buildManifest },
-            ),
-         ),
-      );
+
+      const pool = getPackageLoadPool();
+      let outcome;
+      try {
+         outcome = await pool.loadPackage({
+            packagePath: this.packagePath,
+            packageName: this.packageName,
+            malloyConfig: this.malloyConfig,
+            defaultConnectionName: "duckdb",
+            buildManifest,
+         });
+      } catch (err) {
+         const realError =
+            err instanceof Error
+               ? err
+               : new Error(`Package-load worker pool failure: ${String(err)}`);
+         if (
+            realError instanceof MalloyError ||
+            realError instanceof ModelCompilationError
+         ) {
+            throw realError;
+         }
+         throw new ServiceUnavailableError(
+            `Package-load worker pool unavailable: ${realError.message}`,
+         );
+      }
+
       const nextModels = new Map<string, Model>();
-      for (const model of reloaded) {
-         nextModels.set(model.getPath(), model);
+      for (const sm of outcome.models) {
+         if (sm.compilationError) {
+            const err = Model.deserializeCompilationError(sm.compilationError);
+            logger.warn("Model compilation failed during reload", {
+               packageName: this.packageName,
+               modelPath: sm.modelPath,
+               error: err.message,
+            });
+            nextModels.set(
+               sm.modelPath,
+               Model.fromCompilationError(
+                  this.packageName,
+                  sm.modelPath,
+                  sm.modelType,
+                  err,
+               ),
+            );
+         } else {
+            nextModels.set(
+               sm.modelPath,
+               Model.fromSerialized(
+                  this.packageName,
+                  this.packagePath,
+                  this.malloyConfig,
+                  sm,
+               ),
+            );
+         }
       }
       this.models = nextModels;
    }
@@ -317,20 +463,6 @@ export class Package {
       );
    }
 
-   private static async loadModels(
-      packageName: string,
-      packagePath: string,
-      malloyConfig: MalloyConfig,
-   ): Promise<Map<string, Model>> {
-      const modelPaths = await Package.getModelPaths(packagePath);
-      const models = await Promise.all(
-         modelPaths.map((modelPath) =>
-            Model.create(packageName, packagePath, modelPath, malloyConfig),
-         ),
-      );
-      return new Map(models.map((model) => [model.getPath(), model]));
-   }
-
    private static buildPackageMalloyConfig(
       packagePath: string,
       getEnvironmentMalloyConfig: () => MalloyConfig,
@@ -380,27 +512,6 @@ export class Package {
       return malloyConfig;
    }
 
-   private static async getModelPaths(packagePath: string): Promise<string[]> {
-      let files = undefined;
-      try {
-         files = await recursive(packagePath, [ignoreDotfiles]);
-      } catch (error) {
-         logger.error(error);
-         throw new PackageNotFoundError(
-            `Package config for ${packagePath} does not exist.`,
-         );
-      }
-      return files
-         .map((fullPath: string) => {
-            return path.relative(packagePath, fullPath).replace(/\\/g, "/");
-         })
-         .filter(
-            (modelPath: string) =>
-               modelPath.endsWith(MODEL_FILE_SUFFIX) ||
-               modelPath.endsWith(NOTEBOOK_FILE_SUFFIX),
-         );
-   }
-
    private static async validatePackageManifestExistsOrThrowError(
       packagePath: string,
    ) {
@@ -415,37 +526,32 @@ export class Package {
       }
    }
 
-   private static async readPackageConfig(
-      packagePath: string,
-   ): Promise<ApiPackage> {
-      const packageConfigPath = path.join(packagePath, PACKAGE_MANIFEST_NAME);
-      const packageConfigContents = await fs.readFile(packageConfigPath);
-      // TODO: Validate package manifest.  Define manifest type in public API.
-      const packageManifest = JSON.parse(packageConfigContents.toString());
-      return {
-         name: packageManifest.name,
-         description: packageManifest.description,
-      };
-   }
-
    private static async readDatabases(
       packagePath: string,
+      malloyConfig: MalloyConfig,
    ): Promise<ApiDatabase[]> {
+      const databasePaths = await Package.getDatabasePaths(packagePath);
+      if (databasePaths.length === 0) {
+         return [];
+      }
+      // Resolve the package's duckdb connection ONCE and reuse it for
+      // every schema/row-count probe in this package. Malloy caches the
+      // materialized connection on the MalloyConfig so the same instance
+      // will be returned to model compiles later in `Package.create`.
+      // This is the substantive optimization over the previous code:
+      // we go from `databasePaths.length` separate DuckDBConnections
+      // (each doing its own native init + extension load) to one.
+      const conn = await malloyConfig.connections.lookupConnection("duckdb");
       return await Promise.all(
-         (await Package.getDatabasePaths(packagePath)).map(
-            async (databasePath) => {
-               const databaseInfo = await Package.getDatabaseInfo(
-                  packagePath,
-                  databasePath,
-               );
-
-               return {
-                  path: databasePath,
-                  info: databaseInfo,
-                  type: "embedded",
-               };
-            },
-         ),
+         databasePaths.map(async (databasePath) => ({
+            path: databasePath,
+            info: await Package.getDatabaseInfo(
+               packagePath,
+               databasePath,
+               conn,
+            ),
+            type: "embedded" as const,
+         })),
       );
    }
 
@@ -466,135 +572,38 @@ export class Package {
    private static async getDatabaseInfo(
       packagePath: string,
       databasePath: string,
+      conn: Connection,
    ): Promise<ApiTableDescription> {
       const fullPath = path.join(packagePath, databasePath);
-      // Normalize path to use forward slashes for cross-platform compatibility.
-      // DuckDB on Windows supports forward slashes, which avoids escaping issues.
-      const normalizedPath = fullPath.replace(/\\/g, "/");
-
-      // One DuckDB connection per file (matches the historical
-      // ConnectionRuntime shape). Reused for both the schema probe and
-      // the row-count SQL so we only pay native init once per call.
-      const conn = new DuckDBConnection("duckdb");
-
-      // Schema probe. We need Malloy's view of the column types
-      // (consumers of the API rely on Malloy type strings, not DuckDB
-      // native types), so this stays a Malloy compile. The compile is
-      // CPU-heavy on the main thread relative to the work it produces,
-      // so when the worker pool is enabled we ship the synthetic
-      // `source: temp is duckdb.table(…)` snippet to a worker and read
-      // the resulting modelDef back here. Schema-fetch RPCs from the
-      // worker proxy through the pool against `conn` below.
-      const pool = getCompilePool();
-      let schema: ApiColumn[];
-      if (pool.enabled) {
-         schema = await Package.getSchemaViaPool(
-            pool,
-            packagePath,
-            normalizedPath,
-            conn,
-         );
-      } else {
-         schema = await Package.getSchemaInProcess(normalizedPath, conn);
-      }
 
-      // Row count. We previously compiled
-      // `run: temp->{aggregate: row_count is count()}` via Malloy, which
-      // is another full parse + type-check + SQL gen on the main thread
-      // just to produce a literal `SELECT count(*)`. Skip Malloy entirely
-      // and call DuckDB directly — same wire result, no main-thread CPU.
-      // Single quotes in the path are doubled to be safe against
-      // SQL-injection from filenames (DuckDB uses ANSI quoting rules).
-      const escapedPath = normalizedPath.replace(/'/g, "''");
-      const sqlResult = await conn.runSQL(
-         `SELECT count(*)::BIGINT AS row_count FROM '${escapedPath}'`,
-      );
-      const firstRow = sqlResult.rows[0] as { row_count?: bigint | number };
-      const rowCount = Number(firstRow.row_count ?? 0);
-
-      return { name: databasePath, rowCount, columns: schema };
-   }
-
-   /**
-    * In-process schema probe (legacy / kill-switch path). Builds a
-    * minimal ConnectionRuntime + compiles the synthetic snippet on the
-    * main thread. Same behaviour as the pre-worker-pool implementation.
-    */
-   private static async getSchemaInProcess(
-      normalizedPath: string,
-      conn: DuckDBConnection,
-   ): Promise<ApiColumn[]> {
+      // Create a DuckDB source then:
+      // 1. Load the model and get the table schema from model
+      // 2. Run a query to get the row count from the table
+      // ConnectionRuntime is cheap (just a wrapper), and creating one
+      // per call keeps each probe's compile state isolated. The
+      // expensive piece — the underlying DuckDBConnection — is shared
+      // across all probes via `conn` (resolved once in readDatabases).
       const runtime = new ConnectionRuntime({
          urlReader: new EmptyURLReader(),
          connections: [conn],
       });
+      // Normalize path to use forward slashes for cross-platform compatibility
+      // DuckDB on Windows supports forward slashes, and this avoids escaping issues
+      const normalizedPath = fullPath.replace(/\\/g, "/");
       const model = runtime.loadModel(
          `source: temp is duckdb.table('${normalizedPath}')`,
       );
       const modelDef = await model.getModel();
       const fields = (modelDef._modelDef.contents["temp"] as SourceDef).fields;
-      return fields.map((field): ApiColumn => {
+      const schema = fields.map((field): ApiColumn => {
          return { type: field.type, name: field.name };
       });
-   }
-
-   /**
-    * Worker-pool schema probe. The synthetic Malloy snippet is compiled
-    * in a worker_threads worker; the worker's schema-fetch RPC bounces
-    * back to the main thread, which services it against the
-    * `MalloyConfig` we hold on `conn` below. The returned modelDef has
-    * the resolved field list ready to read.
-    */
-   private static async getSchemaViaPool(
-      pool: ReturnType<typeof getCompilePool>,
-      packagePath: string,
-      normalizedPath: string,
-      conn: DuckDBConnection,
-   ): Promise<ApiColumn[]> {
-      // Keep schema-fetch RPCs from the worker routed to *this* conn
-      // (same instance used by the row-count SQL) by wrapping it in a
-      // MalloyConfig. We don't ship this config across the worker
-      // boundary; the pool holds it on the main side.
-      const malloyConfig = new MalloyConfig(
-         { connections: {} },
-         { config: contextOverlay({ rootDirectory: packagePath }) },
+      const runner = model.loadQuery(
+         "run: temp->{aggregate: row_count is count()}",
       );
-      malloyConfig.wrapConnections(() => ({
-         lookupConnection: async (_name?: string) =>
-            conn as unknown as Connection,
-      }));
-
-      try {
-         const outcome = await pool.compileInline({
-            packagePath,
-            source: `source: temp is duckdb.table('${normalizedPath}')`,
-            malloyConfig,
-            defaultConnectionName: "duckdb",
-         });
-         const modelDef = outcome.modelDef as unknown as {
-            contents: Record<string, SourceDef>;
-         };
-         const fields = modelDef.contents["temp"].fields;
-         return fields.map((field): ApiColumn => {
-            return { type: field.type, name: field.name };
-         });
-      } catch (error) {
-         // Transient pool issues (worker exit, RPC timeout) shouldn't
-         // break package loading. Compile errors here would mean the
-         // file isn't a readable parquet/csv, which in-process would
-         // also throw — so let those propagate.
-         if (
-            error instanceof Error &&
-            !/timed out|exited unexpectedly|shutting down/i.test(error.message)
-         ) {
-            throw error;
-         }
-         logger.warn(
-            "Compile worker failed for database probe; falling back to in-process",
-            { normalizedPath, error: (error as Error).message },
-         );
-         return Package.getSchemaInProcess(normalizedPath, conn);
-      }
+      const result = await runner.run();
+      const rowCount = result.data.value[0].row_count?.valueOf() as number;
+      return { name: databasePath, rowCount, columns: schema };
    }
 
    public setName(name: string) {