first test with op_msg

Author: Michael Maier

lib/bson/decoder.ex

@@ -169,18 +169,11 @@ defmodule BSON.Decoder do
     {string, rest}
   end
 
-  defp subtype(0x00),
-    do: :generic
-  defp subtype(0x01),
-    do: :function
-  defp subtype(0x02),
-    do: :binary_old
-  defp subtype(0x03),
-    do: :uuid_old
-  defp subtype(0x04),
-    do: :uuid
-  defp subtype(0x05),
-    do: :md5
-  defp subtype(int) when is_integer(int) and int in 0x80..0xFF,
-    do: int
+  defp subtype(0x00), do: :generic
+  defp subtype(0x01), do: :function
+  defp subtype(0x02), do: :binary_old
+  defp subtype(0x03), do: :uuid_old
+  defp subtype(0x04), do: :uuid
+  defp subtype(0x05),  do: :md5
+  defp subtype(int) when is_integer(int) and int in 0x80..0xFF, do: int
 end

lib/mongo.ex

@@ -548,6 +548,26 @@ defmodule Mongo do
     end
   end
 
+  @doc false
+  @spec direct_command_msg(pid, BSON.document, Keyword.t) :: {:ok, BSON.document | nil} | {:error, Mongo.Error.t}
+  def direct_command_msg(conn, doc, opts \\ []) do
+    docs = [doc]
+    query = %Query{action: :msg}
+
+    with {:ok, response} <- DBConnection.execute(conn, query, docs, defaults(opts)) do
+      case response do
+        op_msg(docs: [%{"ok" => ok} = doc]) when ok == 1 -> {:ok, doc}
+        op_reply(flags: flags, docs: [%{"$err" => reason, "code" => code}]) when (@reply_query_failure &&& flags) != 0  -> {:error, Mongo.Error.exception(message: reason, code: code)}
+        op_reply(flags: flags) when (@reply_cursor_not_found &&& flags) != 0 ->  {:error, Mongo.Error.exception(message: "cursor not found")}
+        op_reply(docs: [%{"ok" => 0.0, "errmsg" => reason} = error]) ->  {:error, %Mongo.Error{message: "command failed: #{reason}", code: error["code"]}}
+        op_reply(docs: [%{"ok" => ok} = doc]) when ok == 1 ->   {:ok, doc}
+        # TODO: Check if needed
+        op_reply(docs: []) ->  {:ok, nil}
+      end
+    end
+  end
+
+
   @doc """
   Similar to `command/3` but unwraps the result and raises on error.
   """
@@ -579,14 +599,15 @@ defmodule Mongo do
 
     query = [
       insert: coll,
+      "$db": "test",
       documents: [doc],
       ordered: Keyword.get(opts, :ordered),
       writeConcern: write_concern,
       bypassDocumentValidation: Keyword.get(opts, :bypass_document_validation)
     ] |> filter_nils()
 
     with {:ok, conn, _, _} <- select_server(topology_pid, :write, opts),
-         {:ok, doc} <- direct_command(conn, query, opts) do
+         {:ok, doc} <- direct_command_msg(conn, query, opts) do
       case doc do
         %{"writeErrors" => _} -> {:error, %Mongo.WriteError{n: doc["n"], ok: doc["ok"], write_errors: doc["writeErrors"]}}
         _ ->
@@ -644,7 +665,7 @@ defmodule Mongo do
     ] |> filter_nils()
 
     with {:ok, conn, _, _} <- select_server(topology_pid, :write, opts),
-         {:ok, doc} <- direct_command(conn, query, opts) do
+         {:ok, doc} <- direct_command_msg(conn, query, opts) do
       case doc do
         %{"writeErrors" => _} ->
           {:error, %Mongo.WriteError{n: doc["n"], ok: doc["ok"], write_errors: doc["writeErrors"]}}

lib/mongo/messages.ex

@@ -20,6 +20,7 @@ defmodule Mongo.Messages do
   @op_get_more      2005
   @op_delete        2006
   @op_kill_cursors  2007
