From bbb45f3622b7cc51d3dbb6e73476feed2f6f5173 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Thu, 9 Feb 2023 17:49:37 -0500
Subject: [PATCH] fix: sorting on optimized first aggregates

---
 lib/aggregate.ex               | 14 ++++++++++----
 lib/join.ex                    |  5 +++--
 lib/sort.ex                    | 13 +++++++++++++
 test/aggregate_test.exs        | 22 +++++++++++++++++++++-
 test/support/resources/post.ex |  2 ++
 5 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/lib/aggregate.ex b/lib/aggregate.ex
index c815896..63c18f4 100644
--- a/lib/aggregate.ex
+++ b/lib/aggregate.ex
@@ -657,7 +657,7 @@ defmodule AshPostgres.Aggregate do
     ref = %Ash.Query.Ref{
       attribute: Ash.Resource.Info.field(resource, aggregate.field),
       relationship_path: relationship_path,
-      resource: resource
+      resource: query.__ash_bindings__.resource
     }
 
     type = AshPostgres.Types.parameterized_type(aggregate.type, aggregate.constraints)
@@ -678,7 +678,10 @@ defmodule AshPostgres.Aggregate do
           AshPostgres.Sort.sort(
             query,
             aggregate.query.sort,
-            Ash.Resource.Info.related(resource, relationship_path),
+            Ash.Resource.Info.related(
+              query.__ash_bindings__.resource,
+              relationship_path
+            ),
             relationship_path,
             binding,
             true
@@ -745,7 +748,7 @@ defmodule AshPostgres.Aggregate do
     ref = %Ash.Query.Ref{
       attribute: Ash.Resource.Info.field(resource, aggregate.field),
       relationship_path: relationship_path,
-      resource: resource
+      resource: query.__ash_bindings__.resource
     }
 
     field = AshPostgres.Expr.dynamic_expr(query, ref, query.__ash_bindings__, false)
@@ -756,7 +759,10 @@ defmodule AshPostgres.Aggregate do
           AshPostgres.Sort.sort(
             query,
             aggregate.query.sort,
-            Ash.Resource.Info.related(resource, relationship_path),
+            Ash.Resource.Info.related(
+              query.__ash_bindings__.resource,
+              relationship_path
+            ),
             relationship_path,
             binding,
             true
diff --git a/lib/join.ex b/lib/join.ex
index 55a6016..0c18f64 100644
--- a/lib/join.ex
+++ b/lib/join.ex
@@ -346,7 +346,8 @@ defmodule AshPostgres.Join do
 
   defp can_inner_join?(_, _, _), do: false
 
-  defp get_binding(resource, candidate_path, %{__ash_bindings__: _} = query, types) do
+  @doc false
+  def get_binding(resource, candidate_path, %{__ash_bindings__: _} = query, types) do
     types = List.wrap(types)
 
     Enum.find_value(query.__ash_bindings__.bindings, fn
@@ -361,7 +362,7 @@ defmodule AshPostgres.Join do
     end)
   end
 
-  defp get_binding(_, _, _, _), do: nil
+  def get_binding(_, _, _, _), do: nil
 
   defp add_distinct(relationship, _join_type, joined_query) do
     if !joined_query.__ash_bindings__.in_group? && relationship.cardinality == :many &&
diff --git a/lib/sort.ex b/lib/sort.ex
index 9e0b5de..d70f418 100644
--- a/lib/sort.ex
+++ b/lib/sort.ex
@@ -84,6 +84,19 @@ defmodule AshPostgres.Sort do
               end
 
             :error ->
+              aggregate = Ash.Resource.Info.aggregate(resource, sort)
+
+              {binding, sort} =
+                if aggregate && aggregate.kind == :first &&
+                     AshPostgres.Aggregate.single_path?(resource, aggregate.relationship_path) do
+                  {AshPostgres.Join.get_binding(resource, aggregate.relationship_path, query, [
+                     :left,
+                     :inner
+                   ]), aggregate.field}
+                else
+                  {binding, sort}
+                end
+
               Ecto.Query.dynamic(field(as(^binding), ^sort))
           end
 
diff --git a/test/aggregate_test.exs b/test/aggregate_test.exs
index a913ebf..3a5aebe 100644
--- a/test/aggregate_test.exs
+++ b/test/aggregate_test.exs
@@ -1,6 +1,6 @@
 defmodule AshPostgres.AggregateTest do
   use AshPostgres.RepoCase, async: false
-  alias AshPostgres.Test.{Api, Comment, Organization, Post, Rating, User}
+  alias AshPostgres.Test.{Api, Author, Comment, Organization, Post, Rating, User}
 
   require Ash.Query
 
@@ -316,6 +316,26 @@ defmodule AshPostgres.AggregateTest do
                |> Api.read_one!()
     end
 
+    test "first aggregates can be sorted on" do
+      author =
+        Author
+        |> Ash.Changeset.new(%{first_name: "first name"})
+        |> Api.create!()
+
+      post =
+        Post
+        |> Ash.Changeset.new(%{title: "title"})
+        |> Ash.Changeset.manage_relationship(:author, author, type: :append_and_remove)
+        |> Api.create!()
+
+      assert %{author_first_name: "first name"} =
+               Post
+               |> Ash.Query.filter(id == ^post.id)
+               |> Ash.Query.load(:author_first_name)
+               |> Ash.Query.sort(author_first_name: :asc)
+               |> Api.read_one!()
+    end
+
     test "it can be sorted on and produces the appropriate order" do
       post1 =
         Post
diff --git a/test/support/resources/post.ex b/test/support/resources/post.ex
index 5239c53..068824c 100644
--- a/test/support/resources/post.ex
+++ b/test/support/resources/post.ex
@@ -184,6 +184,8 @@ defmodule AshPostgres.Test.Post do
       sort(title: :desc)
     end
 
+    first(:author_first_name, :author, :first_name)
+
     max(:highest_comment_rating, [:comments, :ratings], :score)
     min(:lowest_comment_rating, [:comments, :ratings], :score)
     avg(:avg_comment_rating, [:comments, :ratings], :score)