Fix total entries when query has combinations

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

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()