Merge pull request #29848 from kamipo/fix_distinct_count_with_order_and_limit

Fix `COUNT(DISTINCT ...)` with `ORDER BY` and `LIMIT`

7 files changed, 87 insertions, 45 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index ea9f860b95..afc459ef68 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,7 @@
+*   Fix `COUNT(DISTINCT ...)` with `ORDER BY` and `LIMIT` to keep the existing select list.
+
+    *Ryuta Kamizono*
+
 *   When a `has_one` association is destroyed by `dependent: destroy`,
     `destroyed_by_association` will now be set to the reflection, matching the
     behaviour of `has_many` associations.
diff --git a/activerecord/lib/active_record/collection_cache_key.rb b/activerecord/lib/active_record/collection_cache_key.rb
index 0b29ef2502..b1937a3c68 100644
--- a/activerecord/lib/active_record/collection_cache_key.rb
+++ b/activerecord/lib/active_record/collection_cache_key.rb
@@ -16,7 +16,7 @@ module ActiveRecord
         column = "#{connection.quote_table_name(collection.table_name)}.#{connection.quote_column_name(timestamp_column)}"
         select_values = "COUNT(*) AS #{connection.quote_column_name("size")}, MAX(%s) AS timestamp"
 
-        if collection.limit_value || collection.offset_value
+        if collection.has_limit_or_offset?
           query = collection.spawn
           query.select_values = [column]
           subquery_alias = "subquery_for_cache_key"
diff --git a/activerecord/lib/active_record/relation.rb b/activerecord/lib/active_record/relation.rb
index 92ecdcae92..1c4011b75e 100644
--- a/activerecord/lib/active_record/relation.rb
+++ b/activerecord/lib/active_record/relation.rb
@@ -642,6 +642,10 @@ module ActiveRecord
       @values == klass.unscoped.values
     end
 
+    def has_limit_or_offset? # :nodoc:
+      limit_value || offset_value
+    end
+
     protected
 
       def load_records(records)
diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
index 66a9a3c9c9..d3b5be6bce 100644
--- a/activerecord/lib/active_record/relation/calculations.rb
+++ b/activerecord/lib/active_record/relation/calculations.rb
@@ -131,7 +131,7 @@ module ActiveRecord
     def calculate(operation, column_name)
       if has_include?(column_name)
         relation = construct_relation_for_association_calculations
-        relation = relation.distinct if operation.to_s.downcase == "count"
+        relation.distinct! if operation.to_s.downcase == "count"
 
         relation.calculate(operation, column_name)
       else
@@ -214,8 +214,13 @@ module ActiveRecord
 
         if operation == "count"
           column_name ||= select_for_count
