Support queries yielding heterogeneous results

CHANGELOG.md

@@ -1,6 +1,8 @@
 Unreleased Changes
 ------------------
 
+* Feature - `Aws::Record::BuildableSearch` - Support queries yielding heterogeneous results using `multi_model_filter` (#107)
+
 2.4.1 (2020-05-29)
 ------------------
 

features/searching/search.feature

@@ -190,6 +190,10 @@ Feature: Amazon DynamoDB Querying and Scanning
       ]
       """
 
+  Scenario: Heterogeneous query
+    When we run a heterogeneous query
+    Then we should receive an aws-record collection with multiple model classes
+
   @smart_query
   Scenario: Build a Smart Scan
     When we run the following search:

features/searching/step_definitions.rb

@@ -63,3 +63,21 @@
   SearchTestModel = @model
   @collection = eval(code)
 end
+
+When(/^we run a heterogeneous query$/) do
+  @model_1 = @model.dup
+  @model_2 = @model.dup
+  scan = @model.build_scan.multi_model_filter do |raw_item_attributes|
+    if raw_item_attributes['id'] == "1"
+      @model_1
+    elsif raw_item_attributes['id'] == "2"
+      @model_2
+    end
+  end
+  @collection = scan.complete!
+end
+
+Then(/^we should receive an aws-record collection with multiple model classes/) do
+  result_classes = @collection.map(&:class)
+  expect(result_classes).to include(@model_1, @model_2)
+end

lib/aws-record/record/buildable_search.rb

@@ -178,6 +178,50 @@ def limit(size)
         self
       end
 
+      # Allows you to define a callback that will determine the model class
+      # to be used for each item, allowing queries to return an ItemCollection
+      # with mixed models.  The provided block must return the model class based on
+      # any logic on the raw item attributes or `nil` if no model applies and
+      # the item should be skipped.  Note: The block only has access to raw item
+      # data so attributes must be accessed using their names as defined in the
+      # table, not as the symbols defined in the model class(s).
+      #
+      # @example Scan with heterogeneous results:
+      #   # Example model classes
+      #   class Model_A
+      #     include Aws::Record
+      #     set_table_name(TABLE_NAME)
+      #
+      #     string_attr :uuid, hash_key: true
+      #     string_attr :class_name, range_key: true
+      #
+      #     string_attr :attr_a
+      #   end
+      #
+      #   class Model_B
+      #     include Aws::Record
+      #     set_table_name(TABLE_NAME)
+      #
+      #     string_attr :uuid, hash_key: true
+      #     string_attr :class_name, range_key: true
+      #
+      #     string_attr :attr_b
+      #   end
+      #
+      #   # use multi_model_filter to create a query on TABLE_NAME
+      #   items = Model_A.build_scan.multi_model_filter do |raw_item_attributes|
+      #     case raw_item_attributes['class_name']
+      #     when "A" then Model_A
+      #     when "B" then Model_B
+      #     else
+      #       nil
+      #     end
+      #   end.complete!
+      def multi_model_filter(proc = nil, &block)
+        @params[:model_filter] = proc || block
+        self
+      end
+
       # You must call this method at the end of any query or scan you build.
       #
       # @return [Aws::Record::ItemCollection] The item collection lazy

lib/aws-record/record/item_collection.rb

@@ -19,6 +19,7 @@ class ItemCollection
       def initialize(search_method, search_params, model, client)
         @search_method = search_method
         @search_params = search_params
+        @model_filter = @search_params.delete(:model_filter)
         @model = model
         @client = client
       end
@@ -91,9 +92,11 @@ def empty?
       def _build_items_from_response(items, model)
         ret = []
         items.each do |item|
-          record = model.new
+          model_class = @model_filter ? @model_filter.call(item) : model
+          next unless model_class
+          record = model_class.new
           data = record.instance_variable_get("@data")
-          model.attributes.attributes.each do |name, attr|
+          model_class.attributes.attributes.each do |name, attr|
             data.set_attribute(name, attr.extract(item))
           end
           data.clean!

spec/aws-record/record/item_collection_spec.rb

@@ -214,6 +214,80 @@ module Record
           expect(actual).to eq(expected)
           expect(api_requests.size).to eq(1)
         end
+
+        context 'model_filter is set' do
+
+          let(:model_a) do
+            Class.new do
+              include(Aws::Record)
+              set_table_name("TestTable")
+              integer_attr(:id, hash_key: true)
+              string_attr(:class_name)
+              string_attr(:attr_a)
+            end
+          end
+
+          let(:model_b) do
+            Class.new do
+              include(Aws::Record)
+              set_table_name("TestTable")
+              integer_attr(:id, hash_key: true)
+              string_attr(:class_name)
+              string_attr(:attr_b)
+            end
+          end
+
+          let(:resp) do
+            {
+              items: [
+                { 'id' => 1, 'class_name' => 'A', 'attr_a' => 'a' },
+                { 'id' => 2, 'class_name' => 'B', 'attr_b' => 'b' },
+                { 'id' => 3 }
+              ],
+              count: 3
+            }
+          end
+
+          let(:model_filter) do
+            Proc.new do |raw_item_attributes|
+              case raw_item_attributes['class_name']
+              when "A" then model_a
+              when "B" then model_b
+              else
+                nil
+              end
+            end
+          end
+
+          before { stub_client.stub_responses(:scan, resp) }
+
+          let(:c) do
+            ItemCollection.new(
+              :scan,
+              {table_name: "TestTable", model_filter: model_filter },
+              model_a,
+              stub_client
+            )
+          end
+
+          it 'uses the model proc to determine the returned model classes' do
+            expected = [model_a, model_b]
+            actual = c.map { |item| item.class }
+            expect(actual).to eq(expected)
+          end
+
+          it 'maps class specific attributes' do
+            actual = c.page
+            expect(actual[0].attr_a).to eq('a')
+            expect(actual[1].attr_b).to eq('b')
+          end
+
+          it 'skips items when model_filter returns nil' do
+            actual = c.page
+            expect(actual.size).to eq(2)
+          end
+        end
+
       end
 
       describe "#empty?" do