From 4b23dd179f249f6e14200e275e5e6cae8b477e90 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Thu, 1 Jul 2021 15:22:48 -0400
Subject: [PATCH] fix: allow sorting on aggs, w/o loading

---
 lib/ash/actions/sort.ex | 45 ++++++++++++++++++++++++++++++++---------
 lib/ash/query/query.ex  | 25 +++++++++++++++--------
 2 files changed, 52 insertions(+), 18 deletions(-)

diff --git a/lib/ash/actions/sort.ex b/lib/ash/actions/sort.ex
index 6041ae17..db6f8a15 100644
--- a/lib/ash/actions/sort.ex
+++ b/lib/ash/actions/sort.ex
@@ -67,6 +67,9 @@ defmodule Ash.Actions.Sort do
         attribute = Ash.Resource.Info.attribute(resource, field)
 
         cond do
+          aggregate = Ash.Resource.Info.aggregate(resource, field) ->
+            aggregate_sort(aggregates, aggregate, order, resource, sorts, errors)
+
           Map.has_key?(aggregates, field) ->
             aggregate_sort(aggregates, field, order, resource, sorts, errors)
 
@@ -94,7 +97,7 @@ defmodule Ash.Actions.Sort do
             end
 
           !attribute ->
-            {sorts, [NoSuchAttribute.exception(attribute: field) | errors]}
+            {sorts, [NoSuchAttribute.exception(name: field) | errors]}
 
           Ash.Type.embedded_type?(attribute.type) ->
             {sorts, ["Cannot sort on embedded types" | errors]}
@@ -141,16 +144,38 @@ defmodule Ash.Actions.Sort do
   end
 
   defp aggregate_sort(aggregates, field, order, resource, sorts, errors) do
-    aggregate = Map.get(aggregates, field)
+    {field, type} =
+      case field do
+        field when is_atom(field) ->
+          aggregate = Map.get(aggregates, field)
 
-    if Ash.DataLayer.data_layer_can?(resource, :aggregate_sort) &&
-         Ash.DataLayer.data_layer_can?(
-           resource,
-           {:sort, Ash.Type.storage_type(aggregate.type)}
-         ) do
-      {sorts ++ [{field, order}], errors}
-    else
-      {sorts, [AggregatesNotSupported.exception(resource: resource, feature: "sorting") | errors]}
+          {field, {:ok, aggregate.type}}
+
+        %Ash.Resource.Aggregate{} = agg ->
+          field_type =
+            if agg.field do
+              related = Ash.Resource.Info.related(resource, agg.relationship_path)
+              Ash.Resource.Info.attribute(related, agg.field).type
+            end
+
+          {agg.name, Ash.Query.Aggregate.kind_to_type(agg.kind, field_type)}
+      end
+
+    case type do
+      {:ok, type} ->
+        if Ash.DataLayer.data_layer_can?(resource, :aggregate_sort) &&
+             Ash.DataLayer.data_layer_can?(
+               resource,
+               {:sort, Ash.Type.storage_type(type)}
+             ) do
+          {sorts ++ [{field, order}], errors}
+        else
+          {sorts,
+           [AggregatesNotSupported.exception(resource: resource, feature: "sorting") | errors]}
+        end
+
+      {:error, error} ->
+        {sorts, [error | errors]}
     end
   end
 
diff --git a/lib/ash/query/query.ex b/lib/ash/query/query.ex
index be92729b..b0ed7d37 100644
--- a/lib/ash/query/query.ex
+++ b/lib/ash/query/query.ex
@@ -1541,16 +1541,25 @@ defmodule Ash.Query do
     query = to_query(query)
 
     if Ash.DataLayer.data_layer_can?(query.resource, :sort) do
-      sorts
-      |> List.wrap()
-      |> Enum.reduce(query, fn
-        {sort, direction}, query ->
-          %{query | sort: query.sort ++ [{sort, direction}]}
+      query_with_sort =
+        sorts
+        |> List.wrap()
+        |> Enum.reduce(query, fn
+          {sort, direction}, query ->
+            %{query | sort: query.sort ++ [{sort, direction}]}
 
-        sort, query ->
-          %{query | sort: query.sort ++ [{sort, :asc}]}
+          sort, query ->
+            %{query | sort: query.sort ++ [{sort, :asc}]}
+        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
       end)
-      |> validate_sort()
     else
       add_error(query, :sort, "Data layer does not support sorting")
     end