Have some simple way to create connection pool

2 files changed, 72 insertions, 0 deletions

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index dc2537d36aa..e6612bc3aad 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 b142b3a2781..c9be832b1a0 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)