feat: adds initial findAndModify command support (#51)

* feat: add findAndModify command (wip)

* refact: extracted update parser login

* feat: added update_from_operation

* refact: uses update_from_operation on updates

* feat: findAndModify initial implementation

* feat: added initial upsert support to findAndModify

* feat: supports sort param on findAndModify

Author: fcoury

src/commands/find_and_modify.rs

@@ -0,0 +1,109 @@
+use crate::handler::{CommandExecutionError, Request};
+use crate::parser::parse_update;
+use crate::pg::UpdateResult;
+use crate::{commands::Handler, pg::SqlParam};
+use bson::{doc, Bson, Document};
+
+pub struct FindAndModify {}
+
+impl Handler for FindAndModify {
+    fn new() -> Self {
+        FindAndModify {}
+    }
+
+    fn handle(
+        &self,
+        request: &Request,
+        docs: &Vec<Document>,
+    ) -> Result<Document, CommandExecutionError> {
+        let doc = &docs[0];
+        let db = doc.get_str("$db").unwrap();
+        let collection = doc.get_str("findAndModify").unwrap();
+        let sp = SqlParam::new(db, collection);
+        let query = doc.get_document("query").unwrap();
+        let sort = match doc.get_document("sort") {
+            Ok(sort_doc) => Some(sort_doc),
+            _ => None,
+        };
+        let raw_update = doc.get_document("update").unwrap();
+        let update_doc = parse_update(raw_update);
+        let upsert = doc.get_bool("upsert").unwrap_or(false);
+
+        let mut client = request.get_client();
+        client.create_table_if_not_exists(db, collection).unwrap();
+
+        let res = client
+            .update(
+                &sp,
+                Some(query),
+                sort,
+                update_doc.unwrap(),
+                false,
+                false,
+                true,
+            )
+            .unwrap();
+
+        match res {
+            UpdateResult::Count(total) => {
+                if total == 0 {
+                    if upsert {
+                        let mut obj = query.clone();
+                        obj.extend(extract_operator_values(&raw_update));
+
+                        let res = client.insert_doc(sp, &obj).unwrap();
+                        return Ok(doc! {
+                            "value": null,
+                            "lastErrorObject": {
+                                "updatedExisting": false,
+                                "upserted": res.get_object_id("_id").unwrap().to_string(),
+                                "n": 1,
+                            },
+                            "ok": 1.0,
+                        });
+                    } else {
+                        return Ok(doc! {
+                            "value": null,
+                            "ok": Bson::Double(1.0),
+                        });
+                    }
+                } else {
+                    unreachable!(
+                        "Unexpected numeric result for a findAndUpdate command: {:#?}",
+                        doc
+                    );
+                }
+            }
+            UpdateResult::Document(value) => Ok(doc! {
+                "n": Bson::Int64(1),
+                "value": value,
+                "ok": Bson::Double(1.0),
+            }),
+        }
+    }
+}
+
+fn extract_operator_values(doc: &Document) -> Document {
+    let mut res = Document::new();
+    for (key, value) in doc {
+        if key.starts_with("$") {
+            if let Some(value) = value.as_document() {
+                res.extend(value.clone());
+            }
+        }
+    }
+    res
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extract_operator_values() {
+        assert_eq!(
+            extract_operator_values(&doc! { "$inc": { "score": 1 }, "$set": { "name": "abc" } }),
+            doc! { "score": 1, "name": "abc" }
+        );
+    }
+}

src/commands/mod.rs

@@ -12,6 +12,7 @@ mod delete;
 mod drop;
 mod drop_database;
 mod find;
+mod find_and_modify;
 mod get_cmd_line_opts;
 mod get_parameter;
 mod hello;
@@ -36,6 +37,7 @@ pub use self::delete::Delete;
 pub use self::drop::Drop;
 pub use self::drop_database::DropDatabase;
 pub use self::find::Find;
+pub use self::find_and_modify::FindAndModify;
 pub use self::get_cmd_line_opts::GetCmdLineOpts;
 pub use self::get_parameter::GetParameter;
 pub use self::hello::Hello;
@@ -45,10 +47,7 @@ pub use self::list_collections::ListCollections;
 pub use self::list_databases::ListDatabases;
 pub use self::list_indexes::ListIndexes;
 pub use self::ping::Ping;
-pub use self::update::InvalidUpdateError;
 pub use self::update::Update;
-pub use self::update::UpdateDoc;
-pub use self::update::UpdateOper;
 pub use self::whats_my_uri::WhatsMyUri;
 
 pub trait Handler {

src/commands/update.rs

@@ -1,59 +1,11 @@
 #![allow(dead_code)]
 use crate::handler::{CommandExecutionError, Request};
-use crate::utils::expand_fields;
-use crate::{commands::Handler, pg::SqlParam};
+use crate::parser::parse_update;
+use crate::{commands::Handler, pg::SqlParam, pg::UpdateResult};
 use bson::{doc, Bson, Document};
 
 pub struct Update {}
 
-#[derive(Debug, Clone, PartialEq)]
-pub enum UpdateOper {
-    Update(Vec<UpdateDoc>),
-    Replace(Document),
-}
-
-#[derive(Debug, Clone, PartialEq)]
-pub enum UpdateDoc {
-    Inc(Document),
-    Set(Document),
-    Unset(Document),
-}
-
-#[derive(Debug, Clone, PartialEq)]
-pub struct InvalidUpdateError {
-    reason: String,
-}
-
-impl InvalidUpdateError {
-    pub fn new(reason: String) -> Self {
-        InvalidUpdateError { reason }
-    }
-}
-
-impl UpdateDoc {
-    fn validate(&self) -> Result<UpdateDoc, InvalidUpdateError> {
-        match self {
-            UpdateDoc::Set(doc) => match expand_fields(doc) {
-                Ok(u) => Ok(UpdateDoc::Set(u)),
-                Err(e) => {
-                    return Err(InvalidUpdateError::new(format!(
-                        "Cannot update '{}' and '{}' at the same time",
-                        e.target, e.source
-                    )));
-                }
-            },
-            UpdateDoc::Unset(doc) => Ok(UpdateDoc::Unset(doc.clone())),
-            UpdateDoc::Inc(u) => Ok(UpdateDoc::Inc(u.clone())),
-            // _ => {
-            //     return Err(InvalidUpdateError::new(format!(
-            //         "Unhandled update operation: {:?}",
-            //         self
-            //     )));
-            // }
-        }
-    }
-}
-
 impl Handler for Update {
     fn new() -> Self {
         Update {}
@@ -85,9 +37,22 @@ impl Handler for Update {
                 return Err(CommandExecutionError::new(format!("{:?}", update_doc)));
             }
 
-            n += client
-                .update(&sp, Some(q), update_doc.unwrap(), upsert, multi)
+            let result = client
+                .update(
+                    &sp,
+                    Some(q),
+                    None,
+                    update_doc.unwrap(),
+                    upsert,
+                    multi,
+                    false,
+                )
                 .unwrap();
+
+            match result {
+                UpdateResult::Count(total) => n += total,
+                UpdateResult::Document(_) => n += 1,
+            }
         }
 
         Ok(doc! {
@@ -97,104 +62,3 @@ impl Handler for Update {
         })
     }
 }
-
-fn parse_update(doc: &Document) -> Result<UpdateOper, InvalidUpdateError> {
-    let mut res: Vec<UpdateDoc> = vec![];
-    if !doc.keys().any(|k| k.starts_with("$")) {
-        return Ok(UpdateOper::Replace(doc.clone()));
-    }
-    for (key, value) in doc.iter() {
-        match key.as_str() {
-            "$set" => {
-                let expanded_doc = match expand_fields(value.as_document().unwrap()) {
-                    Ok(doc) => doc,
-                    Err(e) => {
-                        return Err(InvalidUpdateError::new(format!(
-                            "Cannot update '{}' and '{}' at the same time",
-                            e.target, e.source
-                        )));
-                    }
-                };
-                match UpdateDoc::Set(expanded_doc).validate() {
-                    Ok(update_doc) => res.push(update_doc),
-                    Err(e) => {
-                        return Err(InvalidUpdateError::new(format!("{:?}", e)));
-                    }
-                }
-            }
-            "$unset" => {
-                let expanded_doc = match expand_fields(value.as_document().unwrap()) {
-                    Ok(doc) => doc,
-                    Err(e) => {
-                        return Err(InvalidUpdateError::new(format!(
-                            "Cannot update '{}' and '{}' at the same time",
-                            e.target, e.source
-                        )));
-                    }
-                };
-                match UpdateDoc::Unset(expanded_doc).validate() {
-                    Ok(update_doc) => res.push(update_doc),
-                    Err(e) => {
-                        return Err(InvalidUpdateError::new(format!("{:?}", e)));
-                    }
-                }
-            }
-            "$inc" => {
-                let expanded_doc = match expand_fields(value.as_document().unwrap()) {
-                    Ok(doc) => doc,
-                    Err(e) => {
-                        return Err(InvalidUpdateError::new(format!(
-                            "Cannot update '{}' and '{}' at the same time",
-                            e.target, e.source
-                        )));
-                    }
-                };
-                match UpdateDoc::Inc(expanded_doc).validate() {
-                    Ok(update_doc) => res.push(update_doc),
-                    Err(e) => {
-                        return Err(InvalidUpdateError::new(format!("{:?}", e)));
-                    }
-                }
-            }
-            _ => {
-                if key.starts_with("$") || res.len() > 0 {
-                    return Err(InvalidUpdateError::new(format!(
-                        "Unknown modifier: {}",
-                        key
-                    )));
-                }
-            }
-        }
-    }
-    Ok(UpdateOper::Update(res))
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_parse_update() {
-        let set_doc = doc! { "$set": { "a": 1 } };
-        let repl_doc = doc! { "b": 2, "c": 8, "d": 9 };
-        let unknown_doc = doc! { "$xyz": { "a": 1 } };
-        let mixed_doc = doc! { "$set": { "x": 1 }, "b": 2 };
-
-        assert_eq!(
-            parse_update(&set_doc).unwrap(),
-            UpdateOper::Update(vec![UpdateDoc::Set(doc! { "a": 1 })])
-        );
-        assert_eq!(
-            parse_update(&repl_doc).unwrap(),
-            UpdateOper::Replace(repl_doc)
-        );
-        assert_eq!(
-            parse_update(&unknown_doc).unwrap_err(),
-            InvalidUpdateError::new("Unknown modifier: $xyz".to_string())
-        );
-        assert_eq!(
-            parse_update(&mixed_doc).unwrap_err(),
-            InvalidUpdateError::new("Unknown modifier: b".to_string())
-        );
-    }
-}

src/handler.rs

@@ -1,8 +1,8 @@
 #![allow(dead_code)]
 use crate::commands::{
     Aggregate, BuildInfo, CollStats, ConnectionStatus, Create, CreateIndexes, DbStats, Delete,
-    Drop, DropDatabase, Find, GetCmdLineOpts, GetParameter, Handler, Hello, Insert, IsMaster,
-    ListCollections, ListDatabases, ListIndexes, Ping, Update, WhatsMyUri,
+    Drop, DropDatabase, Find, FindAndModify, GetCmdLineOpts, GetParameter, Handler, Hello, Insert,
+    IsMaster, ListCollections, ListDatabases, ListIndexes, Ping, Update, WhatsMyUri,
 };
 use crate::pg::PgDb;
 use crate::wire::{OpCode, OpMsg};
@@ -120,6 +120,8 @@ fn run(request: &Request, docs: &Vec<Document>) -> Result<Document, CommandExecu
 
     if command == "find" {
         Find::new().handle(request, docs)
+    } else if command == "findAndModify" {
+        FindAndModify::new().handle(request, docs)
     } else if command == "aggregate" {
         Aggregate::new().handle(request, docs)
     } else if command == "insert" {

src/parser.rs src/parser/mod.rs

@@ -10,6 +10,13 @@ use mongodb_language_model::{
 use serde_json::Map;
 use std::fmt;
 
+pub use self::update_parser::parse_update;
+pub use self::update_parser::InvalidUpdateError;
+pub use self::update_parser::UpdateDoc;
+pub use self::update_parser::UpdateOper;
+
+mod update_parser;
+
 #[derive(Debug)]
 struct UnimplementedError {
     pub kind: String,