Merge pull request rails#43372 from eileencodes/add-option-to-lazy-load-schema-cache

Add ability to lazily load the schema cache on connection

Author: eileencodes

activerecord/CHANGELOG.md

@@ -1,3 +1,11 @@
+*   Add option to lazily load the schema cache on the connection.
+
+    Previously, the only way to load the schema cache in Active Record was through the Railtie on boot. This option provides the ability to load the schema cache on the connection after it's been established. Loading the cache lazily on the connection can be beneficial for Rails applications that use multiple databases because it will load the cache at the time the connection is established. Currently Railties doesn't have access to the connections before boot.
+
+    To use the cache, set `config.active_record.lazily_load_schema_cache = true` in your application configuration. In addition a `schema_cache_path` should be set in your database configuration if you don't want to use the default "db/schema_cache.yml" path.
+
+    *Eileen M. Uchitelle*
+
 *   Allow automatic `inverse_of` detection for associations with scopes.
 
     Automatic `inverse_of` detection now works for associations with scopes. For

activerecord/lib/active_record.rb

@@ -170,6 +170,12 @@ module Tasks
   autoload :TestDatabases, "active_record/test_databases"
   autoload :TestFixtures, "active_record/fixtures"
 
+  # Lazily load the schema cache. This option will load the schema cache
+  # when a connection is established rather than on boot. If set,
+  # +config.active_record.use_schema_cache_dump+ will be set to false.
+  singleton_class.attr_accessor :lazily_load_schema_cache
+  self.lazily_load_schema_cache = false
+
   # A list of tables or regex's to match tables to ignore when
   # dumping the schema cache. For example if this is set to +[/^_/]+
   # the schema cache will not dump tables named with an underscore.

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -19,6 +19,13 @@ def get_schema_cache(connection)
       def set_schema_cache(cache)
         self.schema_cache = cache
       end
+
+      def lazily_set_schema_cache
+        return unless ActiveRecord.lazily_load_schema_cache
+
+        cache = SchemaCache.load_from(db_config.lazy_schema_cache_path)
+        set_schema_cache(cache)
+      end
     end
 
     class NullPool # :nodoc:
@@ -147,6 +154,8 @@ def initialize(pool_config)
 
         @async_executor = build_async_executor
 
+        lazily_set_schema_cache
+
         @reaper = Reaper.new(self, db_config.reaping_frequency)
         @reaper.run
       end

activerecord/lib/active_record/database_configurations/hash_config.rb

@@ -109,6 +109,18 @@ def schema_cache_path
         configuration_hash[:schema_cache_path]
       end
 
+      def default_schema_cache_path
+        "db/schema_cache.yml"
+      end
+
+      def lazy_schema_cache_path
+        schema_cache_path || default_schema_cache_path
+      end
+
+      def primary? # :nodoc:
+        Base.configurations.primary?(name)
+      end
+
       # Determines whether to dump the schema for a database.
       def schema_dump
         configuration_hash.fetch(:schema_dump, true)

activerecord/lib/active_record/railtie.rb

@@ -130,7 +130,7 @@ class Railtie < Rails::Railtie # :nodoc:
     initializer "active_record.check_schema_cache_dump" do
       check_schema_cache_dump_version = config.active_record.check_schema_cache_dump_version
 
-      if config.active_record.use_schema_cache_dump
+      if config.active_record.use_schema_cache_dump && !config.active_record.lazily_load_schema_cache
         config.after_initialize do |app|
           ActiveSupport.on_load(:active_record) do
             db_config = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env).first

activerecord/test/cases/connection_adapters/schema_cache_test.rb

@@ -319,6 +319,54 @@ def test_clear_data_source_cache
         assert_not @cache.columns_hash?("posts")
       end
 
+      unless in_memory_db?
+        def test_when_lazily_load_schema_cache_is_set_cache_is_lazily_populated_when_est_connection
+          tempfile = Tempfile.new(["schema_cache-", ".yml"])
+          original_config = ActiveRecord::Base.connection_db_config
+          new_config = original_config.configuration_hash.merge(schema_cache_path: tempfile.path)
+
+          ActiveRecord::Base.establish_connection(new_config)
+
+          # cache is empty
+          assert_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@primary_keys)
+          assert_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@data_sources)
+          assert_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@indexes)
+
+          # calling dump_to will load data sources, but not the rest of the cache
+          # so we need to set the cache manually. This essentially mimics the behavior
+          # of the Railtie.
+          cache = SchemaCache.new(ActiveRecord::Base.connection)
+          cache.dump_to(tempfile.path)
+          ActiveRecord::Base.connection.schema_cache = cache
+
+          assert File.exist?(tempfile)
+          assert_not_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@primary_keys)
+          assert_not_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@data_sources)
+          assert_not_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@indexes)
+
+          # assert cache is empty on new connection
+          ActiveRecord::Base.establish_connection(new_config)
+
+          assert File.exist?(tempfile)
+          assert_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@primary_keys)
+          assert_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@data_sources)
+          assert_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@indexes)
+
+          # cache is lazily loaded when lazily loading is on
+          old_config = ActiveRecord.lazily_load_schema_cache
+          ActiveRecord.lazily_load_schema_cache = true
+          ActiveRecord::Base.establish_connection(new_config)
+
+          assert File.exist?(tempfile)
+          assert_not_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@primary_keys)
+          assert_not_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@data_sources)
+          assert_not_empty ActiveRecord::Base.connection.schema_cache.instance_variable_get(:@indexes)
+        ensure
+          ActiveRecord.lazily_load_schema_cache = old_config
+          ActiveRecord::Base.establish_connection(:arunit)
+        end
+      end
+
       private
         def schema_dump_path
           "#{ASSETS_ROOT}/schema_dump_5_1.yml"

activerecord/test/cases/database_configurations/hash_config_test.rb

@@ -135,6 +135,16 @@ def test_database_tasks_overrides_with_value
         config = HashConfig.new("default_env", "primary", database_tasks: "str")
         assert_equal true, config.database_tasks?
       end
+
+      def test_schema_cache_path_default_for_primary
+        config = HashConfig.new("default_env", "primary", {})
+        assert_equal "db/schema_cache.yml", config.default_schema_cache_path
+      end
+
+      def test_schema_cache_path_configuration_hash
+        config = HashConfig.new("default_env", "primary", { schema_cache_path: "db/config_schema_cache.yml" })
+        assert_equal "db/config_schema_cache.yml", config.schema_cache_path
+      end
     end
   end
 end