@senzops/apm-node 1.1.18 → 1.2.0

package/src/instrumentation/mongo.ts

@@ -1,105 +1,202 @@
-import { Context } from '../core/context';
-
-export const instrumentMongo = (debug = false) => {
-  try {
-    const mongodb = require('mongodb');
-    const Collection = mongodb.Collection;
-
-    // Attempt to get Cursor classes
-    // Note: The location of these classes varies by driver version, 
-    // checking common locations
-    const FindCursor = mongodb.FindCursor || require('mongodb/lib/cursor/find_cursor').FindCursor;
-    const AggregationCursor = mongodb.AggregationCursor || require('mongodb/lib/cursor/aggregation_cursor').AggregationCursor;
-
-    if (debug) console.log('[Senzor] Instrumenting MongoDB (Collection + Cursors)...');
-
-    // --- Helper to Record Span ---
-    const recordSpan = (name: string, operation: string, collection: string, startAbs: number, traceStart: number, err?: Error) => {
-      const duration = performance.now() - startAbs;
-      Context.addSpan({
-        name: `MongoDB ${name}`,
-        type: 'db',
-        startTime: performance.now() - traceStart - duration, // Adjust start time to when op actually started
-        duration,
-        status: err ? 500 : 0,
-        meta: { collection, operation, error: err ? err.message : undefined }
-      });
-      if (debug) console.log(`[Senzor] Captured Mongo: ${name} (${duration.toFixed(2)}ms)`);
-    };
-
-    // --- 1. Instrument Immediate Operations (Insert/Update/Delete) ---
-    const immediateMethods = ['insertOne', 'insertMany', 'updateOne', 'updateMany', 'deleteOne', 'deleteMany', 'countDocuments'];
-
-    immediateMethods.forEach((method) => {
-      if (!Collection.prototype[method]) return;
-      const original = Collection.prototype[method];
-
-      Collection.prototype[method] = function (...args: any[]) {
-        const trace = Context.current();
-        if (!trace) return original.apply(this, args);
-
-        const spanStartAbs = performance.now();
-        const traceStart = trace.startTime;
-        const collName = this.collectionName;
-
-        try {
-          const result = original.apply(this, args);
-          if (result && typeof result.then === 'function') {
-            return result.then(
-              (res: any) => { recordSpan(method, method, collName, spanStartAbs, traceStart); return res; },
-              (err: any) => { recordSpan(method, method, collName, spanStartAbs, traceStart, err); throw err; }
-            );
-          }
-          return result;
-        } catch (err: any) {
-          recordSpan(method, method, collName, spanStartAbs, traceStart, err);
-          throw err;
-        }
-      };
-    });
-
-    // --- 2. Instrument Cursor Execution (find -> toArray) ---
-    const patchCursor = (CursorClass: any, label: string) => {
-      if (!CursorClass || !CursorClass.prototype.toArray) return;
-
-      const originalToArray = CursorClass.prototype.toArray;
-
-      CursorClass.prototype.toArray = function (...args: any[]) {
-        const trace = Context.current();
-        // Cursors are often created in context but executed later. 
-        // We check context at execution time.
-        if (!trace) return originalToArray.apply(this, args);
-
-        const spanStartAbs = performance.now();
-        const traceStart = trace.startTime;
-        // Attempt to get collection name from cursor internal state
-        const collName = this.namespace?.collection || 'unknown';
-
-        const onSuccess = (res: any) => {
-          recordSpan(label, label, collName, spanStartAbs, traceStart);
-          return res;
-        };
-        const onError = (err: any) => {
-          recordSpan(label, label, collName, spanStartAbs, traceStart, err);
-          throw err;
-        };
-
-        try {
-          const result = originalToArray.apply(this, args);
-          if (result && typeof result.then === 'function') {
-            return result.then(onSuccess, onError);
-          }
-          return onSuccess(result);
-        } catch (e) {
-          onError(e);
-        }
-      };
-    };
-
-    patchCursor(FindCursor, 'find');
-    patchCursor(AggregationCursor, 'aggregate');
-
-  } catch (e: any) {
-    if (debug) console.warn('[Senzor] MongoDB instrumentation warning:', e.message);
-  }
-};
+import { SenzorOptions } from '../core/types';
+import { hookRequire } from './hook';
+import { patchMethod } from './patch';
+import { runWithCapturedSpan, startCapturedSpan } from './span';
+
+const collectionName = (collection: any): string =>
+  collection?.collectionName ||
+  collection?.s?.namespace?.collection ||
+  collection?.namespace?.collection ||
+  'unknown';
+
+const databaseName = (collection: any): string | undefined =>
+  collection?.dbName ||
+  collection?.s?.namespace?.db ||
+  collection?.namespace?.db;
+
+const cursorCollectionName = (cursor: any): string =>
+  cursor?.namespace?.collection ||
+  cursor?.ns?.collection ||
+  cursor?.cursorNamespace?.collection ||
+  'unknown';
+
+const patchCollectionMethod = (
+  proto: any,
+  method: string,
+  options?: SenzorOptions
+) => {
+  patchMethod(
+    proto,
+    method,
+    `senzor.mongodb.collection.${method}`,
+    (original) =>
+      function patchedMongoCollection(this: any, ...args: any[]) {
+        const collection = collectionName(this);
+        const span = startCapturedSpan(
+          `MongoDB ${method}`,
+          'db',
+          {
+            collection,
+            operation: method,
+            'db.system.name': 'mongodb',
+            'db.collection.name': collection,
+            'db.namespace': databaseName(this)
+              ? `${databaseName(this)}.${collection}`
+              : collection,
+            'db.operation.name': method,
+            library: 'mongodb'
+          },
+          options
+        );
+
+        if (!span) return original.apply(this, args);
+
+        return runWithCapturedSpan(span, () => {
+          try {
+            const result = original.apply(this, args);
+            if (result && typeof result.then === 'function') {
+              return result.then(
+                (value: any) => {
+                  span.end(0, {
+                    matchedCount: value?.matchedCount,
+                    modifiedCount: value?.modifiedCount,
+                    deletedCount: value?.deletedCount,
+                    insertedCount: value?.insertedCount
+                  });
+                  return value;
+                },
+                (error: any) => {
+                  span.end(500, {
+                    error: error?.message,
+                    'error.type': error?.name || 'Error'
+                  });
+                  throw error;
+                }
+              );
+            }
+
+            span.end(0);
+            return result;
+          } catch (error: any) {
+            span.end(500, {
+              error: error?.message,
+              'error.type': error?.name || 'Error'
+            });
+            throw error;
+          }
+        });
+      }
+  );
+};
+
+const patchCursorMethod = (
+  proto: any,
+  method: string,
+  operation: string,
+  options?: SenzorOptions
+) => {
+  patchMethod(
+    proto,
+    method,
+    `senzor.mongodb.cursor.${operation}.${method}`,
+    (original) =>
+      function patchedMongoCursor(this: any, ...args: any[]) {
+        const collection = cursorCollectionName(this);
+        const span = startCapturedSpan(
+          `MongoDB ${operation}`,
+          'db',
+          {
+            collection,
+            operation,
+            'db.system.name': 'mongodb',
+            'db.collection.name': collection,
+            'db.operation.name': operation,
+            library: 'mongodb'
+          },
+          options
+        );
+
+        if (!span) return original.apply(this, args);
+
+        return runWithCapturedSpan(span, () => {
+          try {
+            const result = original.apply(this, args);
+            if (result && typeof result.then === 'function') {
+              return result.then(
+                (value: any) => {
+                  span.end(0, {
+                    resultCount: Array.isArray(value) ? value.length : undefined
+                  });
+                  return value;
+                },
+                (error: any) => {
+                  span.end(500, {
+                    error: error?.message,
+                    'error.type': error?.name || 'Error'
+                  });
+                  throw error;
+                }
+              );
+            }
+
+            span.end(0);
+            return result;
+          } catch (error: any) {
+            span.end(500, {
+              error: error?.message,
+              'error.type': error?.name || 'Error'
+            });
+            throw error;
+          }
+        });
+      }
+  );
+};
+
+const patchMongo = (mongodb: any, options?: SenzorOptions) => {
+  const Collection = mongodb?.Collection || mongodb?.default?.Collection;
+  const collectionProto = Collection?.prototype;
+
+  [
+    'insertOne',
+    'insertMany',
+    'updateOne',
+    'updateMany',
+    'replaceOne',
+    'deleteOne',
+    'deleteMany',
+    'findOne',
+    'findOneAndUpdate',
+    'findOneAndDelete',
+    'findOneAndReplace',
+    'countDocuments',
+    'estimatedDocumentCount',
+    'distinct',
+    'bulkWrite',
+    'createIndex',
+    'dropIndex'
+  ].forEach((method) =>
+    patchCollectionMethod(collectionProto, method, options)
+  );
+
+  const FindCursor =
+    mongodb?.FindCursor || mongodb?.default?.FindCursor;
+  const AggregationCursor =
+    mongodb?.AggregationCursor || mongodb?.default?.AggregationCursor;
+
+  ['toArray', 'next', 'forEach'].forEach((method) =>
+    patchCursorMethod(FindCursor?.prototype, method, 'find', options)
+  );
+  ['toArray', 'next', 'forEach'].forEach((method) =>
+    patchCursorMethod(
+      AggregationCursor?.prototype,
+      method,
+      'aggregate',
+      options
+    )
+  );
+};
+
+export const instrumentMongo = (options?: SenzorOptions) => {
+  hookRequire('mongodb', (exports: any) => patchMongo(exports, options));
+};

