Fix count queries on includes+references for models with composite primary keys

Author: fatkodima

activerecord/lib/active_record/relation/calculations.rb

@@ -234,7 +234,7 @@ def calculate(operation, column_name)
         if operation == "count"
           unless distinct_value || distinct_select?(column_name || select_for_count)
             relation.distinct!
-            relation.select_values = [ klass.primary_key || table[Arel.star] ]
+            relation.select_values = Array(klass.primary_key || table[Arel.star])
           end
           # PostgreSQL: ORDER BY expressions must appear in SELECT list when using DISTINCT
           relation.order_values = [] if group_values.empty?
@@ -459,7 +459,7 @@ def operation_over_aggregate_column(column, operation, distinct)
       end
 
       def execute_simple_calculation(operation, column_name, distinct) # :nodoc:
-        if operation == "count" && (column_name == :all && distinct || has_limit_or_offset?)
+        if build_count_subquery?(operation, column_name, distinct)
           # Shortcut when limit is zero.
           return 0 if limit_value == 0
 
@@ -632,12 +632,30 @@ def type_cast_calculated_value(value, operation, type)
       def select_for_count
         if select_values.present?
           return select_values.first if select_values.one?
-          select_values.join(", ")
+
+          select_values.map do |field|
+            column = arel_column(field.to_s) do |attr_name|
+              Arel.sql(attr_name)
+            end
+
+            if column.is_a?(Arel::Nodes::SqlLiteral)
+              column
+            else
+              "#{adapter_class.quote_table_name(column.relation.name)}.#{adapter_class.quote_column_name(column.name)}"
+            end
+          end.join(", ")
         else
           :all
         end
       end
 
+      def build_count_subquery?(operation, column_name, distinct)
+        # SQLite and older MySQL does not support `COUNT DISTINCT` with `*` or
+        # multiple columns, so we need to use subquery for this.
+        operation == "count" &&
+          (((column_name == :all || select_values.many?) && distinct) || has_limit_or_offset?)
+      end
+
       def build_count_subquery(relation, column_name, distinct)
         if column_name == :all
           column_alias = Arel.star

activerecord/test/cases/calculations_test.rb

@@ -403,6 +403,10 @@ def test_group_by_count_for_a_composite_primary_key_model
     assert_equal(expected, Cpk::Book.where(author_id: book.author_id).group(:author_id).count)
   end
 
+  def test_count_for_a_composite_primary_key_model_with_includes_and_references
+    assert_equal Cpk::Book.count, Cpk::Book.includes(:chapters).references(:chapters).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]