-          column_name = primary_key if column_name == :all && distinct
-          distinct = nil if column_name =~ /\s*DISTINCT[\s(]+/i
+          if column_name == :all
+            if distinct && !(has_limit_or_offset? && order_values.any?)
+              column_name = primary_key
+            end
+          elsif column_name =~ /\s*DISTINCT[\s(]+/i
+            distinct = nil
+          end
         end
 
         if group_values.any?
@@ -242,7 +247,7 @@ module ActiveRecord
       def execute_simple_calculation(operation, column_name, distinct) #:nodoc:
         column_alias = column_name
 
-        if operation == "count" && (limit_value || offset_value)
+        if operation == "count" && has_limit_or_offset?
           # Shortcut when limit is zero.
           return 0 if limit_value == 0
 
@@ -381,14 +386,18 @@ module ActiveRecord
       end
 
       def build_count_subquery(relation, column_name, distinct)
-        column_alias = Arel.sql("count_column")
-        subquery_alias = Arel.sql("subquery_for_count")
+        relation.select_values = [
+          if column_name == :all
+            distinct ? table[Arel.star] : Arel.sql("1")
+          else
+            column_alias = Arel.sql("count_column")
+            aggregate_column(column_name).as(column_alias)
+          end
+        ]
 
-        aliased_column = aggregate_column(column_name == :all ? 1 : column_name).as(column_alias)
-        relation.select_values = [aliased_column]
-        subquery = relation.arel.as(subquery_alias)
+        subquery = relation.arel.as(Arel.sql("subquery_for_count"))
+        select_value = operation_over_aggregate_column(column_alias || Arel.star, "count", false)
 
-        select_value = operation_over_aggregate_column(column_alias, "count", distinct)
         Arel::SelectManager.new(subquery).project(select_value)
       end
   end
diff --git a/activerecord/test/cases/calculations_test.rb b/activerecord/test/cases/calculations_test.rb
index dd6815b149..92d071187d 100644
--- a/activerecord/test/cases/calculations_test.rb
+++ b/activerecord/test/cases/calculations_test.rb
@@ -236,6 +236,30 @@ class CalculationsTest < ActiveRecord::TestCase
     end
   end
 
+  def test_distinct_count_with_order_and_limit
+    assert_equal 4, Account.distinct.order(:firm_id).limit(4).count
+  end
+
+  def test_distinct_count_with_order_and_offset
+    assert_equal 4, Account.distinct.order(:firm_id).offset(2).count
+  end
+
+  def test_distinct_count_with_order_and_limit_and_offset
+    assert_equal 4, Account.distinct.order(:firm_id).limit(4).offset(2).count
+  end
+
+  def test_distinct_joins_count_with_order_and_limit
+    assert_equal 3, Account.joins(:firm).distinct.order(:firm_id).limit(3).count
+  end
+
+  def test_distinct_joins_count_with_order_and_offset
+    assert_equal 3, Account.joins(:firm).distinct.order(:firm_id).offset(2).count
+  end
+
+  def test_distinct_joins_count_with_order_and_limit_and_offset
+    assert_equal 3, Account.joins(:firm).distinct.order(:firm_id).limit(3).offset(2).count
+  end
+
   def test_should_group_by_summed_field_having_condition
     c = Account.group(:firm_id).having("sum(credit_limit) > 50").sum(:credit_limit)
     assert_nil        c[1]
diff --git a/activerecord/test/models/account.rb b/activerecord/test/models/account.rb
new file mode 100644
index 0000000000..0c3cd45a81
--- /dev/null
+++ b/activerecord/test/models/account.rb
@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+class Account < ActiveRecord::Base
+  belongs_to :firm, class_name: "Company"
+  belongs_to :unautosaved_firm, foreign_key: "firm_id", class_name: "Firm", autosave: false
+
+  alias_attribute :available_credit, :credit_limit
+
+  def self.destroyed_account_ids
+    @destroyed_account_ids ||= Hash.new { |h, k| h[k] = [] }
+  end
+
+  # Test private kernel method through collection proxy using has_many.
+  def self.open
+    where("firm_name = ?", "37signals")
+  end
+
+  before_destroy do |account|
+    if account.firm
+      Account.destroyed_account_ids[account.firm.id] << account.id
+    end
+  end
+
+  validate :check_empty_credit_limit
+
+  private
+    def check_empty_credit_limit
+      errors.add("credit_limit", :blank) if credit_limit.blank?
+    end
+
+    def private_method
+      "Sir, yes sir!"
+    end
+end
diff --git a/activerecord/test/models/company.rb b/activerecord/test/models/company.rb
index 50e9553bfe..bbc5fc2b2d 100644
--- a/activerecord/test/models/company.rb
+++ b/activerecord/test/models/company.rb
@@ -187,37 +187,4 @@ end
 class VerySpecialClient < SpecialClient
 end
 
-class Account < ActiveRecord::Base
-  belongs_to :firm, class_name: "Company"
-  belongs_to :unautosaved_firm, foreign_key: "firm_id", class_name: "Firm", autosave: false
-
-  alias_attribute :available_credit, :credit_limit
-
-  def self.destroyed_account_ids
-    @destroyed_account_ids ||= Hash.new { |h, k| h[k] = [] }
-  end
-
-  # Test private kernel method through collection proxy using has_many.
-  def self.open
-    where("firm_name = ?", "37signals")
-  end
-
-  before_destroy do |account|
-    if account.firm
-      Account.destroyed_account_ids[account.firm.id] << account.id
-    end
-    true
-  end
-
-  validate :check_empty_credit_limit
-
-  private
-
-    def check_empty_credit_limit
-      errors.add("credit_limit", :blank) if credit_limit.blank?
-    end
-
-    def private_method
-      "Sir, yes sir!"
-    end
-end
+require "models/account"