package/src/instrumentation/mongoose.ts

@@ -0,0 +1,156 @@
+import { SenzorOptions } from '../core/types';
+import { hookRequire } from './hook';
+import { patchMethod } from './patch';
+import { runWithCapturedSpan, startCapturedSpan } from './span';
+
+const modelName = (target: any): string =>
+  target?.model?.modelName ||
+  target?.constructor?.modelName ||
+  target?.modelName ||
+  'unknown';
+
+const collectionName = (target: any): string =>
+  target?.mongooseCollection?.name ||
+  target?.collection?.name ||
+  target?.model?.collection?.name ||
+  'unknown';
+
+const patchExec = (
+  proto: any,
+  label: 'query' | 'aggregate',
+  options?: SenzorOptions
+) => {
+  patchMethod(
+    proto,
+    'exec',
+    `senzor.mongoose.${label}.exec`,
+    (original) =>
+      function patchedMongooseExec(this: any, ...args: any[]) {
+        const operation =
+          String(this?.op || this?._op || label).toUpperCase();
+        const collection = collectionName(this);
+        const span = startCapturedSpan(
+          `Mongoose ${operation}`,
+          'db',
+          {
+            collection,
+            model: modelName(this),
+            operation,
+            'db.system.name': 'mongodb',
+            'db.collection.name': collection,
+            'db.operation.name': operation,
+            library: 'mongoose'
+          },
+          options
+        );
+
+        if (!span) return original.apply(this, args);
+
+        return runWithCapturedSpan(span, () => {
+          try {
+            const result = original.apply(this, args);
+            if (result && typeof result.then === 'function') {
+              return result.then(
+                (value: any) => {
+                  span.end(0, {
+                    resultCount: Array.isArray(value) ? value.length : undefined
+                  });
+                  return value;
+                },
+                (error: any) => {
+                  span.end(500, {
+                    error: error?.message,
+                    'error.type': error?.name || 'Error'
+                  });
+                  throw error;
+                }
+              );
+            }
+
+            span.end(0);
+            return result;
+          } catch (error: any) {
+            span.end(500, {
+              error: error?.message,
+              'error.type': error?.name || 'Error'
+            });
+            throw error;
+          }
+        });
+      }
+  );
+};
+
+const patchSave = (modelProto: any, options?: SenzorOptions) => {
+  patchMethod(
+    modelProto,
+    'save',
+    'senzor.mongoose.model.save',
+    (original) =>
+      function patchedMongooseSave(this: any, ...args: any[]) {
+        const collection = collectionName(this);
+        const span = startCapturedSpan(
+          'Mongoose SAVE',
+          'db',
+          {
+            collection,
+            model: modelName(this),
+            operation: 'SAVE',
+            'db.system.name': 'mongodb',
+            'db.collection.name': collection,
+            'db.operation.name': 'SAVE',
+            library: 'mongoose'
+          },
+          options
+        );
+
+        if (!span) return original.apply(this, args);
+
+        return runWithCapturedSpan(span, () => {
+          try {
+            const result = original.apply(this, args);
+            if (result && typeof result.then === 'function') {
+              return result.then(
+                (value: any) => {
+                  span.end(0);
+                  return value;
+                },
+                (error: any) => {
+                  span.end(500, {
+                    error: error?.message,
+                    'error.type': error?.name || 'Error'
+                  });
+                  throw error;
+                }
+              );
+            }
+
+            span.end(0);
+            return result;
+          } catch (error: any) {
+            span.end(500, {
+              error: error?.message,
+              'error.type': error?.name || 'Error'
+            });
+            throw error;
+          }
+        });
+      }
+  );
+};
+
+const patchMongoose = (mongoose: any, options?: SenzorOptions) => {
+  patchExec(mongoose?.Query?.prototype, 'query', options);
+  patchExec(mongoose?.Aggregate?.prototype, 'aggregate', options);
+  patchSave(mongoose?.Model?.prototype, options);
+
+  if (mongoose?.default) {
+    patchExec(mongoose.default?.Query?.prototype, 'query', options);
+    patchExec(mongoose.default?.Aggregate?.prototype, 'aggregate', options);
+    patchSave(mongoose.default?.Model?.prototype, options);
+  }
+};
+
+export const instrumentMongoose = (options?: SenzorOptions) => {
+  hookRequire('mongoose', (exports: any) => patchMongoose(exports, options));
+};

