Merge pull request #35847 from kamipo/optimizer_hints_with_count_subquery

Optimizer hints should be applied on Top level query as much as possible

5 files changed, 30 insertions, 5 deletions

diff --git a/activerecord/lib/active_record/collection_cache_key.rb b/activerecord/lib/active_record/collection_cache_key.rb
index 4b6db8a96c..bfcb3e54be 100644
--- a/activerecord/lib/active_record/collection_cache_key.rb
+++ b/activerecord/lib/active_record/collection_cache_key.rb
@@ -23,8 +23,7 @@ module ActiveRecord
           query = collection.select("#{column} AS collection_cache_key_timestamp")
           subquery_alias = "subquery_for_cache_key"
           subquery_column = "#{subquery_alias}.collection_cache_key_timestamp"
-          subquery = query.arel.as(subquery_alias)
-          arel = Arel::SelectManager.new(subquery).project(select_values % subquery_column)
+          arel = query.send(:build_subquery, subquery_alias, select_values % subquery_column)
         else
           query = collection.unscope(:order)
           query.select_values = [select_values % column]
diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
index 4f9ddf302e..0384023a17 100644
--- a/activerecord/lib/active_record/relation/calculations.rb
+++ b/activerecord/lib/active_record/relation/calculations.rb
@@ -416,16 +416,17 @@ module ActiveRecord
 
       def build_count_subquery(relation, column_name, distinct)
         if column_name == :all
+          column_alias = Arel.star
           relation.select_values = [ Arel.sql(FinderMethods::ONE_AS_ONE) ] unless distinct
         else
           column_alias = Arel.sql("count_column")
           relation.select_values = [ aggregate_column(column_name).as(column_alias) ]
         end
 
-        subquery = relation.arel.as(Arel.sql("subquery_for_count"))
-        select_value = operation_over_aggregate_column(column_alias || Arel.star, "count", false)
+        subquery_alias = Arel.sql("subquery_for_count")
+        select_value = operation_over_aggregate_column(column_alias, "count", false)
 
-        Arel::SelectManager.new(subquery).project(select_value)
+        relation.build_subquery(subquery_alias, select_value)
       end
   end
 end
diff --git a/activerecord/lib/active_record/relation/query_methods.rb b/activerecord/lib/active_record/relation/query_methods.rb
index b0535cfff5..5f728f2263 100644
--- a/activerecord/lib/active_record/relation/query_methods.rb
+++ b/activerecord/lib/active_record/relation/query_methods.rb
@@ -992,6 +992,15 @@ module ActiveRecord
       @arel ||= build_arel(aliases)
     end
 
+    protected
+      def build_subquery(subquery_alias, select_value) # :nodoc:
+        subquery = except(:optimizer_hints).arel.as(subquery_alias)
+
+        Arel::SelectManager.new(subquery).project(select_value).tap do |arel|
+          arel.optimizer_hints(*optimizer_hints_values) unless optimizer_hints_values.empty?
+        end
+      end
+
     private
       # Returns a relation value with a given name
       def get_value(name)
diff --git a/activerecord/test/cases/adapters/mysql2/optimizer_hints_test.rb b/activerecord/test/cases/adapters/mysql2/optimizer_hints_test.rb
index b9794c5710..628802b216 100644
--- a/activerecord/test/cases/adapters/mysql2/optimizer_hints_test.rb
+++ b/activerecord/test/cases/adapters/mysql2/optimizer_hints_test.rb
@@ -15,6 +15,14 @@ if supports_optimizer_hints?
       end
     end
 
+    def test_optimizer_hints_with_count_subquery
+      assert_sql(%r{\ASELECT /\*\+ NO_RANGE_OPTIMIZATION\(posts index_posts_on_author_id\) \*/}) do
+        posts = Post.optimizer_hints("NO_RANGE_OPTIMIZATION(posts index_posts_on_author_id)")
+        posts = posts.select(:id).where(author_id: [0, 1]).limit(5)
+        assert_equal 5, posts.count
+      end
+    end
+
     def test_optimizer_hints_is_sanitized
       assert_sql(%r{\ASELECT /\*\+ NO_RANGE_OPTIMIZATION\(posts index_posts_on_author_id\) \*/}) do
         posts = Post.optimizer_hints("/*+ NO_RANGE_OPTIMIZATION(posts index_posts_on_author_id) */")
diff --git a/activerecord/test/cases/adapters/postgresql/optimizer_hints_test.rb b/activerecord/test/cases/adapters/postgresql/optimizer_hints_test.rb
index 5e4bf232e1..5b9f5e0832 100644
--- a/activerecord/test/cases/adapters/postgresql/optimizer_hints_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/optimizer_hints_test.rb
@@ -19,6 +19,14 @@ if supports_optimizer_hints?
       end
     end
 
+    def test_optimizer_hints_with_count_subquery
+      assert_sql(%r{\ASELECT /\*\+ SeqScan\(posts\) \*/}) do
+        posts = Post.optimizer_hints("SeqScan(posts)")
+        posts = posts.select(:id).where(author_id: [0, 1]).limit(5)
+        assert_equal 5, posts.count
+      end
+    end
+
     def test_optimizer_hints_is_sanitized
       assert_sql(%r{\ASELECT /\*\+ SeqScan\(posts\) \*/}) do
         posts = Post.optimizer_hints("/*+ SeqScan(posts) */")