Fixed Promise-CLS shim for 'each' and other methods.

lib/promise.js

 'use strict';
 'use strict';
 
 
 var Promise = require('bluebird')
 var Promise = require('bluebird')
-  , _then = Promise.prototype._then;
-
-Promise.prototype._then = function (didFulfill, didReject, didProgress, receiver, internalData) {
-  if (Promise.Sequelize.cls) {
-    var ns = Promise.Sequelize.cls;
-    if (typeof didFulfill === 'function') didFulfill = ns.bind(didFulfill);
-    if (typeof didReject === 'function') didReject = ns.bind(didReject);
-    if (typeof didProgress === 'function') didProgress = ns.bind(didProgress);
-  }
-  
-  return _then.call(this, didFulfill, didReject, didProgress, receiver, internalData);
-};
-
-module.exports = Promise;
+  , shimmer = require('shimmer');
+
+// functionName: The Promise function that should be shimmed
+// fnArgs: The arguments index that should be CLS enabled (typically all callbacks). Offset from last if negative
+function shimCLS(object, functionName, fnArgs){
+  shimmer.wrap(object, functionName, function(fn) {
+    return function () {
+      if (Promise.Sequelize && Promise.Sequelize.cls) {
+        var ns = Promise.Sequelize.cls;
+        for(var x=0; x<fnArgs.length; x++) {
+          var argIndex = x < 0 ? arguments.length + x : x;
+          if ( argIndex < arguments.length && typeof arguments[argIndex] === 'function' ) arguments[argIndex] = ns.bind( arguments[argIndex] );
+        }
+      }
\ No newline at end of file
+
+      return fn.apply(this, arguments);
+    };
+  });
+}
+
+// Core
+shimCLS(Promise, 'join', [-1]);
+shimCLS(Promise, 'try', [0]);
+shimCLS(Promise, 'method', [0]);
+shimCLS(Promise.prototype, 'then', [0, 1, 2]);
+shimCLS(Promise.prototype, 'spread', [0, 1]);
+shimCLS(Promise.prototype, 'catch', [-1]);
+shimCLS(Promise.prototype, 'error', [0]);
+shimCLS(Promise.prototype, 'finally', [0]);
+
+// Collections
+shimCLS(Promise.prototype, 'map', [0]);
+shimCLS(Promise.prototype, 'reduce', [0]);
+shimCLS(Promise.prototype, 'filter', [0]);
+shimCLS(Promise.prototype, 'each', [0]);
+
+// Utility
+shimCLS(Promise.prototype, 'tap', [0]);
+
+
+module.exports = Promise;

package.json

@@ -40,7 +40,8 @@
     "moment": "^2.9.0",
     "moment": "^2.9.0",
     "node-uuid": "~1.4.1",
     "node-uuid": "~1.4.1",
     "toposort-class": "~0.3.0",
     "toposort-class": "~0.3.0",
-    "validator": "^3.34.0"
+    "validator": "^3.34.0",
+    "shimmer": "1.0.0"
   },
   },
   "devDependencies": {
   "devDependencies": {
     "chai": "^2.1.2",
     "chai": "^2.1.2",