refactor(cli): replace deprecated vm2 with native node:vm runner

cli/vm/BUILD

@@ -1,5 +1,6 @@
 package(default_visibility = ["//visibility:public"])
 
+load("//testing:index.bzl", "ts_test_suite")
 load("//tools:ts_library.bzl", "ts_library")
 
 ts_library(
@@ -10,14 +11,14 @@ ts_library(
     ],
     deps = [
         "//common/protos",
+        "//common/vm:vm_runner",
         "//core",
         "//protos:ts",
         "@npm//@types/glob",
         "@npm//@types/node",
         "@npm//@types/semver",
         "@npm//glob",
         "@npm//semver",
-        "@npm//vm2",
     ],
 )
 

cli/vm/compile.ts

@@ -2,22 +2,26 @@ import * as fs from "fs";
 import * as glob from "glob";
 import * as path from "path";
 import * as semver from "semver";
-import { CompilerFunction, NodeVM } from "vm2";
 
 import { encode64 } from "df/common/protos";
+import { CompilerFunction, VmRunner } from "df/common/vm/vm_runner";
 import { dataform } from "df/protos/ts";
 
 export function compile(compileConfig: dataform.ICompileConfig) {
   compileConfig.projectDir = fs.realpathSync(path.resolve(compileConfig.projectDir));
   const coreBundlePath = path.join(
-    compileConfig.projectDir, "node_modules", "@dataform", "core", "bundle.js"
+    compileConfig.projectDir,
+    "node_modules",
+    "@dataform",
+    "core",
+    "bundle.js",
   );
   if (!fs.existsSync(coreBundlePath)) {
     throw new Error(
       "Could not find a recent installed version of @dataform/core in the project. Check that " +
         "either `dataformCoreVersion` is specified in `workflow_settings.yaml`, or " +
         "`@dataform/core` is specified in `package.json`. If using `package.json`, then run " +
-        "`dataform install`."
+        "`dataform install`.",
     );
   }
 
@@ -27,22 +31,17 @@ export function compile(compileConfig: dataform.ICompileConfig) {
   // through Node's resolver inside the vm covers every install layout
   // (package.json, workflow_settings.yaml, JiT) and matches what the user's
   // code will see. require() caches the bundle so the second call is free.
-  const indexGeneratorVm = new NodeVM({
-    wrapper: "none",
-    require: {
-      context: "sandbox",
-      root: compileConfig.projectDir,
-      external: true,
-      builtin: ["path"]
-    }
+  const indexGeneratorVm = new VmRunner({
+    projectDir: compileConfig.projectDir,
+    builtinModules: ["path"],
   });
   const compiler: CompilerFunction = indexGeneratorVm.run(
     'return require("@dataform/core").compiler',
-    vmIndexFileName
+    vmIndexFileName,
   );
   const dataformCoreVersion: string = indexGeneratorVm.run(
     'return require("@dataform/core").version || "0.0.0"',
-    vmIndexFileName
+    vmIndexFileName,
   );
 
   const cliVersion = readCliVersion();
@@ -61,35 +60,32 @@ export function compile(compileConfig: dataform.ICompileConfig) {
           `${cliVersion}. The CLI requires @dataform/core >= ${minCoreVersion} ` +
           `(matching major.minor). Set \`dataformCoreVersion: ${cliVersion}\` in ` +
           `workflow_settings.yaml (or pin @dataform/core in package.json), then run ` +
-          `\`dataform install\`.`
+          `\`dataform install\`.`,
       );
     }
   }
   const needsCallerFileShim = semver.lt(dataformCoreVersion, "3.0.57");
 
-  // vm2 strips file paths from V8 CallSite objects inside the sandbox, so
-  // getCallerFile() in @dataform/core needs a fallback. Track the currently
+  // While VmRunner preserves V8 CallSite file paths natively, older @dataform/core
+  // versions check global.__dataform_current_file as a fallback. Track the currently
   // executing file via a host-side stack exposed through sandbox helpers, and
   // expose it as a getter on `global.__dataform_current_file`.
   const fileStack: string[] = [];
 
-  // Then use vm2's native compiler integration to apply the compiler to files.
-  const userCodeVm = new NodeVM({
-    wrapper: "none",
+  // Then use VmRunner to apply the compiler to files.
+  const userCodeVm = new VmRunner({
+    projectDir: compileConfig.projectDir,
     sandbox: {
-      __df_enter: (p: string) => { fileStack.push(p); },
-      __df_exit: () => { fileStack.pop(); },
-      __df_current: () => fileStack.length > 0 ? fileStack[fileStack.length - 1] : null
-    },
-    require: {
-      builtin: ["path"],
-      context: "sandbox",
-      external: true,
-      root: compileConfig.projectDir,
-      resolve: (moduleName, parentDirName) =>
-        path.join(parentDirName, path.relative(parentDirName, compileConfig.projectDir), moduleName)
+      __df_enter: (p: string) => {
+        fileStack.push(p);
+      },
+      __df_exit: () => {
+        fileStack.pop();
+      },
+      __df_current: () => (fileStack.length > 0 ? fileStack[fileStack.length - 1] : null),
     },
-    sourceExtensions: ["js", "sql", "sqlx", "yaml", "yml"],
+    builtinModules: ["path"],
+    sourceExtensions: ["js", "sql", "sqlx", "yaml", "yml", "ipynb", "md"],
     compiler: (code, filePath) => {
       let source = code;
       if (needsCallerFileShim && filePath === coreBundlePath) {
@@ -104,31 +100,29 @@ export function compile(compileConfig: dataform.ICompileConfig) {
           __df_exit();
         }
       `;
-    }
+    },
   });
 
   const hasWorkflowSettingsYaml = fs.existsSync(
-    path.join(compileConfig.projectDir, "workflow_settings.yaml")
-  );
-  const hasDataformJson = fs.existsSync(
-    path.join(compileConfig.projectDir, "dataform.json")
+    path.join(compileConfig.projectDir, "workflow_settings.yaml"),
   );
+  const hasDataformJson = fs.existsSync(path.join(compileConfig.projectDir, "dataform.json"));
 
   return userCodeVm.run(
     `
       Object.defineProperty(global, '__dataform_current_file', {
         configurable: true,
         get: function() { return __df_current(); }
       });
-      ${hasWorkflowSettingsYaml
-        ? 'global.workflowSettingsYaml = require("./workflow_settings.yaml");'
-        : ''}
-      ${hasDataformJson
-        ? 'global.dataformJson = require("./dataform.json");'
-        : ''}
+      ${
+        hasWorkflowSettingsYaml
+          ? 'global.workflowSettingsYaml = require("./workflow_settings.yaml");'
+          : ""
+      }
+      ${hasDataformJson ? 'global.dataformJson = require("./dataform.json");' : ""}
       return require("@dataform/core").main("${createCoreExecutionRequest(compileConfig)}")
     `,
-    vmIndexFileName
+    vmIndexFileName,
   );
 }
 
@@ -162,9 +156,7 @@ if (require.main === module) {
 // by pkg_json(version = DF_VERSION). Returns "0.0.0" when unreadable.
 function readCliVersion(): string {
   try {
-    const pkg = JSON.parse(
-      fs.readFileSync(path.join(__dirname, "package.json"), "utf8")
-    );
+    const pkg = JSON.parse(fs.readFileSync(path.join(__dirname, "package.json"), "utf8"));
     return pkg.version || "0.0.0";
   } catch {
     return "0.0.0";
@@ -179,7 +171,7 @@ function readCliVersion(): string {
 const OLD_CORE_THROW =
   'if(!t)throw new Error("Unable to find valid caller file; please report this issue.")';
 const OLD_CORE_WITH_FALLBACK =
-  'if(!t){if(global.__dataform_current_file){t=global.__dataform_current_file}' +
+  "if(!t){if(global.__dataform_current_file){t=global.__dataform_current_file}" +
   'else{throw new Error("Unable to find valid caller file; please report this issue.")}}';
 
 function patchOldCoreCallerFile(source: string): string {
@@ -191,11 +183,11 @@ function patchOldCoreCallerFile(source: string): string {
  */
 function createCoreExecutionRequest(compileConfig: dataform.ICompileConfig): string {
   const filePaths = Array.from(
-    new Set<string>(glob.sync("!(node_modules)/**/*.*", { cwd: compileConfig.projectDir }))
+    new Set<string>(glob.sync("!(node_modules)/**/*.*", { cwd: compileConfig.projectDir })),
   );
 
   return encode64(dataform.CoreExecutionRequest, {
     // Add the list of file paths to the compile config if not already set.
-    compile: { compileConfig: { filePaths, ...compileConfig } }
+    compile: { compileConfig: { filePaths, ...compileConfig } },
   });
 }

cli/vm/jit_worker.ts

@@ -1,10 +1,13 @@
 import * as fs from "fs";
 import * as path from "path";
-import { NodeVM } from "vm2";
 
+import { VmRunner } from "df/common/vm/vm_runner";
 import { dataform } from "df/protos/ts";
 
-const pendingRpcCallbacks = new Map<string, (err: string | null, resBytes: Uint8Array | null) => void>();
+const pendingRpcCallbacks = new Map<
+  string,
+  (err: string | null, resBytes: Uint8Array | null) => void
+>();
 
 export function registerRpcResponseHandler() {
   process.on("message", (res: any) => {
@@ -25,7 +28,8 @@ export function registerJitCompileHandler() {
       if (hasStartedProcessing) {
         process.send({
           type: "jit_error",
-          error: "Worker process received multiple JiT compilation requests. Subsequent requests are rejected."
+          error:
+            "Worker process received multiple JiT compilation requests. Subsequent requests are rejected.",
         });
         return;
       }
@@ -35,34 +39,44 @@ export function registerJitCompileHandler() {
   });
 }
 
-export async function handleJitRequest(message: {
-  request: any;
-  projectDir: string;
-}) {
+export async function handleJitRequest(message: { request: any; projectDir: string }) {
   try {
     const { request, projectDir } = message;
 
-    const projectLocalCorePath = path.join(projectDir, "node_modules", "@dataform", "core", "bundle.js");
+    const projectLocalCorePath = path.join(
+      projectDir,
+      "node_modules",
+      "@dataform",
+      "core",
+      "bundle.js",
+    );
     const hasProjectLocalCore = fs.existsSync(projectLocalCorePath);
 
-    if (!hasProjectLocalCore && !fs.existsSync(path.join(projectDir, "node_modules", "@dataform", "core", "package.json"))) {
+    if (
+      !hasProjectLocalCore &&
+      !fs.existsSync(path.join(projectDir, "node_modules", "@dataform", "core", "package.json"))
+    ) {
       throw new Error(
         "Could not find a recent installed version of @dataform/core in the project. Check that " +
           "either `dataformCoreVersion` is specified in `workflow_settings.yaml`, or " +
           "`@dataform/core` is specified in `package.json`. If using `package.json`, then run " +
-          "`dataform install`."
+          "`dataform install`.",
       );
     }
 
-    const rpcCallback = (method: string, reqBytes: Uint8Array, callback: (err: string | null, resBytes: Uint8Array | null) => void) => {
+    const rpcCallback = (
+      method: string,
+      reqBytes: Uint8Array,
+      callback: (err: string | null, resBytes: Uint8Array | null) => void,
+    ) => {
       const correlationId = Math.random().toString(36).substring(7);
       pendingRpcCallbacks.set(correlationId, callback);
 
       process.send({
         type: "rpc_request",
         method,
         request: reqBytes,
-        correlationId
+        correlationId,
       });
     };
 
@@ -71,20 +85,19 @@ export async function handleJitRequest(message: {
 
     const vmFileName = path.resolve(projectDir, "index.js");
 
-    const vm = new NodeVM({
-      require: {
-        builtin: [],
-        context: "sandbox",
-        external: { modules: ["@dataform/*"], transitive: false },
-        root: projectDir,
-        mock: hasProjectLocalCore ? {} : {
-          "@dataform/core": require("@dataform/core")
-        }
-      },
-      sourceExtensions: ["js", "json", "yaml", "yml"]
+    const vm = new VmRunner({
+      projectDir,
+      builtinModules: [],
+      mockModules: hasProjectLocalCore
+        ? {}
+        : {
+            "@dataform/core": require("@dataform/core"),
+          },
+      sourceExtensions: ["js", "json", "yaml", "yml"],
     });
 
-    const jitCompileInVm = vm.run(`
+    const jitCompileInVm = vm.run(
+      `
       const { jitCompiler } = require("@dataform/core");
 
       global.require = require;
@@ -103,10 +116,14 @@ export async function handleJitRequest(message: {
         const compilerInstance = jitCompiler(internalRpcCallback);
         return await compilerInstance.compile(requestBytesTyped);
       };
-    `, vmFileName);
+    `,
+      vmFileName,
+    );
 
     const responseBytes = await jitCompileInVm(requestBytes, rpcCallback);
-    const response = dataform.JitCompilationResponse.decode(new Uint8Array(responseBytes as number[]));
+    const response = dataform.JitCompilationResponse.decode(
+      new Uint8Array(responseBytes as number[]),
+    );
 
     process.send({ type: "jit_response", response: response.toJSON() });
   } catch (e) {

common/vm/BUILD

@@ -0,0 +1,50 @@
+package(default_visibility = ["//visibility:public"])
+
+load("//testing:index.bzl", "ts_test_suite")
+load("//tools:ts_library.bzl", "ts_library")
+
+ts_library(
+    name = "vm_runner",
+    srcs = [
+        "vm_runner.ts",
+    ],
+    deps = [
+        "@npm//@types/node",
+    ],
+)
+
+ts_test_suite(
+    name = "tests",
+    srcs = ["vm_runner_test.ts"],
+    deps = [
+        ":vm_runner",
+        "//testing",
+        "@npm//@types/chai",
+        "@npm//@types/node",
+        "@npm//chai",
+    ],
+)
+
+load("@build_bazel_rules_nodejs//:index.bzl", "nodejs_binary")
+
+ts_library(
+    name = "benchmark_lib",
+    srcs = ["vm_runner_benchmark.ts"],
+    deps = [
+        ":vm_runner",
+        "@npm//@types/node",
+    ],
+)
+
+nodejs_binary(
+    name = "benchmark",
+    data = [
+        ":benchmark_lib",
+        "@npm//source-map-support",
+    ],
+    entry_point = ":vm_runner_benchmark.ts",
+    templated_args = [
+        "--node_options=--require=source-map-support/register",
+        "--bazel_patch_module_resolver",
+    ],
+)