RUST-1588: Add RunCursorCommand

src/db.rs

@@ -4,6 +4,7 @@ use std::{fmt::Debug, sync::Arc};
 
 #[cfg(feature = "in-use-encryption-unstable")]
 use bson::doc;
+use bson::RawDocumentBuf;
 use futures_util::stream::TryStreamExt;
 
 use crate::{
@@ -15,19 +16,28 @@ use crate::{
         ChangeStream,
     },
     client::session::TransactionState,
-    cmap::conn::PinnedConnectionHandle,
+    cmap::{conn::PinnedConnectionHandle, Command, RawCommandResponse, StreamDescription},
     concern::{ReadConcern, WriteConcern},
-    cursor::Cursor,
+    cursor::{Cursor, CursorSpecification},
     error::{Error, ErrorKind, Result},
     gridfs::{options::GridFsBucketOptions, GridFsBucket},
-    operation::{Aggregate, AggregateTarget, Create, DropDatabase, ListCollections, RunCommand},
+    operation::{
+        Aggregate,
+        AggregateTarget,
+        Create,
+        DropDatabase,
+        ListCollections,
+        Operation,
+        RunCommand,
+    },
     options::{
         AggregateOptions,
         CollectionOptions,
         CreateCollectionOptions,
         DatabaseOptions,
         DropDatabaseOptions,
         ListCollectionsOptions,
+        RunCursorCommandOptions,
     },
     results::CollectionSpecification,
     selection_criteria::SelectionCriteria,
@@ -469,6 +479,53 @@ impl Database {
             .await
     }
 
+    /// Runs a database-level command and returns a cursor to the response.
+    pub async fn run_cursor_command(
+        &self,
+        command: Document,
+        options: RunCursorCommandOptions,
+    ) -> Result<Cursor<Document>> {
+        let rcc = RunCommand::new(self.name().to_string(), command, None, None)?;
+        let rc_command = RunCursorCommand {
+            run_command: rcc,
+            options,
+        };
+        let client = self.client();
+        client.execute_cursor_operation(rc_command).await
+    }
+
+    /// Runs a database-level command and returns a cursor to the response.
+    pub async fn run_cursor_command_with_session(
+        &self,
+        command: Document,
+        options: RunCursorCommandOptions,
+        session: &mut ClientSession,
+    ) -> Result<SessionCursor<Document>> {
+        let mut selection_criteria = SelectionCriteria::ReadPreference(options.read_preference.clone().unwrap()).into();
+        match session.transaction.state {
+            TransactionState::Starting | TransactionState::InProgress => {
+                selection_criteria = match selection_criteria {
+                    Some(selection_criteria) => Some(selection_criteria),
+                    None => {
+                        if let Some(ref options) = session.transaction.options {
+                            options.selection_criteria.clone()
+                        } else {
+                            None
+                        }
+                    }
+                };     
+            }
+            _ => {}
+        }
+        let rcc = RunCommand::new(self.name().to_string(), command, selection_criteria, None)?;
+        let rc_command = RunCursorCommand {
+            run_command: rcc,
+            options,
+        };
+        let client = self.client();
+        client.execute_session_cursor_operation(rc_command, session).await
+    }
+
     /// Runs a database-level command using the provided `ClientSession`.
     ///
     /// If the `ClientSession` provided is currently in a transaction, `command` must not specify a
@@ -607,3 +664,86 @@ impl Database {
         GridFsBucket::new(self.clone(), options.into().unwrap_or_default())
     }
 }
