fix: support tracing awaited mongoose queries (#1007)

package.json

@@ -60,6 +60,7 @@
     "@types/knex": "^0.15.1",
     "@types/methods": "^1.1.0",
     "@types/mocha": "^5.2.5",
+    "@types/mongoose": "^5.3.26",
     "@types/ncp": "^2.0.1",
     "@types/nock": "^10.0.0",
     "@types/node": "~10.7.2",
@@ -81,6 +82,7 @@
     "intelli-espower-loader": "^1.0.1",
     "js-green-licenses": "^0.5.0",
     "jshint": "^2.9.1",
+    "linkinator": "^1.1.2",
     "mocha": "^6.0.0",
     "ncp": "^2.0.0",
     "nock": "^10.0.0",
@@ -95,8 +97,7 @@
     "timekeeper": "^2.0.0",
     "tmp": "0.1.0",
     "ts-node": "^8.0.0",
-    "typescript": "~3.4.0",
-    "linkinator": "^1.1.2"
+    "typescript": "~3.4.0"
   },
   "dependencies": {
     "@google-cloud/common": "^0.32.1",

src/config.ts

@@ -265,6 +265,7 @@ export const defaultConfig = {
     'http2': path.join(pluginDirectory, 'plugin-http2.js'),
     'koa': path.join(pluginDirectory, 'plugin-koa.js'),
     'mongodb-core': path.join(pluginDirectory, 'plugin-mongodb-core.js'),
+    'mongoose': path.join(pluginDirectory, 'plugin-mongoose.js'),
     'mysql': path.join(pluginDirectory, 'plugin-mysql.js'),
     'mysql2': path.join(pluginDirectory, 'plugin-mysql2.js'),
     'pg': path.join(pluginDirectory, 'plugin-pg.js'),

src/plugins/plugin-mongoose.ts

@@ -0,0 +1,41 @@
+/**
+ * Copyright 2019 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as mongooseTypes from 'mongoose';
+import { PluginTypes } from '..';
+
+const plugin: PluginTypes.Plugin = [
+  {
+    versions: '4 - 5',
+    file: 'lib/query.js',
+    intercept: (Query: typeof mongooseTypes.Query, api) => {
+      // Assume that the context desired at Query execution time should be the
+      // context where the Query object was constructed. In most (if not all)
+      // Mongoose read APIs, both of these appear to happen as part of the same
+      // API call.
+      return new Proxy(Query, {
+        apply(target, thisArg, args) {
+          // result is expected to be undefined.
+          const result = target.apply(thisArg, args);
+          thisArg.exec = api.wrap(thisArg.exec);
+          return result;
+        }
+      });
+    }
+  }
+];
+
+export = plugin;

test/plugins/test-trace-mongoose-async-await.ts

@@ -0,0 +1,87 @@
+/**
+ * Copyright 2019 Google Inc. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import * as assert from 'assert';
+import * as mongooseTypes from 'mongoose';
+
+import * as traceTestModule from '../trace';
+import {describeInterop} from '../utils';
+
+describeInterop<typeof mongooseTypes>('mongoose', (fixture) => {
+  let mongoose: typeof mongooseTypes;
+  // Simple will be treated as a class constructor.
+  // tslint:disable-next-line:variable-name
+  let Simple: mongooseTypes.Model<mongooseTypes.Document>;
+
+  /**
+   * Common logic used in multiple tests -- inserts an object into the database.
+   * @param doc
+   */
+  async function insertTestData(doc: {f1: string, f2: boolean, f3: number}) {
+    const data = new Simple(doc);
+    const tracer = traceTestModule.get();
+    await tracer.runInRootSpan({name: 'insert-test-data'}, async (span) => {
+      assert.ok(tracer.isRealSpan(span));
+      await data.save();
+      span.endSpan();
+    });
+  }
+
+  before(async () => {
+    traceTestModule.setCLSForTest();
+    traceTestModule.setPluginLoaderForTest();
+    traceTestModule.start();
+    mongoose = fixture.require();
+    await mongoose.connect('mongodb://localhost:27017/testdb');
+
+    const {Schema} = mongoose;
+    const simpleSchema = new Schema({f1: String, f2: Boolean, f3: Number});
+    Simple = mongoose.model('Simple', simpleSchema);
+  });
+
+  after(async () => {
+    traceTestModule.setCLSForTest(traceTestModule.TestCLS);
+    traceTestModule.setPluginLoaderForTest(traceTestModule.TestPluginLoader);
+    await mongoose.connection.db.dropDatabase();
+    await mongoose.disconnect();
+  });
+
+  afterEach(() => {
+    traceTestModule.clearTraceData();
+  });
+
+  it('Traces creates with async/await', async () => {
+    await insertTestData({f1: 'val', f2: false, f3: 1729});
+    const trace = traceTestModule.getOneTrace(
+        trace => trace.spans.some(span => span.name === 'insert-test-data'));
+    assert.strictEqual(trace.spans.length, 2);
+    assert.strictEqual(trace.spans[1].name, 'mongo-insert');
+  });
+
+  it('Traces queries with async/await', async () => {
+    await insertTestData({f1: 'sim', f2: false, f3: 1729});
+    const tracer = traceTestModule.get();
+    await tracer.runInRootSpan({name: 'query-test-data'}, async (span) => {
+      assert.ok(tracer.isRealSpan(span));
+      await Simple.findOne({f1: 'sim'});
+      span.endSpan();
+    });
+    const trace = traceTestModule.getOneTrace(
+        trace => trace.spans.some(span => span.name === 'query-test-data'));
+    assert.strictEqual(trace.spans.length, 2);
+    assert.strictEqual(trace.spans[1].name, 'mongo-cursor');
+  });
+});

tsconfig.json

@@ -25,6 +25,7 @@
     "test/plugins/test-trace-http.ts",
     "test/plugins/test-trace-http2.ts",
     "test/plugins/test-trace-knex.ts",
+    "test/plugins/test-trace-mongoose-async-await.ts",
     "test/logger.ts",
     "test/nocks.ts",
     "test/test-cls.ts",