RUST-576 Fix count_documents fails when writeConcern set on Collection (#265)

Author: patrickfreed

src/coll/mod.rs

@@ -20,6 +20,7 @@ use crate::{
     operation::{
         Aggregate,
         Count,
+        CountDocuments,
         Delete,
         Distinct,
         DropCollection,
@@ -232,70 +233,9 @@ where
         options: impl Into<Option<CountOptions>>,
     ) -> Result<i64> {
         let options = options.into();
-
-        let mut pipeline = vec![doc! {
-            "$match": filter.into().unwrap_or_default(),
-        }];
-
-        if let Some(skip) = options.as_ref().and_then(|opts| opts.skip) {
-            pipeline.push(doc! {
-                "$skip": skip
-            });
-        }
-
-        if let Some(limit) = options.as_ref().and_then(|opts| opts.limit) {
-            pipeline.push(doc! {
-                "$limit": limit
-            });
-        }
-
-        pipeline.push(doc! {
-            "$group": {
-                "_id": 1,
-                "n": { "$sum": 1 },
-            }
-        });
-
-        let aggregate_options = options.map(|opts| {
-            AggregateOptions::builder()
-                .hint(opts.hint)
-                .max_time(opts.max_time)
-                .collation(opts.collation)
-                .build()
-        });
-
-        let result = match self
-            .aggregate(pipeline, aggregate_options)
-            .await?
-            .next()
-            .await
-        {
-            Some(doc) => doc?,
-            None => return Ok(0),
-        };
-
-        let n = match result.get("n") {
-            Some(n) => n,
-            None => {
-                return Err(ErrorKind::ResponseError {
-                    message: "server response to count_documents aggregate did not contain the \
-                              'n' field"
-                        .into(),
-                }
-                .into())
-            }
-        };
-
-        bson_util::get_int(n).ok_or_else(|| {
-            ErrorKind::ResponseError {
-                message: format!(
-                    "server response to count_documents aggregate should have contained integer \
-                     'n', but instead had {:?}",
-                    n
-                ),
-            }
-            .into()
-        })
+        let filter = filter.into();
+        let op = CountDocuments::new(self.namespace(), filter, options);
+        self.client().execute_operation(op).await
     }
 
     /// Deletes all documents stored in the collection matching `query`.

src/coll/options.rs

@@ -550,6 +550,17 @@ pub struct CountOptions {
     /// information on how to use this option.
     #[builder(default)]
     pub collation: Option<Collation>,
+
+    /// The criteria used to select a server for this operation.
+    ///
+    /// If none specified, the default set on the collection will be used.
+    #[builder(default)]
+    #[serde(skip_serializing)]
+    pub selection_criteria: Option<SelectionCriteria>,
+
+    /// The level of the read concern.
+    #[builder(default)]
+    pub read_concern: Option<ReadConcern>,
 }
 
 // rustfmt tries to split the link up when it's all on one line, which breaks the link, so we wrap

src/operation/count_documents/mod.rs

@@ -0,0 +1,116 @@
+#[cfg(test)]
+mod test;
+
+use bson::{doc, Document};
+
+use super::{Operation, Retryability};
+use crate::{
+    bson_util,
+    cmap::{Command, CommandResponse, StreamDescription},
+    error::{ErrorKind, Result},
+    operation::aggregate::Aggregate,
+    options::{AggregateOptions, CountOptions},
+    selection_criteria::SelectionCriteria,
+    Namespace,
+};
+
+pub(crate) struct CountDocuments {
+    aggregate: Aggregate,
+}
+
+impl CountDocuments {
+    pub(crate) fn new(
+        namespace: Namespace,
+        filter: Option<Document>,
+        options: Option<CountOptions>,
+    ) -> Self {
+        let mut pipeline = vec![doc! {
+            "$match": filter.unwrap_or_default(),
+        }];
+
+        if let Some(skip) = options.as_ref().and_then(|opts| opts.skip) {
+            pipeline.push(doc! {
+                "$skip": skip
+            });
+        }
+
+        if let Some(limit) = options.as_ref().and_then(|opts| opts.limit) {
+            pipeline.push(doc! {
+                "$limit": limit
+            });
+        }
+
+        pipeline.push(doc! {
+            "$group": {
+                "_id": 1,
+                "n": { "$sum": 1 },
+            }
+        });
+
+        let aggregate_options = options.map(|opts| {
+            AggregateOptions::builder()
+                .hint(opts.hint)
+                .max_time(opts.max_time)
+                .collation(opts.collation)
+                .selection_criteria(opts.selection_criteria)
+                .read_concern(opts.read_concern)
+                .build()
+        });
+
+        Self {
+            aggregate: Aggregate::new(namespace, pipeline, aggregate_options),
+        }
+    }
+}
+
+impl Operation for CountDocuments {
+    type O = i64;
+    const NAME: &'static str = Aggregate::NAME;
+
+    fn build(&self, description: &StreamDescription) -> Result<Command> {
+        self.aggregate.build(description)
+    }
+
+    fn handle_response(&self, response: CommandResponse) -> Result<Self::O> {
+        let result = self
+            .aggregate
+            .handle_response(response)
+            .map(|mut spec| spec.initial_buffer.pop_front())?;
+
+        let result_doc = match result {
+            Some(doc) => doc,
+            None => return Ok(0),
+        };
+
+        let n = match result_doc.get("n") {
+            Some(n) => n,
+            None => {
+                return Err(ErrorKind::ResponseError {
+                    message: "server response to count_documents aggregate did not contain the \
+                              'n' field"
+                        .into(),
+                }
+                .into())
+            }
+        };
+
+        bson_util::get_int(n).ok_or_else(|| {
+            ErrorKind::ResponseError {
+                message: format!(
+                    "server response to count_documents aggregate should have contained integer \
+                     'n', but instead had {:?}",
+                    n
+                ),
+            }
+            .into()
+        })
+    }
+
+    fn selection_criteria(&self) -> Option<&SelectionCriteria> {
+        self.aggregate.selection_criteria()
+    }
+
+    fn retryability(&self) -> Retryability {
+        Retryability::Read
+    }
+}

src/operation/count_documents/test.rs

@@ -0,0 +1,124 @@
+use crate::{
+    bson::doc,
+    bson_util,
+    cmap::{CommandResponse, StreamDescription},
+    coll::Namespace,
+    concern::ReadConcern,
+    operation::{test, Operation},
+    options::{CountOptions, Hint},
+};
+
+use super::CountDocuments;
+
+#[cfg_attr(feature = "tokio-runtime", tokio::test)]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn build() {
+    let ns = Namespace {
+        db: "test_db".to_string(),
+        coll: "test_coll".to_string(),
+    };
+    let count_op = CountDocuments::new(ns, Some(doc! { "x": 1 }), None);
+    let mut count_command = count_op
+        .build(&StreamDescription::new_testing())
+        .expect("error on build");
+
+    let mut expected_body = doc! {
+        "aggregate": "test_coll",
+        "pipeline": [
+            { "$match": { "x": 1 } },
+            { "$group": { "_id": 1, "n": { "$sum": 1 } } },
+        ],
+        "cursor": { }
+    };
+
+    bson_util::sort_document(&mut expected_body);
+    bson_util::sort_document(&mut count_command.body);
+
+    assert_eq!(count_command.body, expected_body);
+    assert_eq!(count_command.target_db, "test_db");
+    assert_eq!(count_command.read_pref, None);
+}
+
+#[cfg_attr(feature = "tokio-runtime", tokio::test)]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn build_with_options() {
+    let skip = 2;
+    let limit = 5;
+    let options = CountOptions::builder()
+        .skip(skip)
+        .limit(limit)
+        .hint(Hint::Name("_id_1".to_string()))
+        .read_concern(ReadConcern::available())
+        .build();
+    let ns = Namespace {
+        db: "test_db".to_string(),
+        coll: "test_coll".to_string(),
+    };
+    let count_op = CountDocuments::new(ns, None, Some(options));
+    let mut count_command = count_op
+        .build(&StreamDescription::new_testing())
+        .expect("error on build");
+
+    let mut expected_body = doc! {
+        "aggregate": "test_coll",
+        "pipeline": [
+            { "$match": {} },
+            { "$skip": skip },
+            { "$limit": limit },
+            { "$group": { "_id": 1, "n": { "$sum": 1 } } },
+        ],
+        "hint": "_id_1",
+        "cursor": { },
+        "readConcern": { "level": "available" },
+    };
+
+    bson_util::sort_document(&mut expected_body);
+    bson_util::sort_document(&mut count_command.body);
+
+    assert_eq!(count_command.body, expected_body);
+    assert_eq!(count_command.target_db, "test_db");
+    assert_eq!(count_command.read_pref, None);
+}
+
+#[cfg_attr(feature = "tokio-runtime", tokio::test)]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn op_selection_criteria() {
+    test::op_selection_criteria(|selection_criteria| {
+        let options = CountOptions {
+            selection_criteria,
+            ..Default::default()
+        };
+        CountDocuments::new(Namespace::empty(), None, Some(options))
+    });
+}
+
+#[cfg_attr(feature = "tokio-runtime", tokio::test)]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn handle_success() {
+    let ns = Namespace {
+        db: "test_db".to_string(),
+        coll: "test_coll".to_string(),
+    };
+    let count_op = CountDocuments::new(ns, None, None);
+
+    let n = 26;
+    let response = CommandResponse::with_document(doc! {
+        "cursor" : {
+            "firstBatch" : [
+                {
+                    "_id" : 1,
+                    "n" : n
+                }
+            ],
+            "id" : 0,
+            "ns" : "test_db.test_coll"
+        },
+        "ok" : 1
+    });
+
+    let actual_values = count_op
+        .handle_response(response)
+        .expect("supposed to succeed");
+
+    assert_eq!(actual_values, n);
+}

src/operation/mod.rs

@@ -1,5 +1,6 @@
 mod aggregate;
 mod count;
+mod count_documents;
 mod create;
 mod delete;
 mod distinct;
@@ -37,6 +38,7 @@ use crate::{
 
 pub(crate) use aggregate::Aggregate;
 pub(crate) use count::Count;
+pub(crate) use count_documents::CountDocuments;
 pub(crate) use create::Create;
 pub(crate) use delete::Delete;
 pub(crate) use distinct::Distinct;