package/src/instrumentation/mysql.ts

@@ -0,0 +1,169 @@
+import { getSqlOperation, normalizeSql } from '../core/sanitizer';
+import { SenzorOptions } from '../core/types';
+import { hookRequire } from './hook';
+import { patchMethod } from './patch';
+import { runWithCapturedSpan, startCapturedSpan } from './span';
+
+const extractSql = (args: any[]): string | undefined => {
+  const first = args[0];
+  if (typeof first === 'string') return first;
+  if (first && typeof first.sql === 'string') return first.sql;
+  return undefined;
+};
+
+const patchSqlMethod = (
+  proto: any,
+  method: 'query' | 'execute',
+  library: string,
+  options?: SenzorOptions
+) => {
+  patchMethod(
+    proto,
+    method,
+    `senzor.${library}.${method}`,
+    (original) =>
+      function patchedMysqlMethod(this: any, ...args: any[]) {
+        const sql = extractSql(args);
+        const operation = getSqlOperation(sql) || method.toUpperCase();
+        const span = startCapturedSpan(
+          `MySQL ${operation}`,
+          'db',
+          {
+            query: normalizeSql(sql, options),
+            operation,
+            'db.system.name': 'mysql',
+            'db.operation.name': operation,
+            'db.query.text': normalizeSql(sql, options),
+            library
+          },
+          options
+        );
+
+        if (!span) return original.apply(this, args);
+
+        const callbackIndex = args.findIndex(
+          (arg) => typeof arg === 'function'
+        );
+        if (callbackIndex >= 0) {
+          const originalCallback = args[callbackIndex];
+          args[callbackIndex] = function wrappedMysqlCallback(
+            this: unknown,
+            err: any,
+            rows: any
+          ) {
+            span.end(err ? 500 : 0, {
+              error: err?.message,
+              'error.type': err?.name,
+              rowCount: Array.isArray(rows) ? rows.length : undefined
+            });
+            return originalCallback.apply(this, arguments as any);
+          };
+        }
+
+        return runWithCapturedSpan(span, () => {
+          try {
+            const result = original.apply(this, args);
+
+            if (result && typeof result.then === 'function') {
+              return result.then(
+                (value: any) => {
+                  const rows = Array.isArray(value) ? value[0] : value;
+                  span.end(0, {
+                    rowCount: Array.isArray(rows) ? rows.length : undefined
+                  });
+                  return value;
+                },
+                (error: any) => {
+                  span.end(500, {
+                    error: error?.message,
+                    'error.type': error?.name || 'Error'
+                  });
+                  throw error;
+                }
+              );
+            }
+
+            if (
+              callbackIndex < 0 &&
+              result &&
+              typeof result.once === 'function'
+            ) {
+              result.once('end', () => span.end(0));
+              result.once('error', (error: Error) =>
+                span.end(500, {
+                  error: error.message,
+                  'error.type': error.name
+                })
+              );
+            } else if (callbackIndex < 0) {
+              span.end(0);
+            }
+
+            return result;
+          } catch (error: any) {
+            span.end(500, {
+              error: error?.message,
+              'error.type': error?.name || 'Error'
+            });
+            throw error;
+          }
+        });
+      }
+  );
+};
+
+const patchKnownPrototypes = (
+  mysql: any,
+  library: string,
+  options?: SenzorOptions
+) => {
+  [
+    mysql?.Connection?.prototype,
+    mysql?.Pool?.prototype,
+    mysql?.PoolConnection?.prototype,
+    mysql?.PromiseConnection?.prototype,
+    mysql?.PromisePool?.prototype,
+    mysql?.default?.Connection?.prototype,
+    mysql?.default?.Pool?.prototype
+  ].forEach((proto) => {
+    patchSqlMethod(proto, 'query', library, options);
+    patchSqlMethod(proto, 'execute', library, options);
+  });
+};
+
+const patchFactories = (
+  mysql: any,
+  library: string,
+  options?: SenzorOptions
+) => {
+  ['createConnection', 'createPool'].forEach((factory) => {
+    patchMethod(
+      mysql,
+      factory,
+      `senzor.${library}.${factory}`,
+      (original) =>
+        function patchedMysqlFactory(this: any, ...args: any[]) {
+          const client = original.apply(this, args);
+          patchSqlMethod(client, 'query', library, options);
+          patchSqlMethod(client, 'execute', library, options);
+          patchSqlMethod(Object.getPrototypeOf(client), 'query', library, options);
+          patchSqlMethod(Object.getPrototypeOf(client), 'execute', library, options);
+          return client;
+        }
+    );
+  });
+};
+
+const patchMysql = (
+  mysql: any,
+  library: string,
+  options?: SenzorOptions
+) => {
+  patchKnownPrototypes(mysql, library, options);
+  patchFactories(mysql, library, options);
+};
+
+export const instrumentMysql = (options?: SenzorOptions) => {
+  hookRequire('mysql', (exports: any) => patchMysql(exports, 'mysql', options));
+  hookRequire('mysql2', (exports: any) => patchMysql(exports, 'mysql2', options));
+};

