Add: CounterCache feature to hasMany models

lib/associations/has-many.js

@@ -5,7 +5,8 @@ var Utils = require('./../utils')
   , _ = require('lodash')
   , Association = require('./base')
   , Transaction = require('../transaction')
-  , Model = require('../model');
+  , Model = require('../model')
+  , CounterCache = require('../plugins/counter-cache');
 
 var HasManySingleLinked = require('./has-many-single-linked')
   , HasManyDoubleLinked = require('./has-many-double-linked');
@@ -177,6 +178,10 @@ module.exports = (function() {
       hasSingle: 'has' + singular,
       hasAll: 'has' + plural
     };
+
+    if (this.options.counterCache) {
+      new CounterCache(this, this.options.counterCache !== true ? this.options.counterCache : {});
+    }
   };
 
   // the id is in the target table

lib/plugins/counter-cache.js

+'use strict';
+
+var Utils = require('./../utils')
+  , Helpers = require('../associations/helpers')
+  , Transaction = require('../transaction')
+  , util = require('util')
+  , DataTypes = require('../data-types')
+  , Promise = require('bluebird');
+
+module.exports = (function() {
+  var CounterCache = function(association, options) {
+    this.association = association;
+    this.source = association.source;
+    this.target = association.target;
+    this.options = options || {};
+
+    this.sequelize = this.source.daoFactoryManager.sequelize;
+    this.as = this.options.as;
+
+    if (association.associationType !== 'HasMany') {
+      throw new Error('Can only have CounterCache on HasMany association');
+    }
+
+    if (this.as) {
+      this.isAliased = true;
+      this.columnName = this.as;
+    } else {
+      this.as = 'count_' + this.target.options.name.plural;
+      this.columnName = Utils._.camelizeIf(
+        this.as,
+        !this.source.options.underscored
+      );
+    }
+
+    this.injectAttributes();
+    this.injectHooks();
+  };
+
+  // Add countAssociation attribute to source model
+  CounterCache.prototype.injectAttributes = function() {
+    // Do not try to use a column that's already taken
+    Helpers.checkNamingCollision(this);
+
+    var newAttributes = {};
+
+    newAttributes[this.columnName] = {
+      type: DataTypes.INTEGER,
+      allowNull: false
+    };
+
+    // apparently you can't set defaultValues after model definition
+    this.source._defaultValues[this.columnName] = Utils._.partial(
+      Utils.toDefaultValue,
+      0
+    );
+    this.source._hasDefaultValues = true;
+
+    Utils.mergeDefaults(this.source.rawAttributes, newAttributes);
+
+    // Sync attributes and setters/getters to DAO prototype
+    this.source.refreshAttributes();
+  };
+
+  // Add setAssociaton method to the prototype of the model instance
+  CounterCache.prototype.injectHooks = function() {
+    var association = this.association,
+      counterCacheInstance = this,
+      CounterUtil,
+      fullUpdateHook,
+      atomicHooks,
+      previousTargetId;
+
+    CounterUtil = {
+      update: function (targetId) {
+        var query = CounterUtil._targetQuery(targetId);
+
+        return association.target.count({ where: query }).then(function (count) {
+          var newValues = {};
+
+          query = CounterUtil._sourceQuery(targetId);
+
+          newValues[counterCacheInstance.columnName] = count;
+
+          return association.source.update(newValues, { where: query });
+        });
+      },
+      increment: function (targetId) {
+        var query = CounterUtil._sourceQuery(targetId);
+
+        return association.source.find({ where: query }).then(function (instance) {
+          return instance.increment(counterCacheInstance.columnName, { by: 1 });
+        });
+      },
+      decrement: function (targetId) {
+        var query = CounterUtil._sourceQuery(targetId);
+
+        return association.source.find({ where: query }).then(function (instance) {
+          return instance.decrement(counterCacheInstance.columnName, { by: 1 });
+        });
+      },
+      // helpers
+      _targetQuery: function (id) {
+        var query = {};
+
+        query[association.identifier] = id;
+
+        return query;
+      },
+      _sourceQuery: function (id) {
+        var query = {};
+
+        query[association.source.primaryKeyAttribute] = id;
+
+        return query;
+      }
+    };
+
+    fullUpdateHook = function (target) {
+      var targetId = target.get(association.identifier)
+        , promises = [];
+
+      if (targetId) {
+        promises.push(CounterUtil.update(targetId));
+      }
+
+      if (previousTargetId && previousTargetId !== targetId) {
+        promises.push(CounterUtil.update(previousTargetId));
+      }
+
+      return Promise.all(promises).return(undefined);
+    };
+
+    atomicHooks = {
+      create: function (target) {
+        var targetId = target.get(association.identifier);
+
+        if (targetId) {
+          return CounterUtil.increment(targetId);
+        }
+      },
+      update: function (target) {
+        var targetId = target.get(association.identifier)
+          , promises = [];
+
+        if (targetId && !previousTargetId) {
+          promises.push(CounterUtil.increment(targetId));
+        }
+        if (!targetId && previousTargetId) {
+          promises.push(CounterUtil.decrement(targetId));
+        }
+        if (previousTargetId && targetId && previousTargetId !== targetId) {
+          promises.push(CounterUtil.increment(targetId));
+          promises.push(CounterUtil.decrement(previousTargetId));
+        }
+
+        return Promise.all(promises);
+      },
+      destroy: function (target) {
+        var targetId = target.get(association.identifier);
+
+        if (targetId) {
+          return CounterUtil.decrement(targetId);
+        }
+      }
+    };
+
+    // previousDataValues are cleared before afterUpdate, so we need to save this here
+    association.target.addHook('beforeUpdate', function (target) {
+      previousTargetId = target.previous(association.identifier);
+    });
+
+    if (this.options.atomic === false) {
+      association.target.addHook('afterCreate', fullUpdateHook);
+      association.target.addHook('afterUpdate', fullUpdateHook);
+      association.target.addHook('afterDestroy', fullUpdateHook);
+    } else {
+      association.target.addHook('afterCreate', atomicHooks.create);
+      association.target.addHook('afterUpdate', atomicHooks.update);
+      association.target.addHook('afterDestroy', atomicHooks.destroy);
+    }
+  };
+
+  return CounterCache;
+})();

