Schema Cache: cache table indexes

Useful to not query for indexes when an application uses schema cache.

Ref https://github.com/rails/rails/pull/35546

2 files changed, 43 insertions, 14 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/schema_cache.rb b/activerecord/lib/active_record/connection_adapters/schema_cache.rb
index 69b7a6e690..5b598620d3 100644
--- a/activerecord/lib/active_record/connection_adapters/schema_cache.rb
+++ b/activerecord/lib/active_record/connection_adapters/schema_cache.rb
@@ -13,6 +13,7 @@ module ActiveRecord
         @columns_hash = {}
         @primary_keys = {}
         @data_sources = {}
+        @indexes      = {}
       end
 
       def initialize_dup(other)
@@ -21,22 +22,25 @@ module ActiveRecord
         @columns_hash = @columns_hash.dup
         @primary_keys = @primary_keys.dup
         @data_sources = @data_sources.dup
+        @indexes      = @indexes.dup
       end
 
       def encode_with(coder)
-        coder["columns"] = @columns
+        coder["columns"]      = @columns
         coder["columns_hash"] = @columns_hash
         coder["primary_keys"] = @primary_keys
         coder["data_sources"] = @data_sources
-        coder["version"] = connection.migration_context.current_version
+        coder["indexes"]      = @indexes
+        coder["version"]      = connection.migration_context.current_version
       end
 
       def init_with(coder)
-        @columns = coder["columns"]
+        @columns      = coder["columns"]
         @columns_hash = coder["columns_hash"]
         @primary_keys = coder["primary_keys"]
         @data_sources = coder["data_sources"]
-        @version = coder["version"]
+        @indexes      = coder["indexes"] || {}
+        @version      = coder["version"]
       end
 
       def primary_keys(table_name)
@@ -57,6 +61,7 @@ module ActiveRecord
           primary_keys(table_name)
           columns(table_name)
           columns_hash(table_name)
+          indexes(table_name)
         end
       end
 
@@ -82,12 +87,17 @@ module ActiveRecord
         @columns_hash.key?(table_name)
       end
 
+      def indexes(table_name)
+        @indexes[table_name] ||= connection.indexes(table_name)
+      end
+
       # Clears out internal caches
       def clear!
         @columns.clear
         @columns_hash.clear
         @primary_keys.clear
         @data_sources.clear
+        @indexes.clear
         @version = nil
       end
 
@@ -101,10 +111,10 @@ module ActiveRecord
         @columns_hash.delete name
         @primary_keys.delete name
         @data_sources.delete name
+        @indexes.delete name
       end
 
       private
-
         def prepare_data_sources
           connection.data_sources.each { |source| @data_sources[source] = true }
         end
diff --git a/activerecord/test/cases/connection_adapters/schema_cache_test.rb b/activerecord/test/cases/connection_adapters/schema_cache_test.rb
index 79067c0f8a..722db1f69f 100644
--- a/activerecord/test/cases/connection_adapters/schema_cache_test.rb
+++ b/activerecord/test/cases/connection_adapters/schema_cache_test.rb
@@ -6,8 +6,8 @@ module ActiveRecord
   module ConnectionAdapters
     class SchemaCacheTest < ActiveRecord::TestCase
       def setup
-        connection = ActiveRecord::Base.connection
-        @cache     = SchemaCache.new connection
+        @connection = ActiveRecord::Base.connection
+        @cache      = SchemaCache.new @connection
       end
 
       def test_primary_key
@@ -19,6 +19,7 @@ module ActiveRecord
         @cache.columns_hash("posts")
         @cache.data_sources("posts")
         @cache.primary_keys("posts")
+        @cache.indexes("posts")
 
         new_cache = YAML.load(YAML.dump(@cache))
         assert_no_queries do
@@ -26,21 +27,34 @@ module ActiveRecord
           assert_equal 12, new_cache.columns_hash("posts").size
           assert new_cache.data_sources("posts")
           assert_equal "id", new_cache.primary_keys("posts")
+          assert_equal 1, new_cache.indexes("posts").size
         end
       end
 
       def test_yaml_loads_5_1_dump
-        body = File.open(schema_dump_path).read
-        cache = YAML.load(body)
+        @cache = YAML.load(File.read(schema_dump_path))
 
         assert_no_queries do
-          assert_equal 11, cache.columns("posts").size
-          assert_equal 11, cache.columns_hash("posts").size
-          assert cache.data_sources("posts")
-          assert_equal "id", cache.primary_keys("posts")
+          assert_equal 11, @cache.columns("posts").size
+          assert_equal 11, @cache.columns_hash("posts").size
+          assert @cache.data_sources("posts")
+          assert_equal "id", @cache.primary_keys("posts")
         end
       end
 
+      def test_yaml_loads_5_1_dump_without_indexes_still_queries_for_indexes
+        @cache = YAML.load(File.read(schema_dump_path))
+
+        # Simulate assignment in railtie after loading the cache.
+        old_cache, @connection.schema_cache = @connection.schema_cache, @cache
+
+        assert_queries :any, ignore_none: true do
+          assert_equal 1, @cache.indexes("posts").size
+        end
+      ensure
+        @connection.schema_cache = old_cache
+      end
+
       def test_primary_key_for_non_existent_table
         assert_nil @cache.primary_keys("omgponies")
       end
@@ -55,11 +69,17 @@ module ActiveRecord
         assert_equal columns_hash, @cache.columns_hash("posts")
       end
 
+      def test_caches_indexes
+        indexes = @cache.indexes("posts")
+        assert_equal indexes, @cache.indexes("posts")
+      end
+
       def test_clearing
         @cache.columns("posts")
         @cache.columns_hash("posts")
         @cache.data_sources("posts")
         @cache.primary_keys("posts")
+        @cache.indexes("posts")
 
         @cache.clear!
 
@@ -92,7 +112,6 @@ module ActiveRecord
       end
 
       private
-
         def schema_dump_path
           "test/assets/schema_dump_5_1.yml"
         end