RUST-1800: Modularize unified_runner/operation.rs (#1232)

RUST-1800 Modularize unified_runner/operation.rs

src/test/spec/unified_runner/operation/collection.rs

@@ -0,0 +1,228 @@
+use crate::{
+    error::Result,
+    options::{AggregateOptions, CreateCollectionOptions, DropCollectionOptions},
+    test::spec::unified_runner::{
+        operation::{with_mut_session, with_opt_session, TestOperation},
+        Entity,
+        TestRunner,
+    },
+    Collection,
+    Database,
+};
+use bson::{doc, Bson, Document};
+use futures::{future::BoxFuture, TryStreamExt};
+use futures_util::FutureExt;
+use serde::Deserialize;
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub(super) struct AssertCollectionExists {
+    collection_name: String,
+    database_name: String,
+}
+
+impl TestOperation for AssertCollectionExists {
+    fn execute_test_runner_operation<'a>(
+        &'a self,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, ()> {
+        async move {
+            let db = test_runner.internal_client.database(&self.database_name);
+            let names = db.list_collection_names().await.unwrap();
+            assert!(names.contains(&self.collection_name));
+        }
+        .boxed()
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub(super) struct AssertCollectionNotExists {
+    collection_name: String,
+    database_name: String,
+}
+
+impl TestOperation for AssertCollectionNotExists {
+    fn execute_test_runner_operation<'a>(
+        &'a self,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, ()> {
+        async move {
+            let db = test_runner.internal_client.database(&self.database_name);
+            let names = db.list_collection_names().await.unwrap();
+            assert!(!names.contains(&self.collection_name));
+        }
+        .boxed()
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub(super) struct CreateCollection {
+    collection: String,
+    #[serde(flatten)]
+    options: CreateCollectionOptions,
+    session: Option<String>,
+}
+
+impl TestOperation for CreateCollection {
+    fn execute_entity_operation<'a>(
+        &'a self,
+        id: &'a str,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, Result<Option<Entity>>> {
+        async move {
+            let database = test_runner.get_database(id).await;
+            with_opt_session!(
+                test_runner,
+                &self.session,
+                database
+                    .create_collection(&self.collection)
+                    .with_options(self.options.clone()),
+            )
+            .await?;
+            Ok(Some(Entity::Collection(
+                database.collection(&self.collection),
+            )))
+        }
+        .boxed()
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub(super) struct DropCollection {
+    collection: String,
+    #[serde(flatten)]
+    options: DropCollectionOptions,
+    session: Option<String>,
+}
+
+impl TestOperation for DropCollection {
+    fn execute_entity_operation<'a>(
+        &'a self,
+        id: &'a str,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, Result<Option<Entity>>> {
+        async move {
+            let database = test_runner.get_database(id).await;
+            let collection = database.collection::<Document>(&self.collection).clone();
+            with_opt_session!(
+                test_runner,
+                &self.session,
+                collection.drop().with_options(self.options.clone()),
+            )
+            .await?;
+            Ok(None)
+        }
+        .boxed()
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub(super) struct RenameCollection {
+    to: String,
+}
+
+impl TestOperation for RenameCollection {
+    fn execute_entity_operation<'a>(
+        &'a self,
+        id: &'a str,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, Result<Option<Entity>>> {
+        async move {
+            let target = test_runner.get_collection(id).await;
+            let ns = target.namespace();
+            let mut to_ns = ns.clone();
+            to_ns.coll.clone_from(&self.to);
+            let cmd = doc! {
+                "renameCollection": crate::bson::to_bson(&ns)?,
+                "to": crate::bson::to_bson(&to_ns)?,
+            };
+            let admin = test_runner.internal_client.database("admin");
+            admin.run_command(cmd).await?;
+            Ok(None)
+        }
+        .boxed()
+    }
+}
+
+#[derive(Debug, Deserialize)]
+pub(super) struct Aggregate {
+    pipeline: Vec<Document>,
+    session: Option<String>,
+    #[serde(flatten)]
+    options: AggregateOptions,
+}
+
+impl TestOperation for Aggregate {
+    fn execute_entity_operation<'a>(
+        &'a self,
+        id: &'a str,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, Result<Option<Entity>>> {
+        async move {
+            let result = match &self.session {
+                Some(session_id) => {
+                    enum AggregateEntity {
+                        Collection(Collection<Document>),
+                        Database(Database),
+                        Other(String),
+                    }
+                    let entity = match test_runner.entities.read().await.get(id).unwrap() {
+                        Entity::Collection(c) => AggregateEntity::Collection(c.clone()),
+                        Entity::Database(d) => AggregateEntity::Database(d.clone()),
+                        other => AggregateEntity::Other(format!("{:?}", other)),
+                    };
+                    with_mut_session!(test_runner, session_id, |session| async {
+                        let mut cursor = match entity {
+                            AggregateEntity::Collection(collection) => {
+                                collection
+                                    .aggregate(self.pipeline.clone())
+                                    .with_options(self.options.clone())
+                                    .session(&mut *session)
+                                    .await?
+                            }
+                            AggregateEntity::Database(db) => {
+                                db.aggregate(self.pipeline.clone())
+                                    .with_options(self.options.clone())
+                                    .session(&mut *session)
+                                    .await?
+                            }
+                            AggregateEntity::Other(debug) => {
+                                panic!("Cannot execute aggregate on {}", &debug)
+                            }
+                        };
+                        cursor.stream(session).try_collect::<Vec<Document>>().await
+                    })
+                    .await?
+                }
+                None => {
+                    let entities = test_runner.entities.read().await;
+                    let cursor = match entities.get(id).unwrap() {
+                        Entity::Collection(collection) => {
+                            collection
+                                .aggregate(self.pipeline.clone())
+                                .with_options(self.options.clone())
+                                .await?
+                        }
+                        Entity::Database(db) => {
+                            db.aggregate(self.pipeline.clone())
+                                .with_options(self.options.clone())
+                                .await?
+                        }
+                        other => panic!("Cannot execute aggregate on {:?}", &other),
+                    };
+                    cursor.try_collect::<Vec<Document>>().await?
+                }
+            };
+            Ok(Some(Bson::from(result).into()))
+        }
+        .boxed()
+    }
+
+    fn returns_root_documents(&self) -> bool {
+        true
+    }
+}

src/test/spec/unified_runner/operation/command.rs

@@ -0,0 +1,154 @@
+use crate::{
+    action::Action,
+    error::Result,
+    options::{RunCursorCommandOptions, SelectionCriteria},
+    test::spec::unified_runner::{
+        operation::{with_mut_session, with_opt_session, TestOperation},
+        Entity,
+        TestCursor,
+        TestRunner,
+    },
+};
+use bson::{to_bson, Document};
+use futures::{future::BoxFuture, TryStreamExt};
+use futures_util::FutureExt;
+use serde::Deserialize;
+use tokio::sync::Mutex;
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub(super) struct RunCommand {
+    command: Document,
+    // We don't need to use this field, but it needs to be included during deserialization so that
+    // we can use the deny_unknown_fields tag.
+    #[serde(rename = "commandName")]
+    _command_name: String,
+    read_preference: Option<SelectionCriteria>,
+    session: Option<String>,
+}
+
+impl TestOperation for RunCommand {
+    fn execute_entity_operation<'a>(
+        &'a self,
+        id: &'a str,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, Result<Option<Entity>>> {
+        async move {
+            let command = self.command.clone();
+
+            let db = test_runner.get_database(id).await;
+            let result = with_opt_session!(
+                test_runner,
+                &self.session,
+                db.run_command(command)
+                    .optional(self.read_preference.clone(), |a, rp| {
+                        a.selection_criteria(rp)
+                    }),
+            )
+            .await?;
+            let result = to_bson(&result)?;
+            Ok(Some(result.into()))
+        }
+        .boxed()
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub(super) struct RunCursorCommand {
+    command: Document,
+    // We don't need to use this field, but it needs to be included during deserialization so that
+    // we can use the deny_unknown_fields tag.
+    #[serde(rename = "commandName")]
+    _command_name: String,
+
+    #[serde(flatten)]
+    options: RunCursorCommandOptions,
+    session: Option<String>,
+}
+
+impl TestOperation for RunCursorCommand {
+    fn execute_entity_operation<'a>(
+        &'a self,
+        id: &'a str,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, Result<Option<Entity>>> {
+        async move {
+            let command = self.command.clone();
+            let db = test_runner.get_database(id).await;
+            let options = self.options.clone();
+
+            let action = db.run_cursor_command(command).with_options(options);
+            let result = match &self.session {
+                Some(session_id) => {
+                    with_mut_session!(test_runner, session_id, |session| async {
+                        let mut cursor = action.session(&mut *session).await?;
+                        cursor.stream(session).try_collect::<Vec<_>>().await
+                    })
+                    .await?
+                }
+                None => {
+                    let cursor = action.await?;
+                    cursor.try_collect::<Vec<_>>().await?
+                }
+            };
+
+            Ok(Some(bson::to_bson(&result)?.into()))
+        }
+        .boxed()
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub struct CreateCommandCursor {
+    command: Document,
+    // We don't need to use this field, but it needs to be included during deserialization so that
+    // we can use the deny_unknown_fields tag.
+    #[serde(rename = "commandName")]
+    _command_name: String,
+
+    #[serde(flatten)]
+    options: RunCursorCommandOptions,
+    session: Option<String>,
+}
+
+impl TestOperation for CreateCommandCursor {
+    fn execute_entity_operation<'a>(
+        &'a self,
+        id: &'a str,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, Result<Option<Entity>>> {
+        async move {
+            let command = self.command.clone();
+            let db = test_runner.get_database(id).await;
+            let options = self.options.clone();
+
+            let action = db.run_cursor_command(command).with_options(options);
+            match &self.session {
+                Some(session_id) => {
+                    let mut ses_cursor = None;
+                    with_mut_session!(test_runner, session_id, |session| async {
+                        ses_cursor = Some(action.session(session).await);
+                    })
+                    .await;
+                    let test_cursor = TestCursor::Session {
+                        cursor: ses_cursor.unwrap().unwrap(),
+                        session_id: session_id.clone(),
+                    };
+                    Ok(Some(Entity::Cursor(test_cursor)))
+                }
+                None => {
+                    let doc_cursor = action.await?;
+                    let test_cursor = TestCursor::Normal(Mutex::new(doc_cursor));
+                    Ok(Some(Entity::Cursor(test_cursor)))
+                }
+            }
+        }
+        .boxed()
+    }
+
+    fn returns_root_documents(&self) -> bool {
+        false
+    }
+}