diff --git a/linux/crates/drivers/mongodb/src/shell_command.rs b/linux/crates/drivers/mongodb/src/shell_command.rs index afee17bdf5..91cf43cb7c 100644 --- a/linux/crates/drivers/mongodb/src/shell_command.rs +++ b/linux/crates/drivers/mongodb/src/shell_command.rs @@ -495,34 +495,56 @@ const AGGREGATE_OPTIONS: [&str; 9] = [ ]; /// An aggregate. Its columns are the last `$project`'s, where it ends on -/// one, and its `explain()` asks the planner instead. -fn aggregate(target: Bson, pipeline: Vec, options: Document, database: String, call: &ShellCall) -> Planned { - let columns = pipeline - .last() - .and_then(Bson::as_document) - .and_then(|stage| stage.get_document("$project").ok()) - .map_or(Columns::Discovered, crate::shell_find::projected_columns); +/// one, its `explain()` asks the planner instead, and its `itcount()` +/// counts what it returns on the server. +fn aggregate(target: Bson, mut pipeline: Vec, options: Document, database: String, call: &ShellCall) -> Planned { let mut explained = false; + let mut counted = false; for modifier in call.modifiers() { match modifier.name.as_str() { "explain" => explained = true, + "itcount" => counted = true, "toArray" | "pretty" => {} _ => return Err(unsupported("this aggregate cursor method")), } } - let mut command = doc! { "aggregate": target, "pipeline": pipeline, "cursor": {} }; + let mut cursor = Document::new(); + let mut fields = Document::new(); for (key, value) in options { match key.as_str() { "batchSize" => { - command.insert("cursor", doc! { "batchSize": value }); + cursor.insert("batchSize", value); } "explain" => explained |= value.as_bool().unwrap_or(false), name if AGGREGATE_OPTIONS.contains(&name) => { - command.insert(key, value); + fields.insert(key, value); } _ => return Err(value_error(ValueError(format!("aggregate() takes no option {key}")))), } } + let mut columns = pipeline + .last() + .and_then(Bson::as_document) + .and_then(|stage| stage.get_document("$project").ok()) + .map_or(Columns::Discovered, crate::shell_find::projected_columns); + if counted && !explained { + // The server takes no stage after one that writes, and such a + // pipeline returns no documents to count. + let writes = pipeline + .last() + .and_then(Bson::as_document) + .and_then(|stage| stage.keys().next()) + .is_some_and(|stage| stage == "$out" || stage == "$merge"); + if writes { + return Err(unsupported("itcount() after $out or $merge")); + } + pipeline.push(Bson::Document( + doc! { "$group": { "_id": 1_i32, "n": { "$sum": 1_i32 } } }, + )); + columns = Columns::Discovered; + } + let mut command = doc! { "aggregate": target, "pipeline": pipeline, "cursor": cursor }; + command.extend(fields); if explained { return Ok(Plan::Command { database, @@ -534,7 +556,7 @@ fn aggregate(target: Bson, pipeline: Vec, options: Document, database: Str database, command, columns, - counted: false, + counted, }) } @@ -570,6 +592,28 @@ mod tests { assert!(planned("db.a.aggregate([{$out: 'b'}]).itcount()").is_err()); } + #[test] + fn an_aggregates_itcount_counts_on_the_server() { + let Ok(Plan::Cursor { command, counted, .. }) = + planned("db.a.aggregate([{$match: {size: {$gt: 1}}}]).itcount()") + else { + panic!("not a cursor"); + }; + assert!(counted); + let pipeline = command.get_array("pipeline").unwrap(); + assert_eq!(pipeline.len(), 2); + assert_eq!( + pipeline[1].as_document().unwrap(), + &doc! { "$group": { "_id": 1_i32, "n": { "$sum": 1_i32 } } } + ); + + assert!(planned("db.a.aggregate([{$merge: {into: 'b'}}]).itcount()").is_err()); + assert!(matches!( + planned("db.a.aggregate([{$out: 'b'}], {explain: true}).itcount()"), + Ok(Plan::Command { .. }) + )); + } + #[test] fn a_page_of_aliased_fields_shows_their_own_names() { let Ok(Plan::Cursor { columns, .. }) = planned( diff --git a/linux/crates/drivers/mongodb/tests/integration/statements.rs b/linux/crates/drivers/mongodb/tests/integration/statements.rs index 45b075e576..3460ed4853 100644 --- a/linux/crates/drivers/mongodb/tests/integration/statements.rs +++ b/linux/crates/drivers/mongodb/tests/integration/statements.rs @@ -109,6 +109,14 @@ async fn a_find_reads_the_fields_it_projects_in_the_order_it_sorts() { .await; assert_eq!(counted.end, RunEnd::Completed); assert_eq!(rows(&events), [vec![Value::Int(2)]]); + + let (iterated, events) = run( + session.as_ref(), + &format!("db.getCollection('{collection}').aggregate([{{$match: {{size: {{$gt: 10}}}}}}]).itcount()"), + ) + .await; + assert_eq!(iterated.end, RunEnd::Completed); + assert_eq!(rows(&events), [vec![Value::Int(2)]]); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)]