Repository: honeycombio/beeline-nodejs
Pull request #155 (2019-07-24 12:55:59)

Add instrumentation for `pg` package
Addresses issue #37.

Our team uses the [pg package](https://www.npmjs.com/package/pg) to manage our Postgres interactions and were looking to extend our tracing to include those queries to the database. Didn't see an easy way to introduce a custom instrumentation from our own codebase so looking to upstream and share.

A lot of this is based off the pre-existing mysql2 (thanks to @toshok and @ashishbista for their work on that one!) and also a little inspiration from [opentracing-auto's pg instrumentation](https://github.com/RisingStack/opentracing-auto/blob/master/src/instrumentation/pg.js).

Main difference to the mysql2 instrumentation is that this enhances the span with error/result information, similarly to [how beeline-python does it.](https://github.com/honeycombio/beeline-python/blob/master/beeline/middleware/django/__init__.py#L17)

[Here is a minimum project to test the instrumentation out with](https://github.com/allyjweir/pg-trace-example). Check the README.md for setup details.


Patch:
diff --git a/CONTRIBUTORS b/CONTRIBUTORS
index 3e6f095..57b94c4 100644
--- a/CONTRIBUTORS
+++ b/CONTRIBUTORS
@@ -5,3 +5,4 @@ Christine Yen
 Erwin van der Koogh
 Ally Weir
 Ashish Bista
+Mordy Tikotzky
diff --git a/lib/instrumentation.js b/lib/instrumentation.js
index df28b8c..bd25459 100644
--- a/lib/instrumentation.js
+++ b/lib/instrumentation.js
@@ -18,6 +18,7 @@ const instrumentations = [
   "mongoose",
   "mpromise",
   "mysql2",
+  "pg",
   "react-dom/server",
   "sequelize",
 ];
@@ -100,7 +101,6 @@ const instrumentLoad = (exports.instrumentLoad = (mod, loadRequest, parent, opts
   }
 
   let packageVersion = getPackageVersion(loadRequest, parent);
-
   debug(`loading instrumentation for ${loadRequest}@${packageVersion}`);
   let instrumentation = require(instrumentationPath(loadRequest));
   let new_mod;
diff --git a/lib/instrumentation/express.js b/lib/instrumentation/express.js
index 0ee43d7..88d5582 100644
--- a/lib/instrumentation/express.js
+++ b/lib/instrumentation/express.js
@@ -87,12 +87,10 @@ const getMagicMiddleware = ({ userContext, traceIdSource, packageVersion }) => (
       [schema.TRACE_SPAN_NAME]: "request",
       [schema.TRACE_ID_SOURCE]: traceContext.source,
       "request.host": req.hostname,
-      "request.base_url": req.baseUrl,
       "request.original_url": req.originalUrl,
       "request.remote_addr": req.ip,
       "request.secure": req.secure,
       "request.method": req.method,
-      "request.route": req.route ? req.route.path : undefined,
       "request.scheme": req.protocol,
       "request.path": req.path,
       "request.query": req.query,
@@ -129,8 +127,14 @@ const getMagicMiddleware = ({ userContext, traceIdSource, packageVersion }) => (
     }
 
     api.addContext({
+      "request.base_url": req.baseUrl,
       "response.status_code": String(response.statusCode),
     });
+    if (req.route) {
+      api.addContext({
+        "request.route": req.route.path,
+      });
+    }
     if (req.params) {
       Object.keys(req.params).forEach(param =>
         api.addContext({
diff --git a/lib/instrumentation/pg.js b/lib/instrumentation/pg.js
new file mode 100644
index 0000000..b3f1fe3
--- /dev/null
+++ b/lib/instrumentation/pg.js
@@ -0,0 +1,76 @@
+/* eslint-env node */
+const shimmer = require("shimmer"),
+  api = require("../api"),
+  schema = require("../schema");
+
+const getQueryString = function(args) {
+  if (typeof args[0] === "string") return args[0];
+  if (typeof args[0] === "object") return args[0].text;
+  return undefined;
+};
+
+const getQueryArgs = function(args) {
+  if (typeof args[0] === "object") return args[0].values;
+  if (Array.isArray(args[1])) return args[1];
+  return undefined;
+};
+
+const getQueryName = function(args) {
+  if (typeof args[0] === "object") return args[0].name;
+  return undefined;
+};
+
+let instrumentPg = function(pg, opts = {}) {
+  shimmer.wrap(pg.Client.prototype, "query", function(query) {
+    return function(...args) {
+      if (args.length < 1) {
+        return query.apply(this, args);
+      }
+      if (!api.traceActive()) {
+        return query.apply(this, args);
+      }
+
+      api.startAsyncSpan(
+        {
+          [schema.EVENT_TYPE]: "pg",
+          [schema.PACKAGE_VERSION]: opts.packageVersion,
+          [schema.TRACE_SPAN_NAME]: "query",
+          "db.query": getQueryString(args),
+          "db.query_args": getQueryArgs(args),
+          "db.query_name": getQueryName(args),
+        },
+        span => {
+          let cb = args[args.length - 1];
+          // XXX(toshok) this bindFunction shouldn't be necessary, but we aren't
+          // finishing the event properly without it.
+          let wrapped_cb = api.bindFunctionToTrace(function(...cb_args) {
+            let error = cb_args[0];
+            let result = cb_args[1];
+
+            if (error !== null && error instanceof Error) {
+              api.addContext({
+                "db.error": error.message,
+                "db.error_stack": error.stack,
+                "db.error_hint": error.hint,
+              });
+            }
+
+            if (result) {
+              api.addContext({
+                "db.rows_affected": result.rowCount,
+              });
+            }
+
+            api.finishSpan(span, "query");
+            return cb(...cb_args);
+          });
+
+          return query.apply(this, args.slice(0, -1).concat(wrapped_cb));
+        }
+      );
+    };
+  });
+  return pg;
+};
+
+module.exports = instrumentPg;
diff --git a/package-lock.json b/package-lock.json
index 33ee7cd..1a14b49 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -5138,9 +5138,9 @@
       }
     },
     "lodash": {
-      "version": "4.17.11",
-      "resolved": "https://registry.npmjs.org/lodash/-/lodash-4.17.11.tgz",
-      "integrity": "sha512-cQKh8igo5QUhZ7lg38DYWAxMvjSAKG0A8wGSVimP07SIUEK2UO+arSRKbRZWtelMtN5V0Hkwh5ryOto/SshYIg==",
+      "version": "4.17.15",
+      "resolved": "https://registry.npmjs.org/lodash/-/lodash-4.17.15.tgz",
+      "integrity": "sha512-8xOcRHvCjnocdS5cpwXQXVzmmh5e5+saE2QGoeQmbKmRS6J3VQppPOIt0MnmE+4xlZoumy0GPG0D0MVIQbNA1A==",
       "dev": true
     },
     "lodash.sortby": {


Test patch:
diff --git a/lib/instrumentation/express.test.js b/lib/instrumentation/express.test.js
index d75ee16..7cd9304 100644
--- a/lib/instrumentation/express.test.js
+++ b/lib/instrumentation/express.test.js
@@ -53,7 +53,7 @@ describe("userContext", () => {
             [schema.BEELINE_VERSION]: pkg.version,
             "request.host": "127.0.0.1",
             "request.base_url": "",
-            "request.route": undefined,
+            "request.route": "/",
             "request.original_url": "/",
             "request.remote_addr": "::ffff:127.0.0.1",
             "request.secure": false,
@@ -133,6 +133,62 @@ describe("userContext as function", () => {
   });
 });
 
+describe("req proprties should only be read after a request matches a route", () => {
+  const express = instrumentExpress(require("express"), {});
+
+  let server;
+  function initializeTestServer() {
+    let app = express();
+    let router = express.Router();
+    app.get("/:name", function(req, res) {
+      res.status(200).send("ok");
+    });
+
+    router.get("/router/:name", function(req, res) {
+      res.status(200).send("ok");
+    });
+
+    app.use("/base", router);
+
+    server = app.listen(3000);
+  }
+
+  beforeEach(() => {
+    api.configure({ impl: "mock" });
+    initializeTestServer();
+  });
+  afterEach(() => {
+    api._resetForTesting();
+    server.close();
+  });
+
+  test("it correctly sends the matched path", done => {
+    request(server)
+      .get("/mordy")
+      .expect(200, () => {
+        setTimeout(() => {
+          expect(api._apiForTesting().sentEvents.length).toBe(1);
+          let ev = api._apiForTesting().sentEvents[0];
+          expect(ev["request.route"]).toBe("/:name");
+          done();
+        }, 1000);
+      });
+  });
+
+  test("it correctly sends the matched base_url", done => {
+    request(server)
+      .get("/base/router/mordy")
+      .expect(200, () => {
+        setTimeout(() => {
+          expect(api._apiForTesting().sentEvents.length).toBe(1);
+          let ev = api._apiForTesting().sentEvents[0];
+          expect(ev["request.base_url"]).toBe("/base");
+          done();
+        }, 1000);
+      });
+  });
+});
+
 describe("request id from http headers", () => {
   const express = instrumentExpress(require("express"), {});