lib/adaptive_alias.rb in adaptive_alias-0.0.3 vs lib/adaptive_alias.rb in adaptive_alias-0.1.0

@@ -1,9 +1,10 @@
 # frozen_string_literal: true
 
 require 'adaptive_alias/version'
 require 'adaptive_alias/active_model_patches/read_attribute'
+require 'adaptive_alias/active_model_patches/write_attribute'
 require 'adaptive_alias/active_model_patches/remove_alias_attribute'
 require 'adaptive_alias/patches/backward_patch'
 require 'adaptive_alias/patches/forward_patch'
 
 require 'adaptive_alias/hooks/association'
@@ -29,30 +30,34 @@
 
       Module.new do
         extend ActiveSupport::Concern
 
         included do
-          if column_names.include?(new_column)
-            Patches::BackwardPatch.new(self, old_column, new_column).apply!
-          else
-            Patches::ForwardPatch.new(self, old_column, new_column).apply!
-          end
+          patch = (column_names.include?(new_column) ? Patches::BackwardPatch : Patches::ForwardPatch).new(self, old_column, new_column)
+          patch.apply!
+          patch.mark_removable
         end
       end
     end
 
-    def rescue_statement_invalid(relation)
+    def rescue_statement_invalid(relation, &block)
       yield
     rescue ActiveRecord::StatementInvalid => error
       raise error if AdaptiveAlias.current_patches.all?{|_key, patch| !patch.fix_association.call(relation, error) }
-      retry
+
+      result = rescue_statement_invalid(relation, &block)
+      AdaptiveAlias.current_patches.each_value(&:mark_removable)
+      return result
     end
 
-    def rescue_missing_attribute
+    def rescue_missing_attribute(klass, &block)
       yield
     rescue ActiveModel::MissingAttributeError => error
-      raise error if AdaptiveAlias.current_patches.all?{|_key, patch| !patch.fix_missing_attribute.call }
-      retry
+      raise error if AdaptiveAlias.current_patches.all?{|_key, patch| !patch.fix_missing_attribute.call(klass, error) }
+
+      result = rescue_missing_attribute(klass, &block)
+      AdaptiveAlias.current_patches.each_value(&:mark_removable)
+      return result
     end
 
     def get_or_create_model_module(klass)
       return @model_modules[klass] if @model_modules[klass]