diff --git a/src/collection.rs b/src/collection.rs index ac6d4a0..04c265e 100644 --- a/src/collection.rs +++ b/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, + pub new: bool, + pub fields: Option +} + +impl FindAndModifyOptions { + pub fn default() -> FindAndModifyOptions { + FindAndModifyOptions { + sort: None, + new: false, + fields: None + } + } + + fn fields_bsonc(&self) -> Option { + 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, 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 { + 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 { let cstr = unsafe { CStr::from_ptr(bindings::mongoc_collection_get_name(self.inner)) diff --git a/tests/collection.rs b/tests/collection.rs index b989172..d37df5b 100644 --- a/tests/collection.rs +++ b/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