Merge pull request #19 from lowjoel/fix-custom-stamper-columns

Fix custom stamper columns

lib/active_record/userstamp/stampable.rb

@@ -96,29 +96,32 @@ def has_stamper?
   end
 
   def set_creator_attribute
-    attribute = ActiveRecord::Userstamp.config.creator_attribute
-    return if !has_stamper? || attribute.nil? || !has_attribute?(attribute)
+    return unless has_stamper?
 
-    current_attribute_value = send(attribute)
-    return if current_attribute_value.present?
+    creator_association = self.class.reflect_on_association(:creator)
+    return unless creator_association
+    return if creator.present?
 
-    ActiveRecord::Userstamp::Utilities.assign_attribute(self, attribute)
+    ActiveRecord::Userstamp::Utilities.assign_stamper(self, creator_association)
   end
 
   def set_updater_attribute
-    attribute = ActiveRecord::Userstamp.config.updater_attribute
-    return if !has_stamper? || attribute.nil? || !has_attribute?(attribute)
+    return unless has_stamper?
 
-    return if !self.new_record? && !self.changed?
+    updater_association = self.class.reflect_on_association(:updater)
+    return unless updater_association
+    return if !new_record? && !changed?
 
-    ActiveRecord::Userstamp::Utilities.assign_attribute(self, attribute)
+    ActiveRecord::Userstamp::Utilities.assign_stamper(self, updater_association)
   end
 
   def set_deleter_attribute
-    attribute = ActiveRecord::Userstamp.config.deleter_attribute
-    return if !has_stamper? || attribute.nil? || !has_attribute?(attribute)
+    return unless has_stamper?
 
-    ActiveRecord::Userstamp::Utilities.assign_attribute(self, attribute)
+    deleter_association = self.class.reflect_on_association(:deleter)
+    return unless deleter_association
+
+    ActiveRecord::Userstamp::Utilities.assign_stamper(self, deleter_association)
     save
   end
 end

lib/active_record/userstamp/utilities.rb

@@ -38,18 +38,22 @@ def self.available_association_columns(model)
     nil
   end
 
-  # Assigns the given attribute to the record, based on the model's stamper.
+  # Assigns the stamper to the given association reflection in the record.
   #
-  # If the stamper is a record, then it is assigned to the attribute; if it is a number, then it
-  # is assigned to the +_id+ attribute
+  # If the stamper is a record, then it is assigned to the association; if it is a number, then it
+  # is assigned to the foreign key.
   #
   # @param [ActiveRecord::Base] record The record to assign.
-  # @param [Symbol] attribute The attribute to assign.
-  def self.assign_attribute(record, attribute)
-    attribute = attribute.to_s
+  # @param [ActiveRecord::Reflection] association The association to assign
+  def self.assign_stamper(record, association)
     stamp_value = record.class.stamper_class.stamper
+    attribute =
+      if stamp_value.is_a?(ActiveRecord::Base)
+        association.name
+      else
+        association.foreign_key
+      end
 
-    attribute = attribute[0..-4] if !stamp_value.is_a?(Fixnum) && attribute.end_with?('_id')
     record.send("#{attribute}=", stamp_value)
   end
 end