diff --git a/lib/quo/eager_query.rb b/lib/quo/eager_query.rb index a44a095..2198688 100644 --- a/lib/quo/eager_query.rb +++ b/lib/quo/eager_query.rb @@ -6,7 +6,11 @@ class EagerQuery < Quo.base_query_class # This is useful when the total count is known and not equal to size # of wrapped collection. def count - options[:total_count] || super + options[:total_count] || underlying_query.count + end + + def page_count + query_with_logging.count end # Is this query object paged? (when no total count) diff --git a/lib/quo/query.rb b/lib/quo/query.rb index 8ffdedf..5b5475a 100644 --- a/lib/quo/query.rb +++ b/lib/quo/query.rb @@ -107,15 +107,18 @@ def select(*options) # Delegate SQL calculation methods to the underlying query delegate :sum, :average, :minimum, :maximum, to: :query_with_logging - # Gets the count of all results ignoring the current page and page size (if set) - delegate :count, to: :underlying_query + # Gets the count of all results ignoring the current page and page size (if set). + def count + count_query(underlying_query) + end + alias_method :total_count, :count alias_method :size, :count # Gets the actual count of elements in the page of results (assuming paging is being used, otherwise the count of # all results) def page_count - query_with_logging.count + count_query(query_with_logging) end # Delegate methods that let us get the model class (available on AR relations) @@ -327,5 +330,22 @@ def test_eager(rel) def test_relation(rel) rel.is_a?(ActiveRecord::Relation) end + + # Note we reselect the query as this prevents query errors if the SELECT clause is not compatible with COUNT + # (SQLException: wrong number of arguments to function COUNT()). We do this in two ways, either with the primary key + # or with Arel.star. The primary key is the most compatible way to count, but if the query does not have a primary + # we fallback. The fallback "*" wont work in certain situations though, specifically if we have a limit() on the query + # which Arel constructs as a subquery. In this case we will get a SQL error as the generated SQL contains + # `SELECT COUNT(count_column) FROM (SELECT * AS count_column FROM ...) subquery_for_count` where the error is: + # `ActiveRecord::StatementInvalid: SQLite3::SQLException: near "AS": syntax error` + # Either way DB engines know how to count efficiently. + def count_query(query) + pk = query.model.primary_key + if pk + query.reselect(pk).count + else + query.reselect(Arel.star).count + end + end end end diff --git a/lib/quo/results.rb b/lib/quo/results.rb index 4ec1f80..e111b57 100644 --- a/lib/quo/results.rb +++ b/lib/quo/results.rb @@ -20,8 +20,9 @@ def initialize(query, transformer: nil) :all?, :any?, :none?, - :one?, - :count + :one? + + def_delegators :query, :count def group_by(&block) debug_callstack @@ -66,7 +67,7 @@ def respond_to_missing?(name, include_private = false) private - attr_reader :transformer, :unwrapped + attr_reader :query, :transformer, :unwrapped def transform_results(results) return results unless transformer diff --git a/test/quo/query_test.rb b/test/quo/query_test.rb index b6bf65f..b5f3669 100644 --- a/test/quo/query_test.rb +++ b/test/quo/query_test.rb @@ -69,6 +69,10 @@ def setup assert_equal 2, UnreadCommentsQuery.new(page_size: 1).count end + test "#count with selects" do + assert_equal 2, Quo::WrappedQuery.wrap(Comment.where(read: false).joins(:post).select(:id, "posts.id")).new.count + end + test "#page_count" do assert_equal 2, UnreadCommentsQuery.new.page_count end