From e394a055299f9722266486b3853215b5b7d0f9d8 Mon Sep 17 00:00:00 2001
From: Milo Lee <leexu.dev@gmail.com>
Date: Thu, 26 Aug 2021 22:20:42 +0800
Subject: [PATCH] Fix total entries when query has combinations

---
 lib/scrivener/paginater/ecto/query.ex        | 39 +++++++-------------
 test/scrivener/paginator/ecto/query_test.exs | 12 ++++++
 2 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/lib/scrivener/paginater/ecto/query.ex b/lib/scrivener/paginater/ecto/query.ex
index 58f01fd..c3f9aac 100644
--- a/lib/scrivener/paginater/ecto/query.ex
+++ b/lib/scrivener/paginater/ecto/query.ex
@@ -47,21 +47,7 @@ defimpl Scrivener.Paginater, for: Ecto.Query do
 
   defp total_entries(query, repo, caller, options) do
     prefix = options[:prefix]
-
-    total_entries =
-      query
-      |> exclude(:preload)
-      |> exclude(:order_by)
-      |> aggregate()
-      |> one(repo, caller, prefix)
-
-    total_entries || 0
-  end
-
-  defp aggregate(%{distinct: %{expr: expr}} = query) when expr == true or is_list(expr) do
-    query
-    |> exclude(:select)
-    |> count()
+    aggregate(query, repo, caller, prefix)
   end
 
   defp aggregate(
@@ -74,24 +60,27 @@ defimpl Scrivener.Paginater, for: Ecto.Query do
              }
              | _
            ]
-         } = query
+         } = query,
+         repo,
+         caller,
+         prefix
        ) do
     query
+    |> exclude(:preload)
+    |> exclude(:order_by)
     |> exclude(:select)
     |> select([{x, source_index}], struct(x, ^[field]))
-    |> count()
+    |> subquery()
+    |> select(count("*"))
+    |> one(repo, caller, prefix)
   end
 
-  defp aggregate(query) do
-    query
-    |> exclude(:select)
-    |> select(count("*"))
+  defp aggregate(query, repo, caller, nil) do
+    repo.aggregate(query, :count, caller: caller)
   end
 
-  defp count(query) do
-    query
-    |> subquery
-    |> select(count("*"))
+  defp aggregate(query, repo, caller, prefix) do
+    repo.aggregate(query, :count, caller: caller, prefix: prefix)
   end
 
   defp total_pages(0, _), do: 1
diff --git a/test/scrivener/paginator/ecto/query_test.exs b/test/scrivener/paginator/ecto/query_test.exs
index 325b24f..b1f9ddc 100644
--- a/test/scrivener/paginator/ecto/query_test.exs
+++ b/test/scrivener/paginator/ecto/query_test.exs
@@ -320,6 +320,18 @@ defmodule Scrivener.Paginator.Ecto.QueryTest do
       assert page.total_entries == 2
     end
 
+    test "can be used with combinations" do
+      create_posts()
+
+      page =
+        Post
+        |> Post.published()
+        |> union(^Post.unpublished(Post))
+        |> Scrivener.Ecto.Repo.paginate()
+
+      assert page.total_entries == 7
+    end
+
     test "can be provided a Scrivener.Config directly" do
       posts = create_posts()