+
+pub(super) struct RunCursorCommand<'conn> {
+    run_command: RunCommand<'conn>,
+    options: RunCursorCommandOptions,
+}
+
+impl<'conn> Operation for RunCursorCommand<'conn> {
+    type O = CursorSpecification;
+    type Command = RawDocumentBuf;
+
+    const NAME: &'static str = "run_cursor_command";
+
+    fn build(&mut self, _description: &StreamDescription) -> Result<Command<Self::Command>> {
+        self.run_command.build(_description)
+    }
+
+    fn serialize_command(&mut self, cmd: Command<Self::Command>) -> Result<Vec<u8>> {
+        self.run_command.serialize_command(cmd)
+    }
+
+    fn extract_at_cluster_time(
+        &self,
+        _response: &bson::RawDocument,
+    ) -> Result<Option<bson::Timestamp>> {
+        self.run_command.extract_at_cluster_time(_response)
+    }
+
+    fn handle_error(&self, error: Error) -> Result<Self::O> {
+        Err(error)
+    }
+
+    fn selection_criteria(&self) -> Option<&SelectionCriteria> {
+        self.run_command.selection_criteria()
+    }
+
+    fn is_acknowledged(&self) -> bool {
+        self.run_command.is_acknowledged()
+    }
+
+    fn write_concern(&self) -> Option<&WriteConcern> {
+        self.run_command.write_concern()
+    }
+
+    fn supports_read_concern(&self, _description: &StreamDescription) -> bool {
+        self.run_command.supports_read_concern(_description)
+    }
+
+    fn supports_sessions(&self) -> bool {
+        self.run_command.supports_sessions()
+    }
+
+    fn retryability(&self) -> crate::operation::Retryability {
+        self.run_command.retryability()
+    }
+
+    fn update_for_retry(&mut self) {
+        self.run_command.update_for_retry()
+    }
+
+    fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> {
+        self.run_command.pinned_connection()
+    }
+
+    fn name(&self) -> &str {
+        self.run_command.name()
+    }
+
+    fn handle_response(
+        &self,
+        response: RawCommandResponse,
+        description: &StreamDescription,
+    ) -> Result<Self::O> {
+        let doc = Operation::handle_response(&self.run_command, response, description)?;
+        let cursor_info = bson::from_document(doc)?;
+        Ok(CursorSpecification::new(
+            cursor_info,
+            description.server_address.clone(),
+            self.options.batch_size,
+            self.options.max_time_ms,
+            self.options.comment.clone(),
+        ))
+    }
+}

src/db/options.rs

@@ -8,8 +8,8 @@ use typed_builder::TypedBuilder;
 use crate::{
     bson::{Bson, Document},
     concern::{ReadConcern, WriteConcern},
-    options::Collation,
-    selection_criteria::SelectionCriteria,
+    options::{Collation, CursorType},
+    selection_criteria::{SelectionCriteria, ReadPreference},
     serde_util,
 };
 
@@ -312,3 +312,35 @@ pub struct ChangeStreamPreAndPostImages {
     /// If `true`, change streams will be able to include pre- and post-images.
     pub enabled: bool,
 }
+
+/// Determines the lifetime of a cursor.
+#[derive(Clone, Debug)]
+#[non_exhaustive]
+pub enum TimeoutMode {
+    /// Times out after one Iteration.
+    Iteration,
+    /// Times out after the end of the Cursor Lifetime.
+    CursorLifetime,
+}
+
+/// Specifies the options to a
+/// [`Database::RunCursorCommand`](../struct.Database.html#method.run_cursor_command) operation.
+#[derive(Clone, Debug, Default, TypedBuilder)]
+#[builder(field_defaults(default, setter(into)))]
+#[non_exhaustive]
+pub struct RunCursorCommandOptions {
+    /// Optional string enum value, one of 'iteration' | 'cursorLifetime'.
+    pub timeout_mode: Option<TimeoutMode>,
+    /// The default read preference for operations.
+    pub read_preference: Option<ReadPreference>,
+    /// Optional string enum value, one of 'tailable' | 'tailableAwait' | 'nonTailable'.
+    pub cursor_type: Option<CursorType>,
+    /// Number of documents to return per batch.
+    pub batch_size: Option<u32>,
+    /// Optional non-negative integer value. Use this value to configure the maxTimeMS option sent
+    /// on subsequent getMore commands.
+    pub max_time_ms: Option<Duration>,
+    /// Optional BSON value. Use this value to configure the comment option sent on subsequent
+    /// getMore commands.
+    pub comment: Option<Bson>,
+}

src/results.rs

@@ -6,8 +6,8 @@ use crate::{
     bson::{serde_helpers, Bson, Document},
     change_stream::event::ResumeToken,
     db::options::CreateCollectionOptions,
-    Namespace,
     serde_util,
+    Namespace,
 };
 
 use bson::{Binary, RawDocumentBuf};