diff --git a/Makefile b/Makefile index 4486b38..f2ad018 100644 --- a/Makefile +++ b/Makefile @@ -7,4 +7,4 @@ test: emake prove t/*.t clean: - rm -rf $(wildcard ebin/*.beam) erl_crash.dump \ No newline at end of file + rm -rf $(wildcard ebin/*.beam) erl_crash.dump diff --git a/src/emongo.erl b/src/emongo.erl index 5775125..adbd6b2 100644 --- a/src/emongo.erl +++ b/src/emongo.erl @@ -30,7 +30,7 @@ auth/3, find/2, find/3, find/4, find_all/2, find_all/3, find_all/4, get_more/4, get_more/5, find_one/3, find_one/4, kill_cursors/2, insert/3, update/4, update/5, update_sync/4, update_sync/5, - delete/2, delete/3, ensure_index/3, count/2, dec2hex/1, + delete/2, delete/3, ensure_index/3, count/2, count/4, dec2hex/1, hex2dec/1]). -include("emongo.hrl"). @@ -158,7 +158,13 @@ find(PoolId, Collection, Query) when is_record(Query, emo_query) -> %% Result = documents() | response() find(PoolId, Collection, Selector, Options) when ?IS_DOCUMENT(Selector), is_list(Options) -> {Pid, Pool} = gen_server:call(?MODULE, {pid, PoolId}, infinity), - Query = create_query(Options, Selector), + Query = case proplists:is_defined(fields, Options) of + true -> + Query1 = create_query(proplists:delete(fields,Options), Selector), + Query1#emo_query{field_selector=[{Field, 1} || Field <- proplists:get_value(fields, Options)]}; + false -> + create_query(Options, Selector) + end, Packet = emongo_packet:do_query(Pool#pool.database, Collection, Pool#pool.req_id, Query), Resp = emongo_conn:send_recv(Pid, Pool#pool.req_id, Packet, proplists:get_value(timeout, Options, ?TIMEOUT)), case lists:member(response_options, Options) of @@ -292,6 +298,18 @@ count(PoolId, Collection) -> undefined end. +count(PoolId, Collection, [], []) -> + count(PoolId, Collection); +count(PoolId, Collection, Selector, Options) -> + {Pid, Pool} = gen_server:call(?MODULE, {pid, PoolId}, infinity), + Query = create_query(Options ++ [{fields, []}, {limit, 1}], #emo_query{}, transform_selector(Selector), [{<<"count">>, Collection}]), + Packet = emongo_packet:do_query(Pool#pool.database, "$cmd", Pool#pool.req_id, Query), + case emongo_conn:send_recv(Pid, Pool#pool.req_id, Packet, ?TIMEOUT) of + #response{documents=[[{<<"n">>,Count}|_]]} -> + round(Count); + _ -> + undefined + end. %drop_collection(PoolId, Collection) when is_atom(PoolId), is_list(Collection) -> %%==================================================================== @@ -507,7 +525,11 @@ create_query([], QueryRec, [], OptDoc) -> QueryRec#emo_query{q=OptDoc}; create_query([], QueryRec, QueryDoc, OptDoc) -> - QueryRec#emo_query{q=(OptDoc ++ [{<<"query">>, QueryDoc}])}; + QueryRec#emo_query{q=(OptDoc ++ [{<<"query">>, QueryDoc}])}; + +create_query([{fields, Fields}], QueryRec, QueryDoc, OptDoc) -> + QueryFields = [{Field, 1} || Field <- Fields], + QueryRec#emo_query{q=(OptDoc ++ [{<<"query">>, QueryDoc}] ++ [{<<"fields">>, QueryFields}])}; create_query([{limit, Limit}|Options], QueryRec, QueryDoc, OptDoc) -> QueryRec1 = QueryRec#emo_query{limit=Limit}, @@ -523,8 +545,9 @@ create_query([{orderby, Orderby}|Options], QueryRec, QueryDoc, OptDoc) -> create_query(Options, QueryRec, QueryDoc, OptDoc1); create_query([{fields, Fields}|Options], QueryRec, QueryDoc, OptDoc) -> - QueryRec1 = QueryRec#emo_query{field_selector=[{Field, 1} || Field <- Fields]}, - create_query(Options, QueryRec1, QueryDoc, OptDoc); + %% The fields document needs to be last + %% so it requires some special handling + create_query(proplists:delete(fields, Options) ++ [{fields, Fields}], QueryRec, QueryDoc, OptDoc); create_query([_|Options], QueryRec, QueryDoc, OptDoc) -> create_query(Options, QueryRec, QueryDoc, OptDoc). diff --git a/src/emongo_bson.erl b/src/emongo_bson.erl index a908c1b..e0f35ff 100644 --- a/src/emongo_bson.erl +++ b/src/emongo_bson.erl @@ -30,6 +30,10 @@ encode([]) -> encode([{_,_}|_]=List) when is_list(List) -> Bin = iolist_to_binary([encode_key_value(Key, Val) || {Key, Val} <- List]), <<(size(Bin)+5):32/little-signed, Bin/binary, 0:8>>. + +encode_document([{_,_}|_]=List) when is_list(List) -> + Bin = iolist_to_binary([encode_key_value(Key, Val) || {Key, Val} <- List]), + <<(size(Bin)+5):32/little-signed, Bin/binary>>. %% FLOAT encode_key_value(Key, Val) when is_float(Val) -> @@ -45,14 +49,20 @@ encode_key_value(Key, Val) when is_binary(Val) orelse Val == [] orelse (is_list( {incomplete, Bin1, Bin2} -> exit({cannot_convert_chars_to_binary, Val, Bin1, Bin2}); Val1 -> - <<2, Key1/binary, 0, (byte_size(Val1)+1):32/little-signed, Val1/binary, 0:8>> + case Key1 of + <<"fields">> -> + FieldsVal = <<0,0,0,0>>, + <<3, Key1/binary, 0, (byte_size(FieldsVal)+1):8/signed, FieldsVal/binary, 0:8>>; + _ -> + <<2, Key1/binary, 0, (byte_size(Val1)+1):32/little-signed, Val1/binary, 0:8>> + end end; %% NESTED OBJECT encode_key_value(Key, [{_,_}|_]=Val) -> Key1 = encode_key(Key), - Val1 = encode(Val), - <<3, Key1/binary, 0, Val1/binary>>; + Val1 = encode_document(Val), + <<3, Key1/binary, 0, Val1/binary, 0:8>>; %% DATA ARRAY encode_key_value(Key, {array, Val}) when is_list(Val) -> @@ -230,4 +240,4 @@ decode_value(18, <>) -> {Int, Tail}; decode_value(_, _) -> - exit(oh_fuck). \ No newline at end of file + exit(oh_fuck). diff --git a/src/emongo_packet.erl b/src/emongo_packet.erl index 8b58fa4..694a7cf 100644 --- a/src/emongo_packet.erl +++ b/src/emongo_packet.erl @@ -51,14 +51,14 @@ do_query(Database, Collection, ReqID, Query) when is_record(Query, emo_query) -> is_binary(Query#emo_query.q) -> Query#emo_query.q; true -> emongo_bson:encode(Query#emo_query.q) end, - EncodedFieldSelector = if - Query#emo_query.field_selector == [] -> <<>>; - true -> emongo_bson:encode(Query#emo_query.field_selector) - end, - Message = <>, + EncodedFieldSelector = if + Query#emo_query.field_selector == [] -> <<>>; + true -> emongo_bson:encode(Query#emo_query.field_selector) + end, + Message = <>, Length = byte_size(Message), <<(Length+16):32/little-signed, ReqID:32/little-signed, 0:32, ?OP_QUERY:32/little-signed, Message/binary>>. @@ -127,4 +127,4 @@ index_name([{Key, Val}|Tail], Bin) -> is_integer(Val) -> list_to_binary(integer_to_list(Val)); true -> <<>> end, - index_name(Tail, <>). \ No newline at end of file + index_name(Tail, <>).