From 349b499f8572906fe5781b34f6ac3d017ea1fb32 Mon Sep 17 00:00:00 2001
From: Brian Gough <bjg@network-theory.co.uk>
Date: Tue, 28 Jul 2015 16:53:55 +0100
Subject: [PATCH] add compatibility with v2 mongo driver

---
 libraries/metrics/mongodb.coffee | 76 ++++++++++++++++++++++++++++----
 1 file changed, 67 insertions(+), 9 deletions(-)

diff --git a/libraries/metrics/mongodb.coffee b/libraries/metrics/mongodb.coffee
index 045382d367..4358985101 100644
--- a/libraries/metrics/mongodb.coffee
+++ b/libraries/metrics/mongodb.coffee
@@ -1,15 +1,25 @@
-_ = require("underscore")
-
 module.exports =
 	monitor: (mongodb_require_path, logger) ->
-		Db = require("#{mongodb_require_path}/lib/mongodb/db").Db
+
+		try
+			# for the v1 driver the methods to wrap are in the mongodb
+			# module in lib/mongodb/db.js
+			mongodb = require("#{mongodb_require_path}")
+
+		try
+			# for the v2 driver the relevant methods are in the mongodb-core
+			# module in lib/topologies/{server,replset,mongos}.js
+			v2_path = mongodb_require_path.replace(/\/mongodb$/, '/mongodb-core')
+			mongodbCore = require(v2_path)
 
 		Metrics = require("./metrics")
 
 		monitorMethod = (base, method, type) ->
-			_method = base[method]
+			return unless base?
+			return unless (_method = base[method])?
+			arglen = _method.length
 
-			base[method] = (db_command, options, callback) ->
+			mongo_driver_v1_wrapper = (db_command, options, callback) ->
 				if (typeof callback == 'undefined')
 					callback = options
 					options = {}
@@ -38,7 +48,55 @@ module.exports =
 						"mongo request"
 					callback.apply this, arguments
 
-		monitorMethod(Db.prototype, "_executeQueryCommand",  "query")
-		monitorMethod(Db.prototype, "_executeRemoveCommand", "remove")
-		monitorMethod(Db.prototype, "_executeInsertCommand", "insert")
-		monitorMethod(Db.prototype, "_executeUpdateCommand", "update")
\ No newline at end of file
+			mongo_driver_v2_wrapper = (ns, ops, options, callback) ->
+				if (typeof callback == 'undefined')
+					callback = options
+					options = {}
+					callback = () ->
+
+				if ns.match(/\$cmd$/)
+					# Ignore noisy command methods like authenticating, ismaster and ping
+					return _method.call this, ns, ops, options, callback
+
+				key = "mongo-requests.#{ns}.#{type}"
+				if ops[0].q?  # ops[0].q
+					query = Object.keys(ops[0].q).sort().join("_")
+					key += "." + query
+
+				timer = new Metrics.Timer(key)
+				start = new Date()
+				_method.call this, ns, ops, options, () ->
+					timer.done()
+					time = new Date() - start
+					logger.log
+						query: ops[0].q
+						query_type: type
+						collection: ns
+						"response-time": new Date() - start
+						"mongo request"
+					callback.apply this, arguments
+
+			if arglen == 3
+				base[method] = mongo_driver_v1_wrapper
+			else if arglen == 4
+				base[method] = mongo_driver_v2_wrapper
+
+		monitorMethod(mongodb?.Db.prototype, "_executeQueryCommand",  "query")
+		monitorMethod(mongodb?.Db.prototype, "_executeRemoveCommand", "remove")
+		monitorMethod(mongodb?.Db.prototype, "_executeInsertCommand", "insert")
+		monitorMethod(mongodb?.Db.prototype, "_executeUpdateCommand", "update")
+
+		monitorMethod(mongodbCore?.Server.prototype, "command", "command")
+		monitorMethod(mongodbCore?.Server.prototype, "remove", "remove")
+		monitorMethod(mongodbCore?.Server.prototype, "insert", "insert")
+		monitorMethod(mongodbCore?.Server.prototype, "update", "update")
+
+		monitorMethod(mongodbCore?.ReplSet.prototype, "command", "command")
+		monitorMethod(mongodbCore?.ReplSet.prototype, "remove", "remove")
+		monitorMethod(mongodbCore?.ReplSet.prototype, "insert", "insert")
+		monitorMethod(mongodbCore?.ReplSet.prototype, "update", "update")
+
+		monitorMethod(mongodbCore?.Mongos.prototype, "command", "command")
+		monitorMethod(mongodbCore?.Mongos.prototype, "remove", "remove")
+		monitorMethod(mongodbCore?.Mongos.prototype, "insert", "insert")
+		monitorMethod(mongodbCore?.Mongos.prototype, "update", "update")