+  @op_msg_code      2013
 
   @update_flags [
     upsert: 0x1,
@@ -54,6 +55,7 @@ defmodule Mongo.Messages do
   defrecord  :op_delete, [:coll, :flags, :query]
   defrecord  :op_kill_cursors, [:cursor_ids]
   defrecord  :op_reply, [:flags, :cursor_id, :from, :num, :docs]
+  defrecord  :op_msg, [:flags, :type, :docs]
 
   def encode(request_id, op) do
     iodata = encode_op(op)
@@ -64,6 +66,25 @@ defmodule Mongo.Messages do
     [encode_header(header)|iodata]
   end
 
+  def decode_response_msg(msg_header(length: length) = header, iolist) when is_list(iolist) do
+    IO.puts "decode_response_msg: #{inspect header}"
+    if IO.iodata_length(iolist) >= length,
+       do: decode_response_msg(header, IO.iodata_to_binary(iolist)),
+       else: :error
+  end
+  def decode_response_msg(msg_header(length: length, response_to: response_to) = header, binary) when byte_size(binary) >= length do
+    IO.puts "decode_response_msg: #{inspect header}"
+    <<response::binary(length), rest::binary>> = binary
+
+    IO.puts inspect decode_msg(response)
+    IO.puts inspect rest
+    {:ok, response_to, decode_msg(response), rest}
+  end
+  def decode_response_msg(_header, _binary) do
+    :error
+  end
+
+
   def decode_response(msg_header(length: length) = header, iolist) when is_list(iolist) do
     if IO.iodata_length(iolist) >= length,
       do: decode_response(header, IO.iodata_to_binary(iolist)),
@@ -82,8 +103,7 @@ defmodule Mongo.Messages do
       do: IO.iodata_to_binary(iolist) |> decode_header,
     else: :error
   end
-  def decode_header(<<length::int32, request_id::int32, response_to::int32,
-                       op_code::int32, rest::binary>>) do
+  def decode_header(<<length::int32, request_id::int32, response_to::int32, op_code::int32, rest::binary>>) do
     header = msg_header(length: length-@header_size, request_id: request_id, response_to: response_to, op_code: op_code)
     {:ok, header, rest}
   end
@@ -106,6 +126,10 @@ defmodule Mongo.Messages do
      docs]
   end
 
