/src/dbdrivers/postgresql/z_db.erl

  1%% @author Marc Worrell <marc@worrell.nl>
  2%% @copyright 2009 Marc Worrell
  3%% Date: 2009-04-07
  4%%
  5%% @doc Interface to database, uses database definition from Context
  6
  7%% Copyright 2009 Marc Worrell
  8%%
  9%% Licensed under the Apache License, Version 2.0 (the "License");
 10%% you may not use this file except in compliance with the License.
 11%% You may obtain a copy of the License at
 12%% 
 13%%     http://www.apache.org/licenses/LICENSE-2.0
 14%% 
 15%% Unless required by applicable law or agreed to in writing, software
 16%% distributed under the License is distributed on an "AS IS" BASIS,
 17%% WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 18%% See the License for the specific language governing permissions and
 19%% limitations under the License.
 20
 21-module(z_db).
 22-author("Marc Worrell <marc@worrell.nl").
 23
 24%% interface functions
 25-export([
 26    has_connection/1,
 27    transaction/2,
 28    transaction_clear/1,
 29    set/3,
 30    get/2,
 31    get_parameter/2,
 32    assoc_row/2,
 33    assoc_row/3,
 34    assoc_props_row/2,
 35    assoc_props_row/3,
 36    assoc/2,
 37    assoc/3,
 38    assoc_props/2,
 39    assoc_props/3,
 40    q/2,
 41    q/3,
 42    q1/2,
 43    q1/3,
 44    q_row/2,
 45    q_row/3,
 46    equery/2,
 47    equery/3,
 48    insert/2,
 49    insert/3,
 50    update/4,
 51    delete/3,
 52    select/3,
 53    columns/2,
 54    column_names/2,
 55    update_sequence/3,
 56    table_exists/2,
 57    ensure_table/3,
 58    drop_table/2,
 59    flush/1,
 60    
 61    assert_table_name/1,
 62    prepare_cols/2
 63]).
 64
 65
 66-include_lib("pgsql.hrl").
 67-include_lib("zotonic.hrl").
 68
 69
 70%% @doc Perform a function inside a transaction, do a rollback on exceptions
 71%% @spec transaction(Function, Context) -> FunctionResult | {error, Reason}
 72transaction(Function, #context{dbc=undefined} = Context) ->
 73    case has_connection(Context) of
 74        true ->
 75            Host     = Context#context.host,
 76            {ok, C}  = pgsql_pool:get_connection(Host),
 77            Context1 = Context#context{dbc=C},
 78            Result = try
 79                        {ok, [], []} = pgsql:squery(C, "BEGIN"),
 80                        R = Function(Context1),
 81                        {ok, [], []} = pgsql:squery(C, "COMMIT"),
 82                        R
 83                     catch
 84                        _:Why ->
 85                            pgsql:squery(C, "ROLLBACK"),
 86                            {rollback, {Why, erlang:get_stacktrace()}}
 87                     end,
 88            pgsql_pool:return_connection(Host, C),
 89            Result;
 90        false ->
 91            {rollback, {no_database_connection, erlang:get_stacktrace()}}
 92    end;
 93transaction(Function, Context) ->
 94    % Nested transaction, only keep the outermost transaction
 95    Function(Context).
 96
 97%% @doc Clear any transaction in the context, useful when starting a thread with this context.
 98transaction_clear(#context{dbc=undefined} = Context) ->
 99    Context;
100transaction_clear(Context) ->
101    Context#context{dbc=undefined}.
102
103
104%% @doc Simple get/set functions for db property lists
105set(Key, Props, Value) ->
106    lists:keystore(Key, 1, Props, {Key, Value}).
107get(Key, Props) ->
108    proplists:get_value(Key, Props).
109
110
111%% @doc Check if we have database connection
112has_connection(#context{host=Host}) ->
113    is_pid(erlang:whereis(Host)).
114
115
116%% @doc Transaction handler safe function for fetching a db connection
117get_connection(#context{dbc=undefined, host=Host} = Context) ->
118    case has_connection(Context) of
119        true ->
120            {ok, C} = pgsql_pool:get_connection(Host),
121            C;
122        false ->
123            none
124    end;
125get_connection(Context) ->
126    Context#context.dbc.
127
128%% @doc Transaction handler safe function for releasing a db connection
129return_connection(C, #context{dbc=undefined, host=Host}) ->
130    pgsql_pool:return_connection(Host, C);
131return_connection(_C, _Context) -> 
132    ok.
133
134
135assoc_row(Sql, Context) ->
136    assoc_row(Sql, [], Context).
137
138assoc_row(Sql, Parameters, Context) ->
139    case assoc(Sql, Parameters, Context) of
140        [Row|_] -> Row;
141        [] -> undefined
142    end.
143
144assoc_props_row(Sql, Context) ->
145    assoc_props_row(Sql, [], Context).
146
147assoc_props_row(Sql, Parameters, Context) ->
148    case assoc_props(Sql, Parameters, Context) of
149        [Row|_] -> Row;
150        [] -> undefined
151    end.
152    
153
154get_parameter(Parameter, Context) ->
155    C = get_connection(Context),
156    try
157        {ok, Result} = pgsql:get_parameter(C, z_convert:to_binary(Parameter)),
158        Result
159    after
160        return_connection(C, Context)
161    end.
162    
163
164%% @doc Return property lists of the results of a query on the database in the Context
165%% @spec assoc(SqlQuery, Context) -> Rows
166assoc(Sql, Context) ->
167    assoc(Sql, [], Context).
168
169assoc(Sql, Parameters, Context) ->
170    case get_connection(Context) of
171        none -> [];
172        C ->
173            try
174                {ok, Result} = pgsql:assoc(C, Sql, Parameters),
175                Result
176            after
177                return_connection(C, Context)
178            end
179    end.
180
181
182assoc_props(Sql, Context) ->
183    assoc_props(Sql, [], Context).
184
185assoc_props(Sql, Parameters, Context) ->
186    case get_connection(Context) of
187        none -> [];
188        C ->
189            try
190                {ok, Result} = pgsql:assoc(C, Sql, Parameters),
191                merge_props(Result)
192            after
193                return_connection(C, Context)
194            end
195    end.
196
197
198q(Sql, Context) ->
199    q(Sql, [], Context).
200
201q(Sql, Parameters, Context) ->
202    case get_connection(Context) of
203        none -> [];
204        C ->
205            try
206                case pgsql:equery(C, Sql, Parameters) of
207                    {ok, _Affected, _Cols, Rows} -> Rows;
208                    {ok, _Cols, Rows} -> Rows;
209                    {ok, Rows} -> Rows
210                end
211            after
212                return_connection(C, Context)
213            end
214    end.
215
216q1(Sql, Context) ->
217    q1(Sql, [], Context).
218
219q1(Sql, Parameters, Context) ->
220    case get_connection(Context) of
221        none -> undefined;
222        C ->
223            try
224                case pgsql:equery1(C, Sql, Parameters) of
225                    {ok, Value} -> Value;
226                    {error, noresult} -> undefined
227                end
228            after
229                return_connection(C, Context)
230            end
231    end.
232
233q_row(Sql, Context) ->
234    q_row(Sql, [], Context).
235
236q_row(Sql, Args, Context) ->
237    case q(Sql, Args, Context) of
238        [Row|_] -> Row;
239        [] -> undefined
240    end.
241
242
243equery(Sql, Context) ->
244    equery(Sql, [], Context).
245    
246equery(Sql, Parameters, Context) ->
247    case get_connection(Context) of
248        none -> 
249            {error, noresult};
250        C ->
251            try
252                pgsql:equery(C, Sql, Parameters)
253            after
254                return_connection(C, Context)
255            end
256    end.
257
258%% @doc Insert a new row in a table, use only default values.
259%% @spec insert(Table, Context) -> {ok, Id}
260insert(Table, Context) when is_atom(Table) ->
261    insert(atom_to_list(Table), Context);
262insert(Table, Context) ->
263    assert_table_name(Table),
264    C = get_connection(Context),
265    try
266        pgsql:equery1(C, "insert into \""++Table++"\" default values returning id")
267    after
268        return_connection(C, Context)
269    end.
270
271
272%% @doc Insert a row, setting the fields to the props.  Unknown columns are serialized in the props column.
273%% When the table has an 'id' column then the new id is returned.
274%% @spec insert(Table::atom(), Props::proplist(), Context) -> {ok, Id} | Error
275insert(Table, [], Context) ->  
276    insert(Table, Context);
277insert(Table, Props, Context) when is_atom(Table) ->
278    insert(atom_to_list(Table), Props, Context);
279insert(Table, Props, Context) ->  
280    assert_table_name(Table),
281    Cols = column_names(Table, Context),
282    InsertProps = prepare_cols(Cols, Props),
283
284    InsertProps1 = case proplists:get_value(props, InsertProps) of
285        undefined ->
286            InsertProps;
287        PropsCol -> 
288            lists:keystore(props, 1, InsertProps, {props, cleanup_props(PropsCol)})
289    end,
290    
291    %% Build the SQL insert statement
292    {ColNames,Parameters} = lists:unzip(InsertProps1),
293    Sql = "insert into \""++Table++"\" (\"" 
294             ++ string:join([ atom_to_list(ColName) || ColName <- ColNames ], "\", \"")
295             ++ "\") values ("
296             ++ string:join([ [$$ | integer_to_list(N)] || N <- lists:seq(1, length(Parameters)) ], ", ")
297             ++ ")",
298
299    FinalSql = case lists:member(id, Cols) of
300        true -> Sql ++ " returning id";
301        false -> Sql
302    end,
303
304    C = get_connection(Context),
305    try
306        Id = case pgsql:equery1(C, FinalSql, Parameters) of
307                {ok, IdVal} -> IdVal;
308                {error, noresult} -> undefined
309             end,
310         {ok, Id}
311    after
312        return_connection(C, Context)
313    end.
314
315%% @doc Update a row in a table, merging the props list with any new props values
316%% @spec update(Table, Id, Parameters, Context) -> {ok, RowsUpdated}
317update(Table, Id, Parameters, Context) when is_atom(Table) ->
318    update(atom_to_list(Table), Id, Parameters, Context);
319update(Table, Id, Parameters, Context) ->
320    assert_table_name(Table),
321    Cols         = column_names(Table, Context),
322    UpdateProps  = prepare_cols(Cols, Parameters),
323    C            = get_connection(Context),
324    try
325        UpdateProps1 = case proplists:is_defined(props, UpdateProps) of
326            true ->
327                % Merge the new props with the props in the database
328                {ok, OldProps} = pgsql:equery1(C, "select props from \""++Table++"\" where id = $1", [Id]),
329                case is_list(OldProps) of
330                    true ->
331                        FReplace = fun ({P,_} = T, L) -> lists:keystore(P, 1, L, T) end,
332                        NewProps = lists:foldl(FReplace, OldProps, proplists:get_value(props, UpdateProps)),
333                        lists:keystore(props, 1, UpdateProps, {props, cleanup_props(NewProps)});
334                    false ->
335                        UpdateProps
336                end;
337            false ->
338                UpdateProps
339        end,
340
341        {ColNames,Params} = lists:unzip(UpdateProps1),
342        ColNamesNr = lists:zip(ColNames, lists:seq(2, length(ColNames)+1)),
343
344        Sql = "update \""++Table++"\" set " 
345                 ++ string:join([ "\"" ++ atom_to_list(ColName) ++ "\" = $" ++ integer_to_list(Nr) || {ColName, Nr} <- ColNamesNr ], ", ")
346                 ++ " where id = $1",
347        {ok, RowsUpdated} = pgsql:equery1(C, Sql, [Id | Params]),
348        {ok, RowsUpdated}
349    after
350        return_connection(C, Context)
351    end.
352
353
354%% @doc Delete a row from a table, the row must have a column with the name 'id'
355%% @spec delete(Table, Id, Context) -> {ok, RowsDeleted}
356delete(Table, Id, Context) when is_atom(Table) ->
357    delete(atom_to_list(Table), Id, Context);
358delete(Table, Id, Context) ->
359    assert_table_name(Table),
360    C = get_connection(Context),
361    try
362        Sql = "delete from \""++Table++"\" where id = $1", 
363        {ok, RowsDeleted} = pgsql:equery1(C, Sql, [Id]),
364        {ok, RowsDeleted}
365    after
366        return_connection(C, Context)
367    end.
368
369
370
371%% @doc Read a row from a table, the row must have a column with the name 'id'.  
372%% The props column contents is merged with the other properties returned.
373%% @spec select(Table, Id, Context) -> {ok, Row}
374select(Table, Id, Context) when is_atom(Table) ->
375    select(atom_to_list(Table), Id, Context);
376select(Table, Id, Context) ->
377    assert_table_name(Table),
378    C = get_connection(Context),
379    {ok, Row} = try
380        Sql = "select * from \""++Table++"\" where id = $1 limit 1", 
381        pgsql:assoc(C, Sql, [Id])
382    after
383        return_connection(C, Context)
384    end,
385    
386    Props = case Row of
387        [R] ->
388            case proplists:get_value(props, R) of
389                PropsCol when is_list(PropsCol) -> 
390                    lists:keydelete(props, 1, R) ++ PropsCol;
391                _ ->
392                    R
393            end;
394        [] ->
395            []
396    end,
397    {ok, Props}.
398
399
400%% @doc Remove all undefined props, translate texts to binaries.
401cleanup_props(Ps) when is_list(Ps) ->
402    [ {K,to_binary_string(V)} || {K,V} <- Ps, V /= undefined ];
403cleanup_props(P) -> 
404    P.
405
406    to_binary_string([]) -> [];
407    to_binary_string(L) when is_list(L) ->
408        case z_string:is_string(L) of
409            true -> list_to_binary(L);
410            false -> L
411        end;
412    to_binary_string({trans, Tr}) ->
413        {trans, [ {Lang,to_binary(V)} || {Lang,V} <- Tr ]};
414    to_binary_string(V) -> 
415        V.
416
417    to_binary(L) when is_list(L) -> list_to_binary(L);
418    to_binary(V) -> V.
419
420
421%% @doc Check if all cols are valid columns in the target table, move unknown properties to the props column (if exists)
422prepare_cols(Cols, Props) ->
423    {CProps, PProps} = split_props(Props, Cols),
424    case PProps of
425        [] ->
426            CProps;
427        _  -> 
428            PPropsMerged = case proplists:is_defined(props, CProps) of
429                            true ->
430                                FReplace = fun ({P,_} = T, L) -> lists:keystore(P, 1, L, T) end,
431                                lists:foldl(FReplace, proplists:get_value(props, CProps), PProps);
432                            false -> 
433                                PProps
434                           end,
435            [{props, PPropsMerged} | proplists:delete(props, CProps)]
436    end.
437
438split_props(Props, Cols) ->
439    {CProps, PProps} = lists:partition(fun ({P,_V}) -> lists:member(P, Cols) end, Props),
440    case PProps of
441        [] -> ok;
442        _  -> z_utils:assert(lists:member(props, Cols), {unknown_column, PProps})
443    end,
444    {CProps, PProps}.
445
446
447%% @doc Return a property list with all columns of the table. (example: [{id,int4,modifier},...])
448%% @spec columns(Table, Context) -> [ #column_def{} ]
449columns(Table, Context) when is_atom(Table) ->
450    columns(atom_to_list(Table), Context);
451columns(Table, Context) ->
452    {ok, Db} = pgsql_pool:get_database(?HOST(Context)),
453    {ok, Schema} = pgsql_pool:get_database_opt(schema, ?HOST(Context)),
454    case z_depcache:get({columns, Db, Schema, Table}, Context) of
455        {ok, Cols} -> 
456            Cols;
457        _ ->
458            Cols = q("  select column_name, data_type, character_maximum_length, is_nullable, column_default
459                        from information_schema.columns
460                        where table_catalog = $1
461                          and table_schema = $2
462                          and table_name = $3
463                        order by ordinal_position", [Db, Schema, Table], Context),
464            Cols1 = [ columns1(Col) || Col <- Cols ],
465            z_depcache:set({columns, Db, Schema, Table}, Cols1, ?YEAR, [{database, Db}], Context),
466            Cols1
467    end.
468    
469
470    columns1({<<"id">>, <<"integer">>, undefined, Nullable, <<"nextval(", _/binary>>}) ->
471        #column_def{
472            name = id,
473            type = "serial",
474            length = undefined,
475            is_nullable = z_convert:to_bool(Nullable),
476            default = undefined
477        };
478    columns1({Name,Type,MaxLength,Nullable,Default}) ->
479        #column_def{
480            name = z_convert:to_atom(Name),
481            type = z_convert:to_list(Type),
482            length = MaxLength,
483            is_nullable = z_convert:to_bool(Nullable),
484            default = column_default(Default)
485        }.
486    
487    column_default(undefined) -> undefined;
488    column_default(<<"nextval(", _/binary>>) -> undefined;
489    column_default(Default) -> binary_to_list(Default).
490
491
492%% @doc Return a list with the column names of a table.  The names are sorted.
493%% @spec column_names(Table, Context) -> [ atom() ]
494column_names(Table, Context) ->
495    Names = [ C#column_def.name || C <- columns(Table, Context)],
496    lists:sort(Names).
497
498
499%% @doc Flush all cached information about the database.
500flush(Context) ->
501    {ok, Db} = pgsql_pool:get_database(?HOST(Context)),
502    z_depcache:flush({database, Db}, Context).
503    
504
505%% @doc Update the sequence of the ids in the table. They will be renumbered according to their position in the id list.
506%% @spec update_sequence(Table, IdList, Context) -> void()
507update_sequence(Table, Ids, Context) when is_atom(Table) ->
508    update_sequence(atom_to_list(Table), Ids, Context);
509update_sequence(Table, Ids, Context) ->
510    assert_table_name(Table),
511    Args = lists:zip(Ids, lists:seq(1, length(Ids))),
512    case get_connection(Context) of
513        none -> [];
514        C ->
515	    try
516		[ {ok, _} = pgsql:equery1(C, "update \""++Table++"\" set seq = $2 where id = $1", Arg) || Arg <- Args ]
517	    after
518		return_connection(C, Context)
519	    end
520    end.
521
522
523
524%% @doc Check the information schema if a certain table exists in the context database.
525%% @spec table_exists(TableName, Context) -> bool()
526table_exists(Table, Context) ->
527    {ok, Db} = pgsql_pool:get_database(?HOST(Context)),
528    {ok, Schema} = pgsql_pool:get_database_opt(schema, ?HOST(Context)),
529    case q1("   select count(*) 
530                from information_schema.tables 
531                where table_catalog = $1 
532                  and table_name = $2 
533                  and table_schema = $3
534                  and table_type = 'BASE TABLE'", [Db, Table, Schema], Context) of
535        1 -> true;
536        0 -> false
537    end.
538
539
540%% @doc Make sure that a table is dropped, only when the table exists
541drop_table(Name, Context) when is_atom(Name) ->
542    drop_table(atom_to_list(Name), Context);
543drop_table(Name, Context) ->
544    case table_exists(Name, Context) of
545        true -> q("drop table \""++Name++"\"", Context);
546        false -> ok
547    end.
548
549
550%% @doc Ensure that a table with the given columns exists, alter any existing table
551%% to add, modify or drop columns.  The 'id' (with type serial) column _must_ be defined
552%% when creating the table.
553ensure_table(Table, Cols, Context) when is_atom(Table) ->
554    ensure_table(atom_to_list(Table), Cols, Context);
555ensure_table(Table, Cols, Context) ->
556    case table_exists(Table, Context) of
557        false ->
558            ensure_table_create(Table, Cols, Context);
559        true ->
560            ExistingCols = lists:sort(columns(Table, Context)),
561            WantedCols = lists:sort(Cols),
562            case ensure_table_alter_cols(WantedCols, ExistingCols, []) of
563                [] -> ok;
564                Diff ->
565                    {ok, Db} = pgsql_pool:get_database(?HOST(Context)),
566                    {ok, Schema} = pgsql_pool:get_database_opt(schema, ?HOST(Context)),
567                    z_db:q("ALTER TABLE \""++Table++"\" " ++ string:join(Diff, ","), Context),
568                    z_depcache:flush({columns, Db, Schema, Table}, Context),
569                    ok
570            end
571    end.
572
573    ensure_table_create(Name, Cols, Context) ->
574        ColsSQL = ensure_table_create_cols(Cols, []),
575        z_db:q("CREATE TABLE \""++Name++"\" ("++string:join(ColsSQL, ",") ++ table_create_primary_key(Cols) ++ ")", Context),
576        ok.
577
578    table_create_primary_key([]) -> [];
579    table_create_primary_key([#column_def{name=id, type="serial"}|_]) -> ", primary key(id)";
580    table_create_primary_key([_|Cols]) -> table_create_primary_key(Cols).
581
582    ensure_table_create_cols([], Acc) ->
583        lists:reverse(Acc);
584    ensure_table_create_cols([C|Cols], Acc) ->
585        M = lists:flatten([$", atom_to_list(C#column_def.name), $", 32, column_spec(C)]),
586        ensure_table_create_cols(Cols, [M|Acc]).
587
588
589    ensure_table_alter_cols([], [], Acc) ->
590        lists:reverse(Acc);
591    ensure_table_alter_cols([N|Ns], [N|Es], Acc) ->
592        ensure_table_alter_cols(Ns, Es, Acc);
593    ensure_table_alter_cols([N|Ns], [E|Es], Acc) when N#column_def.name == E#column_def.name ->
594        M = lists:flatten(["ALTER COLUMN \"", atom_to_list(N#column_def.name), "\" TYPE ", column_spec(N)]),
595        M1 = case N#column_def.is_nullable of 
596                true  -> M ++ lists:flatten([", ALTER COLUMN \"", atom_to_list(N#column_def.name), "\" DROP NOT NULL"]);
597                false -> M ++ lists:flatten([", ALTER COLUMN \"", atom_to_list(N#column_def.name), "\" SET NOT NULL"])
598             end,
599        M2 = case N#column_def.default of
600                undefined -> M1 ++ lists:flatten([", ALTER COLUMN \"", atom_to_list(N#column_def.name), "\" DROP DEFAULT"]);
601                Default -> M1 ++ lists:flatten([", ALTER COLUMN \"", atom_to_list(N#column_def.name), "\" SET DEFAULT ", Default])
602             end,
603        ensure_table_alter_cols(Ns, Es, [M2|Acc]);
604    ensure_table_alter_cols([N|Ns], Es, Acc) when Es == [] orelse N < hd(Es) ->
605        M = lists:flatten(["ADD COLUMN \"", atom_to_list(N#column_def.name), "\" ", 
606                            column_spec(N), 
607                            column_spec_nullable(N#column_def.is_nullable), 
608                            column_spec_default(N#column_def.default)]),
609        ensure_table_alter_cols(Ns, Es, [M|Acc]);
610    ensure_table_alter_cols(Ns, [E|Es], Acc) when Ns == [] orelse E < hd(Ns) ->
611        M = lists:flatten(["DROP COLUMN \"", atom_to_list(E#column_def.name), "\""]),
612        ensure_table_alter_cols(Ns, Es, [M|Acc]).
613
614    column_spec(#column_def{type=Type, length=undefined}) ->
615        Type;
616    column_spec(#column_def{type=Type, length=Length}) ->
617        lists:flatten([Type, $(, integer_to_list(Length), $)]).
618    
619    column_spec_nullable(true) -> "";
620    column_spec_nullable(false) -> " not null".
621    
622    column_spec_default(undefined) -> "";
623    column_spec_default(Default) -> [32, Default].
624
625
626%% @doc Check if a name is a valid SQL table name. Crashes when invalid
627%% @spec assert_table_name(String) -> true
628assert_table_name([H|T]) when (H >= $a andalso H =< $z) orelse H == $_ ->
629    assert_table_name1(T).
630assert_table_name1([]) ->
631    true;
632assert_table_name1([H|T]) when (H >= $a andalso H =< $z) orelse (H >= $0 andalso H =< $9) orelse H == $_ ->
633    assert_table_name1(T).
634
635
636
637%% @doc Merge the contents of the props column into the result rows
638%% @spec merge_props(list()) -> list()
639merge_props(undefined) ->
640    undefined;
641merge_props(List) ->
642    merge_props(List, []).
643    
644merge_props([], Acc) ->
645    lists:reverse(Acc);
646merge_props([R|Rest], Acc) ->
647    case proplists:get_value(props, R) of
648        undefined ->
649            merge_props(Rest, [R|Acc]);
650        <<>> ->
651            merge_props(Rest, [R|Acc]);
652        List ->
653            merge_props(Rest, [lists:keydelete(props, 1, R)++List|Acc])
654    end.