Count with selectors

Makefile

@@ -7,4 +7,4 @@ test: emake
 	prove t/*.t
 
 clean:
-	rm -rf $(wildcard ebin/*.beam) erl_crash.dump
+	rm -rf $(wildcard ebin/*.beam) erl_crash.dump

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).

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:64/little-signed, Tail/binary>>) ->
 	{Int, Tail};
 
 decode_value(_, _) ->
-	exit(oh_fuck).
+	exit(oh_fuck).

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 = <<OptsSum:32/little-signed, FullName/binary, 0:8,
-				(Query#emo_query.offset):32/little-signed, 
-				(Query#emo_query.limit):32/little-signed, 
-				EncodedDocument/binary, EncodedFieldSelector/binary>>,
+    EncodedFieldSelector = if 
+        Query#emo_query.field_selector == [] -> <<>>; 
+        true -> emongo_bson:encode(Query#emo_query.field_selector) 
+    end,
+    Message = <<OptsSum:32/little-signed, FullName/binary, 0:8,
+                (Query#emo_query.offset):32/little-signed, 
+                (Query#emo_query.limit):32/little-signed, 
+                EncodedDocument/binary, EncodedFieldSelector/binary>>,
 	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, <<Bin/binary, KeyBin/binary, "_", ValBin/binary>>).
+	index_name(Tail, <<Bin/binary, KeyBin/binary, "_", ValBin/binary>>).