instrumentation: fixed issue with pg "native" getter instrumentation

Author: Martin Kuba

lib/instrumentation/pg.js

@@ -36,6 +36,8 @@ function initializeSegment(tracer, segment, config) {
 }
 
 module.exports = function initialize(agent, pgsql) {
+  if (!pgsql) return
+
   var tracer = agent.tracer
 
   // allows for native wrapping to not happen if not necessary
@@ -45,15 +47,20 @@ module.exports = function initialize(agent, pgsql) {
     return instrumentPGNative('pg', pgsql)
   }
 
-  // using ('pg').native in their require
+  // The pg module defines "native" getter which sets up the native client lazily
+  // (only when called).  We replace the getter, so that we can instrument the native
+  // client.  The original getter replaces itself with the instance of the native
+  // client, so only instrument if the getter exists (otherwise assume already
+  // instrumented).
   var origGetter = pgsql.__lookupGetter__('native')
-  delete pgsql.native
-  pgsql.__defineGetter__('native', function getNative() {
-    var temp = origGetter()
-    instrumentPGNative('pg.native', pgsql.native)
-    return temp
-  })
-
+  if (origGetter) {
+    delete pgsql.native
+    pgsql.__defineGetter__('native', function getNative() {
+      var temp = origGetter()
+      instrumentPGNative('pg.native', temp)
+      return temp
+    })
+  }
 
   // wrapping for native
   function instrumentPGNative(eng, pg) {

test/unit/instrumentation/postgresql.test.js

@@ -1,6 +1,112 @@
 'use strict'
 
+var path   = require('path')
+  , chai   = require('chai')
+  , expect = chai.expect
+  , helper = require('../../lib/agent_helper')
+
 describe("agent instrumentation of PostgreSQL", function () {
+  var agent
+    , initialize
+
+  before(function () {
+    agent = helper.loadMockedAgent()
+    initialize = require('../../../lib/instrumentation/pg')
+  })
+
+  after(function () {
+    helper.unloadAgent(agent)
+  })
+
+  describe("shouldn't cause bootstrapping to fail", function () {
+    it("when passed no module", function () {
+      expect(function () { initialize(agent); }).not.throws()
+    })
+
+    it("when passed an empty module", function () {
+      expect(function () { initialize(agent, {}); }).not.throws()
+    })
+  })
+
+  describe("lazy loading of native PG client", function() {
+
+    function getMockModule() {
+      function PG(clientConstructor) {
+        this.Client = clientConstructor
+      }
+
+      function DefaultClient() {}
+      DefaultClient.prototype.query = function() {}
+      function NativeClient() {}
+      NativeClient.prototype.query = function() {}
+
+      var mockPg = new PG(DefaultClient)
+      mockPg.__defineGetter__("native", function() {
+        delete mockPg.native;
+        mockPg.native = new PG(NativeClient);
+        return mockPg.native;
+      })
+      return mockPg
+    }
+
+    it("instruments when native getter is called", function() {
+      var mockPg = getMockModule()
+
+      initialize(agent, mockPg)
+
+      var pg = mockPg.native
+      expect(pg.Client['__NR_original'].name).equal('NativeClient')
+
+      var pg = mockPg
+      expect(pg.Client.name).equal('DefaultClient')
+    })
+
+    it("does not fail when getter is called multiple times", function() {
+      var mockPg = getMockModule()
+
+      initialize(agent, mockPg)
+      var pg1 = mockPg.native
+
+      initialize(agent, mockPg)
+      var pg2 = mockPg.native
+
+      expect(pg1).equal(pg2)
+    })
+
+    it("does not interfere with non-native instrumentation", function() {
+      var mockPg = getMockModule()
+
+      initialize(agent, mockPg)
+      var nativeClient = mockPg.native
+      expect(nativeClient.Client['__NR_original'].name).equal('NativeClient')
+      var defaultClient = mockPg
+      expect(defaultClient.Client.name).equal('DefaultClient')
+
+      initialize(agent, mockPg)
+      var nativeClient = mockPg.native
+      expect(nativeClient.Client['__NR_original'].name).equal('NativeClient')
+      var defaultClient = mockPg
+      expect(defaultClient.Client.name).equal('DefaultClient')
+    })
+
+    it("when pg modules is refreshed in cache", function() {
+      var mockPg = getMockModule()
+
+      // instrument once
+      initialize(agent, mockPg)
+      var pg1 = mockPg.native
+      expect(pg1.Client['__NR_original'].name).equal('NativeClient')
+
+      // simulate deleting from module cache
+      mockPg = getMockModule()
+      initialize(agent, mockPg)
+      var pg2 = mockPg.native
+      expect(pg2.Client['__NR_original'].name).equal('NativeClient')
+
+      expect(pg1).not.equal(pg2)
+    })
+  })
+
   describe("for each operation", function () {
     it("should update the global database aggregate statistics")
     it("should also update the global web aggregate statistics")