Allow applications to register custom database configurations

activerecord/CHANGELOG.md

@@ -1,3 +1,47 @@
+*   Allow applications to register a custom database configuration handler.
+
+    Adds a mechanism for registering a custom handler for cases where you want database configurations to respond to custom methods. This is useful for non-Rails database adapters or tools like Vitess that you may want to configure differently from a standard `HashConfig` or `UrlConfig`.
+
+    Given the following database YAML we want the `animals` db to create a `CustomConfig` object instead while the `primary` database will be a `UrlConfig`:
+
+    ```yaml
+    development:
+      primary:
+        url: postgres://localhost/primary
+      animals:
+        url: postgres://localhost/animals
+        custom_config:
+          sharded: 1
+    ```
+
+    To register a custom handler first make a class that has your custom methods:
+
+    ```ruby
+    class CustomConfig < ActiveRecord::DatabaseConfigurations::UrlConfig
+      def sharded?
+        custom_config.fetch("sharded", false)
+      end
+
+      private
+        def custom_config
+          configuration_hash.fetch(:custom_config)
+        end
+    end
+    ```
+
+    Then register the config in an inializer:
+
+    ```ruby
+    ActiveRecord::DatabaseConfigurations.register_db_config_handler do |env_name, name, url, config|
+      next unless config.key?(:custom_config)
+      CustomConfig.new(env_name, name, url, config)
+    end
+    ```
+
+    When the application is booted, configuration hashes with the `:custom_config` key will be `CustomConfig` objects and respond to `sharded?`. Applications must handle the condition in which Active Record should use their custom handler.
+
+    *Eileen M. Uchitelle and John Crepezzi*
+
 *   `ActiveRecord::Base.serialize` no longer uses YAML by default.
 
     YAML isn't particularly performant and can lead to security issues

activerecord/lib/active_record/database_configurations.rb

@@ -8,14 +8,62 @@
 
 module ActiveRecord
   # ActiveRecord::DatabaseConfigurations returns an array of DatabaseConfig
-  # objects (either a HashConfig or UrlConfig) that are constructed from the
-  # application's database configuration hash or URL string.
+  # objects that are constructed from the application's database
+  # configuration hash or URL string.
+  #
+  # The array of +DatabaseConfig+ objects in an application default to
+  # either a +HashConfig+ or +UrlConfig+. If you register a custom handler,
+  # objects will be created according to the conditions of the handler.
   class DatabaseConfigurations
     class InvalidConfigurationError < StandardError; end
 
     attr_reader :configurations
     delegate :any?, to: :configurations
 
+    singleton_class.attr_accessor :db_config_handlers # :nodoc:
+    self.db_config_handlers = [] # :nodoc:
+
+    # Allows an application to register a custom handler for database configuration
+    # objects. This is useful for creating a custom handler that responds to
+    # methods your application needs but Active Record doesn't implement. For
+    # example if you are using Vitess, you may want your Vitess configurations
+    # to respond to `sharded?`. To implement this define the following in an
+    # initializer:
+    #
+    #   ActiveRecord::DatabaseConfigurations.register_db_config_handler do |env_name, name, url, config|
+    #     next unless config.key?(:vitess)
+    #     VitessConfig.new(env_name, name, config)
+    #   end
+    #
+    # Note: applications must handle the condition in which custom config should be
+    # created in your handler registration otherwise all objects will use the custom
+    # handler.
+    #
+    # Then define your +VitessConfig+ to respond to the methods your application
+    # needs. It is recommended that you inherit from one of the existing
+    # database config classes to avoid having to reimplement all methods. Custom
+    # config handlers should only implement methods Active Record does not.
+    #
+    #   class VitessConfig < ActiveRecord::DatabaseConfigurations::UrlConfig
+    #     def sharded?
+    #       configuration_hash.fetch("sharded", false)
+    #     end
+    #   end
+    #
+    # For configs that have a +:vitess+ key, a +VitessConfig+ object will be
+    # created instead of a +UrlConfig+.
+    def self.register_db_config_handler(&block)
+      db_config_handlers << block
+    end
+
+    register_db_config_handler do |env_name, name, url, config|
+      if url
+        UrlConfig.new(env_name, name, url, config)
+      else
+        HashConfig.new(env_name, name, config)
+      end
+    end
+
     def initialize(configurations = {})
       @configurations = build_configs(configurations)
     end
@@ -219,15 +267,16 @@ def build_db_config_from_string(env_name, name, config)
       end
 
       def build_db_config_from_hash(env_name, name, config)
-        if config.has_key?(:url)
-          url = config[:url]
-          config_without_url = config.dup
-          config_without_url.delete :url
+        url = config[:url]
+        config_without_url = config.dup
+        config_without_url.delete :url
 
-          UrlConfig.new(env_name, name, url, config_without_url)
-        else
-          HashConfig.new(env_name, name, config)
+        DatabaseConfigurations.db_config_handlers.reverse_each do |handler|
+          config = handler.call(env_name, name, url, config_without_url)
+          return config if config
         end
+
+        nil
       end
 
       def merge_db_environment_variables(current_env, configs)

activerecord/test/cases/connection_adapters/merge_and_resolve_default_url_config_test.rb

@@ -67,7 +67,7 @@ def test_resolver_with_nil_database_url_and_current_env
         ENV["RAILS_ENV"] = "foo"
         config = { "foo" => { "adapter" => "postgres", "url" => ENV["DATABASE_URL"] } }
         actual = resolve_db_config(:foo, config)
-        expected_config = { adapter: "postgres", url: nil }
+        expected_config = { adapter: "postgres" }
 
         assert_equal expected_config, actual.configuration_hash
       end

activerecord/test/cases/database_configurations_test.rb

@@ -84,6 +84,50 @@ def test_find_db_config_prioritize_db_config_object_for_the_current_env
     assert_equal ActiveRecord::ConnectionHandling::DEFAULT_ENV.call, config.env_name
     assert_equal ":memory:", config.database
   end
+
+  class CustomHashConfig < ActiveRecord::DatabaseConfigurations::HashConfig
+    def sharded?
+      custom_config.fetch("sharded", false)
+    end
+
+    private
+      def custom_config
+        configuration_hash.fetch(:custom_config)
+      end
+  end
+
+  def test_registering_a_custom_config_object
+    previous_handlers = ActiveRecord::DatabaseConfigurations.db_config_handlers
+
+    ActiveRecord::DatabaseConfigurations.register_db_config_handler do |env_name, name, _, config|
+      next unless config.key?(:custom_config)
+      CustomHashConfig.new(env_name, name, config)
+    end
+
+    configs = ActiveRecord::DatabaseConfigurations.new({
+      "test" => {
+        "config_1" => {
+          "database" => "db",
+          "custom_config" => {
+            "sharded" => 1
+          }
+        },
+        "config_2" => {
+          "database" => "db"
+        }
+      }
+    }).configurations
+
+    custom_config = configs.first
+    hash_config = configs.last
+
+    assert custom_config.is_a?(CustomHashConfig)
+    assert hash_config.is_a?(ActiveRecord::DatabaseConfigurations::HashConfig)
+
+    assert_predicate custom_config, :sharded?
+  ensure
+    ActiveRecord::DatabaseConfigurations.db_config_handlers = previous_handlers
+  end
 end
 
 class LegacyDatabaseConfigurationsTest < ActiveRecord::TestCase