Merge branch 'create-connection-pool' into 'master'

Introduce Gitlab::Database.with_connection_pool

See merge request !9192

2 files changed, 73 insertions, 0 deletions

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index a3d6be9959b..a47d7e98a62 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -69,6 +69,31 @@ module Gitlab
       end
     end
 
+    def self.with_connection_pool(pool_size)
+      pool = create_connection_pool(pool_size)
+
+      begin
+        yield(pool)
+      ensure
+        pool.disconnect!
+      end
+    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 41252f31997..f01c42aff91 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -77,6 +77,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)