+  defp encode_op(op_msg(flags: flags, docs: [doc])) do
+    [<<0::int32>>, <<0x00>>, doc]
+  end
+
   defp encode_op(op_query(flags: flags, coll: coll, num_skip: num_skip,
                           num_return: num_return, query: query, select: select)) do
     [<<blit_flags(:query, flags)::int32>>,
@@ -140,11 +164,16 @@ defmodule Mongo.Messages do
   defp op_to_code(op_get_more()),     do: @op_get_more
   defp op_to_code(op_delete()),       do: @op_delete
   defp op_to_code(op_kill_cursors()), do: @op_kill_cursors
+  defp op_to_code(op_msg()),          do: @op_msg_code
 
   defp decode_reply(<<flags::int32, cursor_id::int64, from::int32, num::int32, rest::binary>>) do
     op_reply(flags: flags, cursor_id: cursor_id, from: from, num: num, docs: rest)
   end
 
+  defp decode_msg(<<flags::int32, _type::int8, rest::binary>>) do
+    op_msg(flags: flags, docs: rest)
+  end
+
   defp encode_header(msg_header(length: length, request_id: request_id,
                                 response_to: response_to, op_code: op_code)) do
     <<length::int32, request_id::int32, response_to::int32, op_code::int32>>

lib/mongo/mongo_db_connection.ex

@@ -143,12 +143,24 @@ defmodule Mongo.MongoDBConnection do
     {:ok, state.wire_version, state}
   end
 
+  defp handle_execute(:msg, nil, [doc], opts, state) do
+    op_msg(docs: [doc])
+    |> get_response_msg(state)
+  end
+
+  defp get_response_msg(op, state) do
+    with {:ok, response} <- Utils.post_msg(state.request_id, op, state),
+         state = %{state | request_id: state.request_id + 1},
+         do: {:ok, response, state}
+  end
+
   defp handle_execute(:command, nil, [query], opts, s) do
     flags = Keyword.take(opts, @find_one_flags)
     op_query(coll: Utils.namespace("$cmd", s, opts[:database]), query: query, select: "", num_skip: 0, num_return: 1, flags: flags(flags))
     |> get_response(s)
   end
 
+
   defp get_response(op, state) do
     with {:ok, response} <- Utils.post_request(state.request_id, op, state),
          state = %{state | request_id: state.request_id + 1},

lib/mongo/query.ex

@@ -5,23 +5,27 @@ defmodule Mongo.Query do
 end
 
 defimpl DBConnection.Query, for: Mongo.Query do
-  import Mongo.Messages, only: [op_reply: 1, op_reply: 2]
+  import Mongo.Messages, only: [op_reply: 1, op_reply: 2, op_msg: 1, op_msg: 2]
 
   def parse(query, _opts), do: query
   def describe(query, _opts), do: query
 
   def encode(query, params, _opts) do
-    if query.encoded? do
+    result = if query.encoded? do
       params
     else
       Enum.map(params, fn
         nil -> ""
         doc -> BSON.Encoder.document(doc)
       end)
     end
+
+    IO.puts "Encode message :#{inspect result}"
+    result
   end
 
   def decode(_query, :ok, _opts), do: :ok
   def decode(_query, wire_version, _opts) when is_integer(wire_version), do: wire_version
   def decode(_query, op_reply(docs: docs) = reply, _opts), do: op_reply(reply, docs: BSON.Decoder.documents(docs))
+  def decode(_query, op_msg(docs: docs) = msg, _opts), do: op_msg(msg, docs: BSON.Decoder.documents(docs))
 end

lib/mongo_db_connection/utils.ex

@@ -17,6 +17,20 @@ defmodule Mongo.MongoDBConnection.Utils do
          do: {:ok, response}
   end
 
+  @doc"""
+    Sends a request id and waits for the response with the same id
+  """
+  def post_msg(id, ops, state) when is_list(ops) do
+    with :ok <- send(ops, state),
+         {:ok, ^id, response} <- recv_msg(state),
+         do: {:ok, response}
+  end
+  def post_msg(id, op, state) do
+    with :ok <- send(id, op, state),
+         {:ok, ^id, response} <- recv_msg(state),
+         do: {:ok, response}
+  end
+
   def command(id, command, s) do
     ns =
       if Keyword.get(command, :mechanism) == "MONGODB-X509" && Keyword.get(command, :authenticate) == 1 do
@@ -82,6 +96,30 @@ defmodule Mongo.MongoDBConnection.Utils do
     recv(nil, "", s)
   end
 
+  def recv_msg(s) do
+    recv_msg(nil, "", s)
+  end
+  defp recv_msg(nil, data, %{socket: {mod, sock}} = s) do
+    case decode_header(data) do
+      {:ok, header, rest} -> recv_msg(header, rest, s)
+      :error ->
+        case mod.recv(sock, 0, s.timeout) do
+          {:ok, tail}      -> recv_msg(nil, [data|tail], s)
+          {:error, reason} -> recv_error(reason, s)
+        end
+    end
+  end
+  defp recv_msg(header, data, %{socket: {mod, sock}} = s) do
+    case decode_response_msg(header, data) do
+      {:ok, id, msg, ""} -> {:ok, id, msg}
+      :error ->
+        case mod.recv(sock, 0, s.timeout) do
+          {:ok, tail}      -> recv_msg(header, [data|tail], s)
+          {:error, reason} -> recv_error(reason, s)
+        end
+    end
+  end
+
   # TODO: Optimize to reduce :gen_tcp.recv and decode_response calls
   #       based on message size in header.
   #       :gen.tcp.recv(socket, min(size, max_packet))