psqlcp.app.src

{application, psqlcp, [
    {description, "Postgresql connection pool with poolboy"},
    {vsn, "0.1"},
    {applications, [kernel, stdlib]},
    {modules, [psqlcp, psqlcp_worker]},
    {registered, [psqlcp]},
    {mod, {psqlcp, []}},
    {env, [
        {pools, [
            {fk, [
                {size, 10},
                {max_overflow, 10},
                {hostname, "127.0.0.1"},
                {database, "fk"},
                {username, "root"},
                {password, "111"}
            ]},
            {tvu, [
                {size, 10},
                {max_overflow, 10},
                {hostname, "127.0.0.1"},
                {database, "tvu"},
                {username, "root"},
                {password, "1111"}
            ]}
        ]}
    ]}
]}. 

psqlcp.erl

-module(psqlcp).
-behaviour(application).
-behaviour(supervisor).

-export([start/0, stop/0, squery/2, equery/3, with_transaction/2, with_connection/2]).
-export([start/2, stop/1]).
-export([init/1]).

start() ->
    application:start(?MODULE).

stop() ->
    application:stop(?MODULE).

start(_Type, _Args) ->
    supervisor:start_link({local, psqlcp_sup}, ?MODULE, []).

stop(_State) ->
    ok.

init([]) ->
    {ok, Pools} = application:get_env(psqlcp, pools),
    PoolSpecs = lists:map(fun({Pool, Pool_config}) ->
        Args = [{name, {local, Pool}},
                {worker_module, psqlcp_worker}]
                ++ Pool_config,
        poolboy:child_spec(Pool, Args)
    end, Pools),
    {ok, {{one_for_one, 10, 10}, PoolSpecs}}.

%%%
%%% -----------------------------------------------------------------------
%%%

with_connection(Pool, Function) when erlang:is_function(Function, 1) ->
    poolboy:transaction(Pool, fun(Conn) ->
        erlang:apply(Function, [Conn])
    end).

with_transaction(Pool, Function) when erlang:is_function(Function, 1) ->
    poolboy:transaction(Pool, fun(Conn) ->
        ?MODULE:squery(Conn, <<"begin">>),
        Ret = erlang:apply(Function, [Conn]),
        ?MODULE:squery(Conn, <<"commit">>),
        Ret
    end).

%%%
%%% -----------------------------------------------------------------------
%%%

equery(Conn, Stmt, Params) when erlang:is_pid(Conn) ->
    equery({conn, Conn}, Stmt, Params);

equery(Pool, Stmt, Params) when erlang:is_atom(Pool) ->
    equery({pool, Pool}, Stmt, Params);

equery({conn, Conn}, Stmt, Params) when erlang:is_pid(Conn) ->
    gen_server:call(Conn, {equery, Stmt, Params});

equery({pool, Pool}, Stmt, Params)  ->
    poolboy:transaction(Pool, fun(Conn) ->
        equery({conn, Conn}, Stmt, Params)
    end).

%%%
%%% -----------------------------------------------------------------------
%%%

squery(Conn, Sql) when erlang:is_pid(Conn) ->
    squery({conn, Conn}, Sql);

squery(Pool, Sql) when erlang:is_atom(Pool) ->
    squery({pool, Pool}, Sql);

squery({conn, Conn}, Sql) ->
    gen_server:call(Conn, {equery, Sql});

squery({pool, Pool}, Sql) ->
    poolboy:transaction(Pool, fun(Conn) ->
        gen_server:call(Conn, {squery, Sql})
    end).

%%%
%%% -----------------------------------------------------------------------
%%%

psqlcp_worker.erl

-module(psqlcp_worker).
-behaviour(gen_server).
-behaviour(poolboy_worker).

-export([start_link/1]).
-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2,
         code_change/3]).

-define(RECONNECT_TIMEOUT, 6000).

-record(state, {
    conn,
    hostname,
    database,
    username,
    password,
    timeout=?RECONNECT_TIMEOUT,
    offline=true
}).

start_link(Args) ->
    gen_server:start_link(?MODULE, Args, []).

init(Args) ->
    Hostname = proplists:get_value(hostname, Args),
    Database = proplists:get_value(database, Args),
    Username = proplists:get_value(username, Args),
    Password = proplists:get_value(password, Args),
    State = connect(Hostname, Username, Password, Database),
    {ok, State, ?RECONNECT_TIMEOUT}.

connect(Hostname, Username, Password, Database)->
    connect(Hostname, Username, Password, Database, ?RECONNECT_TIMEOUT).

connect(Hostname, Username, Password, Database, Timeout)->
    case pgsql:connect(Hostname, Username, Password, [
        {database, Database}
    ]) of
    {ok, Conn}  ->
        #state{
            offline=false,
            conn=Conn,
            hostname=Hostname,
            username=Username,
            password=Password,
            database=Database,
            timeout=Timeout
        };
    {error, _} ->
        #state{
            offline=true,
            hostname=Hostname,
            username=Username,
            password=Password,
            database=Database,
            timeout=Timeout
        }
    end.

handle_call(Request,From,
        #state{
            offline=true,
            hostname=Hostname,
            username=Username,
            password=Password,
            database=Database
        }
    ) ->
    handle_call(
        Request,
        From,
        connect(Hostname, Username, Password, Database)
    );

handle_call({squery, Sql}, _From, #state{conn=Conn}=State) ->
    {reply, pgsql:squery(Conn, Sql), State};
handle_call({equery, Stmt, Params}, _From, #state{conn=Conn}=State) ->
    {reply, pgsql:equery(Conn, Stmt, Params), State};
handle_call(_Request, _From, State) ->
    {reply, ok, State}.

handle_cast(_Msg, State) ->
    {noreply, State}.

handle_info(_Info, #state{
            offline=true,
            hostname=Hostname,
            username=Username,
            password=Password,
            database=Database
        }) ->
    Newstate = connect(Hostname, Username, Password, Database),
    {noreply, Newstate, Newstate#state.timeout};

handle_info(_Info, State) ->
    {noreply, State, State#state.timeout}.

terminate(_Reason, #state{offline=true}) ->
    ok;

terminate(_Reason, #state{conn=Conn}) ->
    ok = pgsql:close(Conn),
    ok.

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.