feat: Support MongoDB

Fixes #106

.appmapignore

@@ -0,0 +1,2 @@
+src
+test

.github/workflows/ci.yml

@@ -29,6 +29,10 @@ jobs:
           --health-retries 5
         ports:
           - 5432:5432
+      mongodb:
+        image: mongo
+        ports:
+          - 27017:27017
 
     steps:
       - uses: actions/checkout@v3
@@ -40,6 +44,7 @@ jobs:
       - run: yarn test
         env:
           POSTGRES_URL: postgres://postgres:postgres@localhost:5432
+          MONGODB_URI: mongodb://localhost:27017
 
   windows-test:
     runs-on: windows-latest

package.json

@@ -57,6 +57,7 @@
     "eslint-plugin-prettier": "^5.0.1",
     "fast-glob": "^3.3.1",
     "jest": "^29.6.2",
+    "mongodb": "^6.3.0",
     "next": "^14.0.4",
     "prettier": "^3.0.2",
     "react": "^18",

src/hooks/mongo.ts

@@ -0,0 +1,154 @@
+import { inspect } from "node:util";
+
+import type mongodb from "mongodb";
+
+import { identifier } from "../generate";
+import { fixReturnEventIfPromiseResult, recording } from "../recorder";
+import { FunctionInfo } from "../registry";
+import { getTime } from "../util/getTime";
+import { setCustomInspect } from "../parameter";
+
+export default function mongoHook(mod: typeof mongodb) {
+  const collectionMethods: Partial<Record<MethodLikeKeys<mongodb.Collection>, readonly string[]>> =
+    {
+      insertOne: ["doc", "options"],
+      insertMany: ["docs", "options"],
+      bulkWrite: ["operations", "options"],
+      updateOne: ["filter", "update", "options"],
+      replaceOne: ["filter", "replacement", "options"],
+      updateMany: ["filter", "update", "options"],
+      deleteOne: ["filter", "options"],
+      deleteMany: ["filter", "options"],
+      rename: ["newName", "options"],
+      drop: ["options"],
+      findOne: ["filter", "options"],
+      find: ["filter", "options"],
+      options: ["options"],
+      isCapped: ["options"],
+      createIndex: ["indexSpec", "options"],
+      createIndexes: ["indexSpecs", "options"],
+      dropIndex: ["indexName", "options"],
+      dropIndexes: ["options"],
+      listIndexes: ["options"],
+      indexExists: ["indexes", "options"],
+      indexInformation: ["options"],
+      estimatedDocumentCount: ["options"],
+      countDocuments: ["filter", "options"],
+      distinct: ["key", "filter", "options"],
+      indexes: ["options"],
+      findOneAndDelete: ["filter", "options"],
+      findOneAndReplace: ["filter", "replacement", "options"],
+      findOneAndUpdate: ["filter", "update", "options"],
+      aggregate: ["pipeline", "options"],
+      watch: ["pipeline", "options"],
+      count: ["filter", "options"],
+    };
+  for (const [method, args] of Object.entries(collectionMethods))
+    patchMethod(mod.Collection.prototype, method as MethodLikeKeys<mongodb.Collection>, args);
+
+  setCustomInspect(mod.Collection.prototype, (c) => `[Collection ${c.collectionName}]`);
+  setCustomInspect(
+    mod.AbstractCursor?.prototype,
+    (c) => `[${c.constructor.name} ${c.namespace.toString()}]`,
+  );
+  return mod;
+}
+
+mongoHook.applicable = function (id: string) {
+  return id === "mongodb";
+};
+
+const funInfos = new Map<string, FunctionInfo>();
+
+// the spec requires every method in the same path to have unique line number
+// we're emitting synthetic methods so make sure they all differ
+let lineNo = 1;
+
+function functionInfo(name: string, collection: string, argnames: readonly string[]) {
+  const key = [collection, name].join(":");
+  if (!funInfos.has(key))
+    funInfos.set(key, {
+      async: true,
+      generator: false,
+      id: name,
+      klassOrFile: collection,
+      params: argnames.map(identifier),
+      static: false,
+      location: { lineno: lineNo++, path: `mongodb/${collection}` },
+    });
+  return funInfos.get(key)!;
+}
+
+// use custom inspect so IDs are rendered properly
+const customInspect = (v: unknown) => inspect(v, { customInspect: true });
+
+function patchMethod<K extends MethodLikeKeys<mongodb.Collection>>(
+  obj: typeof mongodb.Collection.prototype,
+  methodName: K,
+  argNames: readonly string[],
+) {
+  const original = obj[methodName];
+
+  if (isPatched(original)) return;
+
+  const patched = function (
+    this: mongodb.Collection,
+    ...args: unknown[]
+  ): ReturnType<typeof original> {
+    if (!recording) return Reflect.apply(original, this, args) as ReturnType<typeof original>;
+
+    const funInfo = functionInfo(methodName, this.collectionName, argNames);
+    const callback = extractOptionalCallback(args);
+    if (callback) {
+      const callEvent = recording.functionCall(funInfo, this, args);
+      const startTime = getTime();
+      args.push((err: unknown, res: unknown) => {
+        setCustomInspect(res, customInspect);
+        if (err) recording.functionException(callEvent.id, err, getTime() - startTime);
+        else recording.functionReturn(callEvent.id, res, getTime() - startTime);
+        return callback(err, res) as unknown;
+      });
+      return Reflect.apply(original, this, args) as ReturnType<typeof original>;
+    }
+
+    const callEvent = recording.functionCall(funInfo, this, args);
+    const startTime = getTime();
+
+    try {
+      const result = Reflect.apply(original, this, args) as unknown;
+      setCustomInspect(result, customInspect);
+      const returnEvent = recording.functionReturn(callEvent.id, result, getTime() - startTime);
+      return fixReturnEventIfPromiseResult(result, returnEvent, callEvent, startTime) as ReturnType<
+        typeof original
+      >;
+    } catch (exn: unknown) {
+      recording.functionException(callEvent.id, exn, getTime() - startTime);
+      throw exn;
+    }
+  };
+
+  markPatched(patched);
+
+  obj[methodName] = patched as typeof original;
+}
+
+function extractOptionalCallback(args: unknown[]): FunctionLike | undefined {
+  if (typeof args.at(-1) === "function") return args.pop() as FunctionLike;
+}
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type FunctionLike = (...args: any) => any;
+
+type MethodLikeKeys<T extends object> = keyof {
+  [K in keyof T as T[K] extends FunctionLike ? K : never]: T[K];
+};
+
+const patchedMarker = Symbol("AppMap-patched");
+
+function markPatched(patched: object) {
+  (patched as { [patchedMarker]: boolean })[patchedMarker] = true;
+}
+
+function isPatched<T extends object>(original: T): boolean {
+  return patchedMarker in original;
+}

