find_by through collection proxy

demo_functionality.rb

@@ -147,6 +147,10 @@ raise unless fa_posts_all.first.creator_id == fa_user.id
 raise unless FaPost.exists?(fa_id: 12_345)
 raise if FaPost.exists?(fa_id: 12_346)
 
+posts = fa_user.created_posts
+fa_post2 = posts.find_by(fa_id: 12_345)
+raise unless fa_post2.id == fa_post.id
+
 e621_user = E621User.create!(username: "bob", e621_id: 67_890)
 raise unless e621_user.persisted?
 raise unless e621_user.username == "bob"

lib/has_aux_table/relation_extensions.rb

@@ -49,12 +49,22 @@ module HasAuxTable
       relation_class =
         main_class.relation_delegate_class(ActiveRecord::Relation)
 
-      Util.hook_method(relation_class, :where!, true) do |original, opts, *rest|
-        if opts.is_a?(Hash)
-          opts_remapped = aux_config.remap_conditions(opts)
-          T.unsafe(original).call(opts_remapped, *rest)
-        else
-          T.unsafe(original).call(opts, *rest)
+      collection_proxy_class =
+        main_class.relation_delegate_class(
+          ActiveRecord::Associations::CollectionProxy
+        )
+
+      [
+        [relation_class, :where!],
+        [collection_proxy_class, :where]
+      ].each do |klass, method_name|
+        Util.hook_method(klass, method_name, true) do |original, opts, *rest|
+          if opts.is_a?(Hash)
+            opts_remapped = aux_config.remap_conditions(opts)
+            T.unsafe(original).call(opts_remapped, *rest)
+          else
+            T.unsafe(original).call(opts, *rest)
+          end
         end
       end
 

spec/active_record/aux_table_spec.rb

@@ -305,6 +305,12 @@ RSpec.describe HasAuxTable do
         found_car = Car.find_by(fuel_type: "diesel")
         expect(found_car).to be_nil
       end
+
+      it "works with find_by!" do
+        expect { Car.find_by!(fuel_type: "diesel") }.to raise_error(
+          ActiveRecord::RecordNotFound
+        )
+      end
     end
 
     describe "where method with automatic joins" do
@@ -722,6 +728,18 @@ RSpec.describe HasAuxTable do
         driver = @car.drivers.create!(name: "John Doe")
         expect { driver.destroy }.to change { @car.reload.drivers.count }.by(-1)
       end
+
+      it "can be queried through the association" do
+        driver = @car.drivers.create!(name: "John Doe", license_number: 123_456)
+        expect(@car.drivers.where(name: "John Doe")).to eq([driver])
+
+        drivers = @car.drivers
+        d = drivers.find_by!(license_number: 123_456)
+        expect(d.id).to eq(driver.id)
+
+        d = drivers.find_by(license_number: 123_456)
+        expect(d.id).to eq(driver.id)
+      end
     end
 
     describe "#reload" do