From 362180190ca8c3bc6cd8a5746daca01d1cd0665b Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Thu, 13 Oct 2022 00:36:29 -0400
Subject: [PATCH] fix: properly load calcs/aggs on manual relationships

It needs to be optimized, but it will do the trick for now.
---
 lib/ash/actions/load.ex          | 26 +++++++++++++++++++-
 lib/ash/flow/step/transaction.ex |  2 +-
 test/actions/async_load_test.exs | 41 ++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/lib/ash/actions/load.ex b/lib/ash/actions/load.ex
index 442deed8..4630a85d 100644
--- a/lib/ash/actions/load.ex
+++ b/lib/ash/actions/load.ex
@@ -510,7 +510,8 @@ defmodule Ash.Actions.Load do
         related_query,
         request_opts,
         fn ->
-          mod.load(data, opts, %{
+          data
+          |> mod.load(opts, %{
             relationship: relationship,
             query: related_query,
             root_query: root_query,
@@ -519,6 +520,29 @@ defmodule Ash.Actions.Load do
             api: related_query.api,
             tenant: related_query.tenant
           })
+          |> case do
+            {:ok, result} ->
+              # TODO: this will result in quite a few requests potentially for aggs/calcs
+              # This should be optimized.
+              Enum.reduce_while(result, {:ok, %{}}, fn {key, records}, {:ok, acc} ->
+                case related_query.api.load(records, %{related_query | load: []},
+                       lazy?: true,
+                       tenant: related_query.tenant,
+                       actor: request_opts[:actor],
+                       authorize?: request_opts[:authorize?],
+                       tracer: request_opts[:tracer]
+                     ) do
+                  {:ok, results} ->
+                    {:cont, {:ok, Map.put(acc, key, results)}}
+
+                  {:error, error} ->
+                    {:halt, {:error, error}}
+                end
+              end)
+
+            {:error, error} ->
+              {:error, error}
+          end
         end
       )
     end)
diff --git a/lib/ash/flow/step/transaction.ex b/lib/ash/flow/step/transaction.ex
index 7338bcc7..f5a0c988 100644
--- a/lib/ash/flow/step/transaction.ex
+++ b/lib/ash/flow/step/transaction.ex
@@ -9,7 +9,7 @@ defmodule Ash.Flow.Step.Transaction do
         output: [
           type: :any,
           doc:
-            "Which step or steps to use when constructing the output list. Defaults to the last step.",
+            "Which step or steps to use when constructing the output. Defaults to the last step.",
           links: []
         ],
         timeout: [
diff --git a/test/actions/async_load_test.exs b/test/actions/async_load_test.exs
index 23af9f9b..0e28b919 100644
--- a/test/actions/async_load_test.exs
+++ b/test/actions/async_load_test.exs
@@ -99,6 +99,10 @@ defmodule Ash.Test.Actions.AsyncLoadTest do
         source_attribute_on_join_resource: :post_id
     end
 
+    calculations do
+      calculate :title_plus_title, :string, expr((title || "foo") <> (title || "bar"))
+    end
+
     policies do
       policy action(:authorized_actor) do
         authorize_if expr(actor_id == ^actor(:id))
@@ -316,6 +320,43 @@ defmodule Ash.Test.Actions.AsyncLoadTest do
                |> Enum.flat_map(&Map.get(&1, :posts_in_same_category))
     end
 
+    test "it allows loading calculations on and through manual relationships" do
+      post1 =
+        Post
+        |> new(%{title: "post1", category: "foo"})
+        |> Api.create!()
+
+      Post
+      |> new(%{title: "post2", category: "bar"})
+      |> Api.create!()
+
+      Post
+      |> new(%{title: "post2", category: "foo"})
+      |> Api.create!()
+
+      assert %Post{
+               title: "post1",
+               title_plus_title: "post1post1",
+               posts_in_same_category: [
+                 %{
+                   title: "post2",
+                   title_plus_title: "post2post2",
+                   posts_in_same_category: [
+                     %{title: "post1", title_plus_title: "post1post1"}
+                   ]
+                 }
+               ]
+             } =
+               post1
+               |> Api.load!([
+                 :title_plus_title,
+                 posts_in_same_category: [
+                   :title_plus_title,
+                   posts_in_same_category: [:title_plus_title]
+                 ]
+               ])
+    end
+
     test "it allows loading related data" do
       author =
         Author