Allow a custom uuid_column name in has_based_uuid

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    based_uuid (0.6.0)
+    based_uuid (0.6.1)
       activerecord (>= 7.0)
       activesupport (>= 7.0)
 

README.md

@@ -73,6 +73,26 @@ def to_param
 end
 ```
 
+### Custom UUID column name
+
+`BasedUUID` will respect the value of `Model.primary_key`, so it supports custom primary key names:
+
+```ruby
+class Transaction < ApplicationRecord
+  self.primary_key = "txid"
+
+  has_based_uuid prefix: :tx
+end
+```
+
+If you want to use a different column, other than the primary key, you can pass it as an option to `has_based_uuid`:
+
+```ruby
+class Session < ApplicationRecord
+  has_based_uuid prefix: :sid, uuid_column: :session_id
+end
+```
+
 ### Use outside ActiveRecord
 
 BasedUUID can be used outside ActiveRecord, too. You can encode any UUID with it:

lib/based_uuid/has_based_uuid.rb

@@ -10,13 +10,16 @@ module HasBasedUUID
 
     included do
       class_attribute :_based_uuid_prefix
+      class_attribute :_based_uuid_column
     end
 
     class_methods do
-      def has_based_uuid(prefix: nil)
+      def has_based_uuid(prefix: nil, uuid_column: primary_key)
         include ModelExtensions
 
         self._based_uuid_prefix = prefix
+        self._based_uuid_column = uuid_column
+
         BasedUUID.register_model_prefix(prefix, self) if prefix
       end
     end
@@ -30,7 +33,7 @@ def find_by_based_uuid(token)
         prefix, uuid_base32 = BasedUUID.split(token)
         raise ArgumentError, "Invalid prefix" if prefix && prefix.to_sym != _based_uuid_prefix
 
-        find_by(primary_key => Base32UUID.decode(uuid_base32))
+        find_by(_based_uuid_column => Base32UUID.decode(uuid_base32))
       end
 
       def find_by_based_uuid!(token)
@@ -39,15 +42,15 @@ def find_by_based_uuid!(token)
     end
 
     def based_uuid(prefix: true)
-      raise ArgumentError, "UUID is empty" if _primary_key_value.blank?
+      raise ArgumentError, "UUID is empty" if _uuid_column_value.blank?
 
-      BasedUUID.encode(uuid: _primary_key_value, prefix: prefix ? self.class._based_uuid_prefix : nil)
+      BasedUUID.encode(uuid: _uuid_column_value, prefix: prefix ? self.class._based_uuid_prefix : nil)
     end
 
     private
 
-    def _primary_key_value
-      self[self.class.primary_key]
+    def _uuid_column_value
+      self[self.class._based_uuid_column]
     end
   end
 end

lib/based_uuid/version.rb

@@ -1,3 +1,3 @@
 module BasedUUID
-  VERSION = "0.6.0".freeze
+  VERSION = "0.6.1".freeze
 end

test/test_has_based_uuid.rb

@@ -37,6 +37,10 @@ class Transaction < FakeActiveRecordBase
   has_based_uuid prefix: :tx
 end
 
+class Session < FakeActiveRecordBase
+  has_based_uuid prefix: :sid, uuid_column: :session_id
+end
+
 class TestHasBasedUUID < Minitest::Test
   def setup
     @user = User.create(id: "59a9608b-bbbf-4d1a-9adc-2dc34875e423")
@@ -57,6 +61,16 @@ def test_based_uuid
     assert_raises(ArgumentError) { User.new.based_uuid }
   end
 
+  def test_based_uuid_for_custom_primary_keys
+    assert_equal "tx_34swbe3m298v4sarc4vpppsvja",
+                 Transaction.new(txid: "64cf16e1-d049-46c9-9561-84ddad6cee4a").based_uuid
+  end
+
+  def test_based_uuid_for_custom_uuid_columns
+    assert_equal "sid_5akm5ysprd930remqpbbhc0pca",
+                 Session.new(session_id: "aa9d0bec-db0d-48c1-8752-f65ae2c0598a").based_uuid
+  end
+
   def test_find_by_based_uuid
     assert_nil User.find_by_based_uuid(random_user.based_uuid)
 
@@ -78,6 +92,14 @@ def test_find_by_based_uuid_for_custom_primary_keys
     end
   end
 
+  def test_find_by_based_uuid_for_custom_uuid_columns
+    session = Session.create(session_id: SecureRandom.uuid)
+
+    Session.find_by_based_uuid(session.based_uuid).tap do |found_session|
+      assert_equal session, found_session
+    end
+  end
+
   private
 
   def random_user