lib/kms_encrypted/model.rb in kms_encrypted-1.0.0 vs lib/kms_encrypted/model.rb in kms_encrypted-1.0.1

@@ -6,23 +6,33 @@
       key_method = name ? "kms_key_#{name}" : "kms_key"
       key_column = "encrypted_#{key_method}"
       context_method = name ? "kms_encryption_context_#{name}" : "kms_encryption_context"
 
       class_eval do
-        class << self
-          def kms_keys
-            @kms_keys ||= {}
-          end unless respond_to?(:kms_keys)
+        @kms_keys ||= {}
+
+        unless respond_to?(:kms_keys)
+          def self.kms_keys
+            parent_keys =
+              if superclass.respond_to?(:kms_keys)
+                superclass.kms_keys
+              else
+                {}
+              end
+
+            parent_keys.merge(@kms_keys || {})
+          end
         end
-        kms_keys[key_method.to_sym] = {
+
+        @kms_keys[key_method.to_sym] = {
           key_id: key_id,
           name: name,
           version: version,
           previous_versions: previous_versions,
           upgrade_context: upgrade_context
         }
 
-        if kms_keys.size == 1
+        if @kms_keys.size == 1
           after_save :encrypt_kms_keys
 
           # fetch all keys together so only need to update database once
           def encrypt_kms_keys
             updates = {}