feat: allow retrieving the count of paginated relationships (#1183)

lib/ash/actions/read/read.ex

@@ -427,6 +427,7 @@ defmodule Ash.Actions.Read do
              ),
            query <- Map.put(query, :filter, filter),
            query <- Ash.Query.unset(query, :calculations),
+           query <- add_relationship_count_aggregates(query),
            {%{valid?: true} = query, before_notifications} <- run_before_action(query),
            {:ok, count} <-
              fetch_count(
@@ -479,6 +480,48 @@ defmodule Ash.Actions.Read do
     end)
   end
 
+  defp add_relationship_count_aggregates(query) do
+    Enum.reduce(query.load, query, fn {relationship_name, related_query}, query ->
+      relationship = Ash.Resource.Info.relationship(query.resource, relationship_name)
+
+      related_query =
+        case related_query do
+          [] -> Ash.Query.new(relationship.destination)
+          query -> query
+        end
+
+      needs_count? = related_query.page && related_query.page[:count] == true
+
+      if needs_count? do
+        related_query =
+          Ash.Query.unset(related_query, [
+            :sort,
+            :distinct,
+            :distinct_sort,
+            :lock,
+            :load,
+            :page,
+            :aggregates
+          ])
+
+        aggregate_name = paginated_relationship_count_aggregate_name(relationship.name)
+
+        query
+        |> Ash.Query.aggregate(aggregate_name, :count, relationship.name,
+          query: related_query,
+          default: 0
+        )
+      else
+        query
+      end
+    end)
+  end
+
+  @doc false
+  def paginated_relationship_count_aggregate_name(relationship_name) do
+    "__paginated_#{relationship_name}_count__"
+  end
+
   @doc false
   def cleanup_field_auth(records, query, top_level? \\ true)
 
@@ -1403,7 +1446,7 @@ defmodule Ash.Actions.Read do
         data
 
       opts[:return_unpaged?] && original_query.page[:limit] ->
-        Ash.Page.Unpaged.new(data, count, opts)
+        Ash.Page.Unpaged.new(data, opts)
 
       original_query.page[:limit] ->
         to_page(data, action, count, sort, original_query, opts)
@@ -1785,7 +1828,8 @@ defmodule Ash.Actions.Read do
 
     cond do
       Map.has_key?(query.context, :accessing_from) and needs_count? ->
-        {:error, "Cannot request count when paginating relationships"}
+        # Relationship count is fetched by the parent using aggregates, just return nil here
+        {:ok, {:ok, nil}}
 
       needs_count? ->
         with {:ok, filter} <-

lib/ash/actions/read/relationships.ex

@@ -560,34 +560,47 @@ defmodule Ash.Actions.Read.Relationships do
        ) do
     %Ash.Page.Unpaged{
       related_records: related_records,
-      count: count,
       opts: opts
     } = unpaged
 
-    to_page_fun =
+    attach_fun =
       if relationship.cardinality == :many do
-        fn value, record ->
+        fn record, relationship_name, value ->
+          count_key =
+            Ash.Actions.Read.paginated_relationship_count_aggregate_name(relationship.name)
+
+          # Retrieve the count (if present) while deleting it from the record aggregates
+          {count, record} = pop_in(record.aggregates[count_key])
+
           # We scope the lateral join to the specific record, so that next runs of rerun
           # just fetch the entries related to this record
           related_query =
             Ash.Query.set_context(related_query, %{
               data_layer: %{lateral_join_source: {[record], lateral_join_source_path}}
             })
 
-          Ash.Actions.Read.to_page(
-            value,
-            related_query.action,
-            count,
-            related_query.sort,
-            related_query,
-            opts
-          )
+          page =
+            Ash.Actions.Read.to_page(
+              value,
+              related_query.action,
+              count,
+              related_query.sort,
+              related_query,
+              opts
+            )
+
+          attach_related(record, relationship_name, page)
         end
       else
-        fn value, _record -> value end
+        &attach_related/3
       end
 
