Merge pull request #342 from arjes/add_metadata_to_page_iterator

Adding metadata to page iteration

Author: andrykonchin

README.md

@@ -590,11 +590,30 @@ Access to the native DynamoDB pages can be obtained via the `find_by_pages`
 method, which yields arrays of records.
 
 ```ruby
-Address.find_by_pages do |addresses|
+Address.find_by_pages do |addresses, metadata|
   # have an array of pages
 end
 ```
 
+Each yielded pages returns metadata as the second argument, which is a hash
+including a key `:last_evaluated_key`. The value of this key can be used for
+the `start` method to fetch the next page of records.
+
+```ruby
+class UserController < ApplicationController
+  def index
+    next_page = params[:next_page_token] ? JSON.parse(Base64.decode64(params[:next_page_token])) : ''
+
+    records, metadata = User.start(next_page).find_by_pages.first
+
+    render json: {
+      records: records,
+      next_page_token: Base64.encode64(metadata[:last_evaluated_key].to_json)
+    }
+  end
+end
+```
+
 #### Sort Conditions and Filters
 
 You are able to optimize query with condition for sort key. Following operators are available: `gt`, `lt`, `gte`, `lte`,

lib/dynamoid/adapter_plugin/aws_sdk_v3.rb

@@ -532,11 +532,15 @@ def put_item(table_name, object, options = {})
       #
       # @todo Provide support for various other options http://docs.aws.amazon.com/sdkforruby/api/Aws/DynamoDB/Client.html#query-instance_method
       def query(table_name, options = {})
-        return enum_for(:query, table_name, options) unless block_given?
-        table = describe_table(table_name)
+        Enumerator.new do |yielder|
+          table = describe_table(table_name)
 
-        Query.new(client, table, options).call.each do |page|
-          yield page.items.map{ |row| result_item_to_hash(row) }
+          Query.new(client, table, options).call.each do |page|
+            yielder.yield(
+              page.items.map { |row| result_item_to_hash(row) },
+              last_evaluated_key: page.last_evaluated_key
+            )
+          end
         end
       end
 
@@ -561,11 +565,15 @@ def query_count(table_name, options = {})
       #
       # @todo: Provide support for various options http://docs.aws.amazon.com/sdkforruby/api/Aws/DynamoDB/Client.html#scan-instance_method
       def scan(table_name, conditions = {}, options = {})
-        return enum_for(:scan, table_name, conditions, options) unless block_given?
-        table = describe_table(table_name)
+        Enumerator.new do |yielder|
+          table = describe_table(table_name)
 
-        Scan.new(client, table, conditions, options).call.each do |page|
-          yield page.items.map{ |row| result_item_to_hash(row) }
+          Scan.new(client, table, conditions, options).call.each do |page|
+            yielder.yield(
+              page.items.map { |row| result_item_to_hash(row) },
+              last_evaluated_key: page.last_evaluated_key
+            )
+          end
         end
       end
 

lib/dynamoid/criteria.rb

@@ -13,12 +13,12 @@ module ClassMethods
         # see Dynamoid::Criteria::Chain.
         #
         # @since 0.2.0
-        define_method(meth) do |*args|
+        define_method(meth) do |*args, &blk|
           chain = Dynamoid::Criteria::Chain.new(self)
           if args
-            chain.send(meth, *args)
+            chain.send(meth, *args, &blk)
           else
-            chain.send(meth)
+            chain.send(meth, &blk)
           end
         end
       end

lib/dynamoid/criteria/chain.rb

@@ -163,10 +163,10 @@ def pages
       #
       # @since 3.1.0
       def pages_via_query
-        return enum_for(:pages_via_query) unless block_given?
-
-        Dynamoid.adapter.query(source.table_name, range_query).each do |items|
-          yield items.map { |hash| source.from_database(hash) }
+        Enumerator.new do |yielder|
+          Dynamoid.adapter.query(source.table_name, range_query).each do |items, metadata|
+            yielder.yield items.map { |hash| source.from_database(hash) }, metadata.slice(:last_evaluated_key)
+          end
         end
       end
 
@@ -176,10 +176,10 @@ def pages_via_query
       #
       # @since 3.1.0
       def pages_via_scan
-        return enum_for(:pages_via_scan) unless block_given?
-
-        Dynamoid.adapter.scan(source.table_name, scan_query, scan_opts).each do |items|
-          yield items.map { |hash| source.from_database(hash) }
+        Enumerator.new do |yielder|
+          Dynamoid.adapter.scan(source.table_name, scan_query, scan_opts).each do |items, metadata|
+            yielder.yield(items.map { |hash| source.from_database(hash) }, metadata.slice(:last_evaluated_key))
+          end
         end
       end
 

spec/dynamoid/adapter_plugin/aws_sdk_v3_spec.rb

@@ -258,27 +258,27 @@ def dynamo_request(table_name, scan_hash = {}, select_opts = {})
     end
 
     it 'performs query on a table with a range and selects items in a range' do
-      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_between: [0.0, 3.0]).to_a).to eq [[{ id: '1', range: BigDecimal(1) }, { id: '1', range: BigDecimal(3) }]]
+      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_between: [0.0, 3.0]).to_a).to eq [[[{ id: '1', range: BigDecimal(1) }, { id: '1', range: BigDecimal(3) }], {last_evaluated_key: nil}]]
     end
 
     it 'performs query on a table with a range and selects items in a range with :select option' do
-      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_between: [0.0, 3.0], select: 'ALL_ATTRIBUTES').to_a).to eq [[{ id: '1', range: BigDecimal(1) }, { id: '1', range: BigDecimal(3) }]]
+      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_between: [0.0, 3.0], select: 'ALL_ATTRIBUTES').to_a).to eq [[[{ id: '1', range: BigDecimal(1) }, { id: '1', range: BigDecimal(3) }], {last_evaluated_key: nil}]]
     end
 
     it 'performs query on a table with a range and selects items greater than' do
