Update to work in Rails 4