RUST-658 Change estimatedDocumentCount() to use the $collStats agg stage

Author: saghm

src/client/executor.rs

@@ -322,7 +322,7 @@ impl Client {
                     handler.handle_command_succeeded_event(command_succeeded_event);
                 });
 
-                op.handle_response(response)
+                op.handle_response(response, connection.stream_description()?)
             }
         }
     }

src/coll/mod.rs

@@ -265,6 +265,7 @@ where
         resolve_options!(self, options, [read_concern, selection_criteria]);
 
         let op = Count::new(self.namespace(), options);
+
         self.client().execute_operation(op, None).await
     }
 

src/operation/aggregate/mod.rs

@@ -63,7 +63,11 @@ impl Operation for Aggregate {
         ))
     }
 
-    fn handle_response(&self, response: CommandResponse) -> Result<Self::O> {
+    fn handle_response(
+        &self,
+        response: CommandResponse,
+        _description: &StreamDescription,
+    ) -> Result<Self::O> {
         let body: CursorBody = response.body()?;
 
         if self.is_out_or_merge() {

src/operation/aggregate/test.rs

@@ -200,10 +200,10 @@ async fn handle_success() {
         "ok": 1.0
     };
 
-    let result = aggregate.handle_response(CommandResponse::with_document_and_address(
-        address.clone(),
-        response.clone(),
-    ));
+    let result = aggregate.handle_response(
+        CommandResponse::with_document_and_address(address.clone(), response.clone()),
+        &Default::default(),
+    );
     assert!(result.is_ok());
 
     let cursor_spec = result.unwrap();
@@ -228,10 +228,10 @@ async fn handle_success() {
                 .build(),
         ),
     );
-    let result = aggregate.handle_response(CommandResponse::with_document_and_address(
-        address.clone(),
-        response,
-    ));
+    let result = aggregate.handle_response(
+        CommandResponse::with_document_and_address(address.clone(), response),
+        &Default::default(),
+    );
     assert!(result.is_ok());
 
     let cursor_spec = result.unwrap();
@@ -266,7 +266,7 @@ async fn handle_max_await_time() {
     let aggregate = Aggregate::empty();
 
     let spec = aggregate
-        .handle_response(response.clone())
+        .handle_response(response.clone(), &Default::default())
         .expect("handle should succeed");
     assert!(spec.max_time().is_none());
 
@@ -276,7 +276,7 @@ async fn handle_max_await_time() {
         .build();
     let aggregate = Aggregate::new(Namespace::empty(), Vec::new(), Some(options));
     let spec = aggregate
-        .handle_response(response)
+        .handle_response(response, &Default::default())
         .expect("handle_should_succeed");
     assert_eq!(spec.max_time(), Some(max_await));
 }
@@ -308,7 +308,7 @@ async fn handle_write_concern_error() {
     );
 
     let error = aggregate
-        .handle_response(response)
+        .handle_response(response, &Default::default())
         .expect_err("should get wc error");
     match error.kind {
         ErrorKind::WriteError(WriteFailure::WriteConcernError(_)) => {}
@@ -323,7 +323,7 @@ async fn handle_invalid_response() {
 
     let garbled = doc! { "asdfasf": "ASdfasdf" };
     assert!(aggregate
-        .handle_response(CommandResponse::with_document(garbled))
+        .handle_response(CommandResponse::with_document(garbled), &Default::default())
         .is_err());
 
     let missing_cursor_field = doc! {
@@ -333,6 +333,9 @@ async fn handle_invalid_response() {
         }
     };
     assert!(aggregate
-        .handle_response(CommandResponse::with_document(missing_cursor_field))
+        .handle_response(
+            CommandResponse::with_document(missing_cursor_field),
+            &Default::default()
+        )
         .is_err());
 }

src/operation/count/mod.rs

@@ -4,14 +4,16 @@ mod test;
 use serde::Deserialize;
 
 use crate::{
-    bson::{doc, Document},
+    bson::doc,
     cmap::{Command, CommandResponse, StreamDescription},
     coll::{options::EstimatedDocumentCountOptions, Namespace},
-    error::Result,
-    operation::{append_options, Operation, Retryability},
+    error::{Error, ErrorKind, Result},
+    operation::{append_options, CursorBody, Operation, Retryability},
     selection_criteria::SelectionCriteria,
 };
 
+const SERVER_4_9_0_WIRE_VERSION: i32 = 12;
+
 pub(crate) struct Count {
     ns: Namespace,
     options: Option<EstimatedDocumentCountOptions>,
@@ -38,9 +40,30 @@ impl Operation for Count {
     type O = i64;
     const NAME: &'static str = "count";
 
-    fn build(&self, _description: &StreamDescription) -> Result<Command> {
-        let mut body: Document = doc! {
-            Self::NAME: self.ns.coll.clone(),
+    fn build(&self, description: &StreamDescription) -> Result<Command> {
+        let mut body = match description.max_wire_version {
+            Some(v) if v >= SERVER_4_9_0_WIRE_VERSION => {
+                doc! {
+                    "aggregate": self.ns.coll.clone(),
+                    "pipeline": [
+                        {
+                            "$collStats": { "count": {} },
+                        },
+                        {
+                            "$group": {
+                                "_id": 1,
+                                "n": { "$sum": "$count" },
+                            },
+                        },
+                    ],
+                    "cursor": {},
+                }
+            }
+            _ => {
+                doc! {
+                    Self::NAME: self.ns.coll.clone(),
+                }
+            }
         };
 
         append_options(&mut body, self.options.as_ref())?;
@@ -52,8 +75,37 @@ impl Operation for Count {
         ))
     }
 
-    fn handle_response(&self, response: CommandResponse) -> Result<Self::O> {
-        response.body::<ResponseBody>().map(|body| body.n)
+    fn handle_response(
+        &self,
+        response: CommandResponse,
+        description: &StreamDescription,
+    ) -> Result<Self::O> {
+        let response_body: ResponseBody = match description.max_wire_version {
+            Some(v) if v >= SERVER_4_9_0_WIRE_VERSION => {
+                let CursorBody { mut cursor } = response.body()?;
+
+                cursor
+                    .first_batch
+                    .pop_front()
+                    .and_then(|doc| bson::from_document(doc).ok())
+                    .ok_or_else(|| {
+                        Error::from(ErrorKind::ResponseError {
+                            message: "invalid server response to count operation".into(),
+                        })
+                    })?
+            }
+            _ => response.body()?,
+        };
+
+        Ok(response_body.n)
+    }
+
+    fn handle_error(&self, error: Error) -> Result<Self::O> {
+        if error.is_ns_not_found() {
+            Ok(0)
+        } else {
+            Err(error)
+        }
     }
 
     fn selection_criteria(&self) -> Option<&SelectionCriteria> {

src/operation/count/test.rs

@@ -82,7 +82,7 @@ async fn handle_success() {
     let response = CommandResponse::with_document(doc! { "n" : n, "ok" : 1 });
 
     let actual_values = count_op
-        .handle_response(response)
+        .handle_response(response, &Default::default())
         .expect("supposed to succeed");
 
     assert_eq!(actual_values, n);
@@ -95,7 +95,7 @@ async fn handle_response_no_n() {
 
     let response = CommandResponse::with_document(doc! { "ok" : 1 });
 
-    let result = count_op.handle_response(response);
+    let result = count_op.handle_response(response, &Default::default());
     match result.as_ref().map_err(|e| &e.kind) {
         Err(ErrorKind::ResponseError { .. }) => {}
         other => panic!("expected response error, but got {:?}", other),

src/operation/count_documents/mod.rs

@@ -71,10 +71,14 @@ impl Operation for CountDocuments {
         self.aggregate.build(description)
     }
 
-    fn handle_response(&self, response: CommandResponse) -> Result<Self::O> {
+    fn handle_response(
+        &self,
+        response: CommandResponse,
+        description: &StreamDescription,
+    ) -> Result<Self::O> {
         let result = self
             .aggregate
-            .handle_response(response)
+            .handle_response(response, description)
             .map(|mut spec| spec.initial_buffer.pop_front())?;
 
         let result_doc = match result {

src/operation/count_documents/test.rs

@@ -117,7 +117,7 @@ async fn handle_success() {
     });
 
     let actual_values = count_op
-        .handle_response(response)
+        .handle_response(response, &Default::default())
         .expect("supposed to succeed");
 
     assert_eq!(actual_values, n);

src/operation/create/mod.rs

@@ -50,7 +50,11 @@ impl Operation for Create {
         ))
     }
 
-    fn handle_response(&self, response: CommandResponse) -> Result<Self::O> {
+    fn handle_response(
+        &self,
+        response: CommandResponse,
+        _description: &StreamDescription,
+    ) -> Result<Self::O> {
         response.body::<WriteConcernOnlyBody>()?.validate()
     }
 

src/operation/create/test.rs

@@ -80,9 +80,9 @@ async fn handle_success() {
     let op = Create::empty();
 
     let ok_response = CommandResponse::with_document(doc! { "ok": 1.0 });
-    assert!(op.handle_response(ok_response).is_ok());
+    assert!(op.handle_response(ok_response, &Default::default()).is_ok());
     let ok_extra = CommandResponse::with_document(doc! { "ok": 1.0, "hello": "world" });
-    assert!(op.handle_response(ok_extra).is_ok());
+    assert!(op.handle_response(ok_extra, &Default::default()).is_ok());
 }
 
 #[cfg_attr(feature = "tokio-runtime", tokio::test)]
@@ -99,7 +99,7 @@ async fn handle_write_concern_error() {
         "ok": 1
     });
 
-    let result = op.handle_response(response);
+    let result = op.handle_response(response, &Default::default());
     assert!(result.is_err());
 
     match result.unwrap_err().kind {