From c642b8654d6ece394aa28495879362348f62fbd7 Mon Sep 17 00:00:00 2001
From: shao <shao@1555.info>
Date: Wed, 18 Apr 2018 18:33:47 +0900
Subject: [PATCH 1/7] add ActiveModel::Datastore.save_all feature

---
 README.md                                     | 12 +++
 lib/active_model/datastore/batch_operation.rb | 43 +++++++++++
 lib/activemodel/datastore.rb                  |  1 +
 test/cases/batch_operation_test.rb            | 76 +++++++++++++++++++
 test/test_helper.rb                           |  1 +
 5 files changed, 133 insertions(+)
 create mode 100644 lib/active_model/datastore/batch_operation.rb
 create mode 100644 test/cases/batch_operation_test.rb

diff --git a/README.md b/README.md
index 1790057..a56b39d 100644
--- a/README.md
+++ b/README.md
@@ -161,6 +161,18 @@ class User
 end
 ```
 
+To improve performance, you can save many entities at once. (Batch operation)
+
+```ruby
+group = Group.find(params[:group_id])
+users = %w{ alice bob charlie }.map{ |name| User.new(name: name) }
+if ActiveModel::Datastore.save_all(users, parent: group)
+  format.html { redirect_to action: :index, notice: "#{users.count} users were successfully created.' }
+else
+  format.html { render :new }
+end
+```
+
 ## <a name="controller"></a>Controller Example
 
 Now on to the controller! A scaffold generated controller works out of the box:
diff --git a/lib/active_model/datastore/batch_operation.rb b/lib/active_model/datastore/batch_operation.rb
new file mode 100644
index 0000000..21d8ca0
--- /dev/null
+++ b/lib/active_model/datastore/batch_operation.rb
@@ -0,0 +1,43 @@
+##
+# Batch operations
+#
+#
+# Such batch calls are faster than making separate calls for each individual entity because they incur the overhead
+# for only one service call.
+#
+# reference: https://cloud.google.com/datastore/docs/concepts/entities#batch_operations
+#
+# group = Group.find(params[:group_id])
+# users = %w{ alice bob charlie }.map{ |name| User.new(name: name) }
+# saved_users = ActiveModel::CloudDatastore.save_all(users, parent: group)
+#
+module ActiveModel::Datastore
+  def self.save_all(entries, parent: nil)
+    invalid_entries = entries.reject{|entry| entry.valid?}
+    return if invalid_entries.present?
+    results = []
+    entries.each_slice(500) do |sliced_entries|
+      entities_to_save = []
+      saved_entities = nil
+      fn = lambda do |n|
+        entry = sliced_entries[n]
+        entry.run_callbacks(:save) do
+          entities_to_save << entry.build_entity(parent)
+          if n + 1 < sliced_entries.count
+            # recursive call
+            fn.call(n + 1)
+          else
+            # batch insert
+            saved_entities = entry.class.retry_on_exception? { CloudDatastore.dataset.save entities_to_save }
+          end
+          sliced_entries[n].id = saved_entities[n].key.id if saved_entities
+          sliced_entries[n].parent_key_id = saved_entities[n].key.parent.id if saved_entities && saved_entities[n].key.parent.present?
+          results << sliced_entries if n == 0 && saved_entities
+          saved_entities.present?
+        end
+      end
+      fn.call(0)
+    end
+    results.flatten
+  end
+end
diff --git a/lib/activemodel/datastore.rb b/lib/activemodel/datastore.rb
index a426a4a..03820b4 100644
--- a/lib/activemodel/datastore.rb
+++ b/lib/activemodel/datastore.rb
@@ -11,4 +11,5 @@
 require 'active_model/datastore/nested_attr'
 require 'active_model/datastore/property_values'
 require 'active_model/datastore/track_changes'
+require 'active_model/datastore/batch_operation'
 require 'active_model/datastore'
diff --git a/test/cases/batch_operation_test.rb b/test/cases/batch_operation_test.rb
new file mode 100644
index 0000000..c4c8d74
--- /dev/null
+++ b/test/cases/batch_operation_test.rb
@@ -0,0 +1,76 @@
+require 'test_helper'
+
+class ActiveModel::BatchOperationTest < ActiveSupport::TestCase
+  test 'batch save' do
+    count = MockModel.count_test_entities
+    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
+    assert ActiveModel::Datastore.save_all(mock_models)
+    assert_equal count + 3, MockModel.count_test_entities
+    mock_models.each do |mock_model|
+      assert_not_nil mock_model.id
+      key = CloudDatastore.dataset.key 'MockModel', mock_model.id
+      key.parent = CloudDatastore.dataset.key('ParentMockModel', MOCK_PARENT_ID)
+      entity = CloudDatastore.dataset.find key
+      assert_equal mock_model.id, entity.key.id
+      assert_equal 'MockModel', entity.key.kind
+      assert_equal 'ParentMockModel', entity.key.parent.kind
+      assert_equal MOCK_PARENT_ID, entity.key.parent.id
+    end
+  end
+
+  test 'before validation callback on batch save' do
+    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
+    mock_models.each do |mock_model|
+      class << mock_model
+        before_validation { self.name = nil }
+      end
+    end
+    refute ActiveModel::Datastore.save_all(mock_models)
+    mock_models.each do |mock_model|
+      assert_nil mock_model.name
+    end
+    assert_equal 0, MockModel.count_test_entities
+  end
+
+  test 'after validation callback on batch save' do
+    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
+    mock_models.each do |mock_model|
+      class << mock_model
+        after_validation { self.name = nil }
+      end
+    end
+    assert ActiveModel::Datastore.save_all(mock_models)
+    mock_models.each do |mock_model|
+      assert_nil mock_model.name
+    end
+    assert_equal mock_models.count, MockModel.count_test_entities
+  end
+
+  test 'before save callback on batch save' do
+    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
+    mock_models.each do |mock_model|
+      class << mock_model
+        before_save { self.name = name.upcase }
+      end
+    end
+    assert ActiveModel::Datastore.save_all(mock_models)
+    mock_models.each_with_index do |mock_model, index|
+      assert_equal mock_model.name, %w{ alice bob charlie }[index].upcase
+      assert_equal MockModel.all[index].name, %w{ alice bob charlie }[index].upcase
+    end
+  end
+
+  test 'after save callback on batch save' do
+    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
+    mock_models.each do |mock_model|
+      class << mock_model
+        after_save { self.name = name.upcase }
+      end
+    end
+    assert ActiveModel::Datastore.save_all(mock_models)
+    mock_models.each_with_index do |mock_model, index|
+      assert_equal mock_model.name, %w{ alice bob charlie }[index].upcase
+      assert_equal MockModel.all[index].name, %w{ alice bob charlie }[index]
+    end
+  end
+end
\ No newline at end of file
diff --git a/test/test_helper.rb b/test/test_helper.rb
index d7b2de8..ea06b81 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -15,6 +15,7 @@
 require 'active_model/datastore/nested_attr'
 require 'active_model/datastore/property_values'
 require 'active_model/datastore/track_changes'
+require 'active_model/datastore/batch_operation'
 require 'active_model/datastore'
 require 'action_controller/metal/strong_parameters'
 

From 0e06ae22bddc6b435858394a897a8bb9f6664c16 Mon Sep 17 00:00:00 2001
From: shao <shao@1555.info>
Date: Wed, 18 Apr 2018 18:51:55 +0900
Subject: [PATCH 2/7] wrong module namespace in document

---
 lib/active_model/datastore/batch_operation.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/active_model/datastore/batch_operation.rb b/lib/active_model/datastore/batch_operation.rb
index 21d8ca0..d35db05 100644
--- a/lib/active_model/datastore/batch_operation.rb
+++ b/lib/active_model/datastore/batch_operation.rb
@@ -9,7 +9,7 @@
 #
 # group = Group.find(params[:group_id])
 # users = %w{ alice bob charlie }.map{ |name| User.new(name: name) }
-# saved_users = ActiveModel::CloudDatastore.save_all(users, parent: group)
+# saved_users = ActiveModel::Datastore.save_all(users, parent: group)
 #
 module ActiveModel::Datastore
   def self.save_all(entries, parent: nil)

From ac3f0fb003e896030f1fc963cc29cc22d065378c Mon Sep 17 00:00:00 2001
From: shao <shao@1555.info>
Date: Wed, 18 Apr 2018 20:04:38 +0900
Subject: [PATCH 3/7] code style improvement

---
 lib/active_model/datastore.rb                 | 10 ++--
 lib/active_model/datastore/batch_operation.rb | 23 ++++----
 test/cases/batch_operation_test.rb            | 54 ++++++++++---------
 3 files changed, 48 insertions(+), 39 deletions(-)

diff --git a/lib/active_model/datastore.rb b/lib/active_model/datastore.rb
index 540c1a3..f54ac0c 100644
--- a/lib/active_model/datastore.rb
+++ b/lib/active_model/datastore.rb
@@ -193,6 +193,12 @@ def destroy
     end
   end
 
+  def fill_id_from_entity(entity)
+    self.id = entity.key.id
+    self.parent_key_id = entity.key.parent.id if entity.key.parent.present?
+    entity
+  end
+
   private
 
   def save_entity(parent = nil)
@@ -200,9 +206,7 @@ def save_entity(parent = nil)
     run_callbacks :save do
       entity = build_entity(parent)
       success = self.class.retry_on_exception? { CloudDatastore.dataset.save entity }
-      self.id = entity.key.id if success
-      self.parent_key_id = entity.key.parent.id if entity.key.parent.present?
-      success
+      success ? fill_id_from_entity(entity) : false
     end
   end
 
diff --git a/lib/active_model/datastore/batch_operation.rb b/lib/active_model/datastore/batch_operation.rb
index d35db05..f84f9f1 100644
--- a/lib/active_model/datastore/batch_operation.rb
+++ b/lib/active_model/datastore/batch_operation.rb
@@ -2,8 +2,8 @@
 # Batch operations
 #
 #
-# Such batch calls are faster than making separate calls for each individual entity because they incur the overhead
-# for only one service call.
+# Such batch calls are faster than making separate calls for each individual entity because they
+# incur the overhead for only one service call.
 #
 # reference: https://cloud.google.com/datastore/docs/concepts/entities#batch_operations
 #
@@ -13,10 +13,9 @@
 #
 module ActiveModel::Datastore
   def self.save_all(entries, parent: nil)
-    invalid_entries = entries.reject{|entry| entry.valid?}
-    return if invalid_entries.present?
+    return if entries.reject(&:valid?).present?
     results = []
-    entries.each_slice(500) do |sliced_entries|
+    entries.each_slice(500).map do |sliced_entries|
       entities_to_save = []
       saved_entities = nil
       fn = lambda do |n|
@@ -28,16 +27,20 @@ def self.save_all(entries, parent: nil)
             fn.call(n + 1)
           else
             # batch insert
-            saved_entities = entry.class.retry_on_exception? { CloudDatastore.dataset.save entities_to_save }
+            saved_entities = entry.class.retry_on_exception? do
+              CloudDatastore.dataset.save entities_to_save
+            end
+          end
+          if saved_entities.present?
+            sliced_entries[n].fill_id_from_entity(saved_entities[n])
+            results << sliced_entries if n.zero?
           end
-          sliced_entries[n].id = saved_entities[n].key.id if saved_entities
-          sliced_entries[n].parent_key_id = saved_entities[n].key.parent.id if saved_entities && saved_entities[n].key.parent.present?
-          results << sliced_entries if n == 0 && saved_entities
           saved_entities.present?
         end
       end
       fn.call(0)
-    end
+      sliced_entries
+    end.flatten
     results.flatten
   end
 end
diff --git a/test/cases/batch_operation_test.rb b/test/cases/batch_operation_test.rb
index c4c8d74..bf2bea4 100644
--- a/test/cases/batch_operation_test.rb
+++ b/test/cases/batch_operation_test.rb
@@ -1,12 +1,18 @@
 require 'test_helper'
 
 class ActiveModel::BatchOperationTest < ActiveSupport::TestCase
+  def setup
+    super
+    @mock_models = %w[alice bob charlie].map do |name|
+      MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)
+    end
+  end
+
   test 'batch save' do
     count = MockModel.count_test_entities
-    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
-    assert ActiveModel::Datastore.save_all(mock_models)
-    assert_equal count + 3, MockModel.count_test_entities
-    mock_models.each do |mock_model|
+    assert ActiveModel::Datastore.save_all(@mock_models)
+    assert_equal count + @mock_models.count, MockModel.count_test_entities
+    @mock_models.each do |mock_model|
       assert_not_nil mock_model.id
       key = CloudDatastore.dataset.key 'MockModel', mock_model.id
       key.parent = CloudDatastore.dataset.key('ParentMockModel', MOCK_PARENT_ID)
@@ -19,58 +25,54 @@ class ActiveModel::BatchOperationTest < ActiveSupport::TestCase
   end
 
   test 'before validation callback on batch save' do
-    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
-    mock_models.each do |mock_model|
+    @mock_models.each do |mock_model|
       class << mock_model
         before_validation { self.name = nil }
       end
     end
-    refute ActiveModel::Datastore.save_all(mock_models)
-    mock_models.each do |mock_model|
+    refute ActiveModel::Datastore.save_all(@mock_models)
+    @mock_models.each do |mock_model|
       assert_nil mock_model.name
     end
     assert_equal 0, MockModel.count_test_entities
   end
 
   test 'after validation callback on batch save' do
-    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
-    mock_models.each do |mock_model|
+    @mock_models.each do |mock_model|
       class << mock_model
         after_validation { self.name = nil }
       end
     end
-    assert ActiveModel::Datastore.save_all(mock_models)
-    mock_models.each do |mock_model|
+    assert ActiveModel::Datastore.save_all(@mock_models)
+    @mock_models.each do |mock_model|
       assert_nil mock_model.name
     end
-    assert_equal mock_models.count, MockModel.count_test_entities
+    assert_equal @mock_models.count, MockModel.count_test_entities
   end
 
   test 'before save callback on batch save' do
-    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
-    mock_models.each do |mock_model|
+    @mock_models.each do |mock_model|
       class << mock_model
         before_save { self.name = name.upcase }
       end
     end
-    assert ActiveModel::Datastore.save_all(mock_models)
-    mock_models.each_with_index do |mock_model, index|
-      assert_equal mock_model.name, %w{ alice bob charlie }[index].upcase
-      assert_equal MockModel.all[index].name, %w{ alice bob charlie }[index].upcase
+    assert ActiveModel::Datastore.save_all(@mock_models)
+    @mock_models.each_with_index do |mock_model, index|
+      assert_equal mock_model.name, %w[alice bob charlie][index].upcase
+      assert_equal MockModel.all[index].name, %w[alice bob charlie][index].upcase
     end
   end
 
   test 'after save callback on batch save' do
-    mock_models = %w{ alice bob charlie }.map{|name| MockModel.new(name: name, parent_key_id: MOCK_PARENT_ID)}
-    mock_models.each do |mock_model|
+    @mock_models.each do |mock_model|
       class << mock_model
         after_save { self.name = name.upcase }
       end
     end
-    assert ActiveModel::Datastore.save_all(mock_models)
-    mock_models.each_with_index do |mock_model, index|
-      assert_equal mock_model.name, %w{ alice bob charlie }[index].upcase
-      assert_equal MockModel.all[index].name, %w{ alice bob charlie }[index]
+    assert ActiveModel::Datastore.save_all(@mock_models)
+    @mock_models.each_with_index do |mock_model, index|
+      assert_equal mock_model.name, %w[alice bob charlie][index].upcase
+      assert_equal MockModel.all[index].name, %w[alice bob charlie][index]
     end
   end
-end
\ No newline at end of file
+end

From eeebc06ebc3806ff604436b79d3cebd68a0ce7c2 Mon Sep 17 00:00:00 2001
From: shao <shao@1555.info>
Date: Wed, 18 Apr 2018 20:16:25 +0900
Subject: [PATCH 4/7] refactor

---
 lib/active_model/datastore.rb                 |  1 +
 lib/active_model/datastore/batch_operation.rb | 18 +++++-------------
 2 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/lib/active_model/datastore.rb b/lib/active_model/datastore.rb
index f54ac0c..533393e 100644
--- a/lib/active_model/datastore.rb
+++ b/lib/active_model/datastore.rb
@@ -194,6 +194,7 @@ def destroy
   end
 
   def fill_id_from_entity(entity)
+    return if entity.nil?
     self.id = entity.key.id
     self.parent_key_id = entity.key.parent.id if entity.key.parent.present?
     entity
diff --git a/lib/active_model/datastore/batch_operation.rb b/lib/active_model/datastore/batch_operation.rb
index f84f9f1..6952f74 100644
--- a/lib/active_model/datastore/batch_operation.rb
+++ b/lib/active_model/datastore/batch_operation.rb
@@ -14,33 +14,25 @@
 module ActiveModel::Datastore
   def self.save_all(entries, parent: nil)
     return if entries.reject(&:valid?).present?
-    results = []
     entries.each_slice(500).map do |sliced_entries|
-      entities_to_save = []
-      saved_entities = nil
+      entities = []
       fn = lambda do |n|
         entry = sliced_entries[n]
         entry.run_callbacks(:save) do
-          entities_to_save << entry.build_entity(parent)
+          entities << entry.build_entity(parent)
           if n + 1 < sliced_entries.count
             # recursive call
             fn.call(n + 1)
           else
             # batch insert
-            saved_entities = entry.class.retry_on_exception? do
-              CloudDatastore.dataset.save entities_to_save
-            end
+            results = entry.class.retry_on_exception? { CloudDatastore.dataset.save entities }
           end
-          if saved_entities.present?
-            sliced_entries[n].fill_id_from_entity(saved_entities[n])
-            results << sliced_entries if n.zero?
-          end
-          saved_entities.present?
+          sliced_entries[n].fill_id_from_entity(results&.fetch(n))
+          results.present?
         end
       end
       fn.call(0)
       sliced_entries
     end.flatten
-    results.flatten
   end
 end

From 04f5bdfcc5a1384c8cef839b1cfbf79fb3ac90b5 Mon Sep 17 00:00:00 2001
From: shao <shao@1555.info>
Date: Wed, 18 Apr 2018 20:31:58 +0900
Subject: [PATCH 5/7] Revert "refactor"

This reverts commit eeebc06ebc3806ff604436b79d3cebd68a0ce7c2.
---
 lib/active_model/datastore.rb                 |  1 -
 lib/active_model/datastore/batch_operation.rb | 18 +++++++++++++-----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/lib/active_model/datastore.rb b/lib/active_model/datastore.rb
index 533393e..f54ac0c 100644
--- a/lib/active_model/datastore.rb
+++ b/lib/active_model/datastore.rb
@@ -194,7 +194,6 @@ def destroy
   end
 
   def fill_id_from_entity(entity)
-    return if entity.nil?
     self.id = entity.key.id
     self.parent_key_id = entity.key.parent.id if entity.key.parent.present?
     entity
diff --git a/lib/active_model/datastore/batch_operation.rb b/lib/active_model/datastore/batch_operation.rb
index 6952f74..f84f9f1 100644
--- a/lib/active_model/datastore/batch_operation.rb
+++ b/lib/active_model/datastore/batch_operation.rb
@@ -14,25 +14,33 @@
 module ActiveModel::Datastore
   def self.save_all(entries, parent: nil)
     return if entries.reject(&:valid?).present?
+    results = []
     entries.each_slice(500).map do |sliced_entries|
-      entities = []
+      entities_to_save = []
+      saved_entities = nil
       fn = lambda do |n|
         entry = sliced_entries[n]
         entry.run_callbacks(:save) do
-          entities << entry.build_entity(parent)
+          entities_to_save << entry.build_entity(parent)
           if n + 1 < sliced_entries.count
             # recursive call
             fn.call(n + 1)
           else
             # batch insert
-            results = entry.class.retry_on_exception? { CloudDatastore.dataset.save entities }
+            saved_entities = entry.class.retry_on_exception? do
+              CloudDatastore.dataset.save entities_to_save
+            end
           end
-          sliced_entries[n].fill_id_from_entity(results&.fetch(n))
-          results.present?
+          if saved_entities.present?
+            sliced_entries[n].fill_id_from_entity(saved_entities[n])
+            results << sliced_entries if n.zero?
+          end
+          saved_entities.present?
         end
       end
       fn.call(0)
       sliced_entries
     end.flatten
+    results.flatten
   end
 end

From 65afb12f345a0d1c479410dd982a7681792a1417 Mon Sep 17 00:00:00 2001
From: shao <shao@1555.info>
Date: Wed, 18 Apr 2018 20:40:24 +0900
Subject: [PATCH 6/7] okay, okay, test goes well

---
 lib/active_model/datastore/batch_operation.rb | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/lib/active_model/datastore/batch_operation.rb b/lib/active_model/datastore/batch_operation.rb
index f84f9f1..54228bc 100644
--- a/lib/active_model/datastore/batch_operation.rb
+++ b/lib/active_model/datastore/batch_operation.rb
@@ -14,33 +14,26 @@
 module ActiveModel::Datastore
   def self.save_all(entries, parent: nil)
     return if entries.reject(&:valid?).present?
-    results = []
     entries.each_slice(500).map do |sliced_entries|
-      entities_to_save = []
-      saved_entities = nil
+      entities = []
+      results = nil
       fn = lambda do |n|
         entry = sliced_entries[n]
         entry.run_callbacks(:save) do
-          entities_to_save << entry.build_entity(parent)
+          entities << entry.build_entity(parent)
           if n + 1 < sliced_entries.count
             # recursive call
             fn.call(n + 1)
           else
             # batch insert
-            saved_entities = entry.class.retry_on_exception? do
-              CloudDatastore.dataset.save entities_to_save
-            end
+            results = entry.class.retry_on_exception? { CloudDatastore.dataset.save entities }
           end
-          if saved_entities.present?
-            sliced_entries[n].fill_id_from_entity(saved_entities[n])
-            results << sliced_entries if n.zero?
-          end
-          saved_entities.present?
+          sliced_entries[n].fill_id_from_entity(results[n]) if results.present?
+          results.present?
         end
       end
       fn.call(0)
       sliced_entries
     end.flatten
-    results.flatten
   end
 end

From 050bfc5a9fe0eff3d48c918402d286a11c73cc9b Mon Sep 17 00:00:00 2001
From: shao <shao@1555.info>
Date: Wed, 18 Apr 2018 20:44:10 +0900
Subject: [PATCH 7/7] finally pass all test

---
 lib/active_model/datastore/batch_operation.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/active_model/datastore/batch_operation.rb b/lib/active_model/datastore/batch_operation.rb
index 54228bc..e254475 100644
--- a/lib/active_model/datastore/batch_operation.rb
+++ b/lib/active_model/datastore/batch_operation.rb
@@ -28,7 +28,7 @@ def self.save_all(entries, parent: nil)
             # batch insert
             results = entry.class.retry_on_exception? { CloudDatastore.dataset.save entities }
           end
-          sliced_entries[n].fill_id_from_entity(results[n]) if results.present?
+          sliced_entries[n].fill_id_from_entity(results&.fetch(n))
           results.present?
         end
       end