module ActiveRecordPartitioning

# File lib/activerecord_partitioning.rb, line 40
def current_connection_pool_config
  Thread.current[:current_connection_pool_config]
end

# File lib/activerecord_partitioning.rb, line 32
def default_config
  @default_config
end

# File lib/activerecord_partitioning.rb, line 36
def default_config=(config)
  @default_config = config
end

# File lib/activerecord_partitioning.rb, line 13
def reset
  self.default_config = nil
  ActiveRecord::Base.clear_active_connections!
  ActiveRecord::Base.connection_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
end

# File lib/activerecord_partitioning.rb, line 6
def setup(key_name, default_config = nil, store = {})
  self.default_config = default_config.try(:symbolize_keys)
  new_pools = ConnectionPools.new(key_name, store)
  new_pools.merge!(ActiveRecord::Base.connection_handler.connection_pools)
  ActiveRecord::Base.connection_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new(new_pools)
end

# File lib/activerecord_partitioning.rb, line 19
def with_connection_pool(config, &block)
  config.try(:symbolize_keys!)
  origin = Thread.current[:current_connection_pool_config]
  Thread.current[:current_connection_pool_config] = (self.default_config || {}).merge(config || {})
  if ActiveRecord::Base.connection_pool.nil?
    ActiveRecord::Base.establish_connection(self.current_connection_pool_config)
  end
  yield if block_given?
ensure
  ActiveRecord::Base.clear_active_connections!
  Thread.current[:current_connection_pool_config] = origin
end