From dd6fd1fbc372c3ed89958ff0fe5d35e7c84070c0 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Tue, 4 Oct 2022 00:22:40 -0400
Subject: [PATCH] improvement: allow select/load callbacks for calcs to return
 irrelevant keys test: add a test case for nested required calculations

---
 lib/ash/query/query.ex    |  6 ++++++
 test/calculation_test.exs | 22 ++++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/lib/ash/query/query.ex b/lib/ash/query/query.ex
index a5ec259a..0de343ee 100644
--- a/lib/ash/query/query.ex
+++ b/lib/ash/query/query.ex
@@ -838,6 +838,7 @@ defmodule Ash.Query do
                 |> Kernel.||([])
                 |> Enum.concat(module.select(query, opts, calculation.context) || [])
                 |> Enum.uniq()
+                |> Enum.filter(&Ash.Resource.Info.attribute(query.resource, &1))
 
               loads =
                 module.load(
@@ -847,6 +848,7 @@ defmodule Ash.Query do
                 )
                 |> Ash.Actions.Helpers.validate_calculation_load!(module)
                 |> Enum.concat(List.wrap(resource_calculation.load))
+                |> Enum.reject(&Ash.Resource.Info.attribute(query.resource, &1))
 
               calculation = %{
                 calculation
@@ -933,6 +935,7 @@ defmodule Ash.Query do
             |> Kernel.||([])
             |> Enum.concat(module.select(query, opts, calculation.context) || [])
             |> Enum.uniq()
+            |> Enum.filter(&Ash.Resource.Info.attribute(query.resource, &1))
 
           loads =
             module.load(
@@ -942,6 +945,7 @@ defmodule Ash.Query do
             )
             |> Ash.Actions.Helpers.validate_calculation_load!(module)
             |> Enum.concat(List.wrap(resource_calculation.load))
+            |> Enum.reject(&Ash.Resource.Info.attribute(query.resource, &1))
 
           calculation = %{
             calculation
@@ -1421,6 +1425,7 @@ defmodule Ash.Query do
           |> module.select(opts, calculation.context)
           |> List.wrap()
           |> Enum.concat(calculation.select)
+          |> Enum.filter(&Ash.Resource.Info.attribute(query.resource, &1))
 
         loads =
           module.load(
@@ -1430,6 +1435,7 @@ defmodule Ash.Query do
           )
           |> Ash.Actions.Helpers.validate_calculation_load!(module)
           |> Enum.concat(List.wrap(calculation.required_loads))
+          |> Enum.reject(&Ash.Resource.Info.attribute(query.resource, &1))
 
         calculation = %{calculation | select: fields_to_select, required_loads: loads}
         %{query | calculations: Map.put(query.calculations, name, calculation)}
diff --git a/test/calculation_test.exs b/test/calculation_test.exs
index f8c89728..f50c820a 100644
--- a/test/calculation_test.exs
+++ b/test/calculation_test.exs
@@ -39,6 +39,10 @@ defmodule Ash.Test.CalculationTest do
       opts[:keys]
     end
 
+    def select(_query, opts, _) do
+      opts[:keys]
+    end
+
     def calculate(records, opts, _) do
       Enum.map(records, fn record ->
         Enum.map_join(opts[:keys], " ", fn key ->
@@ -94,6 +98,10 @@ defmodule Ash.Test.CalculationTest do
                 :string,
                 {ConcatWithLoad, keys: [:full_name, :full_name]}
 
+      calculate :full_name_plus_full_name_plus_full_name,
+                :string,
+                {ConcatWithLoad, keys: [:full_name, :full_name_plus_full_name]}
+
       calculate :slug, :string, expr(full_name <> "123"), load: [:full_name]
 
       calculate :expr_full_name, :string, expr(first_name <> " " <> last_name) do
@@ -202,6 +210,20 @@ defmodule Ash.Test.CalculationTest do
     assert full_names == ["brian cranston brian cranston", "zach daniel zach daniel"]
   end
 
+  test "it loads anything specified by the load callback, even when nested" do
+    full_names =
+      User
+      |> Ash.Query.load(:full_name_plus_full_name_plus_full_name)
+      |> Api.read!()
+      |> Enum.map(& &1.full_name_plus_full_name_plus_full_name)
+      |> Enum.sort()
+
+    assert full_names == [
+             "brian cranston brian cranston brian cranston",
+             "zach daniel zach daniel zach daniel"
+           ]
+  end
+
   test "it doesn't reload anything specified by the load callback if its already been loaded when using `lazy?: true`" do
     full_names =
       User