Add find_and_modify on collection

src/collection.rs

@@ -3,6 +3,7 @@ use std::ffi::CStr;
 use std::borrow::Cow;
 
 use mongoc::bindings;
+use bsonc;
 
 use bson::Document;
 
@@ -43,6 +44,35 @@ impl BulkOperationOptions {
     }
 }
 
+pub struct FindAndModifyOptions {
+    pub sort:   Option<Document>,
+    pub new:    bool,
+    pub fields: Option<Document>
+}
+
+impl FindAndModifyOptions {
+    pub fn default() -> FindAndModifyOptions {
+        FindAndModifyOptions {
+            sort:   None,
+            new:    false,
+            fields: None
+        }
+    }
+
+    fn fields_bsonc(&self) -> Option<bsonc::Bsonc> {
+        match self.fields {
+            Some(ref f) => Some(bsonc::Bsonc::from_document(f).unwrap()),
+            None => None
+        }
+    }
+}
+
+pub enum FindAndModifyOperation<'a> {
+    Update(&'a Document),
+    Upsert(&'a Document),
+    Remove
+}
+
 pub struct CountOptions {
     pub query_flags: Flags<QueryFlag>,
     pub skip:        u32,
@@ -316,6 +346,76 @@ impl<'a> Collection<'a> {
         ))
     }
 
+    // Update and return an object.
+    //
+    // This is a thin wrapper around the findAndModify command. Pass in
+    // an operation that either updates, upserts or removes.
+    pub fn find_and_modify(
+        &'a self,
+        query:     &Document,
+        operation: FindAndModifyOperation<'a>,
+        options:   Option<&FindAndModifyOptions>
+    ) -> Result<Document> {
+        assert!(!self.inner.is_null());
+
+        let default_options = FindAndModifyOptions::default();
+        let options         = options.unwrap_or(&default_options);
+        let fields_bsonc    = options.fields_bsonc();
+
+        // Bsonc to store the reply
+        let mut reply = Bsonc::new();
+        // Empty error that might be filled
+        let mut error = BsoncError::empty();
+
+        // Do these before the mongoc call to make sure we keep
+        // them around long enough.
+        let sort = match options.sort {
+            Some(ref doc) => {
+                try!(Bsonc::from_document(doc)).inner()
+            },
+            None => ptr::null()
+        };
+        let update = match operation {
+            FindAndModifyOperation::Update(ref doc) | FindAndModifyOperation::Upsert(ref doc) => {
+                try!(Bsonc::from_document(doc)).inner()
+            },
+            FindAndModifyOperation::Remove => ptr::null()
+        };
+
+        let success = unsafe {
+            bindings::mongoc_collection_find_and_modify(
+                self.inner,
+                try!(Bsonc::from_document(&query)).inner(),
+                sort,
+                update,
+                match fields_bsonc {
+                    Some(ref f) => f.inner(),
+                    None => ptr::null()
+                },
+                match operation {
+                    FindAndModifyOperation::Remove => true,
+                    _ => false
+                } as u8,
+                match operation {
+                    FindAndModifyOperation::Upsert(_) => true,
+                    _ => false
+                } as u8,
+                options.new as u8,
+                reply.mut_inner(),
+                error.mut_inner()
+            )
+        };
+
+        if success == 1 {
+            match reply.as_document() {
+                Ok(document) => return Ok(document),
+                Err(error)   => return Err(error.into())
+            }
+        } else {
+            Err(error.into())
+        }
+    }
+
     pub fn get_name(&self) -> Cow<str> {
         let cstr = unsafe {
             CStr::from_ptr(bindings::mongoc_collection_get_name(self.inner))

tests/collection.rs

@@ -1,6 +1,7 @@
 use bson;
 
 use mongo_driver::CommandAndFindOptions;
+use mongo_driver::collection::FindAndModifyOperation;
 use mongo_driver::uri::Uri;
 use mongo_driver::client::ClientPool;
 use mongo_driver::flags;
@@ -122,6 +123,53 @@ fn test_mutation_and_finding() {
     assert_eq!(0, collection.count(&query, None).unwrap());
 }
 
+#[test]
+fn test_find_and_modify() {
+    let uri        = Uri::new("mongodb://localhost:27017/").unwrap();
+    let pool       = ClientPool::new(uri, None);
+    let client     = pool.pop();
+    let collection = client.get_collection("rust_driver_test", "find_and_modify");
+
+    // Upsert something, it should now exist
+    let query = doc! {
+        "key_1" => "Value 1"
+    };
+    let update = doc! {
+        "$set" => {"content" => 1}
+    };
+    let result = collection.find_and_modify(
+        &query,
+        FindAndModifyOperation::Upsert(&update),
+        None
+    );
+    assert!(result.is_ok());
+    assert_eq!(update.get("content"), result.unwrap().get("content"));
+    assert_eq!(1, collection.count(&query, None).unwrap());
+
+    // Update this record
+    let update2 = doc! {
+        "$set" => {"content" => 2}
+    };
+    let result = collection.find_and_modify(
+        &query,
+        FindAndModifyOperation::Update(&update2),
+        None
+    );
+    assert!(result.is_ok());
+    assert_eq!(1, collection.count(&query, None).unwrap());
+    let found_document = collection.find(&query, None).unwrap().next().unwrap().unwrap();
+    assert_eq!(update2.get("content"), found_document.get("content"));
+
+    // Remove it
+    let result = collection.find_and_modify(
+        &query,
+        FindAndModifyOperation::Remove,
+        None
+    );
+    assert!(result.is_ok());
+    assert_eq!(0, collection.count(&query, None).unwrap());
+}
+
 #[test]
 fn test_insert_failure() {
     let uri        = Uri::new("mongodb://localhost:27018/").unwrap(); // There should be no mongo server here