Faster way to access prefix IDs for associated models

lib/prefixed_ids.rb

@@ -16,10 +16,10 @@ class Error < StandardError; end
   mattr_accessor :models, default: {}
 
   def self.find(prefix_id)
-    prefix, _ = split_id(prefix_id)
+    prefix, = split_id(prefix_id)
     models.fetch(prefix).find_by_prefix_id(prefix_id)
   rescue KeyError
-    raise Error, "Unable to find model with prefix `#{prefix}`. Available prefixes are: #{models.keys.join(", ")}"
+    raise Error, "Unable to find model with prefix `#{prefix}`. Available prefixes are: #{models.keys.join(', ')}"
   end
 
   # Splits a prefixed ID into its prefix and ID
@@ -42,6 +42,7 @@ def has_prefix_id(prefix, override_find: true, override_param: true, fallback: t
         include Attribute
         include Finder if override_find
         include ToParam if override_param
+
         self._prefix_id = PrefixId.new(self, prefix, **options)
         self._prefix_id_fallback = fallback
 
@@ -97,6 +98,7 @@ def find(*ids)
         prefix_ids = ids.flatten.map do |id|
           prefix_id = _prefix_id.decode(id, fallback: _prefix_id_fallback)
           raise Error, "#{id} is not a valid prefix_id" if !_prefix_id_fallback && prefix_id.nil?
+
           prefix_id
         end
         prefix_ids = [prefix_ids] if ids.first.is_a?(Array)
@@ -108,6 +110,29 @@ def relation
         super.tap { |r| r.extend ClassMethods }
       end
 
+      def belongs_to(*args, **options, &)
+        association = super
+
+        name = args.first
+        reflection = association[name]
+
+        return association if reflection.polymorphic?
+        return association if reflection.klass._prefix_id.blank?
+
+        generated_association_methods.class_eval <<-CODE, __FILE__, __LINE__ + 1
+          def #{name}_prefix_id
+            #{reflection.klass}._prefix_id.encode(#{reflection.foreign_key})
+          end
+
+          def #{name}_prefix_id=(prefix_id)
+            decoded_id = #{reflection.klass}._prefix_id.decode(prefix_id, fallback: #{reflection.klass}._prefix_id_fallback)
+            send("#{reflection.foreign_key}=", decoded_id)
+          end
+        CODE
+
+        association
+      end
+
       def has_many(*args, &block)
         options = args.extract_options!
         options[:extend] = Array(options[:extend]).push(ClassMethods)

test/dummy/app/models/post.rb

@@ -1,4 +1,5 @@
 class Post < ApplicationRecord
   has_prefix_id :post, override_exists: false
   belongs_to :user
+  belongs_to :nonprefixed_item, optional: true
 end

test/dummy/app/models/tag.rb

@@ -0,0 +1,5 @@
+class Tag < ApplicationRecord
+  has_prefix_id :tag
+
+  belongs_to :taggable, polymorphic: true
+end

test/dummy/db/migrate/20210503145247_create_posts.rb

@@ -2,6 +2,7 @@ class CreatePosts < ActiveRecord::Migration[6.1]
   def change
     create_table :posts do |t|
       t.references :user, null: false, foreign_key: true
+      t.references :nonprefixed_item, null: true, foreign_key: true
       t.timestamps
     end
   end

test/dummy/db/migrate/20241113200714_create_tags.rb

@@ -0,0 +1,8 @@
+class CreateTags < ActiveRecord::Migration[7.2]
+  def change
+    create_table :tags do |t|
+      t.references :taggable, polymorphic: true, null: false
+      t.timestamps
+    end
+  end
+end

test/dummy/db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema.define(version: 2024_07_14_120000) do
+ActiveRecord::Schema[7.2].define(version: 2024_11_13_200714) do
   create_table "accounts", force: :cascade do |t|
     t.integer "user_id"
     t.datetime "created_at", null: false
@@ -32,11 +32,21 @@
 
   create_table "posts", force: :cascade do |t|
     t.integer "user_id", null: false
+    t.integer "nonprefixed_item_id"
     t.datetime "created_at", null: false
     t.datetime "updated_at", null: false
+    t.index ["nonprefixed_item_id"], name: "index_posts_on_nonprefixed_item_id"
     t.index ["user_id"], name: "index_posts_on_user_id"
   end
 
+  create_table "tags", force: :cascade do |t|
+    t.string "taggable_type", null: false
+    t.integer "taggable_id", null: false
+    t.datetime "created_at", null: false
+    t.datetime "updated_at", null: false
+    t.index ["taggable_type", "taggable_id"], name: "index_tags_on_taggable"
+  end
+
   create_table "teams", force: :cascade do |t|
     t.datetime "created_at", null: false
     t.datetime "updated_at", null: false
@@ -47,5 +57,6 @@
     t.datetime "updated_at", null: false
   end
 
+  add_foreign_key "posts", "nonprefixed_items"
   add_foreign_key "posts", "users"
 end

test/dummy/test/fixtures/tags.yml

@@ -0,0 +1,11 @@
+# Read about fixtures at https://api.rubyonrails.org/classes/ActiveRecord/FixtureSet.html
+
+# This model initially had no columns defined. If you add columns to the
+# model remove the "{}" from the fixture names and add the columns immediately
+# below each fixture, per the syntax in the comments below
+#
+one:
+  taggable: one (User)
+
+two:
+  taggable: one (Post)

test/dummy/test/models/tag_test.rb

@@ -0,0 +1,7 @@
+require "test_helper"
+
+class TagTest < ActiveSupport::TestCase
+  # test "the truth" do
+  #   assert true
+  # end
+end

test/prefixed_ids_test.rb

@@ -24,7 +24,7 @@ class PrefixedIdsTest < ActiveSupport::TestCase
   test "can get prefix IDs from multiple original IDs" do
     assert_equal(
       [users(:one).prefix_id, users(:two).prefix_id, users(:three).prefix_id],
-      User.prefix_ids([users(:one).id, users(:two).id, users(:three).id])
+      User.prefix_ids([users(:one).id, users(:two).id, users(:three).id]),
     )
   end
 
@@ -35,7 +35,7 @@ class PrefixedIdsTest < ActiveSupport::TestCase
   test "can get original IDs from multiple prefix IDs" do
     assert_equal(
       [users(:one).id, users(:two).id, users(:three).id],
-      User.decode_prefix_ids([users(:one).prefix_id, users(:two).prefix_id, users(:three).prefix_id])
+      User.decode_prefix_ids([users(:one).prefix_id, users(:two).prefix_id, users(:three).prefix_id]),
     )
   end
 
@@ -203,6 +203,47 @@ class PrefixedIdsTest < ActiveSupport::TestCase
     assert_nil Post.new.to_param
   end
 
+  test "helper for getting prefix ID from belongs to models when associated model has prefix ID" do
+    post = posts(:one)
+    assert_equal post.user_prefix_id, post.user.prefix_id
+  end
+
+  test "no helper for getting prefix ID from belongs to models when associated model does not have prefix ID" do
+    post = posts(:one)
+    assert_raises NoMethodError do
+      post.nonprefixed_item_prefix_id
+    end
+  end
+
+  test "setter for using prefix ID while creating models with mass assignment" do
+    post = Post.create!(user_prefix_id: users(:two).prefix_id)
+    assert_equal users(:two), post.user
+  end
+
+  test "no setter for using prefix ID while creating models with mass assignment when associated model does not have prefix ID" do
+    assert_raises ActiveModel::UnknownAttributeError do
+      Post.create!(user: users(:two), nonprefixed_item_prefix_id: "abc123")
+    end
+  end
+
+  test "setter for belongs to models that have prefix IDs" do
+    post = Post.new
+    post.user_prefix_id = users(:two).prefix_id
+    assert_equal users(:two), post.user
+  end
+
+  test "setter not created on models without has_prefix_id" do
+    assert_raises NoMethodError do
+      NonprefixedItem.new.user_prefix_id
+    end
+  end
+
+  test "setter not created on polymorphic belongs to models" do
+    assert_raises NoMethodError do
+      Tag.new.taggable_prefix_id
+    end
+  end
+
   if PrefixedIds::Test.rails71_and_up?
     test "compound primary - can get prefix ID from original ID" do
       assert compound_primary_items(:one).id.is_a?(Array)