How to set a custom repo for an action with relationship loads?

I have a read action in my resource that I want to run in my replica repo instead o the main one, to do that I added this prepare to the action:

read :my_action do
  ...
  prepare query, _ ->
    Ash.Query.set_context(query, %{data_layer: %{repo: Core.ReplicaRepo}})
  end

This work, but now let’s say I call this action like this:

MyResource |> Ash.Query.for_read(:my_action, %{}, load: [:load_1, :load_2]) |> Ash.read!()

This will correctly get the data from :my_action using the Core.ReplicaRepo, but it will still use the primary repo for the :load_1 and :load_2 queries.

Can I somehow set a custom repo for these loads as-well (or set a custom read action for them) without needing to change the relationship default read action?

I’m unsure if this works, but I would try this.

You can pass a initial query to use for the load, there you could set the context.

MyResource |> Ash.Query.for_read(:my_action, %{}, load: [
  load_1: Ash.Query.set_context(LoadResource1, %{data_layer: %{repo: Core.ReplicaRepo}}),
load_2: Ash.Query.set_context(LoadResource2, %{data_layer: %{repo: Core.ReplicaRepo}})
]) |> Ash.read!()

I think this is probably similar to this problem: Can I pass context thru manage_relationship?. We might consider adding support for that in the future.