Merge pull request #5209 from matthewheck/add_beforecount_hook

add a beforeCount hook

lib/hooks.js

@@ -59,6 +59,7 @@ var hookTypes = {
   beforeFindAfterExpandIncludeAll: {params: 1},
   beforeFindAfterOptions: {params: 1},
   afterFind: {params: 2},
+  beforeCount: {params: 1},
   beforeDefine: {params: 2, sync: true},
   afterDefine: {params: 1, sync: true},
   beforeInit: {params: 2, sync: true},
@@ -376,6 +377,13 @@ Hooks.hasHooks = Hooks.hasHook;
  */
 
 /**
+ * A hook that is run before a count query
+ * @param {String}   name
+ * @param {Function} fn   A callback function that is called with options
+ * @name beforeCount
+ */
+
+/**
  * A hook that is run before a define call
  * @param {String}   name
  * @param {Function} fn   A callback function that is called with attributes, options

lib/model.js

@@ -1578,28 +1578,37 @@ Model.prototype.aggregate = function(attribute, aggregateFunction, options) {
  */
 Model.prototype.count = function(options) {
   options = Utils._.clone(options || {});
-  conformOptions(options, this);
-  this.$injectScope(options);
+
+  _.defaults(options, { hooks: true });
 
   var col = '*';
 
-  if (options.include) {
-    col = this.name + '.' + this.primaryKeyField;
-    expandIncludeAll.call(this, options);
-    validateIncludedElements.call(this, options);
-  }
+  return Promise.bind(this).then(function() {
+    conformOptions(options, this);
+    this.$injectScope(options);
 
-  Utils.mapOptionFieldNames(options, this);
+    if (options.hooks) {
+      return this.runHooks('beforeCount', options);
+    }
+  }).then(function() {
+    if (options.include) {
+      col = this.name + '.' + this.primaryKeyField;
+      expandIncludeAll.call(this, options);
+      validateIncludedElements.call(this, options);
+    }
 
-  options.plain = options.group ? false : true;
-  options.dataType = new DataTypes.INTEGER();
-  options.includeIgnoreAttributes = false;
-  options.limit = null;
-  options.offset = null;
-  options.order = null;
-  options.attributes = [];
+    Utils.mapOptionFieldNames(options, this);
 
-  return this.aggregate(col, 'count', options);
+    options.plain = options.group ? false : true;
+    options.dataType = new DataTypes.INTEGER();
+    options.includeIgnoreAttributes = false;
+    options.limit = null;
+    options.offset = null;
+    options.order = null;
+    options.attributes = [];
+
+    return this.aggregate(col, 'count', options);
+  });
 };
 
 /**

test/integration/hooks.test.js

@@ -978,6 +978,49 @@ describe(Support.getTestDialectTeaser('Hooks'), function() {
     });
   });
 
+  describe('#count', function() {
+    beforeEach(function() {
+      return this.User.bulkCreate([
+        {username: 'adam', mood: 'happy'},
+        {username: 'joe', mood: 'sad'},
+        {username: 'joe', mood: 'happy'}
+      ]);
+    });
+
+    describe('on success', function() {
+      it('hook runs', function() {
+        var beforeHook = false;
+
+        this.User.beforeCount(function() {
+          beforeHook = true;
+        });
+
+        return this.User.count().then(function(count) {
+          expect(count).to.equal(3);
+          expect(beforeHook).to.be.true;
+        });
+      });
+
+      it('beforeCount hook can change options', function() {
+        this.User.beforeCount(function(options) {
+          options.where.username = 'adam';
+        });
+
+        return expect(this.User.count({where: {username: 'joe'}})).to.eventually.equal(1);
+      });
+    });
+
+    describe('on error', function() {
+      it('in beforeCount hook returns error', function() {
+        this.User.beforeCount(function() {
+          throw new Error('Oops!');
+        });
+
+        return expect(this.User.count({where: {username: 'adam'}})).to.be.rejectedWith('Oops!');
+      });
+    });
+  });
+
   describe('#define', function() {
     before(function() {
       this.sequelize.addHook('beforeDefine', function(attributes, options) {