improve command serialization

Author: isabelatkinson

src/bson_util.rs

@@ -125,6 +125,17 @@ pub(crate) fn read_document_bytes<R: Read>(mut reader: R) -> Result<Vec<u8>> {
     Ok(bytes)
 }
 
+pub(crate) fn extend_raw_document_buf(
+    this: &mut RawDocumentBuf,
+    other: RawDocumentBuf,
+) -> Result<()> {
+    for result in other.iter() {
+        let (k, v) = result?;
+        this.append(k, v.to_raw_bson());
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod test {
     use crate::bson_util::num_decimal_digits;

src/cmap/conn/command.rs

@@ -3,7 +3,8 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use super::wire::Message;
 use crate::{
-    bson::Document,
+    bson::{rawdoc, Document},
+    bson_util::extend_raw_document_buf,
     client::{options::ServerApi, ClusterTime, HELLO_COMMAND_NAMES, REDACTED_COMMANDS},
     error::{Error, ErrorKind, Result},
     hello::{HelloCommandResponse, HelloReply},
@@ -177,6 +178,19 @@ impl<T> Command<T> {
     }
 }
 
+impl Command<RawDocumentBuf> {
+    pub(crate) fn into_bson_bytes(mut self) -> Result<Vec<u8>> {
+        let mut command = self.body;
+
+        // Clear the body of the command to avoid re-serializing.
+        self.body = rawdoc! {};
+        let rest_of_command = bson::to_raw_document_buf(&self)?;
+
+        extend_raw_document_buf(&mut command, rest_of_command)?;
+        Ok(command.into_bytes())
+    }
+}
+
 #[derive(Debug, Clone)]
 pub(crate) struct RawCommandResponse {
     pub(crate) source: ServerAddress,

src/operation.rs

@@ -32,7 +32,7 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use crate::{
     bson::{self, Bson, Document},
-    bson_util,
+    bson_util::{self, extend_raw_document_buf},
     client::{ClusterTime, HELLO_COMMAND_NAMES, REDACTED_COMMANDS},
     cmap::{conn::PinnedConnectionHandle, Command, RawCommandResponse, StreamDescription},
     error::{
@@ -243,10 +243,7 @@ pub(crate) fn append_options_to_raw_document<T: Serialize>(
 ) -> Result<()> {
     if let Some(options) = options {
         let options_raw_doc = bson::to_raw_document_buf(options)?;
-        for result in options_raw_doc.into_iter() {
-            let (key, value) = result?;
-            doc.append(key, value.to_raw_bson());
-        }
+        extend_raw_document_buf(doc, options_raw_doc)?;
     }
     Ok(())
 }

src/operation/find_and_modify.rs

@@ -133,6 +133,10 @@ impl<'a, R: Serialize, T: DeserializeOwned> OperationWithDefaults for FindAndMod
         ))
     }
 
+    fn serialize_command(&mut self, cmd: Command<Self::Command>) -> Result<Vec<u8>> {
+        cmd.into_bson_bytes()
+    }
+
     fn handle_response(
         &self,
         response: RawCommandResponse,

src/operation/update.rs

@@ -182,6 +182,10 @@ impl<'a, T: Serialize> OperationWithDefaults for Update<'a, T> {
         ))
     }
 
+    fn serialize_command(&mut self, cmd: Command<Self::Command>) -> Result<Vec<u8>> {
+        cmd.into_bson_bytes()
+    }
+
     fn handle_response(
         &self,
         raw_response: RawCommandResponse,