connect callback accepts pool and socket, fix tests

src/epgsql_pool.app.src

@@ -16,10 +16,11 @@
          {min_reconnect_timeout, 100},
          {keep_alive_timeout, 60000},
          %%
-         %% callback is function of arity 1 which accepts pool name
-         %%
+         %% callback is function of arity 2 which accepts pool name and connection socket pid
          %% {on_connect_callback, {Module :: atom(), Function :: atom()}},
          {on_connect_callback, undefined},
+         %%
+         %% callback is function of arity 1 which accepts pool name
          %% {on_disconnect_callback, {Module :: atom(), Function :: atom()}}
          {on_disconnect_callback, undefined}
         ]}

src/epgsql_pool_worker.erl

@@ -107,11 +107,13 @@ handle_info(open_connection, #state{pool_name = PoolName, connection = Connectio
             {noreply, State#state{connection = Connection4}}
     end;
 
-handle_info(on_connect, #state{pool_name = PoolName} = State) ->
+handle_info(on_connect, #state{pool_name = PoolName, connection = Connection} = State) ->
     case application:get_env(epgsql_pool, on_connect_callback) of
         undefined -> ignore;
         {ok, undefined} -> ignore;
-        {ok, {M, F}} -> M:F(PoolName)
+        {ok, {M, F}} ->
+            #epgsql_connection{sock = Sock} = Connection,
+            M:F(PoolName, Sock)
     end,
     {noreply, State};
 

src/epgsql_test.erl

@@ -1,6 +1,6 @@
 -module(epgsql_test).
 
--export([test_run/0, on_connect/1, on_disconnect/1]).
+-export([test_run/0, on_connect/2, on_disconnect/1]).
 
 -spec test_run() -> ok.
 test_run() ->
@@ -16,8 +16,10 @@ test_run() ->
     {ok, _} = epgsql_pool:start(my_pool, 2, 2, Params),
 
     Qs = [
-        "create table category (id int, name varchar)",
-        "insert into category values (1, 'some')",
+        "CREATE TABLE category (id bigserial, title text, PRIMARY KEY (id))",
+        "CREATE TABLE item (id bigserial, category_id bigint, title text, num int, PRIMARY KEY (id), "
+        "FOREIGN KEY (category_id) REFERENCES category (id) ON DELETE SET NULL)",
+        "insert into category (title) values ('some'), ('other')",
         "select * from category"
     ],
 
@@ -30,9 +32,9 @@ test_run() ->
     ok.
 
 
--spec on_connect(term()) -> ok.
-on_connect(PoolName) ->
-    error_logger:info_msg("On Connect ~p", [PoolName]),
+-spec on_connect(term(), pid()) -> ok.
+on_connect(PoolName, Sock) ->
+    error_logger:info_msg("On Connect ~p ~p", [PoolName, Sock]),
     ok.
 
 

test/epgsql_pool_tests.erl

@@ -7,32 +7,36 @@
 get_set_settings_test() ->
      application:ensure_all_started(epgsql_pool),
 
-    ?assertEqual(#{connection_timeout => 10000,
-                   keep_alive_timeout => 60000,
-                   max_reconnect_timeout => 5000,
-                   min_reconnect_timeout => 100,
-                   pooler_get_worker_timeout => 10000,
-                   pooler_max_queue => 1000,
-                   query_timeout => 10000,
-                   transaction_timeout => 20000},
-        epgsql_pool:get_settings()),
+    ?assertMatch(#{
+        connection_timeout := 10000,
+        keep_alive_timeout := 60000,
+        max_reconnect_timeout := 5000,
+        min_reconnect_timeout := 100,
+        pooler_get_worker_timeout := 10000,
+        pooler_max_queue := 1000,
+        query_timeout := 10000,
+        transaction_timeout := 20000
+    }, epgsql_pool:get_settings()),
 
-    ok = epgsql_pool:set_settings(#{aa => bb,
-                                    pooler_max_queue => 500,
-                                    cc => dd,
-                                    max_reconnect_timeout => 777,
-                                    min_reconnect_timeout => 42,
-                                    dd => 42,
-                                    query_timeout => 555}),
+    ok = epgsql_pool:set_settings(#{
+        aa => bb,
+        pooler_max_queue => 500,
+        cc => dd,
+        max_reconnect_timeout => 777,
+        min_reconnect_timeout => 42,
+        dd => 42,
+        query_timeout => 555
+    }),
 
-    ?assertEqual(#{connection_timeout => 10000,
-                   keep_alive_timeout => 60000,
-                   max_reconnect_timeout => 777,
-                   min_reconnect_timeout => 42,
-                   pooler_get_worker_timeout => 10000,
-                   pooler_max_queue => 500,
-                   query_timeout => 555,
-                   transaction_timeout => 20000},
-                 epgsql_pool:get_settings()),
+    ?assertMatch(#{
+        connection_timeout := 10000,
+        keep_alive_timeout := 60000,
+        max_reconnect_timeout := 777,
+        min_reconnect_timeout := 42,
+        pooler_get_worker_timeout := 10000,
+        pooler_max_queue := 500,
+        query_timeout := 555,
+        transaction_timeout := 20000
+    }, epgsql_pool:get_settings()),
 
     ok.