Support activerecord associations preloader

- Add new `:preload` exposure option
- Add `Grape::Entity.preload_and_represent` helper method
- Add `Grape::Entity::Preloader` class for preload activerecord associations

Author: OuYangJinTing

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 #### Features
 
+* Add activerecord associations preloader to avoid N+1 queries for : new `:preload` exposure option and `Grape::Entity.preload_and_represent` helper. Requires ActiveRecord >= 7.0; otherwise a warning is emitted and no preload is performed.
 * Your contribution here.
 
 #### Fixes

Gemfile

@@ -17,6 +17,7 @@ group :development, :test do
 end
 
 group :test do
+  gem 'activerecord'
   gem 'coveralls_reborn', require: false
   gem 'growl'
   gem 'guard'
@@ -25,4 +26,5 @@ group :test do
   gem 'rb-fsevent'
   gem 'ruby-grape-danger', '~> 0.2', require: false
   gem 'simplecov', require: false
+  gem 'sqlite3'
 end

README.md

@@ -25,6 +25,7 @@
       - [Format Before Exposing](#format-before-exposing)
       - [Expose Nil](#expose-nil)
       - [Default Value](#default-value)
+      - [ActiveRecord Associations Preloader](#activerecord-associations-preloader)
       - [Documentation](#documentation)
     - [Options Hash](#options-hash)
       - [Passing Additional Option To Nested Exposure](#passing-additional-option-to-nested-exposure)
@@ -510,6 +511,77 @@ module  Entities
 end
 ```
 
+#### ActiveRecord Associations Preloader
+
+Avoid N+1 queries by preload ActiveRecord associations. You can declare which association to preload per exposure via the `:preload` option, and perform representation with `preload_and_represent`.
+
+- Requirements: ActiveRecord >= 7.0. On lower versions, a warning is emitted and no preload is performed.
+- The `:preload` option value must be a Symbol.
+- Respects `only`/`except` options when deciding which associations to preload.
+
+Example entity definitions with `:preload`:
+
+```ruby
+class Tag < ApplicationRecord
+  include Grape::Entity::DSL
+
+  belongs_to :target, polymorphic: true
+
+  entity do
+    # ...
+  end
+end
+
+class Book < ApplicationRecord
+  include Grape::Entity::DSL
+
+  belongs_to :author, foreign_key: :author_id, class_name: 'User'
+  has_many :tags, as: :target
+
+  entity do
+    # ...
+    expose :tags, using: Tag::Entity, preload: :tags
+  end
+end
+
+class User < ApplicationRecord
+  include Grape::Entity::DSL
+
+  has_many :books, foreign_key: :author_id
+  has_many :tags, as: :target
+
+  entity do
+    # ...
+    expose :books, using: Book::Entity, preload: :books
+    expose :tags, using: Tag::Entity, preload: :tags
+  end
+end
+```
+
+Preload and represent in one call:
+
+```ruby
+# Preload all declared associations that will be exposed
+User::Entity.preload_and_represent(users)
+
+# Only preload what will be exposed (respects :only / :except)
+User::Entity.preload_and_represent(users, only: [:tags])
+User::Entity.preload_and_represent(users, except: [:tags])
+```
+
+Nested exposure preloading also works, for example:
+
+```ruby
+class UserWithNest < User::Entity
+  unexpose :books
+  expose :nesting do
+    expose :books, using: Book::Entity, preload: :books
+  end
+end
+
+UserWithNest.preload_and_represent(users)
+```
+
 #### Documentation
 
 Expose documentation with the field. Gets bubbled up when used with Grape and various API documentation systems.

lib/grape_entity.rb

@@ -10,3 +10,4 @@
 require 'grape_entity/exposure'
 require 'grape_entity/options'
 require 'grape_entity/deprecated'
+require 'grape_entity/preloader'

lib/grape_entity/entity.rb

@@ -187,11 +187,12 @@ def self.inherited(subclass)
     # @option options :merge This option allows you to merge an exposed field to the root
     #
     # rubocop:disable Layout/LineLength
-    def self.expose(*args, &block)
+    def self.expose(*args, &block) # rubocop:disable Metrics/AbcSize
       options = merge_options(args.last.is_a?(Hash) ? args.pop : {})
 
-      if args.size > 1
+      raise ArgumentError, 'The :preload option must be a Symbol.' if options.key?(:preload) && !options[:preload].is_a?(Symbol)
 
+      if args.size > 1
         raise ArgumentError, 'You may not use the :as option on multi-attribute exposures.' if options[:as]
         raise ArgumentError, 'You may not use the :expose_nil on multi-attribute exposures.' if options.key?(:expose_nil)
         raise ArgumentError, 'You may not use block-setting on multi-attribute exposures.' if block_given?
@@ -451,6 +452,14 @@ def self.represent(objects, options = {})
       root_element ? { root_element => inner } : inner
     end
 
+    # Same as the represent method, but the activerecord associations declared by the preload option will be preloaded,
+    # Therefore, it can avoid n+1 queries.
+    def self.preload_and_represent(objects, options = {})
+      options = Options.new(options) unless options.is_a?(Options)
+      Preloader.new(self, objects, options).call
+      represent(objects, options)
+    end
+
     # This method returns the entity's root or collection root node, or its parent's
     # @param root_type: either :collection_root or just :root
     def self.root_element(root_type)
@@ -618,6 +627,7 @@ def to_xml(options = {})
       expose_nil
       override
       default
+      preload
     ].to_set.freeze
 
     # Merges the given options with current block options.

lib/grape_entity/exposure/base.rb

@@ -7,7 +7,7 @@ module Grape
   class Entity
     module Exposure
       class Base
-        attr_reader :attribute, :is_safe, :documentation, :override, :conditions, :for_merge
+        attr_reader :attribute, :is_safe, :documentation, :override, :conditions, :for_merge, :preload
 
         def self.new(attribute, options, conditions, ...)
           super(attribute, options, conditions).tap { |e| e.setup(...) }
@@ -24,9 +24,14 @@ def initialize(attribute, options, conditions)
           @attr_path_proc = options[:attr_path]
           @documentation = options[:documentation]
           @override = options[:override]
+          @preload = options[:preload]
           @conditions = conditions
         end
 
+        def preload?
+          !preload.nil?
+        end
+
         def dup(&block)
           self.class.new(*dup_args, &block)
         end
@@ -116,8 +121,12 @@ def attr_path(entity, options)
           end
         end
 
+        def proc_key?
+          @key.respond_to?(:call)
+        end
+
         def key(entity = nil)
-          @key.respond_to?(:call) ? entity.exec_with_object(@options, &@key) : @key
+          proc_key? ? entity.exec_with_object(@options, &@key) : @key
         end
 
         def with_attr_path(entity, options, &block)

lib/grape_entity/preloader.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module Grape
+  class Entity
+    class Preloader
+      attr_reader :entity_class, :objects, :options
+
+      def initialize(entity_class, objects, options)
+        @entity_class = entity_class
+        @objects = Array.wrap(objects)
+        @options = options
+      end
+
+      if defined?(ActiveRecord) && ActiveRecord.respond_to?(:version) && ActiveRecord.version >= Gem::Version.new('7.0')
+        def call
+          associations = {}
+          collect_associations(entity_class.root_exposures, associations, options)
+          ActiveRecord::Associations::Preloader.new(records: objects, associations: associations).call
+        end
+      else
+        def call
+          warn 'The Preloader work normally requires activerecord(>= 7.0) gem'
+        end
+      end
+
+      private
+
+      def collect_associations(exposures, associations, options)
+        exposures.each do |exposure|
+          next unless exposure.should_return_key?(options)
+
+          new_associations = associations[exposure.preload] ||= {} if exposure.preload?
+          next if exposure.proc_key?
+
+          if exposure.is_a?(Exposure::NestingExposure)
+            collect_associations(exposure.nested_exposures, associations, subexposure_options_for(exposure, options))
+          elsif exposure.is_a?(Exposure::RepresentExposure) && new_associations
+            collect_associations(exposure.using_class.root_exposures, new_associations, subexposure_options_for(exposure, options)) # rubocop:disable Layout/LineLength
+          end
+        end
+      end
+
+      def subexposure_options_for(exposure, options)
+        options.for_nesting(exposure.instance_variable_get(:@key))
+      end
+    end
+  end
+end

spec/grape_entity/entity_spec.rb

@@ -40,6 +40,11 @@
         it 'makes sure unknown options are not silently ignored' do
           expect { subject.expose :name, unknown: nil }.to raise_error ArgumentError
         end
+
+        it 'ensures :preload option must be a Symbol' do
+          expect { subject.expose :name, preload: 'author' }.to raise_error ArgumentError
+          expect { subject.expose :name, preload: :author }.not_to raise_error
+        end
       end
 
       context 'with a :merge option' do

spec/grape_entity/preloader_spec.rb

@@ -0,0 +1,143 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Grape::Entity::Preloader do
+  class SQLCounter
+    class << self
+      attr_accessor :ignored_sql, :log, :log_all
+
+      def clear_log
+        self.log = []
+        self.log_all = []
+      end
+    end
+    clear_log
+
+    def call(_name, _start, _finish, _message_id, values)
+      return if values[:cached]
+
+      sql = values[:sql]
+      self.class.log_all << sql
+      self.class.log << sql unless %w[SCHEMA TRANSACTION].include? values[:name]
+    end
+
+    ActiveSupport::Notifications.subscribe('sql.active_record', new)
+  end
+
+  class ApplicationRecord < ActiveRecord::Base
+    self.abstract_class = true
+  end
+
+  class Tag < ApplicationRecord
+    include Grape::Entity::DSL
+
+    connection.create_table(:tags) do |t|
+      t.string :name
+      t.references :target, polymorphic: true
+    end
+
+    belongs_to :target, polymorphic: true
+
+    entity do
+      expose :name
+    end
+  end
+
+  class Book < ApplicationRecord
+    include Grape::Entity::DSL
+
+    connection.create_table(:books) do |t|
+      t.string :name
+      t.references :author
+    end
+
+    belongs_to :author, foreign_key: :author_id, class_name: 'User'
+    has_many :tags, as: :target, dependent: :destroy
+
+    entity do
+      expose :name
+      expose :tags, using: Tag::Entity, preload: :tags
+    end
+  end
+
+  class User < ApplicationRecord
+    include Grape::Entity::DSL
+
+    connection.create_table(:users) do |t|
+      t.string :name
+    end
+
+    has_many :books, foreign_key: :author_id, dependent: :destroy
+    has_many :tags, as: :target, dependent: :destroy
+
+    entity do
+      expose :name
+      expose :books, using: Book::Entity, preload: :books
+      expose :tags, using: Tag::Entity, preload: :tags
+    end
+  end
+
+  let!(:users) { [User.create(name: 'User1'), User.create(name: 'User2')] }
+  let!(:user_tags) { [Tag.create(name: 'Tag1', target: users[0]), Tag.create(name: 'Tag2', target: users[1])] }
+  let!(:books) { [Book.create(name: 'Book1', author: users[0]), Book.create(name: 'Book2', author: users[1])] }
+  let!(:book_tags) { [Tag.create(name: 'Tag1', target: books[0]), Tag.create(name: 'Tag2', target: books[1])] }
+
+  before { SQLCounter.clear_log }
+
+  it 'preload associations through RepresentExposure' do
+    User::Entity.preload_and_represent(users)
+
+    expect(SQLCounter.log).to eq([
+                                   'SELECT "books".* FROM "books" WHERE "books"."author_id" IN (?, ?)',
+                                   'SELECT "tags".* FROM "tags" WHERE "tags"."target_type" = ? AND "tags"."target_id" IN (?, ?)',
+                                   'SELECT "tags".* FROM "tags" WHERE "tags"."target_type" = ? AND "tags"."target_id" IN (?, ?)'
+                                 ])
+  end
+
+  it 'preload associations through NestingExposure' do
+    Class.new(User::Entity) do
+      unexpose :books
+      expose :nesting do
+        expose :books, using: Book::Entity, preload: :books
+      end
+    end.preload_and_represent(users)
+
+    expect(SQLCounter.log).to eq([
+                                   'SELECT "books".* FROM "books" WHERE "books"."author_id" IN (?, ?)',
+                                   'SELECT "tags".* FROM "tags" WHERE "tags"."target_type" = ? AND "tags"."target_id" IN (?, ?)',
+                                   'SELECT "tags".* FROM "tags" WHERE "tags"."target_type" = ? AND "tags"."target_id" IN (?, ?)'
+                                 ])
+  end
+
+  it 'preload same associations multiple times' do
+    Class.new(User::Entity) do
+      expose(:other_books, preload: :books) { :other_books }
+    end.preload_and_represent(users)
+
+    expect(SQLCounter.log).to eq([
+                                   'SELECT "books".* FROM "books" WHERE "books"."author_id" IN (?, ?)',
+                                   'SELECT "tags".* FROM "tags" WHERE "tags"."target_type" = ? AND "tags"."target_id" IN (?, ?)',
+                                   'SELECT "tags".* FROM "tags" WHERE "tags"."target_type" = ? AND "tags"."target_id" IN (?, ?)'
+                                 ])
+  end
+
+  describe 'only preload associations that are specified in the options' do
+    it 'through :only option' do
+      User::Entity.preload_and_represent(users, only: [:tags])
+
+      expect(SQLCounter.log).to eq([
+                                     'SELECT "tags".* FROM "tags" WHERE "tags"."target_type" = ? AND "tags"."target_id" IN (?, ?)'
+                                   ])
+    end
+
+    it 'through :except option' do
+      User::Entity.preload_and_represent(users, except: [:tags])
+
+      expect(SQLCounter.log).to eq([
+                                     'SELECT "books".* FROM "books" WHERE "books"."author_id" IN (?, ?)',
+                                     'SELECT "tags".* FROM "tags" WHERE "tags"."target_type" = ? AND "tags"."target_id" IN (?, ?)'
+                                   ])
+    end
+  end
+end