Просмотр исходного кода

Minor fixes and optimization from review. GH-137

* Assume CommandComplete for COPY always returns number of rows
* Move "format" conversion from int to atom to a separate function
* Use lists:zipwith for the iteration over Types + Values
* Fix typespecs
Sergey Prokhorov 4 лет назад
Родитель
Сommit
37f8797973

+ 2 - 6
src/commands/epgsql_cmd_copy_done.erl

@@ -35,12 +35,8 @@ execute(Sock0, St) ->
     end.
 
 handle_message(?COMMAND_COMPLETE, Bin, Sock, St) ->
-    Complete = epgsql_wire:decode_complete(Bin),
-    Res = case Complete of
-        {copy, Count} -> {ok, Count};
-        copy -> ok
-    end,
-    {add_result, Res, {complete, Complete}, Sock, St};
+    Complete = {copy, Count} = epgsql_wire:decode_complete(Bin),
+    {add_result, {ok, Count}, {complete, Complete}, Sock, St};
 handle_message(?ERROR, Error, Sock, St) ->
     Result = {error, Error},
     {add_result, Result, Result, Sock, St};

+ 9 - 10
src/commands/epgsql_cmd_copy_from_stdin.erl

@@ -32,7 +32,7 @@
 -export([init/1, execute/2, handle_message/4]).
 -export_type([response/0]).
 
--type response() :: ok | {error, epgsql:query_error()}.
+-type response() :: {ok, [text | binary]} | {error, epgsql:query_error()}.
 
 -include("epgsql.hrl").
 -include("protocol.hrl").
@@ -61,13 +61,9 @@ execute(Sock, #copy_stdin{query = SQL, format = Format} = St) ->
 %% CopyBothResponses
 handle_message(?COPY_IN_RESPONSE, <<BinOrText, NumColumns:?int16, Formats/binary>>, Sock,
                #copy_stdin{initiator = Initiator, format = RequestedFormat}) ->
-    ColumnFormats =
-        [case Format of
-             0 -> text;
-             1 -> binary
-         end || <<Format:?int16>> <= Formats],
+    ColumnFormats = [format_to_atom(Format) || <<Format:?int16>> <= Formats],
     length(ColumnFormats) =:= NumColumns orelse error(invalid_copy_in_response),
-    CopyState = init_copy_state(BinOrText, RequestedFormat, ColumnFormats, Initiator),
+    CopyState = init_copy_state(format_to_atom(BinOrText), RequestedFormat, ColumnFormats, Initiator),
     Sock1 = epgsql_sock:set_attr(subproto_state, CopyState, Sock),
     Res = {ok, ColumnFormats},
     {finish, Res, Res, epgsql_sock:set_packet_handler(on_copy_from_stdin, Sock1)};
@@ -77,14 +73,14 @@ handle_message(?ERROR, Error, _Sock, _State) ->
 handle_message(_, _, _, _) ->
     unknown.
 
-init_copy_state(0, text, ColumnFormats, Initiator) ->
-    %% When BinOrText is 0, all "columns" should be 0 format as well.
+init_copy_state(text, text, ColumnFormats, Initiator) ->
+    %% When BinOrText is `text', all "columns" should be `text' format as well.
     %% See https://www.postgresql.org/docs/current/protocol-message-formats.html
     %% CopyInResponse
     (lists:member(binary, ColumnFormats) == false)
         orelse error(invalid_copy_in_response),
     #copy{initiator = Initiator, format = text};
-init_copy_state(1, {binary, ColumnTypes}, ColumnFormats, Initiator) ->
+init_copy_state(binary, {binary, ColumnTypes}, ColumnFormats, Initiator) ->
     %% https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-COPY
     %% "As of the present implementation, all columns in a given COPY operation will use the same
     %% format, but the message design does not assume this."
@@ -101,3 +97,6 @@ init_copy_state(ServerExpectedFormat, RequestedFormat, _, _Initiator) ->
     %% `epgsql:copy_from_stdin(C, "COPY ... WITH (FORMAT binary)", text)' or maybe PostgreSQL
     %% got some new format epgsql is not aware of
     error({format_mismatch, RequestedFormat, ServerExpectedFormat}).
+
+format_to_atom(0) -> text;
+format_to_atom(1) -> binary.

+ 4 - 0
src/epgsql.erl

@@ -472,6 +472,10 @@ copy_from_stdin(C, SQL) ->
 %% @param SQL have to be `COPY ... FROM STDIN ...' statement
 %% @param Format data transfer format specification: `text' or `{binary, epgsql_type()}'. Have to
 %%        match `WHERE (FORMAT ???)' from SQL (`text' for `text'/`csv' OR `{binary, ..}' for `binary').
+%% @returns in case of success, `{ok, [text | binary]}' tuple is returned. List describes the expected
+%%        payload format for each column of input. In current implementation all the atoms in a list
+%%        will be the same and will match the atom in `Format' parameter. It may change in the future
+%%        if PostgreSQL will introduce alternative payload formats.
 -spec copy_from_stdin(connection(), sql_query(), text | {binary, [epgsql_type()]}) ->
           epgsql_cmd_copy_from_stdin:response().
 copy_from_stdin(C, SQL, Format) ->

+ 10 - 8
src/epgsql_wire.erl

@@ -334,14 +334,16 @@ encode_copy_row(ValuesTuple, Types, Codec) when is_tuple(ValuesTuple) ->
 encode_copy_row(Values, Types, Codec) ->
     NumCols = length(Types),
     [<<NumCols:?int16>>
-    | [
-       case epgsql_binary:is_null(Value, Codec) of
-           true ->
-               <<-1:?int32>>;
-           false ->
-               epgsql_binary:encode(Type, Value, Codec)
-       end || {Type, Value} <- lists:zip(Types, Values) % TODO: parallel iteration ninstead
-      ]].
+    | lists:zipwith(
+        fun(Type, Value) ->
+                case epgsql_binary:is_null(Value, Codec) of
+                    true ->
+                        <<-1:?int32>>;
+                    false ->
+                        epgsql_binary:encode(Type, Value, Codec)
+                end
+        end, Types, Values)
+    ].
 
 %% @doc encode binary copy data file header
 %%