From 53b603b777147f7012004f33022cd3e65c068cc1 Mon Sep 17 00:00:00 2001
From: Nate Matykiewicz <natematykiewicz@gmail.com>
Date: Thu, 9 May 2024 12:52:41 -0500
Subject: [PATCH] Added option to validate which models are allowed to be
 associated to a polymorphic belongs_to

---
 activerecord/CHANGELOG.md                     |  9 ++++++++
 .../associations/builder/belongs_to.rb        |  6 +++++
 .../belongs_to_associations_test.rb           | 22 +++++++++++++++++++
 guides/source/association_basics.md           | 10 +++++++++
 4 files changed, 47 insertions(+)

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index aa47f9c759373..8b75f4c201c48 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,2 +1,11 @@
+*   Added option to validate which models are allowed to be associated to a polymorphic `belongs_to`.
+
+    Ex:
+    ```ruby
+    belongs_to :commentable, polymorphic: ["Post", "Comment"]
+    ```
+    This automatically adds an inclusion validator to commentable_type, to ensure that it's either Post or Comment.
+
+    *Nate Matykiewicz*
 
 Please check [7-2-stable](https://github.com/rails/rails/blob/7-2-stable/activerecord/CHANGELOG.md) for previous changes.
diff --git a/activerecord/lib/active_record/associations/builder/belongs_to.rb b/activerecord/lib/active_record/associations/builder/belongs_to.rb
index aa5a7c3dace8c..908d0551bc4a4 100644
--- a/activerecord/lib/active_record/associations/builder/belongs_to.rb
+++ b/activerecord/lib/active_record/associations/builder/belongs_to.rb
@@ -121,6 +121,12 @@ def self.define_validations(model, reflection)
         required = !reflection.options[:optional]
       end
 
+      if reflection.options[:polymorphic].is_a?(Array)
+        model.validates reflection.foreign_type,
+          inclusion: { in: reflection.options[:polymorphic].map { |klass| klass.to_s.freeze } },
+          if: ->(record) { record.read_attribute(reflection.foreign_type) }
+      end
+
       super
 
       if required
diff --git a/activerecord/test/cases/associations/belongs_to_associations_test.rb b/activerecord/test/cases/associations/belongs_to_associations_test.rb
index 834a55c181bc8..23fc4d8278a68 100644
--- a/activerecord/test/cases/associations/belongs_to_associations_test.rb
+++ b/activerecord/test/cases/associations/belongs_to_associations_test.rb
@@ -1374,6 +1374,28 @@ def test_polymorphic_counter_cache
     end
   end
 
+  class TestPolymorphicArrayAllowed < ActiveRecord::Base
+    self.table_name = "wheels"
+    belongs_to :wheelable, polymorphic: ["Book", :Essay], optional: false
+  end
+
+  def test_polymorphic_with_allowed_class_string_array
+    wheel = TestPolymorphicArrayAllowed.new
+
+    assert_not_predicate wheel, :valid?
+    assert_equal ["Wheelable must exist"], wheel.errors.full_messages
+
+    wheel.wheelable = Book.create!
+    assert_predicate wheel, :valid?
+
+    wheel.wheelable = Essay.create!
+    assert_predicate wheel, :valid?
+
+    wheel.wheelable = Citation.create!
+    assert_not_predicate wheel, :valid?
+    assert_equal ["Wheelable type is not included in the list"], wheel.errors.full_messages
+  end
+
   def test_polymorphic_with_custom_foreign_type
     sponsor = sponsors(:moustache_club_sponsor_for_groucho)
     groucho = members(:groucho)
diff --git a/guides/source/association_basics.md b/guides/source/association_basics.md
index 13de54f91d883..91aeaf8e6c9f2 100644
--- a/guides/source/association_basics.md
+++ b/guides/source/association_basics.md
@@ -538,6 +538,16 @@ class CreatePictures < ActiveRecord::Migration[7.2]
 end
 ```
 
+In this example, `imageable` can have any model assigned to it. If you wanted to ensure that the `imageable_type` is only `Employee` or `Product`, you can set the `polymorphic` option to an array of allowed class names. This will [add an inclusion validator][inclusion validator] to the `imageable_type` column, which adds a validation error if `imageable` is not one of the expected classes:
+
+[inclusion validator]: active_record_validations.html#inclusion
+
+```ruby
+class Picture < ApplicationRecord
+  belongs_to :imageable, polymorphic: ["Employee", "Product"]
+end
+```
+
 NOTE: Since polymorphic associations rely on storing class names in the
 database, that data must remain synchronized with the class name used by the
 Ruby code. When renaming a class, make sure to update the data in the