Basic support for OP_MSG, fixes #2

Author: zookzook

lib/mongo.ex

@@ -281,9 +281,13 @@ defmodule Mongo do
     opts = Keyword.drop(opts, ~w(bypass_document_validation max_time projection return_document sort upsert collation)a)
 
     with {:ok, conn, _, _} <- select_server(topology_pid, :write, opts),
-         {:ok, doc} <- direct_command(conn, query, opts), do: {:ok, doc["value"]}
+         {:ok, doc} <- direct_command(conn, query, opts) do
+        {:ok, doc["value"]}
+    end
+
   end
 
+
   @doc """
   Finds a document and replaces it.
 
@@ -487,30 +491,29 @@ defmodule Mongo do
       other -> other
     end
 
-    cmd = [
-              {"find", coll},
-              {"filter", filter},
-              {"limit", opts[:limit]},
-              {"hint", opts[:hint]},
-              {"singleBatch", opts[:single_batch]},
-              {"readConcern", opts[:read_concern]},
-              {"max", opts[:max]},
-              {"min", opts[:min]},
-              {"collation", opts[:collation]},
-              {"returnKey", opts[:return_key]},
-              {"showRecordId", opts[:show_record_id]},
-              {"tailable", opts[:tailable]},
-              {"oplogReplay", opts[:oplog_replay]},
-              {"tailable", opts[:tailable]},
-              {"noCursorTimeout", opts[:no_cursor_timeout]},
-              {"awaitData", opts[:await_data]},
-              {"batchSize", opts[:batch_size]},
-              {"projection", opts[:projection]},
-              {"comment", opts[:comment]},
-              {"maxTimeMS", opts[:max_time]},
-              {"skip", opts[:skip]},
-              {"sort", opts[:sort]}
-            ]
+    cmd = [find: coll,
+           filter: filter,
+           limit: opts[:limit],
+           hint: opts[:hint],
+           singleBatch: opts[:single_batch],
+           readConcern: opts[:read_concern],
+           max: opts[:max],
+           min: opts[:min],
+           collation: opts[:collation],
+           returnKey: opts[:return_key],
+           showRecordId: opts[:show_record_id],
+           tailable: opts[:tailable],
+           oplogReplay: opts[:oplog_replay],
+           tailable: opts[:tailable],
+           noCursorTimeout: opts[:no_cursor_timeout],
+           awaitData: opts[:await_data],
+           batchSize: opts[:batch_size],
+           projection: opts[:projection],
+           comment: opts[:comment],
+           maxTimeMS: opts[:max_time],
+           skip: opts[:skip],
+           sort: opts[:sort]
+          ]
 
     cmd = filter_nils(cmd)
 
@@ -564,18 +567,18 @@ defmodule Mongo do
   def direct_command(conn, cmd, opts \\ []) do
     action = %Query{action: :command}
 
-    with {:ok, _query, response} <- DBConnection.execute(conn, action, [cmd], defaults(opts)) do
-      case response do
-        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
+    with {:ok, _query, doc} <- DBConnection.execute(conn, action, [cmd], defaults(opts)),
+         {:ok, doc} <- check_for_error(doc) do
+      {:ok, doc}
     end
   end
 
+  defp check_for_error(%{"ok" => ok} = response) when ok == 1 do
+    {:ok, response}
+  end
+  defp check_for_error(%{"code" => code, "codeName" => _name, "errmsg" => msg}) do
+    {:error, Mongo.Error.exception(message: msg, code: code)}
+  end
 
   @doc """
   Returns the current wire version.
@@ -600,7 +603,7 @@ defmodule Mongo do
   """
   @spec ping(GenServer.server) :: result(BSON.document)
   def ping(topology_pid) do
-    command(topology_pid, %{ping: 1}, [batch_size: 1])
+    command(topology_pid, [ping: 1], [batch_size: 1])
   end
 
   @doc """
@@ -977,7 +980,7 @@ defmodule Mongo do
   defp select_servers(topology_pid, type, opts, start_time) do
     topology = Topology.topology(topology_pid)
     with {:ok, servers, slave_ok, mongos?} <- TopologyDescription.select_servers(topology, type, opts) do
-      case Enum.empty? servers do
+      case Enum.empty?(servers) do
         true ->
           case Topology.wait_for_connection(topology_pid, @sel_timeout, start_time) do
             {:ok, _servers} -> select_servers(topology_pid, type, opts, start_time)
