Skip to content

Commit

Permalink
WIP
Browse files Browse the repository at this point in the history
  • Loading branch information
byroot authored and hcmaATshopify committed Sep 3, 2024
1 parent e19d0d4 commit 2fb0930
Show file tree
Hide file tree
Showing 7 changed files with 140 additions and 45 deletions.
10 changes: 8 additions & 2 deletions activerecord/lib/active_record/associations/association.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 }
Expand Down
10 changes: 10 additions & 0 deletions activerecord/lib/active_record/relation.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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?
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
31 changes: 31 additions & 0 deletions activerecord/test/cases/associations/has_many_associations_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
32 changes: 32 additions & 0 deletions activerecord/test/cases/associations/has_one_associations_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
65 changes: 41 additions & 24 deletions activerecord/test/cases/helper.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 0 additions & 15 deletions activerecord/test/cases/relation/load_async_test.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit 2fb0930

Please sign in to comment.