From 2abd8e311921069d319b72e6eb145f7eb4e5e777 Mon Sep 17 00:00:00 2001
From: Ray Zane <ray@promptworks.com>
Date: Fri, 1 Sep 2017 15:15:34 -0400
Subject: [PATCH] Add includes, preload, and eager_load to compat

---
 lib/baby_squeel/compat.rb       | 67 +++++++++++++++++++++++++++++++++
 spec/baby_squeel/compat_spec.rb | 33 ++++++++++++++++
 2 files changed, 100 insertions(+)

diff --git a/lib/baby_squeel/compat.rb b/lib/baby_squeel/compat.rb
index 02971a2..2b2d6d9 100644
--- a/lib/baby_squeel/compat.rb
+++ b/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.
diff --git a/spec/baby_squeel/compat_spec.rb b/spec/baby_squeel/compat_spec.rb
index 814cdf9..95a98e4 100644
--- a/spec/baby_squeel/compat_spec.rb
+++ b/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