@@ -1054,24 +1057,6 @@ defmodule Mongo do
     Keyword.put_new(opts, :timeout, @timeout)
   end
 
-  defp get_last_error(:ok) do
-    :ok
-  end
-  defp get_last_error(op_reply(docs: [%{"ok" => ok, "err" => nil} = doc])) when ok == 1 do
-    {:ok, doc}
-  end
-  defp get_last_error(op_reply(docs: [%{"ok" => ok, "err" => message, "code" => code}])) when ok == 1 do
-    # If a batch insert (OP_INSERT) fails some documents may still have been
-    # inserted, but mongo always returns {n: 0}
-    # When we support the 2.6 bulk write API we will get number of inserted
-    # documents and should change the return value to be something like:
-    # {:error, %WriteResult{}, %Error{}}
-    {:error, Mongo.Error.exception(message: message, code: code)}
-  end
-  defp get_last_error(op_reply(docs: [%{"ok" => 0.0, "errmsg" => message, "code" => code}])) do
-    {:error, Mongo.Error.exception(message: message, code: code)}
-  end
-
   defp assign_ids(list) when is_list(list) do
     Enum.map(list, &assign_id/1)
     |> Enum.unzip
@@ -1101,7 +1086,4 @@ defmodule Mongo do
   defp add_id([{key, _}|_] = list, id) when is_binary(key), do: [{"_id", id}|list]
   defp add_id([], id), do: [{"_id", id}]
 
-  defp maybe_failure(op_reply(flags: flags, docs: [%{"$err" => reason, "code" => code}])) when (@reply_query_failure &&& flags) != 0, do: {:error, Mongo.Error.exception(message: reason, code: code)}
-  defp maybe_failure(op_reply(flags: flags)) when (@reply_cursor_not_found &&& flags) != 0, do: {:error, Mongo.Error.exception(message: "cursor not found")}
-  defp maybe_failure(_op_reply),  do: :ok
 end

lib/mongo/cursor.ex

@@ -132,10 +132,10 @@ defmodule Mongo.Cursor do
     def get_more(conn, coll, cursor, nil, opts) do
 
       cmd = [
-        {"getMore", cursor},
-        {"collection", coll},
-        {"batchSize", opts[:batch_size]},
-        {"maxTimeMS", opts[:max_time]}
+        getMore: cursor,
+        collection: coll,
+        batchSize: opts[:batch_size],
+        maxTimeMS: opts[:max_time]
       ] |> filter_nils()
 
       with {:ok, %{"cursor" => %{ "id" => cursor_id, "nextBatch" => docs}, "ok" => ok}} when ok == 1 <- Mongo.direct_command(conn, cmd, opts) do
