active-record-6.1-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_specification_name
        # if !defined?(@connection_specification_name) || @connection_specification_name.nil?
        #   return self == Base ? Base.name : superclass.connection_specification_name
        # end
        # @connection_specification_name
        connection_pool_name(DbConfig.db_key)
      end

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

        # we want to have multiple connection pools, one per db
        # owner_name = primary_class? ? Base.name : name
        owner_name = primary_class? ? Base.name : connection_pool_name(@db_key)
        self.connection_specification_name = owner_name

        db_config = ActiveRecord::Base.configurations.resolve(config_or_env)
        # key the db config by connection_pool_name instead of the class
        # [db_config, self]
        [db_config, owner_name]
      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