package/src/instrumentation/patch.ts

@@ -0,0 +1,56 @@
+const PATCHES = Symbol.for('senzor.patch.keys');
+const ORIGINAL = Symbol.for('senzor.patch.original');
+
+type WrappedFunction = Function & {
+  [PATCHES]?: Set<string>;
+  [ORIGINAL]?: Function;
+};
+
+export const patchMethod = (
+  target: any,
+  methodName: string,
+  patchKey: string,
+  wrapper: (original: Function) => Function
+): boolean => {
+  if (!target) return false;
+
+  const current = target[methodName] as WrappedFunction | undefined;
+  if (typeof current !== 'function') return false;
+
+  const existingPatches = current[PATCHES];
+  if (existingPatches?.has(patchKey)) return false;
+
+  const original = current[ORIGINAL] || current;
+  const wrapped = wrapper(current) as WrappedFunction;
+  const patches = new Set(existingPatches || []);
+  patches.add(patchKey);
+
+  try {
+    Object.defineProperty(wrapped, PATCHES, {
+      value: patches,
+      enumerable: false
+    });
+    Object.defineProperty(wrapped, ORIGINAL, {
+      value: original,
+      enumerable: false
+    });
+  } catch {
+    return false;
+  }
+
+  try {
+    target[methodName] = wrapped;
+    return true;
+  } catch {
+    return false;
+  }
+};
+
+export const isPatched = (
+  target: any,
+  methodName: string,
+  patchKey: string
+): boolean => {
+  const current = target?.[methodName] as WrappedFunction | undefined;
+  return Boolean(current?.[PATCHES]?.has(patchKey));
+};