@@ -149,10 +149,10 @@ defmodule Mongo.Cursor do
                                on_resume_token: fun, topology_pid: topology_pid) = change_stream, opts) do
 
       get_more = [
-                  {"getMore", cursor_id},
-                  {"collection", coll},
-                  {"batchSize", opts[:batch_size]},
-                  {"maxTimeMS", opts[:max_time]}
+                   getMore: cursor_id,
+                   collection: coll,
+                   batchSize: opts[:batch_size],
+                   maxTimeMS: opts[:max_time]
                  ] |> filter_nils()
 
       with {:ok, %{"operationTime" => op_time,
@@ -266,8 +266,8 @@ defmodule Mongo.Cursor do
     def kill_cursors(conn, coll, cursor_ids, opts) do
 
       cmd = [
-        {"killCursors", coll},
-        {"cursors", cursor_ids}
+        killCursors: coll,
+        cursors: cursor_ids
       ]
 
       cmd = filter_nils(cmd)

lib/mongo/grid_fs/bucket.ex

@@ -94,8 +94,8 @@ defmodule Mongo.GridFs.Bucket do
   """
   @spec drop(Bucket.t) :: Mongo.result(BSON.document)
   def drop(%Bucket{topology_pid: topology_pid, opts: opts} = bucket) do
-    {:ok, _} = Mongo.command(topology_pid, %{drop: files_collection_name(bucket)}, opts)
-    {:ok, _} = Mongo.command(topology_pid, %{drop: chunks_collection_name(bucket)}, opts)
+    {:ok, _} = Mongo.command(topology_pid, [drop: files_collection_name(bucket)], opts)
+    {:ok, _} = Mongo.command(topology_pid, [drop: chunks_collection_name(bucket)], opts)
   end
 
   @doc """

lib/mongo/messages.ex

@@ -8,10 +8,6 @@ defmodule Mongo.Messages do
   defmacro __using__(_opts) do
     quote do
       import unquote(__MODULE__)
-      @reply_cursor_not_found   0x1
-      @reply_query_failure      0x2
-      @reply_shard_config_stale 0x4
-      @reply_await_capable      0x8
     end
   end
 
@@ -37,7 +33,10 @@ defmodule Mongo.Messages do
   defrecordp :msg_header, [:length, :request_id, :response_to, :op_code]
   defrecord  :op_query, [:flags, :coll, :num_skip, :num_return, :query, :select]
   defrecord  :op_reply, [:flags, :cursor_id, :from, :num, :docs]
-  defrecord  :op_msg, [:flags, :type, :docs]
+  defrecord  :sequence, [:size, :identifier, :docs]
+  defrecord  :payload, [:doc, :sequence]
+  defrecord  :section, [:payload_type, :payload]
+  defrecord  :op_msg, [:flags, :sections]
 
   @doc """
     Decodes the header from response of a request sent by the mongodb server
@@ -66,8 +65,9 @@ defmodule Mongo.Messages do
   def decode_response(msg_header(length: length, response_to: response_to, op_code: op_code), binary)  when byte_size(binary) >= length do
     <<response::binary(length), rest::binary>> = binary
     case op_code do
-      @op_reply -> {:ok, response_to, decode_reply(response), rest}
-      _         -> :error
+      @op_reply    -> {:ok, response_to, decode_reply(response), rest}
+      @op_msg_code -> {:ok, response_to, decode_msg(response), rest}
+      _            -> :error
     end
   end
   def decode_response(_header, _binary), do: :error
@@ -78,19 +78,69 @@ defmodule Mongo.Messages do
   def 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: BSON.Decoder.documents(rest))
   end
+  def decode_msg(<<flags::int32, rest::binary>>) do
+    op_msg(flags: flags, sections: decode_sections(rest))
+  end
+
+  def decode_sections(binary), do: decode_sections(binary, [])
+  def decode_sections("", acc), do: Enum.reverse(acc)
+
+  def decode_sections(<<0x00::int8, payload::binary>>, acc) do
+    <<size::int32, _rest::binary>>      = payload
+    <<doc::binary(size), rest::binary>> = payload
+    with {doc, ""} <- BSON.Decoder.document(doc) do
+      decode_sections(rest, [section(payload_type: 0, payload: payload(doc: doc)) | acc])
+    end
+  end
+  def decode_sections(<<0x01::int8, payload::binary>>, acc) do
+    <<size::int32, _rest::binary>>           = payload
+    <<sequence::binary(size), rest::binary>> = payload
+    decode_sections(rest, [section(payload_type: 1, payload: payload(sequence: decode_sequence(sequence))) | acc])
+  end
+
+  def decode_sequence(<<size::int32, rest::binary>>) do
+    with {identifier, docs} <- cstring(rest) do
+      sequence(size: size, identifier: identifier, docs: BSON.Decoder.documents(docs))
+    end
+  end
+
+  defp cstring(binary) do
+    [string, rest] = :binary.split(binary, <<0x00>>)
+    {string, rest}
+  end
 
   def encode(request_id, op_query() = op) do
     iodata = encode_op(op)
     header = msg_header(length: IO.iodata_length(iodata) + @header_size,  request_id: request_id, response_to: 0, op_code: @op_query)
     [encode_header(header)|iodata]
   end
 
+  def encode(request_id, op_msg() = op) do
+    iodata = encode_op(op)
+    header = msg_header(length: IO.iodata_length(iodata) + @header_size,  request_id: request_id, response_to: 0, op_code: @op_msg_code)
+    [encode_header(header)|iodata]
+  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>>
   end
 
-  defp encode_op(op_msg(flags: flags, docs: [doc])) do
-    [<<0::int32>>, <<0x00>>, doc]
+  defp encode_op(op_msg(flags: flags, sections: sections)) do
+    # todo: flags encoding
+    [<<0::int32>>, encode_sections(sections)]
+  end
+  defp encode_sections(sections) do
+    Enum.map(sections, fn section -> encode_section(section) end)
+  end
+
+  defp encode_section(section(payload_type: t, payload: payload)) do
+    [<<t::int8>>, encode_payload(payload)]
+  end
+  defp encode_payload(payload(doc: doc, sequence: nil)) do
+    BSON.Encoder.document(doc)
+  end
+  defp encode_payload(payload(doc: nil, sequence: sequence(size: size, identifier: identifier, docs: docs))) do
+    [<<size::int32>>, identifier, <<0x00>>, BSON.Encoder.encode(docs)]
   end
 
   defp encode_op(op_query(flags: flags, coll: coll, num_skip: num_skip,

lib/mongo/mongo_db_connection.ex

@@ -1,6 +1,6 @@
 defmodule Mongo.MongoDBConnection do
   @moduledoc """
-  Implementierung für das DBConnection-Protokoll.
+  Implementation of the DBConnection behaviour module.
   """
 
   use DBConnection
@@ -13,6 +13,7 @@ defmodule Mongo.MongoDBConnection do
 
   @impl true
   def connect(opts) do
+
     {write_concern, opts} = Keyword.split(opts, @write_concern)
     write_concern = Keyword.put_new(write_concern, :w, 1)
 
@@ -54,9 +55,7 @@ defmodule Mongo.MongoDBConnection do
       end
 
     case result do
-      {:ok, state} ->
-        ## IO.puts inspect state
-        {:ok, state}
+      {:ok, state} -> {:ok, state}
 
       {:disconnect, reason, state} ->
         reason = case reason do
@@ -128,6 +127,7 @@ defmodule Mongo.MongoDBConnection do
     end
   end
 
+  @impl true
   def checkout(state), do: {:ok, state}
   @impl true
   def checkin(state), do: {:ok, state}
@@ -171,22 +171,39 @@ defmodule Mongo.MongoDBConnection do
   defp execute_action(:wire_version, _, _, state) do
     {:ok, state.wire_version, state}
   end
+  defp execute_action(:command, [cmd], opts, %{wire_version: version} = state) when version >= 6 do
+
+    cmd = cmd ++ ["$db": opts[:database] || state.database,
+                  "$readPreference": [mode: update_read_preferences(opts[:slave_ok])]]
 
-  defp execute_action(:command, [query], opts, state) do
     timeout = Keyword.get(opts, :max_time, 0)
-    flags   = Keyword.take(opts, @find_one_flags)
-    op      = op_query(coll: Utils.namespace("$cmd", state, opts[:database]), query: query, select: "", num_skip: 0, num_return: 1, flags: flags(flags))
+    op      = op_msg(flags: 0, sections: [section(payload_type: 0, payload: payload(doc: cmd))])
 
-    with {:ok, response} <- Utils.post_request(state.request_id, op, state, timeout),
-         state = %{state | request_id: state.request_id + 1},
-         do: {:ok, response, state}
+    with {:ok, doc} <- Utils.post_request(op, state.request_id, state, timeout),
+         state = %{state | request_id: state.request_id + 1} do
+      {:ok, doc, state}
+    end
   end
+  defp execute_action(:command, [cmd], opts, state) do
 
+    timeout = Keyword.get(opts, :max_time, 0)
+    flags   = Keyword.take(opts, @find_one_flags)
+    op      = op_query(coll: Utils.namespace("$cmd", state, opts[:database]), query: cmd, select: "", num_skip: 0, num_return: 1, flags: flags(flags))
+
+    with {:ok, doc} <- Utils.post_request(op, state.request_id, state, timeout),
+         state = %{state | request_id: state.request_id + 1}  do
+      {:ok, doc, state}
+    end
+  end
   defp execute_action(:error, _query, _opts, state) do
     exception = Mongo.Error.exception("Test-case")
     {:disconnect, exception, state}
   end
 
+  def update_read_preferences(true), do: "primaryPreferred"
+  def update_read_preferences(false), do: "primary"
+  def update_read_preferences(nil), do: "primary"
+
   defp flags(flags) do
     Enum.reduce(flags, [], fn
       {flag, true},   acc -> [flag|acc]

lib/mongo/monitor.ex

@@ -104,7 +104,7 @@ defmodule Mongo.Monitor do
   defp call_is_master(conn_pid, opts) do
     start_time = System.monotonic_time
     result = try do
-      Mongo.direct_command(conn_pid, %{isMaster: 1}, opts)
+      Mongo.direct_command(conn_pid, [isMaster: 1], opts)
     rescue
       e -> {:error, e}
     end