Add includes, preload, and eager_load to compat

lib/baby_squeel/compat.rb

@@ -8,6 +8,46 @@ def self.enable!
       ::ActiveRecord::Relation.prepend QueryMethods
     end
 
+    class KeyPath
+      def self.evaluate(&block)
+        if block.arity.zero?
+          unwrap new.instance_eval(&block)
+        else
+          unwrap yield(new)
+        end
+      end
+
+      def self.unwrap(path)
+        if path.kind_of? self
+          path.value
+        elsif path.respond_to? :map
+          path.map(&method(:unwrap))
+        else
+          []
+        end
+      end
+
+      def initialize(*path)
+        @path = path
+      end
+
+      def value
+        @path.reverse.reduce({}) do |acc, name|
+          { name => acc }
+        end
+      end
+
+      private
+
+      def respond_to_missing?(*)
+        true
+      end
+
+      def method_missing(name, *)
+        self.class.new(*@path, name)
+      end
+    end
+
     module DSL
       # An alias for BabySqueel::DSL#sql
       def `(str)
@@ -36,6 +76,33 @@ def joins(*args, &block)
         end
       end
 
+      # Overrides ActiveRecord::QueryMethods#includes
+      def includes(*args, &block)
+        if block_given? && args.empty?
+          super KeyPath.evaluate(&block)
+        else
+          super
+        end
+      end
+
+      # Overrides ActiveRecord::QueryMethods#eager_load
+      def eager_load(*args, &block)
+        if block_given? && args.empty?
+          super KeyPath.evaluate(&block)
+        else
+          super
+        end
+      end
+
+      # Overrides ActiveRecord::QueryMethods#preload
+      def preload(*args, &block)
+        if block_given? && args.empty?
+          super KeyPath.evaluate(&block)
+        else
+          super
+        end
+      end
+
       # Heads up, Array#select conflicts with
       # ActiveRecord::QueryMethods#select. So, if arity
       # is given to the block, we'll use Array#select.

spec/baby_squeel/compat_spec.rb

@@ -23,3 +23,36 @@
     end
   end
 end
+
+describe 'BabySqueel::Compat::QueryMethods', :compat do
+  describe '#includes' do
+    it 'accepts a block' do
+      posts = Post.includes { author.posts }
+      expect(posts.includes_values).to eq([ author: { posts: {} } ])
+    end
+
+    it 'handles arrays' do
+      posts = Post.includes { [comments, author.posts] }
+      expect(posts.includes_values).to eq([ { comments: {} }, author: { posts: {} } ])
+    end
+
+    it 'handles nil' do
+      posts = Post.includes { nil }
+      expect(posts.includes_values).to eq([])
+    end
+  end
+
+  describe '#eager_load' do
+    it 'accepts a block' do
+      posts = Post.eager_load { author.posts }
+      expect(posts.eager_load_values).to eq([ author: { posts: {} } ])
+    end
+  end
+
+  describe '#preload' do
+    it 'accepts a block' do
+      posts = Post.preload { author.posts }
+      expect(posts.preload_values).to eq([ author: { posts: {} } ])
+    end
+  end
+end