allow connection initialization

spec/database_spec.cr

@@ -0,0 +1,42 @@
+require "./spec_helper"
+
+describe DB::Database do
+  it "allows connection initialization" do
+    cnn_setup = 0
+    DB.open "dummy://localhost:1027?initial_pool_size=2&max_pool_size=4&max_idle_pool_size=1" do |db|
+      cnn_setup.should eq(0)
+
+      db.setup_connection do |cnn|
+        cnn_setup += 1
+      end
+
+      cnn_setup.should eq(2)
+
+      db.using_connection do
+        cnn_setup.should eq(2)
+        db.using_connection do
+          cnn_setup.should eq(2)
+          db.using_connection do
+            cnn_setup.should eq(3)
+            db.using_connection do
+              cnn_setup.should eq(4)
+            end
+            # the pool didn't shrink no new initialization should be done next
+            db.using_connection do
+              cnn_setup.should eq(4)
+            end
+          end
+          # the pool shrink 1. max_idle_pool_size=1
+          # after the previous end there where 2.
+          db.using_connection do
+            cnn_setup.should eq(4)
+            # so now there will be a new connection created
+            db.using_connection do
+              cnn_setup.should eq(5)
+            end
+          end
+        end
+      end
+    end
+  end
+end

src/db/database.cr

@@ -21,6 +21,7 @@ module DB
     getter uri
 
     @pool : Pool(Connection)
+    @setup_connection : Connection -> Nil
 
     # :nodoc:
     def initialize(@driver : Driver, @uri : URI)
@@ -28,8 +29,20 @@ module DB
       params = (query = uri.query) ? HTTP::Params.parse(query) : HTTP::Params.new(Hash(String, Array(String)).new)
       pool_options = @driver.connection_pool_options(params)
 
+      @setup_connection = ->(conn : Connection) {}
       @pool = uninitialized Pool(Connection) # in order to use self in the factory proc
-      @pool = Pool.new(**pool_options) { @driver.build_connection(self).as(Connection) }
+      @pool = Pool.new(**pool_options) {
+        conn = @driver.build_connection(self).as(Connection)
+        @setup_connection.call conn
+        conn
+      }
+    end
+
+    def setup_connection(&proc : Connection -> Nil)
+      @setup_connection = proc
+      @pool.each_resource do |conn|
+        @setup_connection.call conn
+      end
     end
 
     # Closes all connection to the database.

src/db/pool.cr

@@ -46,6 +46,13 @@ module DB
       end
     end
 
+    # :nodoc:
+    def each_resource
+      @available.each do |resource|
+        yield resource
+      end
+    end
+
     private def build_resource : T
       resource = @factory.call
       @total << resource