active-record-6.0-horizontal-sharding.rb

module MyGem
  class BaseRecord < ActiveRecord::Base
    self.abstract_class = true

    class << self
      # What we are doing here is creating a connection pool for each db on class load
      # (see bottom of class) and we override 2 active record base class methods for
      # establishing and retrieving each db connection

      # This method overrides the one in ActiveRecord::ConnectionHandling
      def connection
        pool_name = connection_pool_name(DbConfig.db_key)
        connection_handler.retrieve_connection(pool_name)
      end

      # This method overrides the one in ActiveRecord::ConnectionHandling
      def resolve_config_for_connection(config_or_env) # :nodoc:
        raise 'Anonymous class is not allowed.' unless name

        config_or_env ||= DEFAULT_ENV.call.to_sym
        # we are only changing the commented out line below since for this model
        # we want to have multiple connection pools, one per db
        # pool_name = primary_class? ? "primary" : name
        pool_name = primary_class? ? 'primary' : connection_pool_name(@db_key)
        self.connection_specification_name = pool_name

        resolver = ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new(
          ActiveRecord::Base.configurations
        )
        config_hash = resolver.resolve(config_or_env, pool_name).symbolize_keys
        config_hash[:name] = pool_name

        config_hash
      end

      def connection_pool_name(db_key)
        "MyGem::BaseRecord_#{db_key}"
      end
    end

    # setup db connections on class load
    DbConfig.configurations.each do |db_key, config|
      @db_key = db_key
      establish_connection(config)
    end
  end
end