src/parameter.ts

@@ -16,6 +16,24 @@ export function parameter(value: unknown): AppMap.Parameter {
   });
 }
 
+const kCustomInspect = Symbol("AppMap.customInspect");
+
+export function setCustomInspect<T>(value: T, customInspect: (v: T) => string) {
+  if (value == null || typeof value !== "object") return;
+  if (kCustomInspect in value) return;
+  Object.defineProperty(value, kCustomInspect, {
+    value: customInspect,
+    enumerable: false,
+    writable: true,
+  });
+}
+
+function doInspect(value: unknown): string {
+  if (typeof value === "object" && value && kCustomInspect in value)
+    return (value as { [kCustomInspect]: (v: typeof value) => string })[kCustomInspect](value);
+  return inspect(value, { depth: 1, customInspect: false });
+}
+
 export function stringify(value: unknown): string {
   if (value instanceof IncomingMessage)
     return format("[IncomingMessage: %s %s]", value.method, value.url);
@@ -26,7 +44,7 @@ export function stringify(value: unknown): string {
     );
   // Pause recorder to prevent potential recursive calls by inspect()
   pauseRecorder();
-  const result = inspect(value, { depth: 1, customInspect: false });
+  const result = doInspect(value);
   resumeRecorder();
   return result;
 }

src/requireHook.ts

@@ -1,4 +1,5 @@
 import httpHook from "./hooks/http";
+import mongoHook from "./hooks/mongo";
 import mysqlHook from "./hooks/mysql";
 import pgHook from "./hooks/pg";
 import prismaHook from "./hooks/prisma";
@@ -10,7 +11,7 @@ interface Hook {
   applicable(id: string): boolean;
 }
 
-const hooks: Hook[] = [httpHook, mysqlHook, pgHook, sqliteHook, prismaHook];
+const hooks: Hook[] = [httpHook, mongoHook, mysqlHook, pgHook, sqliteHook, prismaHook];
 
 export default function requireHook(
   original: NodeJS.Require,