From d95a0e1662d437f374d48705c56458b64ca6ccdf Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Thu, 22 Jul 2021 15:22:53 -0400
Subject: [PATCH] fix: load calculations from sorts properly

---
 lib/ash/query/aggregate.ex | 18 +++++++++++-------
 lib/ash/query/query.ex     | 39 ++++++++++++++++++++++++++++++++------
 2 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/lib/ash/query/aggregate.ex b/lib/ash/query/aggregate.ex
index aefa3ab6..8efe494b 100644
--- a/lib/ash/query/aggregate.ex
+++ b/lib/ash/query/aggregate.ex
@@ -466,14 +466,18 @@ defmodule Ash.Query.Aggregate do
         end)
 
       sort_calculations =
-        Enum.flat_map(query.sort, fn {field, _} ->
-          case Map.fetch(query.calculations, field) do
-            :error ->
-              []
+        Enum.flat_map(query.sort, fn
+          {%Ash.Query.Calculation{} = calc, _} ->
+            [calc]
 
-            {:ok, calc} ->
-              [calc]
-          end
+          {field, _} ->
+            case Map.fetch(query.calculations, field) do
+              :error ->
+                []
+
+              {:ok, calc} ->
+                [calc]
+            end
         end)
 
       sort_calc_aggregates =
diff --git a/lib/ash/query/query.ex b/lib/ash/query/query.ex
index 5cc050c2..59c1ac53 100644
--- a/lib/ash/query/query.ex
+++ b/lib/ash/query/query.ex
@@ -1554,12 +1554,39 @@ defmodule Ash.Query do
         end)
         |> validate_sort()
 
-      Enum.reduce(query_with_sort.sort || [], query_with_sort, fn {key, _value}, query ->
-        if Ash.Resource.Info.aggregate(query.resource, key) do
-          Ash.Query.load(query, key)
-        else
-          query
-        end
+      Enum.reduce(query_with_sort.sort || [], query_with_sort, fn
+        {%Ash.Query.Calculation{name: name, module: module, opts: opts} = calculation, _},
+        query ->
+          {resource_load, resource_select} =
+            if resource_calculation = Ash.Resource.Info.calculation(query.resource, name) do
+              {resource_calculation.load, resource_calculation.select}
+            else
+              {[], []}
+            end
+
+          fields_to_select =
+            resource_select
+            |> Kernel.||([])
+            |> Enum.concat(module.select(query, opts, calculation.context) || [])
+
+          calculation = %{calculation | load: name, select: fields_to_select}
+
+          query =
+            query
+            |> module.load(
+              opts,
+              calculation.context
+              |> Map.put(:context, query.context)
+            )
+
+          Ash.Query.load(query, resource_load)
+
+        {key, _value}, query ->
+          if Ash.Resource.Info.aggregate(query.resource, key) do
+            Ash.Query.load(query, key)
+          else
+            query
+          end
       end)
     else
       add_error(query, :sort, "Data layer does not support sorting")