From 3889fa1d9779ebd14afa4c85810006f66fb67b44 Mon Sep 17 00:00:00 2001
From: Danny Radev <danny@dradev.net>
Date: Wed, 15 Aug 2018 17:46:55 +0300
Subject: [PATCH] Persist attributes before save

Persisting Vault attributes on an `after_save` uses two separate queries:
one for the model `INSERT/UPDATE`, and another to `UPDATE` the ciphertext
for the encrypted attributes. Encrypting the attributes with a `before_save`
avoids the second query. In some cases users might _not_ want to have two
queries when saving a single record. This would be necessary for example,
when one has an auditing table and/or stored procedures that take some action
when a record is changed.
---
 lib/vault/encrypted_model.rb      | 30 ++++++++++++++-------
 spec/unit/encrypted_model_spec.rb | 45 +++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/lib/vault/encrypted_model.rb b/lib/vault/encrypted_model.rb
index 64441ed1..924e44e3 100644
--- a/lib/vault/encrypted_model.rb
+++ b/lib/vault/encrypted_model.rb
@@ -116,6 +116,12 @@ def vault_attribute(attribute, options = {})
         self
       end
 
+      # Encrypt Vault attribures before saving them
+      def vault_persist_before_save!
+        skip_callback :save, :after, :__vault_persist_attributes!
+        before_save :__vault_encrypt_attributes!
+      end
+
       # The list of Vault attributes.
       #
       # @return [Hash]
@@ -219,13 +225,7 @@ def __vault_load_attribute!(attribute, options)
       # on this model.
       # @return [true]
       def __vault_persist_attributes!
-        changes = {}
-
-        self.class.__vault_attributes.each do |attribute, options|
-          if c = self.__vault_persist_attribute!(attribute, options)
-            changes.merge!(c)
-          end
-        end
+        changes = __vault_encrypt_attributes!
 
         # If there are any changes to the model, update them all at once,
         # skipping any callbacks and validation. This is okay, because we are
@@ -234,12 +234,24 @@ def __vault_persist_attributes!
           self.update_columns(changes)
         end
 
-        return true
+        true
+      end
+
+      def __vault_encrypt_attributes!
+        changes = {}
+
+        self.class.__vault_attributes.each do |attribute, options|
+          if c = self.__vault_encrypt_attribute!(attribute, options)
+            changes.merge!(c)
+          end
+        end
+
+        changes
       end
 
       # Encrypt a single attribute using Vault and persist back onto the
       # encrypted attribute value.
-      def __vault_persist_attribute!(attribute, options)
+      def __vault_encrypt_attribute!(attribute, options)
         key        = options[:key]
         path       = options[:path]
         serializer = options[:serializer]
diff --git a/spec/unit/encrypted_model_spec.rb b/spec/unit/encrypted_model_spec.rb
index 46c37f07..ffe52ada 100644
--- a/spec/unit/encrypted_model_spec.rb
+++ b/spec/unit/encrypted_model_spec.rb
@@ -42,4 +42,49 @@
       expect(klass.instance_methods).to include(:foo_was)
     end
   end
+
+  describe '#vault_persist_before_save!' do
+    let(:after_save_dummy_class) do
+      Class.new(ActiveRecord::Base) do
+        include Vault::EncryptedModel
+      end
+    end
+
+    let(:before_save_dummy_class) do
+      Class.new(ActiveRecord::Base) do
+        include Vault::EncryptedModel
+        vault_persist_before_save!
+      end
+    end
+
+    context "when not used" do
+      it "the model has an after_save callback" do
+        save_callbacks = after_save_dummy_class._save_callbacks.select do |cb|
+          cb.filter == :__vault_persist_attributes!
+        end
+
+        expect(save_callbacks.length).to eq 1
+        persist_callback = save_callbacks.first
+
+        expect(persist_callback).to be_a ActiveSupport::Callbacks::Callback
+
+        expect(persist_callback.kind).to eq :after
+      end
+    end
+
+    context "when used" do
+      it "the model has a before_save callback" do
+        save_callbacks = before_save_dummy_class._save_callbacks.select do |cb|
+          cb.filter == :__vault_encrypt_attributes!
+        end
+
+        expect(save_callbacks.length).to eq 1
+        persist_callback = save_callbacks.first
+
+        expect(persist_callback).to be_a ActiveSupport::Callbacks::Callback
+
+        expect(persist_callback.kind).to eq :before
+      end
+    end
+  end
 end