-    attach_lateral_join_related_records(records, relationship, related_records, to_page_fun)
+    attach_lateral_join_related_records(
+      records,
+      relationship,
+      related_records,
+      attach_fun
+    )
   end
 
   defp do_attach_related_records(
@@ -813,11 +826,15 @@ defmodule Ash.Actions.Read.Relationships do
     Map.put(record, key, default)
   end
 
+  defp attach_related(record, relationship_name, value) do
+    Map.put(record, relationship_name, value)
+  end
+
   defp attach_lateral_join_related_records(
          [%resource{} | _] = records,
          relationship,
          related_records,
-         maybe_to_page_fun \\ fn related_value, _record -> related_value end
+         attach_fun \\ &attach_related/3
        ) do
     source_attribute =
       Ash.Resource.Info.attribute(relationship.source, relationship.source_attribute)
@@ -844,10 +861,10 @@ defmodule Ash.Actions.Read.Relationships do
       Enum.map(records, fn record ->
         with :error <- Map.fetch(values, Map.take(record, primary_key)),
              :error <- Map.fetch(values, Map.get(record, relationship.source_attribute)) do
-          Map.put(record, relationship.name, maybe_to_page_fun.(default, record))
+          attach_fun.(record, relationship.name, default)
         else
           {:ok, value} ->
-            Map.put(record, relationship.name, maybe_to_page_fun.(value, record))
+            attach_fun.(record, relationship.name, value)
         end
       end)
     else
@@ -875,7 +892,7 @@ defmodule Ash.Actions.Read.Relationships do
             end
           ])
 
-        Map.put(record, relationship.name, maybe_to_page_fun.(related, record))
+        attach_fun.(record, relationship.name, related)
       end)
     end
   end

lib/ash/page/unpaged.ex

@@ -3,12 +3,11 @@ defmodule Ash.Page.Unpaged do
   # related records and then paged
   @moduledoc false
 
-  defstruct [:related_records, :count, :opts]
+  defstruct [:related_records, :opts]
 
-  def new(related_records, count, opts) do
+  def new(related_records, opts) do
     %__MODULE__{
       related_records: related_records,
-      count: count,
       opts: Keyword.delete(opts, :return_unpaged?)
     }
   end

test/actions/load_test.exs

@@ -1420,59 +1420,171 @@ defmodule Ash.Test.Actions.LoadTest do
              } = author1.posts
     end
 
-    test "returns error when requesting count" do
-      Author
-      |> Ash.Changeset.for_create(:create, %{name: "a"})
+    test "doesn't honor required? pagination to maintain backwards compatibility" do
+      author =
+        Author
+        |> Ash.Changeset.for_create(:create, %{name: "a"})
+        |> Ash.create!()
+
+      Post
+      |> Ash.Changeset.for_create(:create, %{title: "b"})
+      |> Ash.Changeset.manage_relationship(:author, author, type: :append_and_remove)
       |> Ash.create!()
 
-      paginated_posts =
+      posts =
         Post
-        |> Ash.Query.page(limit: 1, count: true)
+        |> Ash.Query.for_read(:required_pagination)
 
-      assert {:error,
-              %Ash.Error.Unknown{
-                errors: [
-                  %Ash.Error.Unknown.UnknownError{
-                    error: "Cannot request count when paginating relationships"
-                  }
-                ]
-              }} =
-               Author
-               |> Ash.Query.load(posts: paginated_posts)
-               |> Ash.read()
-
-      assert {:error,
-              %Ash.Error.Unknown{
-                errors: [
-                  %Ash.Error.Unknown.UnknownError{
-                    error: "Cannot request count when paginating relationships"
-                  }
-                ]
-              }} =
+      assert [_post] =
                Author
                |> Ash.read!()
-               |> Ash.load(posts: paginated_posts)
+               |> Ash.load!(posts: posts)
     end
 
-    test "doesn't honor required? pagination to maintain backwards compatibility" do
-      author =
+    test "it allows counting has_many relationships" do
+      author1 =
         Author
         |> Ash.Changeset.for_create(:create, %{name: "a"})
         |> Ash.create!()
 
