diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb index dc2537d36aa9a2d5b826ee5af68dcda6ec621b17..e6612bc3aadeda45d4cdf82e09c797699f0d3c9f 100644 --- a/lib/gitlab/database.rb +++ b/lib/gitlab/database.rb @@ -69,6 +69,30 @@ module Gitlab end end + def self.with_connection_pool(pool_size) + pool = create_connection_pool(pool_size) + + yield(pool) + + ensure + pool.disconnect! + end + + def self.create_connection_pool(pool_size) + # See activerecord-4.2.7.1/lib/active_record/connection_adapters/connection_specification.rb + env = Rails.env + original_config = ActiveRecord::Base.configurations + env_config = original_config[env].merge('pool' => pool_size) + config = original_config.merge(env => env_config) + + spec = + ActiveRecord:: + ConnectionAdapters:: + ConnectionSpecification::Resolver.new(config).spec(env.to_sym) + + ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec) + end + def self.connection ActiveRecord::Base.connection end diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb index b142b3a27817fd020de44b78a3ddc11328e8578d..c9be832b1a091747da83ac41ee5027d7127b827d 100644 --- a/spec/lib/gitlab/database_spec.rb +++ b/spec/lib/gitlab/database_spec.rb @@ -71,6 +71,54 @@ describe Gitlab::Database, lib: true do end end + describe '.with_connection_pool' do + it 'creates a new connection pool and disconnect it after used' do + closed_pool = nil + + described_class.with_connection_pool(1) do |pool| + pool.with_connection do |connection| + connection.execute('SELECT 1 AS value') + end + + expect(pool).to be_connected + + closed_pool = pool + end + + expect(closed_pool).not_to be_connected + end + + it 'disconnects the pool even an exception was raised' do + error = Class.new(RuntimeError) + closed_pool = nil + + begin + described_class.with_connection_pool(1) do |pool| + pool.with_connection do |connection| + connection.execute('SELECT 1 AS value') + end + + closed_pool = pool + + raise error.new('boom') + end + rescue error + end + + expect(closed_pool).not_to be_connected + end + end + + describe '.create_connection_pool' do + it 'creates a new connection pool with specific pool size' do + pool = described_class.create_connection_pool(5) + + expect(pool) + .to be_kind_of(ActiveRecord::ConnectionAdapters::ConnectionPool) + expect(pool.spec.config[:pool]).to eq(5) + end + end + describe '#true_value' do it 'returns correct value for PostgreSQL' do expect(described_class).to receive(:postgresql?).and_return(true)