Fix restore with polymorphic has_one relationships

lib/paranoia.rb

@@ -143,8 +143,16 @@ def restore_associated_records
 
       if association_data.nil? && association.macro.to_s == "has_one"
         association_class_name = association.options[:class_name].present? ? association.options[:class_name] : association.name.to_s.camelize
-        association_foreign_key = association.options[:foreign_key].present? ? association.options[:foreign_key] : "#{self.class.name.to_s.underscore}_id"
-        Object.const_get(association_class_name).only_deleted.where(association_foreign_key => self.id).first.try(:restore, recursive: true)
+        association_foreign_key = association.foreign_key.present? ? association.foreign_key : "#{self.class.name.to_s.underscore}_id"
+
+        if association.type
+          association_polymorphic_type = association.type
+          association_find_conditions = { association_polymorphic_type => self.class.name.to_s, association_foreign_key => self.id }
+        else
+          association_find_conditions = { association_foreign_key => self.id }
+        end
+
+        Object.const_get(association_class_name).only_deleted.where(association_find_conditions).first.try(:restore, recursive: true)
       end
     end
 

test/paranoia_test.rb

@@ -32,6 +32,7 @@ def setup!
   ActiveRecord::Base.connection.execute 'CREATE TABLE custom_column_models (id INTEGER NOT NULL PRIMARY KEY, destroyed_at DATETIME)'
   ActiveRecord::Base.connection.execute 'CREATE TABLE custom_sentinel_models (id INTEGER NOT NULL PRIMARY KEY, deleted_at DATETIME NOT NULL)'
   ActiveRecord::Base.connection.execute 'CREATE TABLE non_paranoid_models (id INTEGER NOT NULL PRIMARY KEY, parent_model_id INTEGER)'
+  ActiveRecord::Base.connection.execute 'CREATE TABLE polymorphic_models (id INTEGER NOT NULL PRIMARY KEY, parent_id INTEGER, parent_type STRING, deleted_at DATETIME)'
 end
 
 class WithDifferentConnection < ActiveRecord::Base
@@ -654,6 +655,19 @@ def test_model_without_db_connection
     setup!
   end
 
+  def test_restore_recursive_on_polymorphic_has_one_association
+    parent = ParentModel.create
+    polymorphic = PolymorphicModel.create(parent: parent)
+
+    parent.destroy
+
+    assert_equal 0, polymorphic.class.count
+
+    parent.restore(recursive: true)
+
+    assert_equal 1, polymorphic.class.count
+  end
+
   private
   def get_featureful_model
     FeaturefulModel.new(:name => "not empty")
@@ -706,6 +720,7 @@ class ParentModel < ActiveRecord::Base
   has_many :very_related_models, :class_name => 'RelatedModel', dependent: :destroy
   has_many :non_paranoid_models, dependent: :destroy
   has_many :asplode_models, dependent: :destroy
+  has_one :polymorphic_model, as: :parent, dependent: :destroy
 end
 
 class RelatedModel < ActiveRecord::Base
@@ -795,3 +810,8 @@ class AsplodeModel < ActiveRecord::Base
 
 class NoConnectionModel < ActiveRecord::Base
 end
+
+class PolymorphicModel < ActiveRecord::Base
+  acts_as_paranoid
+  belongs_to :parent, polymorphic: true
+end