WIP

Author: byroot

activerecord/lib/active_record/associations/association.rb

@@ -191,7 +191,7 @@ def async_load_target
         @target = find_target(async: true) if (@stale_state && stale_target?) || find_target?
 
         loaded! unless loaded?
-        target
+        @target
       end
 
       # We can't dump @reflection and @through_reflection since it contains the scope proc
@@ -243,7 +243,13 @@ def find_target(async: false)
           end
 
           scope = self.scope
-          return scope.to_a if skip_statement_cache?(scope)
+          if skip_statement_cache?(scope)
+            if async
+              return scope.load_async.then(&:to_a)
+            else
+              return scope.to_a 
+            end
+          end
 
           sc = reflection.association_scope_cache(klass, owner) do |params|
             as = AssociationScope.create { params.bind }

activerecord/lib/active_record/relation.rb

@@ -1150,6 +1150,16 @@ def load_async
       self
     end
 
+    def then(&block)
+      if @future_result
+        @future_result.then do
+          yield self
+        end
+      else
+        super
+      end
+    end
+
     # Returns <tt>true</tt> if the relation was scheduled on the background
     # thread pool.
     def scheduled?

activerecord/test/cases/associations/belongs_to_associations_test.rb

@@ -1841,19 +1841,33 @@ def test_destroy_linked_models
 end
 
 class AsyncBelongsToAssociationsTest < ActiveRecord::TestCase
+  include WaitForAsyncTestHelper
+
   fixtures :companies
 
   self.use_transactional_tests = false
 
-  def test_temp_async_load_belongs_to
-    # TODO: proper test?
+  def test_async_load_belongs_to
     client = Client.find(3)
     first_firm = companies(:first_firm)
-    assert_queries_match(/LIMIT|ROWNUM <=|FETCH FIRST/) do
-      client.association(:firm).async_load_target
 
+    promise = client.association(:firm).async_load_target
+    wait_for_async_query
+
+    events = []
+    callback = -> (event) do
+      events << event unless event.payload[:name] == "SCHEMA"
+    end
+    ActiveSupport::Notifications.subscribed(callback, "sql.active_record") do
+      client.firm
+    end
+
+    assert_no_queries do
       assert_equal first_firm, client.firm
       assert_equal first_firm.name, client.firm.name
     end
+
+    assert_equal 1, events.size
+    assert_equal true, events.first.payload[:async]
   end
 end

activerecord/test/cases/associations/has_many_associations_test.rb

@@ -3252,3 +3252,34 @@ def force_signal37_to_load_all_clients_of_firm
       companies(:first_firm).clients_of_firm.load_target
     end
 end
+
+class AsyncHasOneAssociationsTest < ActiveRecord::TestCase
+  include WaitForAsyncTestHelper
+
+  fixtures :companies
+
+  self.use_transactional_tests = false
+
+  def test_async_load_has_many
+    firm = companies(:first_firm)
+
+    promise = firm.association(:clients).async_load_target
+    wait_for_async_query
+
+    events = []
+    callback = -> (event) do
+      events << event unless event.payload[:name] == "SCHEMA"
+    end
+
+    ActiveSupport::Notifications.subscribed(callback, "sql.active_record") do
+      assert_equal 3, firm.clients.size
+    end
+
+    assert_no_queries do
+      assert_not_nil firm.clients[2]
+    end
+
+    assert_equal 1, events.size
+    assert_equal true, events.first.payload[:async]
+  end
+end

activerecord/test/cases/associations/has_one_associations_test.rb

@@ -943,3 +943,35 @@ def test_has_one_with_touch_option_on_nonpersisted_built_associations_doesnt_upd
     MESSAGE
   end
 end
+
+class AsyncHasOneAssociationsTest < ActiveRecord::TestCase
+  include WaitForAsyncTestHelper
+
+  fixtures :companies, :accounts
+
+  self.use_transactional_tests = false
+
+  def test_async_load_has_one
+    firm = companies(:first_firm)
+    first_account = Account.find(1)
+
+    promise = firm.association(:account).async_load_target
+    wait_for_async_query
+
+    events = []
+    callback = -> (event) do
+      events << event unless event.payload[:name] == "SCHEMA"
+    end
+    ActiveSupport::Notifications.subscribed(callback, "sql.active_record") do
+      firm.account
+    end
+
+    assert_no_queries do
+      assert_equal first_account, firm.account
+      assert_equal first_account.credit_limit, firm.account.credit_limit
+    end
+
+    assert_equal 1, events.size
+    assert_equal true, events.first.payload[:async]
+  end
+end

activerecord/test/cases/helper.rb

@@ -45,36 +45,53 @@
 ActiveRecord::ConnectionAdapters.register("abstract", "ActiveRecord::ConnectionAdapters::AbstractAdapter", "active_record/connection_adapters/abstract_adapter")
 ActiveRecord::ConnectionAdapters.register("fake", "FakeActiveRecordAdapter", File.expand_path("../support/fake_adapter.rb", __dir__))
 
-class SQLSubscriber
-  attr_reader :logged
-  attr_reader :payloads
+class ActiveRecord::TestCase
+  class SQLSubscriber
+    attr_reader :logged
+    attr_reader :payloads
+
+    def initialize
+      @logged = []
+      @payloads = []
+    end
+
+    def start(name, id, payload)
+      @payloads << payload
+      @logged << [payload[:sql].squish, payload[:name], payload[:binds]]
+    end
 
-  def initialize
-    @logged = []
-    @payloads = []
+    def finish(name, id, payload); end
   end
 
-  def start(name, id, payload)
-    @payloads << payload
-    @logged << [payload[:sql].squish, payload[:name], payload[:binds]]
+  module InTimeZone
+    private
+      def in_time_zone(zone)
+        old_zone  = Time.zone
+        old_tz    = ActiveRecord::Base.time_zone_aware_attributes
+
+        Time.zone = zone ? ActiveSupport::TimeZone[zone] : nil
+        ActiveRecord::Base.time_zone_aware_attributes = !zone.nil?
+        yield
+      ensure
+        Time.zone = old_zone
+        ActiveRecord::Base.time_zone_aware_attributes = old_tz
+      end
   end
 
-  def finish(name, id, payload); end
-end
+  module WaitForAsyncTestHelper
+    private
+      def wait_for_async_query(connection = ActiveRecord::Base.lease_connection, timeout: 5)
+        return unless connection.async_enabled?
 
-module InTimeZone
-  private
-    def in_time_zone(zone)
-      old_zone  = Time.zone
-      old_tz    = ActiveRecord::Base.time_zone_aware_attributes
-
-      Time.zone = zone ? ActiveSupport::TimeZone[zone] : nil
-      ActiveRecord::Base.time_zone_aware_attributes = !zone.nil?
-      yield
-    ensure
-      Time.zone = old_zone
-      ActiveRecord::Base.time_zone_aware_attributes = old_tz
-    end
+        executor = connection.pool.async_executor
+        (timeout * 100).times do
+          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
 end
 
 # Encryption

activerecord/test/cases/relation/load_async_test.rb

@@ -7,21 +7,6 @@
 require "models/other_dog"
 
 module ActiveRecord
-  module WaitForAsyncTestHelper
-    private
-      def wait_for_async_query(connection = ActiveRecord::Base.lease_connection, timeout: 5)
-        return unless connection.async_enabled?
-
-        executor = connection.pool.async_executor
-        (timeout * 100).times do
-          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
-
   class LoadAsyncTest < ActiveRecord::TestCase
     include WaitForAsyncTestHelper