feat: Add support for table optimizer hints

Add two new methods to `ActiveRecord::Relation::QueryMethods`.

- `#force_index(index_name, direction: )`
- `#index_hint(hint)`

Fixes #266

Author: BuonOmo

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## Ongoing
+
+- Add support for table optimize hints (#266).
+
 ## 7.0.2 - 2023-05-23
 
 - Fix default numbers test to expect the correct result after

lib/active_record/connection_adapters/cockroachdb_adapter.rb

@@ -23,6 +23,8 @@
 require_relative "../migration/cockroachdb/compatibility"
 require_relative "../../version"
 
+require_relative "../relation/query_methods_ext"
+
 # Run to ignore spatial tables that will break schemna dumper.
 # Defined in ./setup.rb
 ActiveRecord::ConnectionAdapters::CockroachDB.initial_setup

lib/active_record/relation/query_methods_ext.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  class Relation
+    module QueryMethodsExt
+      def from!(...) # :nodoc:
+        @force_index = nil
+        @index_hint = nil
+        super
+      end
+
+      # Set table index hint for the query to the
+      # given `index_name`, and `direction` (either
+      # `ASC` or `DESC`).
+      #
+      # Any call to `ActiveRecord::QueryMethods#from`
+      # will reset the index hint. Index hints are
+      # not set if the `from` clause is not a table
+      # name.
+      #
+      # @see https://www.cockroachlabs.com/docs/v22.2/table-expressions#force-index-selection
+      def force_index(index_name, direction: nil)
+        spawn.force_index!(index_name, direction: direction)
+      end
+
+      def force_index!(index_name, direction: nil)
+        return self unless from_clause_is_a_table_name?
+
+        index_name = sanitize_sql(index_name.to_s)
+        direction = direction.to_s.upcase
+        direction = %w[ASC DESC].include?(direction) ? ",#{direction}" : ""
+
+        @force_index = "FORCE_INDEX=#{index_name}#{direction}"
+        self.from_clause = build_from_clause_with_hints
+        self
+      end
+
+      # Set table index hint for the query with the
+      # given `hint`. This allows more control over
+      # the hint than `ActiveRecord::Relation#force_index`.
+      # For instance, you could set it to `NO_FULL_SCAN`.
+      #
+      # Any call to `ActiveRecord::QueryMethods#from`
+      # will reset the index hint. Index hints are
+      # not set if the `from` clause is not a table
+      # name.
+      #
+      # @see https://www.cockroachlabs.com/docs/v22.2/table-expressions#force-index-selection
+      def index_hint(hint)
+        spawn.index_hint!(hint)
+      end
+
+      def index_hint!(hint)
+        return self unless from_clause_is_a_table_name?
+
+        hint = sanitize_sql(hint.to_s)
+        @index_hint = hint.to_s
+        self.from_clause = build_from_clause_with_hints
+        self
+      end
+
+      private
+
+      def from_clause_is_a_table_name?
+        # if empty, we are just dealing with the current table.
+        return true if from_clause.empty?
+        # `from_clause` can be a subquery.
+        return false unless from_clause.value.is_a?(String)
+        # `from_clause` can be a list of tables or a function.
+        # A simple way to check is to see if the string
+        # contains special characters. But we have to
+        # not check against an existing table hint.
+        return !from_clause.value.gsub(/\@{.*?\}/, "").match?(/[,\(]/)
+      end
+
+      def build_from_clause_with_hints
+        table_hints = [@index_hint, @force_index].compact.join(",")
+
+        table_name =
+          if from_clause.empty?
+            quoted_table_name
+          else
+            # Remove previous table hints if any. And spaces.
+            from_clause.value.partition("@").first.strip
+          end
+        Relation::FromClause.new("#{table_name}@{#{table_hints}}", nil)
+      end
+    end
+
+    QueryMethods.prepend(QueryMethodsExt)
+  end
+  # `ActiveRecord::Base` ancestors do not include `QueryMethods`.
+  # But the `#all` method returns a relation, which has `QueryMethods`
+  # as ancestor. That is how active_record is doing is as well.
+  #
+  # @see https://github.com/rails/rails/blob/914130a9f/activerecord/lib/active_record/querying.rb#L23
+  Querying.delegate(:force_index, :index_hint, to: :all)
+end

test/cases/relation/table_hints_test.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+require "cases/helper_cockroachdb"
+require "models/post"
+
+module CockroachDB
+  class TableHintsTest < ActiveRecord::TestCase
+    fixtures :posts
+
+    def test_add_hint
+      force_index = ->(q) { q.force_index("index_posts_on_author_id") }
+      index_hint = ->(q) { q.index_hint("NO_FULL_SCAN") }
+
+      assert_sql(/"posts"@\{FORCE_INDEX=index_posts_on_author_id\}/) do
+        Post.then(&force_index).take
+      end
+
+      assert_sql(/"posts"@\{NO_FULL_SCAN\}/) do
+        Post.then(&index_hint).take
+      end
+
+      [force_index, index_hint].permutation do |procs|
+        assert_sql(/"posts"@\{NO_FULL_SCAN,FORCE_INDEX=index_posts_on_author_id\}/) do
+          Post.then(&procs.reduce(:<<)).take
+        end
+      end
+    end
+
+    def test_choose_index_order
+      idx = "index_posts_on_author_id"
+      assert_sql(/"posts"@\{FORCE_INDEX=#{idx},ASC\}/) do
+        Post.force_index(idx, direction: "ASC").take
+      end
+      assert_sql(/"posts"@\{FORCE_INDEX=#{idx},DESC\}/) do
+        Post.force_index(idx, direction: "DESC").take
+      end
+      assert_sql(/"posts"@\{NO_FULL_SCAN,FORCE_INDEX=#{idx},DESC\}/) do
+        Post.force_index(idx, direction: "DESC").index_hint("NO_FULL_SCAN").take
+      end
+    end
+
+    def test_use_other_table
+      force_index = ->(q) { q.force_index("index_subscribers_on_nick") }
+      index_hint = ->(q) { q.index_hint("NO_FULL_SCAN") }
+      post_from = Post.select(:id).from("subscribers")
+
+      assert_sql(/subscribers@\{FORCE_INDEX=index_subscribers_on_nick\}/) do
+        post_from.then(&force_index).take
+      end
+
+      assert_sql(/subscribers@\{NO_FULL_SCAN\}/) do
+        post_from.then(&index_hint).take
+      end
+
+      [force_index, index_hint].permutation do |procs|
+        assert_sql(/subscribers@\{NO_FULL_SCAN,FORCE_INDEX=index_subscribers_on_nick\}/) do
+          post_from.then(&procs.reduce(:<<)).take
+        end
+      end
+    end
+
+    def test_from_with_space
+      ["foo\t", "foo  "].each do
+        assert_match(/FROM foo@\{H\}/, Post.from(_1).index_hint("H").to_sql)
+        assert_match(
+          /FROM foo@\{H,FORCE_INDEX=i\}/,
+          Post.from(_1).index_hint("H").force_index("i").to_sql
+        )
+      end
+    end
+
+    def test_reset_with_from
+      force_index = ->(q) { q.force_index("type") }
+      index_hint = ->(q) { q.index_hint("NO_FULL_SCAN") }
+
+      [force_index, index_hint].permutation do
+        assert_match(
+          /FROM age\z/,
+          Post.then(&_1).then(&_2).from("age").to_sql
+        )
+      end
+      [force_index, index_hint].each do
+        assert_match(/FROM age\z/, Post.then(&_1).from("age").to_sql)
+      end
+
+    end
+
+    def test_ignore_from_multiple_tables
+      tables = "ta, tb"
+      from = ->(q) { q.from(tables) }
+      force_index = ->(q) { q.force_index("type") }
+      index_hint = ->(q) { q.index_hint("NO_FULL_SCAN") }
+
+      (
+        [from, force_index, index_hint].permutation +
+        [from, force_index].permutation +
+        [from, index_hint].permutation
+      ).each do |procs|
+        assert_match(
+          /FROM #{tables}\z/,
+          Post.then(&procs.reduce(:<<)).to_sql
+        )
+      end
+    end
+
+    def test_ignore_from_subquery
+      subquery = Post.where("created_at < ?", 1.year.ago)
+      from = ->(q) { q.from(subquery) }
+      force_index = ->(q) { q.force_index("safeword") }
+      index_hint = ->(q) { q.index_hint("NO_FULL_SCAN") }
+
+      (
+        [from, force_index, index_hint].permutation +
+        [from, force_index].permutation +
+        [from, index_hint].permutation
+      ).each do |procs|
+        refute_match(
+          /safeword|NO_FULL_SCAN|@/,
+          Post.then(&procs.reduce(:<<)).to_sql
+        )
+      end
+    end
+  end
+end