Merge pull request rails#45440 from Shopify/association-proxy-load-async

Fix CollectionProxy#load_async

Author: byroot

activerecord/lib/active_record/associations/collection_proxy.rb

@@ -1108,7 +1108,7 @@ def inspect # :nodoc:
       ].flat_map { |klass|
         klass.public_instance_methods(false)
       } - self.public_instance_methods(false) - [:select] + [
-        :scoping, :values, :insert, :insert_all, :insert!, :insert_all!, :upsert, :upsert_all
+        :scoping, :values, :insert, :insert_all, :insert!, :insert_all!, :upsert, :upsert_all, :load_async
       ]
 
       delegate(*delegate_methods, to: :scope)

activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb

@@ -364,7 +364,7 @@ def shared_cache? # :nodoc:
       end
 
       def get_database_version # :nodoc:
-        SQLite3Adapter::Version.new(query_value("SELECT sqlite_version(*)"))
+        SQLite3Adapter::Version.new(query_value("SELECT sqlite_version(*)", "SCHEMA"))
       end
 
       def check_version # :nodoc:

activerecord/test/cases/relation/load_async_test.rb

@@ -2,22 +2,22 @@
 
 require "cases/helper"
 require "models/post"
+require "models/category"
 require "models/comment"
 require "models/other_dog"
 
 module ActiveRecord
   module WaitForAsyncTestHelper
     private
-      def wait_for_async_query(relation, timeout: 5)
-        if !relation.connection.async_enabled? || relation.instance_variable_get(:@records)
-          return relation
-        end
+      def wait_for_async_query(connection = ActiveRecord::Base.connection, timeout: 5)
+        return unless connection.async_enabled?
 
-        future_result = relation.instance_variable_get(:@future_result)
+        executor = connection.pool.async_executor
         (timeout * 100).times do
-          return relation unless future_result.pending?
+          return unless executor.scheduled_task_count > executor.completed_task_count
           sleep 0.01
         end
+
         raise Timeout::Error, "The async executor wasn't drained after #{timeout} seconds"
       end
   end
@@ -27,7 +27,7 @@ class LoadAsyncTest < ActiveRecord::TestCase
 
     self.use_transactional_tests = false
 
-    fixtures :posts, :comments
+    fixtures :posts, :comments, :categories, :categories_posts
 
     def test_scheduled?
       deferred_posts = Post.where(author_id: 1).load_async
@@ -52,6 +52,48 @@ def test_reset
       assert_not_predicate deferred_posts, :scheduled?
     end
 
+    unless in_memory_db?
+      def test_load_async_has_many_association
+        post = Post.first
+
+        defered_comments = post.comments.load_async
+        assert_predicate defered_comments, :scheduled?
+
+        events = []
+        callback = -> (event) do
+          events << event unless event.payload[:name] == "SCHEMA"
+        end
+
+        wait_for_async_query
+        ActiveSupport::Notifications.subscribed(callback, "sql.active_record") do
+          defered_comments.to_a
+        end
+
+        assert_equal [["Comment Load", true]], events.map { |e| [e.payload[:name], e.payload[:async]] }
+        assert_not_predicate post.comments, :loaded?
+      end
+
+      def test_load_async_has_many_through_association
+        post = Post.first
+
+        defered_categories = post.scategories.load_async
+        assert_predicate defered_categories, :scheduled?
+
+        events = []
+        callback = -> (event) do
+          events << event unless event.payload[:name] == "SCHEMA"
+        end
+
+        wait_for_async_query
+        ActiveSupport::Notifications.subscribed(callback, "sql.active_record") do
+          defered_categories.to_a
+        end
+
+        assert_equal [["Category Load", true]], events.map { |e| [e.payload[:name], e.payload[:async]] }
+        assert_not_predicate post.scategories, :loaded?
+      end
+    end
+
     def test_notification_forwarding
       expected_records = Post.where(author_id: 1).to_a
 
@@ -66,7 +108,8 @@ def test_notification_forwarding
         end
       end
 
-      deferred_posts = wait_for_async_query(Post.where(author_id: 1).load_async)
+      deferred_posts = Post.where(author_id: 1).load_async
+      wait_for_async_query
 
       assert_equal expected_records, deferred_posts.to_a
       assert_equal Post.connection.supports_concurrent_connections?, status[:async]
@@ -92,7 +135,8 @@ def test_simple_query
         end
       end
 
-      deferred_posts = wait_for_async_query(Post.where(author_id: 1).load_async)
+      deferred_posts = Post.where(author_id: 1).load_async
+      wait_for_async_query
 
       assert_equal expected_records, deferred_posts.to_a
       assert_equal Post.connection.supports_concurrent_connections?, status[:async]
@@ -130,7 +174,8 @@ def test_eager_loading_query
         end
       end
 
-      deferred_posts = wait_for_async_query(Post.where(author_id: 1).eager_load(:comments).load_async)
+      deferred_posts = Post.where(author_id: 1).eager_load(:comments).load_async
+      wait_for_async_query
 
       if in_memory_db?
         assert_not_predicate deferred_posts, :scheduled?
@@ -355,7 +400,8 @@ def test_simple_query
           end
         end
 
-        deferred_posts = wait_for_async_query(Post.where(author_id: 1).load_async)
+        deferred_posts = Post.where(author_id: 1).load_async
+        wait_for_async_query
 
         assert_equal expected_records, deferred_posts.to_a
         assert_equal Post.connection.supports_concurrent_connections?, status[:async]
@@ -388,7 +434,8 @@ def test_eager_loading_query
           end
         end
 
-        deferred_posts = wait_for_async_query(Post.where(author_id: 1).eager_load(:comments).load_async)
+        deferred_posts = Post.where(author_id: 1).eager_load(:comments).load_async
+        wait_for_async_query
 
         assert_predicate deferred_posts, :scheduled?
 
@@ -497,8 +544,8 @@ def test_simple_query
         deferred_posts = Post.where(author_id: 1).load_async
         deferred_dogs = OtherDog.where(id: 1).load_async
 
-        wait_for_async_query(deferred_posts)
-        wait_for_async_query(deferred_dogs)
+        wait_for_async_query
+        wait_for_async_query
 
         assert_equal expected_records, deferred_posts.to_a
         assert_equal expected_dogs, deferred_dogs.to_a