Skip to content

RUST-1845 Support search index type field #1147

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Jun 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .evergreen/config.yml
Original file line number Diff line number Diff line change
Expand Up @@ -581,6 +581,7 @@ task_groups:
add_expansions_to_env: true
args:
- ${DRIVERS_TOOLS}/.evergreen/atlas/teardown-atlas-cluster.sh
- func: "upload test results"
tasks:
- test-search-index

Expand Down Expand Up @@ -886,6 +887,7 @@ tasks:
- name: test-search-index
commands:
- command: subprocess.exec
type: test
params:
working_dir: src
binary: bash
Expand Down
2 changes: 1 addition & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,7 @@ pub use client::session::ClusterTime;
pub use coll::Namespace;
pub use index::IndexModel;
pub use sdam::public::*;
pub use search_index::SearchIndexModel;
pub use search_index::{SearchIndexModel, SearchIndexType};

/// A boxed future.
pub type BoxFuture<'a, T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;
Expand Down
19 changes: 18 additions & 1 deletion src/search_index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use serde::{Deserialize, Serialize};
use typed_builder::TypedBuilder;

/// Specifies the options for a search index.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, Default, TypedBuilder, Serialize, Deserialize)]
#[builder(field_defaults(default, setter(into)))]
#[non_exhaustive]
Expand All @@ -13,8 +14,24 @@ pub struct SearchIndexModel {
pub definition: Document,

/// The name for this index, if present.
#[serde(skip_serializing_if = "Option::is_none")]
pub name: Option<String>,

/// The type for this index, if present.
#[serde(rename = "type")]
pub index_type: Option<SearchIndexType>,
}

/// Specifies the type of search index.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum SearchIndexType {
/// A regular search index.
Search,
/// A vector search index.
VectorSearch,
/// An unknown type of search index.
#[serde(untagged)]
Other(String),
}

pub mod options {
Expand Down
115 changes: 115 additions & 0 deletions src/test/spec/index_management.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,10 @@ use bson::{doc, oid::ObjectId, Document};
use futures_util::TryStreamExt;

use crate::{
search_index::SearchIndexType,
test::{log_uncaptured, spec::unified_runner::run_unified_tests},
Client,
Collection,
SearchIndexModel,
};

Expand Down Expand Up @@ -272,3 +274,116 @@ async fn search_index_drop_not_found() {

coll0.drop_search_index("test-search-index").await.unwrap();
}

async fn wait_for_index(coll: &Collection<Document>, name: &str) -> Document {
let deadline = Instant::now() + Duration::from_secs(60 * 5);
while Instant::now() < deadline {
let mut cursor = coll.list_search_indexes().name(name).await.unwrap();
while let Some(def) = cursor.try_next().await.unwrap() {
if def.get_str("name") == Ok(name) && def.get_bool("queryable") == Ok(true) {
return def;
}
}
tokio::time::sleep(Duration::from_secs(5)).await;
}
panic!("search index creation timed out");
}

// SearchIndex Case 7: Driver can successfully handle search index types when creating indexes
#[tokio::test]
async fn search_index_create_with_type() {
if env::var("INDEX_MANAGEMENT_TEST_PROSE").is_err() {
log_uncaptured("Skipping index management prose test: INDEX_MANAGEMENT_TEST_PROSE not set");
return;
}
let client = Client::test_builder().build().await;
let coll_name = ObjectId::new().to_hex();
let db = client.database("search_index_test");
db.create_collection(&coll_name).await.unwrap();
let coll0 = db.collection::<Document>(&coll_name);

let name = coll0
.create_search_index(
SearchIndexModel::builder()
.name(String::from("test-search-index-case7-implicit"))
.definition(doc! { "mappings": { "dynamic": false } })
.build(),
)
.await
.unwrap();
assert_eq!(name, "test-search-index-case7-implicit");
let index1 = wait_for_index(&coll0, &name).await;
assert_eq!(index1.get_str("type"), Ok("search"));

let name = coll0
.create_search_index(
SearchIndexModel::builder()
.name(String::from("test-search-index-case7-explicit"))
.index_type(SearchIndexType::Search)
.definition(doc! { "mappings": { "dynamic": false } })
.build(),
)
.await
.unwrap();
assert_eq!(name, "test-search-index-case7-explicit");
let index2 = wait_for_index(&coll0, &name).await;
assert_eq!(index2.get_str("type"), Ok("search"));

let name = coll0
.create_search_index(
SearchIndexModel::builder()
.name(String::from("test-search-index-case7-vector"))
.index_type(SearchIndexType::VectorSearch)
.definition(doc! {
"fields": [{
"type": "vector",
"path": "plot_embedding",
"numDimensions": 1536,
"similarity": "euclidean",
}]
})
.build(),
)
.await
.unwrap();
assert_eq!(name, "test-search-index-case7-vector");
let index3 = wait_for_index(&coll0, &name).await;
assert_eq!(index3.get_str("type"), Ok("vectorSearch"));
}

// SearchIndex Case 8: Driver requires explicit type to create a vector search index
#[tokio::test]
async fn search_index_requires_explicit_vector() {
if env::var("INDEX_MANAGEMENT_TEST_PROSE").is_err() {
log_uncaptured("Skipping index management prose test: INDEX_MANAGEMENT_TEST_PROSE not set");
return;
}
let client = Client::test_builder().build().await;
let coll_name = ObjectId::new().to_hex();
let db = client.database("search_index_test");
db.create_collection(&coll_name).await.unwrap();
let coll0 = db.collection::<Document>(&coll_name);

let result = coll0
.create_search_index(
SearchIndexModel::builder()
.name(String::from("test-search-index-case8-error"))
.definition(doc! {
"fields": [{
"type": "vector",
"path": "plot_embedding",
"numDimensions": 1536,
"similarity": "euclidean",
}]
})
.build(),
)
.await;
assert!(
result
.as_ref()
.is_err_and(|e| e.to_string().contains("Attribute mappings missing")),
"invalid result: {:?}",
result
);
}
Loading