+      author2 =
+        Author
+        |> Ash.Changeset.for_create(:create, %{name: "b"})
+        |> Ash.create!()
+
+      for i <- 1..3 do
+        Post
+        |> Ash.Changeset.for_create(:create, %{title: "author1 post#{i}", author_id: author1.id})
+        |> Ash.create!()
+      end
+
+      for i <- 1..6 do
+        Post
+        |> Ash.Changeset.for_create(:create, %{title: "author2 post#{i}", author_id: author2.id})
+        |> Ash.create!()
+      end
+
+      paginated_posts =
+        Post
+        |> Ash.Query.page(limit: 2, offset: 2, count: true)
+
+      assert [author1, author2] =
+               Author
+               |> Ash.Query.sort(:name)
+               |> Ash.Query.load(posts: paginated_posts)
+               |> Ash.read!()
+
+      assert %Ash.Page.Offset{count: 3} = author1.posts
+      assert %Ash.Page.Offset{count: 6} = author2.posts
+    end
+
+    test "it allows counting many_to_many relationships" do
+      categories =
+        for i <- 1..9 do
+          Category
+          |> Ash.Changeset.for_create(:create, %{name: "category#{i}"})
+          |> Ash.create!()
+        end
+
+      categories_1_to_3 = Enum.take(categories, 3)
+      categories_4_to_9 = Enum.slice(categories, 3..9)
+
+      Post
+      |> Ash.Changeset.for_create(:create, %{title: "a"})
+      |> Ash.Changeset.manage_relationship(:categories, categories_1_to_3,
+        type: :append_and_remove
+      )
+      |> Ash.create!()
+
       Post
       |> Ash.Changeset.for_create(:create, %{title: "b"})
-      |> Ash.Changeset.manage_relationship(:author, author, type: :append_and_remove)
+      |> Ash.Changeset.manage_relationship(:categories, categories_4_to_9,
+        type: :append_and_remove
+      )
       |> Ash.create!()
 
-      posts =
+      paginated_categories =
+        Category
+        |> Ash.Query.page(limit: 2, count: true)
+        |> Ash.Query.sort(:name)
+
+      assert [post1, post2] =
+               Post
+               |> Ash.Query.sort(:title)
+               |> Ash.Query.load(categories: paginated_categories)
+               |> Ash.read!()
+
+      assert %Ash.Page.Offset{count: 3} = post1.categories
+      assert %Ash.Page.Offset{count: 6} = post2.categories
+    end
+
+    test "allows counting nested relationships" do
+      author1 =
+        Author
+        |> Ash.Changeset.for_create(:create, %{name: "a"})
+        |> Ash.create!()
+
+      _author2 =
+        Author
+        |> Ash.Changeset.for_create(:create, %{name: "b"})
+        |> Ash.create!()
+
+      categories =
+        for i <- 1..3 do
+          Category
+          |> Ash.Changeset.for_create(:create, %{name: "category#{i}"})
+          |> Ash.create!()
+        end
+
+      for i <- 1..5 do
         Post
-        |> Ash.Query.for_read(:required_pagination)
+        |> Ash.Changeset.for_create(:create, %{title: "author1 post#{i}", author_id: author1.id})
+        |> Ash.Changeset.manage_relationship(:categories, categories, type: :append_and_remove)
+        |> Ash.create!()
+      end
 
-      assert [_post] =
+      paginated_categories =
+        Category
+        |> Ash.Query.page(limit: 1, count: true)
+
+      paginated_posts =
+        Post
+        |> Ash.Query.load(categories: paginated_categories)
+        |> Ash.Query.page(limit: 1, count: true)
+
+      assert %Ash.Page.Offset{results: [author1], count: 2} =
                Author
+               |> Ash.Query.sort(:name)
+               |> Ash.Query.load(posts: paginated_posts)
+               |> Ash.read!(page: [limit: 1, count: true])
+
+      assert %Ash.Page.Offset{count: 5, results: [%{categories: %Ash.Page.Offset{count: 3}}]} =
+               author1.posts
+    end
+
+    test "doesn't leak the internal count aggregate when counting" do
+      author =
+        Author
+        |> Ash.Changeset.for_create(:create, %{name: "a"})
+        |> Ash.create!()
+
+      for i <- 1..3 do
+        Post
+        |> Ash.Changeset.for_create(:create, %{title: "author1 post#{i}", author_id: author.id})
+        |> Ash.create!()
+      end
+
+      paginated_posts =
+        Post
+        |> Ash.Query.page(limit: 2, offset: 2, count: true)
+
+      assert [author] =
+               Author
+               |> Ash.Query.load(posts: paginated_posts)
                |> Ash.read!()
-               |> Ash.load!(posts: posts)
+
+      assert %Ash.Page.Offset{count: 3} = author.posts
+      assert %{} == author.aggregates
     end
   end
 end