collection.rb in heimdallr-1.0.4

- old
+ new

@@ -17,10 +17,11 @@
     # @option options [Boolean] implicit proxy type
     def initialize(context, scope, options={})
       @context, @scope, @options = context, scope, options
 
       @restrictions = @scope.restrictions(context)
+      @options[:eager_loaded] ||= {}
     end
 
     # Collections cannot be restricted with different context or options.
     #
     # @return self
@@ -38,11 +39,11 @@
     #   A proxy for +$1+ method which adds fixtures to the attribute list and
     #   returns a restricted record.
     def self.delegate_as_constructor(name, method)
       class_eval(<<-EOM, __FILE__, __LINE__)
       def #{name}(attributes={})
-        record = @restrictions.request_scope(:fetch).new.restrict(@context, @options)
+        record = @restrictions.request_scope(:fetch).new.restrict(@context, options_with_escape)
         record.#{method}(attributes.merge(@restrictions.fixtures[:create]))
         record
       end
       EOM
     end
@@ -51,11 +52,11 @@
     # @macro [attach] delegate_as_scope
     #   A proxy for +$1+ method which returns a restricted scope.
     def self.delegate_as_scope(name)
       class_eval(<<-EOM, __FILE__, __LINE__)
       def #{name}(*args)
-        Proxy::Collection.new(@context, @scope.#{name}(*args), @options)
+        Proxy::Collection.new(@context, @scope.#{name}(*args), options_with_escape)
       end
       EOM
     end
 
     # @private
@@ -73,11 +74,11 @@
     # @macro [attach] delegate_as_record
     #   A proxy for +$1+ method which returns a restricted record.
     def self.delegate_as_record(name)
       class_eval(<<-EOM, __FILE__, __LINE__)
       def #{name}(*args)
-        @scope.#{name}(*args).restrict(@context, @options)
+        @scope.#{name}(*args).restrict(@context, options_with_eager_load)
       end
       EOM
     end
 
     # @private
@@ -85,11 +86,11 @@
     #   A proxy for +$1+ method which returns an array of restricted records.
     def self.delegate_as_records(name)
       class_eval(<<-EOM, __FILE__, __LINE__)
       def #{name}(*args)
         @scope.#{name}(*args).map do |element|
-          element.restrict(@context, @options)
+          element.restrict(@context, options_with_eager_load)
         end
       end
       EOM
     end
 
@@ -111,13 +112,10 @@
 
     delegate_as_scope :scoped
     delegate_as_scope :uniq
     delegate_as_scope :where
     delegate_as_scope :joins
-    delegate_as_scope :includes
-    delegate_as_scope :eager_load
-    delegate_as_scope :preload
     delegate_as_scope :lock
     delegate_as_scope :limit
     delegate_as_scope :offset
     delegate_as_scope :order
     delegate_as_scope :reorder
@@ -152,45 +150,96 @@
 
     delegate_as_records :all
     delegate_as_records :to_a
     delegate_as_records :to_ary
 
+    # A proxy for +includes+ which adds Heimdallr conditions for eager loaded
+    # associations.
+    def includes(*associations)
+      # Normalize association list to strict nested hash.
+      normalize = ->(list) {
+        if list.is_a? Array
+          list.map(&normalize).reduce(:merge)
+        elsif list.is_a? Symbol
+          { list => {} }
+        elsif list.is_a? Hash
+          hash = {}
+          list.each do |key, value|
+            hash[key] = normalize.(value)
+          end
+          hash
+        end
+      }
+      associations = normalize.(associations)
+
+      current_scope = @scope.includes(associations)
+
+      add_conditions = ->(associations, scope) {
+        associations.each do |association, nested|
+          reflection = scope.reflect_on_association(association)
+          if reflection && !reflection.options[:polymorphic]
+            associated_klass = reflection.klass
+
+            if associated_klass.respond_to? :restrict
+              nested_scope = associated_klass.restrictions(@context).request_scope(:fetch)
+
+              where_values = nested_scope.where_values
+              if where_values.any?
+                current_scope = current_scope.where(*where_values)
+              end
+
+              add_conditions.(nested, associated_klass)
+            end
+          end
+        end
+      }
+
+      unless Heimdallr.skip_eager_condition_injection
+        add_conditions.(associations, current_scope)
+      end
+
+      options = @options.merge(eager_loaded:
+        @options[:eager_loaded].merge(associations))
+
+      Proxy::Collection.new(@context, current_scope, options)
+    end
+
     # A proxy for +find+ which restricts the returned record or records.
     #
     # @return [Proxy::Record, Array<Proxy::Record>]
     def find(*args)
       result = @scope.find(*args)
 
       if result.is_a? Enumerable
         result.map do |element|
-          element.restrict(@context, @options)
+          element.restrict(@context, options_with_eager_load)
         end
       else
-        result.restrict(@context, @options)
+        result.restrict(@context, options_with_eager_load)
       end
     end
 
     # A proxy for +each+ which restricts the yielded records.
     #
     # @yield [record]
     # @yieldparam [Proxy::Record] record
     def each
       @scope.each do |record|
-        yield record.restrict(@context, @options)
+        yield record.restrict(@context, options_with_eager_load)
       end
     end
 
     # Wraps a scope or a record in a corresponding proxy.
     def method_missing(method, *args)
       if method =~ /^find_all_by/
         @scope.send(method, *args).map do |element|
-          element.restrict(@context, @options)
+          element.restrict(@context, options_with_escape)
         end
       elsif method =~ /^find_by/
-        @scope.send(method, *args).restrict(@context, @options)
+        @scope.send(method, *args).restrict(@context, options_with_escape)
       elsif @scope.heimdallr_scopes && @scope.heimdallr_scopes.include?(method)
-        Proxy::Collection.new(@context, @scope.send(method, *args), @options)
+        Proxy::Collection.new(@context, @scope.send(method, *args), options_with_escape)
       elsif @scope.respond_to? method
         raise InsecureOperationError,
             "Potentially insecure method #{method} was called"
       else
         super
@@ -229,8 +278,21 @@
       }.merge(@restrictions.reflection)
     end
 
     def creatable?
       @restrictions.can? :create
+    end
+
+    private
+
+    # Return options hash to pass to children proxies.
+    # Currently this checks only eagerly loaded collections, which
+    # shouldn't be passed around blindly.
+    def options_with_escape
+      @options.reject { |k,v| k == :eager_loaded }
+    end
+
+    def options_with_eager_load
+      @options
     end
   end
 end
\ No newline at end of file