test/plugins/counter-cache.test.js

+/* jshint camelcase: false, expr: true */
+var chai      = require('chai')
+  , expect    = chai.expect
+  , Support   = require(__dirname + '/../support')
+  , DataTypes = require(__dirname + "/../../lib/data-types")
+  , Sequelize = require('../../index')
+  , Promise   = Sequelize.Promise
+  , assert    = require('assert');
+
+chai.config.includeStack = true;
+
+describe(Support.getTestDialectTeaser("CounterCache"), function() {
+  it('adds an integer column', function() {
+    var User  = this.sequelize.define('User', {})
+      , Group = this.sequelize.define('Group', {});
+
+    User.hasMany(Group, { counterCache: true });
+
+    expect(Object.keys(User.attributes)).to.contain('countGroups');
+    expect(User.attributes.countGroups.type).to.equal(DataTypes.INTEGER);
+  });
+
+  it('supports `as`', function() {
+    var User  = this.sequelize.define('User', {})
+      , Group = this.sequelize.define('Group', {});
+
+    User.hasMany(Group, { counterCache: { as: 'countDemGroups' } });
+
+    expect(Object.keys(User.attributes)).to.contain('countDemGroups');
+  });
+
+  it('inits at 0', function() {
+    var User  = this.sequelize.define('User', {})
+      , Group = this.sequelize.define('Group', {});
+
+    User.hasMany(Group, { counterCache: true });
+
+    return this.sequelize.sync({ force: true }).then(function () {
+      return User.create();
+    }).then(function (user) {
+      expect(user.countGroups).to.equal(0);
+    });
+  });
+
+  describe('hooks', function () {
+    var User, Group;
+
+    beforeEach(function() {
+      User  = this.sequelize.define('User', {});
+      Group = this.sequelize.define('Group', {});
+
+      User.hasMany(Group, { counterCache: true });
+
+      return this.sequelize.sync({ force: true });
+    });
+
+    it('increments', function() {
+      return User.create().then(function (user) {
+        expect(user.countGroups).to.equal(0);
+
+        return user.createGroup().return(user);
+      }).then(function (user) {
+        return User.find(user.id);
+      }).then(function (user) {
+        expect(user.countGroups).to.equal(1);
+      });
+    });
+
+    it('decrements', function() {
+      var user;
+
+      return User.create().then(function (tmpUser) {
+        user = tmpUser;
+        return user.createGroup();
+      }).then(function (group) {
+        return group.destroy();
+      }).then(function () {
+        return user.reload();
+      }).then(function () {
+        expect(user.countGroups).to.equal(0);
+      });
+    });
+
+    it('works on update', function () {
+      var user, otherUser;
+
+      return User.create().then(function (tmpUser) {
+        otherUser = tmpUser;
+
+        return User.create();
+      }).then(function (tmpUser) {
+        user = tmpUser;
+        return user.createGroup();
+      }).tap(function (group) {
+        return user.reload();
+      }).tap(function () {
+        expect(user.countGroups).to.equal(1);
+      }).then(function (group) {
+        group.UserId = otherUser.id;
+        return group.save();
+      }).then(function () {
+        return Promise.all([user.reload(), otherUser.reload()]);
+      }).then(function () {
+        expect(user.countGroups).to.equal(0);
+        expect(otherUser.countGroups).to.equal(1);
+      });
+    });
+  });
+})