lib/lockbox/migrator.rb in lockbox-0.3.2 vs lib/lockbox/migrator.rb in lockbox-0.3.3

@@ -11,46 +11,50 @@
     end
 
     def rotate(attributes:)
       fields = {}
       attributes.each do |a|
-        # use key instad of v[:attribute] to make it more intuitive when migrating: true
+        # use key instead of v[:attribute] to make it more intuitive when migrating: true
         field = model.lockbox_attributes[a]
         raise ArgumentError, "Bad attribute: #{a}" unless field
         fields[a] = field
       end
 
-      perform(fields: fields)
+      perform(fields: fields, rotate: true)
     end
 
     # TODO add attributes option
     def migrate(restart:)
       fields = model.lockbox_attributes.select { |k, v| v[:migrating] }
 
+      # need blind indexes for building relation
       blind_indexes = model.respond_to?(:blind_indexes) ? model.blind_indexes.select { |k, v| v[:migrating] } : {}
 
       perform(fields: fields, blind_indexes: blind_indexes, restart: restart)
     end
 
     private
 
-    def perform(fields:, blind_indexes: [], restart: true)
+    def perform(fields:, blind_indexes: [], restart: true, rotate: false)
       relation = @relation
 
-      # remove true condition in 0.4.0
-      if true || (defined?(ActiveRecord::Base) && base_relation.is_a?(ActiveRecord::Base))
+      # unscope if passed a model
+      unless ar_relation?(relation) || mongoid_relation?(relation)
         relation = relation.unscoped
+      else
+        # TODO remove in 0.4.0
+        relation = relation.unscoped
       end
 
       # convert from possible class to ActiveRecord::Relation or Mongoid::Criteria
       relation = relation.all
 
       unless restart
         attributes = fields.map { |_, v| v[:encrypted_attribute] }
         attributes += blind_indexes.map { |_, v| v[:bidx_attribute] }
 
-        if defined?(ActiveRecord::Relation) && relation.is_a?(ActiveRecord::Relation)
+        if ar_relation?(relation)
           base_relation = relation.unscoped
           or_relation = relation.unscoped
 
           attributes.each_with_index do |attribute, i|
             or_relation =
@@ -66,11 +70,11 @@
           relation = relation.or(attributes.map { |a| {a => nil} })
         end
       end
 
       each_batch(relation) do |records|
-        migrate_records(records, fields: fields, blind_indexes: blind_indexes, restart: restart)
+        migrate_records(records, fields: fields, blind_indexes: blind_indexes, restart: restart, rotate: rotate)
       end
     end
 
     def each_batch(relation)
       if relation.respond_to?(:find_in_batches)
@@ -90,28 +94,60 @@
         end
         yield records if records.any?
       end
     end
 
-    def migrate_records(records, fields:, blind_indexes:, restart:)
+    def migrate_records(records, fields:, blind_indexes:, restart:, rotate:)
       # do computation outside of transaction
       # especially expensive blind index computation
-      records.each do |record|
-        fields.each do |k, v|
-          record.send("#{v[:attribute]}=", record.send(k)) if restart || !record.send(v[:encrypted_attribute])
+      if rotate
+        records.each do |record|
+          fields.each do |k, v|
+            # update encrypted attribute directly to skip blind index computation
+            record.send("lockbox_direct_#{k}=", record.send(k))
+          end
         end
-        blind_indexes.each do |k, v|
-          record.send("compute_#{k}_bidx") if restart || !record.send(v[:bidx_attribute])
+      else
+        records.each do |record|
+          if restart
+            fields.each do |k, v|
+              record.send("#{v[:encrypted_attribute]}=", nil)
+            end
+
+            blind_indexes.each do |k, v|
+              record.send("#{v[:bidx_attribute]}=", nil)
+            end
+          end
+
+          fields.each do |k, v|
+            record.send("#{v[:attribute]}=", record.send(k)) unless record.send(v[:encrypted_attribute])
+          end
+
+          # with Blind Index 2.0, bidx_attribute should be already set for each record
+          blind_indexes.each do |k, v|
+            record.send("compute_#{k}_bidx") unless record.send(v[:bidx_attribute])
+          end
         end
       end
 
+      # don't need to save records that went from nil => nil
       records.select! { |r| r.changed? }
 
-      with_transaction do
-        records.each do |record|
-          record.save!(validate: false)
+      if records.any?
+        with_transaction do
+          records.each do |record|
+            record.save!(validate: false)
+          end
         end
       end
+    end
+
+    def ar_relation?(relation)
+      defined?(ActiveRecord::Relation) && relation.is_a?(ActiveRecord::Relation)
+    end
+
+    def mongoid_relation?(relation)
+      defined?(Mongoid::Criteria) && relation.is_a?(Mongoid::Criteria)
     end
 
     def with_transaction
       if @transaction
         @relation.transaction do