RUST-947 Implement Index Management API (#420)

Author: NBSquare

src/coll/mod.rs

@@ -2,7 +2,10 @@ pub mod options;
 
 use std::{borrow::Borrow, collections::HashSet, fmt, fmt::Debug, sync::Arc};
 
-use futures_util::stream::StreamExt;
+use futures_util::{
+    future,
+    stream::{StreamExt, TryStreamExt},
+};
 use serde::{
     de::{DeserializeOwned, Error as DeError},
     Deserialize,
@@ -17,19 +20,30 @@ use crate::{
     client::session::TransactionState,
     concern::{ReadConcern, WriteConcern},
     error::{convert_bulk_errors, BulkWriteError, BulkWriteFailure, Error, ErrorKind, Result},
+    index::IndexModel,
     operation::{
         Aggregate,
         Count,
         CountDocuments,
+        CreateIndexes,
         Delete,
         Distinct,
         DropCollection,
+        DropIndexes,
         Find,
         FindAndModify,
         Insert,
+        ListIndexes,
         Update,
     },
-    results::{DeleteResult, InsertManyResult, InsertOneResult, UpdateResult},
+    results::{
+        CreateIndexResult,
+        CreateIndexesResult,
+        DeleteResult,
+        InsertManyResult,
+        InsertOneResult,
+        UpdateResult,
+    },
     selection_criteria::SelectionCriteria,
     Client,
     ClientSession,
@@ -313,6 +327,46 @@ impl<T> Collection<T> {
         self.client().execute_operation(delete, session).await
     }
 
+    async fn create_indexes_common(
+        &self,
+        indexes: impl IntoIterator<Item = IndexModel>,
+        options: impl Into<Option<CreateIndexOptions>>,
+        session: impl Into<Option<&mut ClientSession>>,
+    ) -> Result<CreateIndexesResult> {
+        let session = session.into();
+
+        let mut options = options.into();
+        resolve_write_concern_with_session!(self, options, session.as_ref())?;
+
+        let indexes: Vec<IndexModel> = indexes.into_iter().collect();
+
+        let create_indexes = CreateIndexes::new(self.namespace(), indexes, options);
+        self.client()
+            .execute_operation(create_indexes, session)
+            .await
+    }
+
+    /// Creates the given index on this collection.
+    pub async fn create_index(
+        &self,
+        index: IndexModel,
+        options: impl Into<Option<CreateIndexOptions>>,
+    ) -> Result<CreateIndexResult> {
+        let response = self
+            .create_indexes_common(vec![index], options, None)
+            .await?;
+        Ok(response.into())
+    }
+
+    /// Creates the given indexes on this collection.
+    pub async fn create_indexes(
+        &self,
+        indexes: impl IntoIterator<Item = IndexModel>,
+        options: impl Into<Option<CreateIndexOptions>>,
+    ) -> Result<CreateIndexesResult> {
+        self.create_indexes_common(indexes, options, None).await
+    }
+
     /// Deletes all documents stored in the collection matching `query`.
     pub async fn delete_many(
         &self,
@@ -423,6 +477,81 @@ impl<T> Collection<T> {
             .await
     }
 
+    async fn drop_index_common(
+        &self,
+        name: impl Into<Option<&str>>,
+        options: impl Into<Option<DropIndexOptions>>,
+        session: impl Into<Option<&mut ClientSession>>,
+    ) -> Result<()> {
+        let session = session.into();
+
+        let mut options = options.into();
+        resolve_write_concern_with_session!(self, options, session.as_ref())?;
+
+        // If there is no provided name, that means we should drop all indexes.
+        let index_name = name.into().unwrap_or("*").to_string();
+
+        let drop_index = DropIndexes::new(self.namespace(), index_name, options);
+        self.client().execute_operation(drop_index, session).await
+    }
+
+    /// Drops the index specified by `name` from this collection.
+    pub async fn drop_index(
+        &self,
+        name: impl AsRef<str>,
+        options: impl Into<Option<DropIndexOptions>>,
+    ) -> Result<()> {
+        let name = name.as_ref();
+        if name == "*" {
+            return Err(ErrorKind::InvalidArgument {
+                message: "Cannot pass name \"*\" to drop_index since more than one index would be \
+                          dropped."
+                    .to_string(),
+            }
+            .into());
+        }
+        self.drop_index_common(name, options, None).await
+    }
+
+    /// Drops all indexes associated with this collection.
+    pub async fn drop_indexes(&self, options: impl Into<Option<DropIndexOptions>>) -> Result<()> {
+        self.drop_index_common(None, options, None).await
+    }
+
+    /// Lists all indexes on this collection.
+    pub async fn list_indexes(
+        &self,
+        options: impl Into<Option<ListIndexOptions>>,
+    ) -> Result<Cursor<IndexModel>> {
+        let list_indexes = ListIndexes::new(self.namespace(), options.into());
+        let client = self.client();
+        client
+            .execute_cursor_operation(list_indexes)
+            .await
+            .map(|(spec, session)| Cursor::new(client.clone(), spec, session))
+    }
+
+    async fn list_index_names_common(
+        &self,
+        cursor: impl TryStreamExt<Ok = IndexModel, Error = Error>,
+    ) -> Result<Vec<String>> {
+        cursor
+            .try_filter_map(|index| future::ok(index.get_name()))
+            .try_collect()
+            .await
+    }
+
+    /// Gets the names of all indexes on the collection.
+    pub async fn list_index_names(&self) -> Result<Vec<String>> {
+        let list_indexes = ListIndexes::new(self.namespace(), None);
+        let client = self.client();
+        let cursor: Cursor<IndexModel> = client
+            .execute_cursor_operation(list_indexes)
+            .await
+            .map(|(spec, session)| Cursor::new(client.clone(), spec, session))?;
+        self.list_index_names_common(cursor).await
+    }
+
     async fn update_many_common(
         &self,
         query: Document,

src/coll/options.rs

@@ -1,5 +1,6 @@
 use std::time::Duration;
 
+use bson::serde_helpers;
 use serde::{de::Error, Deserialize, Deserializer, Serialize, Serializer};
 use serde_with::skip_serializing_none;
 use typed_builder::TypedBuilder;
@@ -831,6 +832,36 @@ pub struct FindOneOptions {
     pub sort: Option<Document>,
 }
 
+/// Specifies the options to a
+/// [`Collection::create_index`](../struct.Collection.html#method.create_index) or [`Collection::
+/// create_indexes`](../struct.Collection.html#method.create_indexes) operation.
+///
+/// For more information, see [`createIndexes`](https://docs.mongodb.com/manual/reference/command/createIndexes/).
+#[serde_with::skip_serializing_none]
+#[derive(Clone, Debug, Default, TypedBuilder, Serialize)]
+#[builder(field_defaults(default, setter(into)))]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+#[non_exhaustive]
+pub struct CreateIndexOptions {
+    /// Specify the commit quorum needed to mark an `index` as ready.
+    pub commit_quorum: Option<CommitQuorum>,
+
+    /// The maximum amount of time to allow the index to build.
+    ///
+    /// This option maps to the `maxTimeMS` MongoDB query option, so the duration will be sent
+    /// across the wire as an integer number of milliseconds.
+    #[serde(
+        rename = "maxTimeMS",
+        default,
+        serialize_with = "bson_util::serialize_duration_option_as_int_millis",
+        deserialize_with = "bson_util::deserialize_duration_option_from_u64_millis"
+    )]
+    pub max_time: Option<Duration>,
+
+    /// The write concern for the operation.
+    pub write_concern: Option<WriteConcern>,
+}
+
 /// Specifies the options to a [`Collection::drop`](../struct.Collection.html#method.drop)
 /// operation.
 #[serde_with::skip_serializing_none]
@@ -842,3 +873,114 @@ pub struct DropCollectionOptions {
     /// The write concern for the operation.
     pub write_concern: Option<WriteConcern>,
 }
+
+/// Specifies the options to a
+/// [`Collection::drop_index`](../struct.Collection.html#method.drop_index) or
+/// [`Collection::drop_indexes`](../struct.Collection.html#method.drop_indexes) operation.
+#[serde_with::skip_serializing_none]
+#[derive(Clone, Debug, Default, Deserialize, TypedBuilder, Serialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+#[builder(field_defaults(default, setter(into)))]
+#[non_exhaustive]
+pub struct DropIndexOptions {
+    /// The maximum amount of time to allow the index to drop.
+    ///
+    /// This option maps to the `maxTimeMS` MongoDB query option, so the duration will be sent
+    /// across the wire as an integer number of milliseconds.
+    #[serde(
+        rename = "maxTimeMS",
+        default,
+        serialize_with = "bson_util::serialize_duration_option_as_int_millis",
+        deserialize_with = "bson_util::deserialize_duration_option_from_u64_millis"
+    )]
+    pub max_time: Option<Duration>,
+
+    /// The write concern for the operation.
+    pub write_concern: Option<WriteConcern>,
+}
+
+/// Specifies the options to a
+/// [`Collection::list_indexes`](../struct.Collection.html#method.list_indexes) operation.
+#[serde_with::skip_serializing_none]
+#[derive(Clone, Debug, Default, Deserialize, TypedBuilder, Serialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+#[builder(field_defaults(default, setter(into)))]
+#[non_exhaustive]
+pub struct ListIndexOptions {
+    /// The maximum amount of time to search for the index.
+    ///
+    /// This option maps to the `maxTimeMS` MongoDB query option, so the duration will be sent
+    /// across the wire as an integer number of milliseconds.
+    #[serde(
+        rename = "maxTimeMS",
+        default,
+        serialize_with = "bson_util::serialize_duration_option_as_int_millis",
+        deserialize_with = "bson_util::deserialize_duration_option_from_u64_millis"
+    )]
+    pub max_time: Option<Duration>,
+
+    /// The number of indexes the server should return per cursor batch.
+    #[serde(default, serialize_with = "bson_util::serialize_u32_option_as_i32")]
+    pub batch_size: Option<u32>,
+}
+
+/// The minimum number of data-bearing voting replica set members (i.e. commit quorum), including
+/// the primary, that must report a successful index build before the primary marks the indexes as
+/// ready.
+///
+/// For more information, see the [documentation](https://docs.mongodb.com/manual/reference/command/createIndexes/#definition)
+#[derive(Clone, Debug, PartialEq)]
+#[non_exhaustive]
+pub enum CommitQuorum {
+    /// A specific number of voting replica set members. When set to 0, disables quorum voting.
+    Nodes(u32),
+
+    /// All data-bearing voting replica set members (default).
+    VotingMembers,
+
+    /// A simple majority of voting members.
+    Majority,
+
+    /// A replica set tag name.
+    Custom(String),
+}
+
+impl Serialize for CommitQuorum {
+    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        match self {
+            CommitQuorum::Nodes(n) => serde_helpers::serialize_u32_as_i32(n, serializer),
+            CommitQuorum::VotingMembers => serializer.serialize_str("votingMembers"),
+            CommitQuorum::Majority => serializer.serialize_str("majority"),
+            CommitQuorum::Custom(s) => serializer.serialize_str(s),
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for CommitQuorum {
+    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        #[derive(Deserialize)]
+        #[serde(untagged)]
+        enum IntOrString {
+            Int(u32),
+            String(String),
+        }
+        match IntOrString::deserialize(deserializer)? {
+            IntOrString::String(s) => {
+                if s == "votingMembers" {
+                    Ok(CommitQuorum::VotingMembers)
+                } else if s == "majority" {
+                    Ok(CommitQuorum::Majority)
+                } else {
+                    Ok(CommitQuorum::Custom(s))
+                }
+            }
+            IntOrString::Int(i) => Ok(CommitQuorum::Nodes(i)),
+        }
+    }
+}

src/index/mod.rs

@@ -0,0 +1,55 @@
+pub mod options;
+
+use crate::bson::Document;
+
+use self::options::*;
+use serde::{Deserialize, Serialize};
+
+use typed_builder::TypedBuilder;
+
+/// Specifies the fields and options for an index. For more information, see the [documentation](https://docs.mongodb.com/manual/indexes/).
+#[derive(Clone, Debug, Default, Deserialize, TypedBuilder, Serialize)]
+#[builder(field_defaults(default, setter(into)))]
+#[serde(rename_all = "camelCase")]
+#[non_exhaustive]
+pub struct IndexModel {
+    /// Specifies the index’s fields. For each field, specify a key-value pair in which the key is
+    /// the name of the field to index and the value is index type.
+    #[serde(rename = "key")]
+    keys: Document,
+
+    /// The options for the index.
+    #[serde(flatten)]
+    options: Option<IndexOptions>,
+}
+
+impl IndexModel {
+    /// If the client did not specify a name, generate and set it. Otherwise, do nothing.
+    pub(crate) fn update_name(&mut self) {
+        if self
+            .options
+            .as_ref()
+            .and_then(|o| o.name.as_ref())
+            .is_none()
+        {
+            let key_names: Vec<String> = self
+                .keys
+                .iter()
+                .map(|(k, v)| format!("{}_{}", k, v.to_string()))
+                .collect();
+            self.options.get_or_insert(IndexOptions::default()).name = Some(key_names.join("_"));
+        }
+    }
+
+    pub(crate) fn get_name(&self) -> Option<String> {
+        self.options.as_ref().and_then(|o| o.name.as_ref()).cloned()
+    }
+
+    #[cfg(test)]
+    pub(crate) fn is_unique(&self) -> bool {
+        self.options
+            .as_ref()
+            .and_then(|o| o.unique)
+            .unwrap_or(false)
+    }
+}