-      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_greater_than: 1.0).to_a).to eq [[{ id: '1', range: BigDecimal(3) }]]
+      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_greater_than: 1.0).to_a).to eq [[[{ id: '1', range: BigDecimal(3) }], {last_evaluated_key: nil}]]
     end
 
     it 'performs query on a table with a range and selects items less than' do
-      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_less_than: 2.0).to_a).to eq [[{ id: '1', range: BigDecimal(1) }]]
+      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_less_than: 2.0).to_a).to eq [[[{ id: '1', range: BigDecimal(1) }], {last_evaluated_key: nil}]]
     end
 
     it 'performs query on a table with a range and selects items gte' do
-      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_gte: 1.0).to_a).to eq [[{ id: '1', range: BigDecimal(1) }, { id: '1', range: BigDecimal(3) }]]
+      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_gte: 1.0).to_a).to eq [[[{ id: '1', range: BigDecimal(1) }, { id: '1', range: BigDecimal(3) }], {last_evaluated_key: nil}]]
     end
 
     it 'performs query on a table with a range and selects items lte' do
-      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_lte: 3.0).to_a).to eq [[{ id: '1', range: BigDecimal(1) }, { id: '1', range: BigDecimal(3) }]]
+      expect(Dynamoid.adapter.query(test_table3, hash_value: '1', range_lte: 3.0).to_a).to eq [[[{ id: '1', range: BigDecimal(1) }, { id: '1', range: BigDecimal(3) }], {last_evaluated_key: nil}]]
     end
 
     it 'performs query on a table and returns items based on returns correct limit' do
@@ -677,11 +677,14 @@ def dynamo_request(table_name, scan_hash = {}, select_opts = {})
                                           ])
 
         results = Dynamoid.adapter.scan(test_table3)
-        expect(results.to_a.first).to contain_exactly(
-          { id: '1', range: 1.0 },
-          { id: '2', range: 2.0 },
-          { id: '3', range: 3.0 }
-        )
+        expect(results.to_a.first).to match [
+          contain_exactly(
+            { id: '1', range: 1.0 },
+            { id: '2', range: 2.0 },
+            { id: '3', range: 3.0 }
+          ),
+          { last_evaluated_key: nil }
+        ]
       end
 
       it 'performs BatchDeleteItem with more than 25 items' do
@@ -790,14 +793,14 @@ def dynamo_request(table_name, scan_hash = {}, select_opts = {})
     it 'performs query on a table and returns items' do
       Dynamoid.adapter.put_item(test_table1, id: '1', name: 'Josh')
 
-      expect(Dynamoid.adapter.query(test_table1, hash_value: '1').first).to eq([id: '1', name: 'Josh'])
+      expect(Dynamoid.adapter.query(test_table1, hash_value: '1').first).to eq([[id: '1', name: 'Josh'], {last_evaluated_key: nil}])
     end
 
     it 'performs query on a table and returns items if there are multiple items' do
       Dynamoid.adapter.put_item(test_table1, id: '1', name: 'Josh')
       Dynamoid.adapter.put_item(test_table1, id: '2', name: 'Justin')
 
-      expect(Dynamoid.adapter.query(test_table1, hash_value: '1').first).to eq([id: '1', name: 'Josh'])
+      expect(Dynamoid.adapter.query(test_table1, hash_value: '1').first).to eq([[id: '1', name: 'Josh'], {last_evaluated_key: nil}])
     end
 
     context 'backoff is specified' do
@@ -834,21 +837,28 @@ def dynamo_request(table_name, scan_hash = {}, select_opts = {})
     it 'performs scan on a table and returns items' do
       Dynamoid.adapter.put_item(test_table1, id: '1', name: 'Josh')
 
-      expect(Dynamoid.adapter.scan(test_table1, name: { eq: 'Josh' }).to_a).to eq [[{ id: '1', name: 'Josh' }]]
+      expect(Dynamoid.adapter.scan(test_table1, name: { eq: 'Josh' }).to_a).to eq [[[{ id: '1', name: 'Josh' }], {last_evaluated_key: nil}]]
     end
 
     it 'performs scan on a table and returns items if there are multiple items but only one match' do
       Dynamoid.adapter.put_item(test_table1, id: '1', name: 'Josh')
       Dynamoid.adapter.put_item(test_table1, id: '2', name: 'Justin')
 
-      expect(Dynamoid.adapter.scan(test_table1, name: { eq: 'Josh' }).to_a).to eq [[{ id: '1', name: 'Josh' }]]
+      expect(Dynamoid.adapter.scan(test_table1, name: { eq: 'Josh' }).to_a).to eq [[[{ id: '1', name: 'Josh' }], {last_evaluated_key: nil}]]
     end
 
     it 'performs scan on a table and returns multiple items if there are multiple matches' do
       Dynamoid.adapter.put_item(test_table1, id: '1', name: 'Josh')
       Dynamoid.adapter.put_item(test_table1, id: '2', name: 'Josh')
 
-      expect(Dynamoid.adapter.scan(test_table1, name: { eq: 'Josh' })).to match_array(include({ name: 'Josh', id: '2' }, name: 'Josh', id: '1'))
+      expect(
+        Dynamoid.adapter.scan(test_table1, name: { eq: 'Josh' }).to_a
+      ).to match([
+        [
+          match_array([{ name: 'Josh', id: '2' }, { name: 'Josh', id: '1' }]),
+          {last_evaluated_key: nil}
+        ]
+      ])
     end
 
     it 'performs scan on a table and returns all items if no criteria are specified' do