/src/dbdrivers/postgresql/epgsql/pgsql.erl

  1%%% Copyright (C) 2008 - Will Glozer.  All rights reserved.
  2
  3%% Modified by Marc Worrell, 2009.
  4%% Added last_id/2, reset_id/2, squery1/2, equery1/2, equery1/3, assoc/2, assoc/3, columns/2
  5%% Adapted with_transaction/2
  6
  7-module(pgsql).
  8
  9-export([connect/2, connect/3, connect/4, close/1]).
 10-export([last_id/2, reset_id/2, squery1/2, equery1/2, equery1/3, assoc/2, assoc/3]).
 11-export([get_parameter/2, squery/2, equery/2, equery/3]).
 12-export([parse/2, parse/3, parse/4, describe/2, describe/3]).
 13-export([bind/3, bind/4, execute/2, execute/3, execute/4]).
 14-export([close/2, close/3, sync/1]).
 15-export([with_transaction/2]).
 16
 17-include("pgsql.hrl").
 18-include("zotonic.hrl").
 19
 20-define(timeout, 5000).
 21
 22%% -- client interface --
 23
 24connect(Host, Opts) ->
 25    connect(Host, os:getenv("USER"), "", Opts).
 26
 27connect(Host, Username, Opts) ->
 28    connect(Host, Username, "", Opts).
 29
 30connect(Host, Username, Password, Opts) ->
 31    {ok, C} = pgsql_connection:start_link(),
 32    pgsql_connection:connect(C, Host, Username, Password, Opts).
 33
 34close(C) when is_pid(C) ->
 35    catch pgsql_connection:stop(C),
 36    ok.
 37
 38get_parameter(C, Name) ->
 39    pgsql_connection:get_parameter(C, Name).
 40
 41last_id(C, Table) when is_atom(Table) ->
 42    last_id(C, atom_to_list(Table));
 43last_id(C, Table) ->
 44    equery1(C, "select currval(pg_get_serial_sequence($1, 'id'))", [Table]).
 45
 46reset_id(C, Table) when is_atom(Table) ->
 47    reset_id(C, atom_to_list(Table));
 48reset_id(C, Table) ->
 49    {ok, Max} = equery1(C, "select max(id) from \""++Table++"\""),
 50    equery1(C, "select setval(pg_get_serial_sequence($1, 'id'), $2)", [Table, Max+1]).
 51
 52assoc(C, Sql) ->
 53    assoc(C, Sql, []).
 54assoc(C, Sql, Parameters) ->
 55    case equery(C, Sql, Parameters) of
 56        {ok, Columns, Rows} ->
 57            Names = [ list_to_atom(binary_to_list(Name)) || #column{name=Name} <- Columns ],
 58            Rows1 = [ lists:zip(Names, tuple_to_list(Row)) || Row <- Rows ],
 59            {ok, Rows1};
 60        Other -> Other
 61    end.
 62
 63squery1(C, Sql) ->
 64    case squery(C,Sql) of
 65        {ok, _Columns, []} -> {error, noresult};
 66        {ok, _RowCount, _Columns, []} -> {error, noresult};
 67        {ok, _Columns, [Row|_]} -> {ok, element(1, Row)};
 68        {ok, _RowCount, _Columns, [Row|_]} -> {ok, element(1, Row)};
 69        Other -> Other
 70    end.
 71
 72equery1(C, Sql) ->
 73    equery1(C, Sql, []).
 74equery1(C, Sql, Parameters) ->
 75    case equery(C, Sql, Parameters) of
 76        {ok, _Columns, []} -> {error, noresult};
 77        {ok, _RowCount, _Columns, []} -> {error, noresult};
 78        {ok, _Columns, [Row|_]} -> {ok, element(1, Row)};
 79        {ok, _RowCount, _Columns, [Row|_]} -> {ok, element(1, Row)};
 80        Other -> Other
 81    end.
 82
 83
 84    
 85squery(C, Sql) ->
 86    ok = pgsql_connection:squery(C, Sql),
 87    case receive_results(C, []) of
 88        [Result] -> Result;
 89        Results  -> Results
 90    end.
 91
 92equery(C, Sql) ->
 93    equery(C, Sql, []).
 94
 95
 96equery(C, Sql, Parameters) when is_tuple(Parameters) ->
 97    equery(C, Sql, tuple_to_list(Parameters));
 98equery(C, Sql, Parameters) ->
 99    case pgsql_connection:parse(C, "", Sql, []) of
100        {ok, #statement{types = Types} = S} ->
101            Typed_Parameters = lists:zip(Types, Parameters),
102            ok = pgsql_connection:equery(C, S, Typed_Parameters),
103            receive_result(C);
104        Error ->
105            ?LOG("SQL error ~p : ~p", [Error, Sql]),
106            Error
107    end.
108
109%% parse
110
111parse(C, Sql) ->
112    parse(C, "", Sql, []).
113
114parse(C, Sql, Types) ->
115    parse(C, "", Sql, Types).
116
117parse(C, Name, Sql, Types) ->
118    pgsql_connection:parse(C, Name, Sql, Types).
119
120%% bind
121
122bind(C, Statement, Parameters) ->
123    bind(C, Statement, "", Parameters).
124
125bind(C, Statement, PortalName, Parameters) ->
126    pgsql_connection:bind(C, Statement, PortalName, Parameters).
127
128%% execute
129
130execute(C, S) ->
131    execute(C, S, "", 0).
132
133execute(C, S, N) ->
134    execute(C, S, "", N).
135
136execute(C, S, PortalName, N) ->
137    pgsql_connection:execute(C, S, PortalName, N),
138    receive_extended_result(C).
139
140%% statement/portal functions
141
142describe(C, #statement{name = Name}) ->
143    pgsql_connection:describe(C, statement, Name).
144
145describe(C, Type, Name) ->
146    pgsql_connection:describe(C, Type, Name).
147
148close(C, #statement{name = Name}) ->
149    pgsql_connection:close(C, statement, Name).
150
151close(C, Type, Name) ->
152    pgsql_connection:close(C, Type, Name).
153
154sync(C) ->
155    pgsql_connection:sync(C).
156
157%% misc helper functions
158with_transaction(C, F) ->
159    try
160        {ok, [], []} = squery(C, "BEGIN"),
161        R = F(C),
162        {ok, [], []} = squery(C, "COMMIT"),
163        R
164    catch
165        E:Why ->
166            ?LOG("Exception in transaction: \"~p,~p\"", [E,Why]),
167            squery(C, "ROLLBACK"),
168            {rollback, Why}
169    end.
170
171%% -- internal functions --
172
173receive_result(C) ->
174    R = receive_result(C, [], []),
175    receive
176        {pgsql, C, done} -> R
177    end.
178
179receive_results(C, Results) ->
180    case receive_result(C, [], []) of
181        done -> lists:reverse(Results);
182        R    -> receive_results(C, [R | Results])
183    end.
184
185receive_result(C, Cols, Rows) ->
186    receive
187        {pgsql, C, {columns, Cols2}} ->
188            receive_result(C, Cols2, Rows);
189        {pgsql, C, {data, Row}} ->
190            receive_result(C, Cols, [Row | Rows]);
191        {pgsql, C, {error, _E} = Error} ->
192            Error;
193        {pgsql, C, {complete, {_Type, Count}}} ->
194            case Rows of
195                [] -> {ok, Count};
196                _L -> {ok, Count, Cols, lists:reverse(Rows)}
197            end;
198        {pgsql, C, {complete, _Type}} ->
199            {ok, Cols, lists:reverse(Rows)};
200        {pgsql, C, {notice, _N}} ->
201            receive_result(C, Cols, Rows);
202        {pgsql, C, done} ->
203            done
204    after
205        ?timeout -> {error, timeout}
206    end.
207
208receive_extended_result(C)->
209    receive_extended_result(C, []).
210
211receive_extended_result(C, Rows) ->
212    receive
213        {pgsql, C, {data, Row}} ->
214            receive_extended_result(C, [Row | Rows]);
215        {pgsql, C, {error, _E} = Error} ->
216            Error;
217        {pgsql, C, suspended} ->
218            {partial, lists:reverse(Rows)};
219        {pgsql, C, {complete, {_Type, Count}}} ->
220            case Rows of
221                [] -> {ok, Count};
222                _L -> {ok, Count, lists:reverse(Rows)}
223            end;
224        {pgsql, C, {complete, _Type}} ->
225            {ok, lists:reverse(Rows)};
226        {pgsql, C, {notice, _N}} ->
227            receive_extended_result(C, Rows)
228    after
229        ?timeout -> {error, timeout}
230    end.