From b06515194ecd73976f57b8c4e38be22d2707956e Mon Sep 17 00:00:00 2001
From: Danny Radev <danny@dradev.net>
Date: Thu, 25 Oct 2018 18:17:10 +0300
Subject: [PATCH] Fix missing Dirty methods not being generated

---
 .../associations/join_dependency/join_part.rb | 24 +++++++++++++++++++
 lib/vault/encrypted_model.rb                  |  7 ------
 lib/vault/rails.rb                            |  1 +
 spec/dummy/app/models/person.rb               |  4 ++++
 4 files changed, 29 insertions(+), 7 deletions(-)
 create mode 100644 lib/vault/active_record/associations/join_dependency/join_part.rb

diff --git a/lib/vault/active_record/associations/join_dependency/join_part.rb b/lib/vault/active_record/associations/join_dependency/join_part.rb
new file mode 100644
index 00000000..f83f714e
--- /dev/null
+++ b/lib/vault/active_record/associations/join_dependency/join_part.rb
@@ -0,0 +1,24 @@
+require 'active_record'
+# Force ther original JoinPart to load
+require 'active_record/associations/join_dependency/join_part'
+
+module ActiveRecord
+  module Associations
+    class JoinDependency
+      class JoinPart
+        # Prevent virtual attributes from being included in JOIN sql queries.
+        # This will work with both ActiveRecord 4 an 5 because in the original
+        # implementation this methos is delegated to the base class - the model.
+        def column_names
+          column_names = base_klass.column_names
+
+          if base_klass.methods.include? :persistable_attribute_names
+            column_names = column_names & base_klass.persistable_attribute_names
+          end
+
+          column_names
+        end
+      end
+    end
+  end
+end
diff --git a/lib/vault/encrypted_model.rb b/lib/vault/encrypted_model.rb
index 85a0d595..aea98b9f 100644
--- a/lib/vault/encrypted_model.rb
+++ b/lib/vault/encrypted_model.rb
@@ -200,13 +200,6 @@ def __vault_load_attributes!
         end
       end
 
-      # In ActiveRecord 4.2, virtual attributes are not excluded from joins aliases
-      if respond_to?(:persistable_attribute_names)
-        def self.column_names
-          super & persistable_attribute_names
-        end
-      end
-
       # Decrypt and load a single attribute from Vault.
       def __vault_load_attribute!(attribute, options)
         # If the user provided a value for the attribute, do not try to load it from Vault
diff --git a/lib/vault/rails.rb b/lib/vault/rails.rb
index d20fa2d3..e1ec5074 100644
--- a/lib/vault/rails.rb
+++ b/lib/vault/rails.rb
@@ -11,6 +11,7 @@
 require_relative 'rails/serializers/integer_serializer'
 require_relative 'rails/serializers/float_serializer'
 require_relative 'rails/version'
+require_relative 'active_record/associations/join_dependency/join_part'
 
 module Vault
   module Rails
diff --git a/spec/dummy/app/models/person.rb b/spec/dummy/app/models/person.rb
index 04892a15..c42d68b1 100644
--- a/spec/dummy/app/models/person.rb
+++ b/spec/dummy/app/models/person.rb
@@ -5,6 +5,10 @@ class Person < ActiveRecord::Base
 
   has_many :problems
 
+  before_validation :format_ssn, if: -> { ssn_changed? }
+
+  def format_ssn; end
+
   vault_attribute :county_plaintext, encrypted_column: :county_encrypted
   vault_attribute_proxy :county, :county_plaintext