Have some simple way to create connection pool

1 files changed, 24 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