Skip to content

Commit 08821ef

Browse files
authored
RUST-1845 Support search index type field (#1147)
1 parent 4153a9f commit 08821ef

File tree

10 files changed

+664
-241
lines changed

10 files changed

+664
-241
lines changed

.evergreen/config.yml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -581,6 +581,7 @@ task_groups:
581581
add_expansions_to_env: true
582582
args:
583583
- ${DRIVERS_TOOLS}/.evergreen/atlas/teardown-atlas-cluster.sh
584+
- func: "upload test results"
584585
tasks:
585586
- test-search-index
586587

@@ -886,6 +887,7 @@ tasks:
886887
- name: test-search-index
887888
commands:
888889
- command: subprocess.exec
890+
type: test
889891
params:
890892
working_dir: src
891893
binary: bash

src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ pub use client::session::ClusterTime;
7171
pub use coll::Namespace;
7272
pub use index::IndexModel;
7373
pub use sdam::public::*;
74-
pub use search_index::SearchIndexModel;
74+
pub use search_index::{SearchIndexModel, SearchIndexType};
7575

7676
/// A boxed future.
7777
pub type BoxFuture<'a, T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;

src/search_index.rs

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use serde::{Deserialize, Serialize};
55
use typed_builder::TypedBuilder;
66

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

1516
/// The name for this index, if present.
16-
#[serde(skip_serializing_if = "Option::is_none")]
1717
pub name: Option<String>,
18+
19+
/// The type for this index, if present.
20+
#[serde(rename = "type")]
21+
pub index_type: Option<SearchIndexType>,
22+
}
23+
24+
/// Specifies the type of search index.
25+
#[derive(Debug, Clone, Serialize, Deserialize)]
26+
#[serde(rename_all = "camelCase")]
27+
pub enum SearchIndexType {
28+
/// A regular search index.
29+
Search,
30+
/// A vector search index.
31+
VectorSearch,
32+
/// An unknown type of search index.
33+
#[serde(untagged)]
34+
Other(String),
1835
}
1936

2037
pub mod options {

src/test/spec/index_management.rs

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@ use bson::{doc, oid::ObjectId, Document};
77
use futures_util::TryStreamExt;
88

99
use crate::{
10+
search_index::SearchIndexType,
1011
test::{log_uncaptured, spec::unified_runner::run_unified_tests},
1112
Client,
13+
Collection,
1214
SearchIndexModel,
1315
};
1416

@@ -272,3 +274,116 @@ async fn search_index_drop_not_found() {
272274

273275
coll0.drop_search_index("test-search-index").await.unwrap();
274276
}
277+
278+
async fn wait_for_index(coll: &Collection<Document>, name: &str) -> Document {
279+
let deadline = Instant::now() + Duration::from_secs(60 * 5);
280+
while Instant::now() < deadline {
281+
let mut cursor = coll.list_search_indexes().name(name).await.unwrap();
282+
while let Some(def) = cursor.try_next().await.unwrap() {
283+
if def.get_str("name") == Ok(name) && def.get_bool("queryable") == Ok(true) {
284+
return def;
285+
}
286+
}
287+
tokio::time::sleep(Duration::from_secs(5)).await;
288+
}
289+
panic!("search index creation timed out");
290+
}
291+
292+
// SearchIndex Case 7: Driver can successfully handle search index types when creating indexes
293+
#[tokio::test]
294+
async fn search_index_create_with_type() {
295+
if env::var("INDEX_MANAGEMENT_TEST_PROSE").is_err() {
296+
log_uncaptured("Skipping index management prose test: INDEX_MANAGEMENT_TEST_PROSE not set");
297+
return;
298+
}
299+
let client = Client::test_builder().build().await;
300+
let coll_name = ObjectId::new().to_hex();
301+
let db = client.database("search_index_test");
302+
db.create_collection(&coll_name).await.unwrap();
303+
let coll0 = db.collection::<Document>(&coll_name);
304+
305+
let name = coll0
306+
.create_search_index(
307+
SearchIndexModel::builder()
308+
.name(String::from("test-search-index-case7-implicit"))
309+
.definition(doc! { "mappings": { "dynamic": false } })
310+
.build(),
311+
)
312+
.await
313+
.unwrap();
314+
assert_eq!(name, "test-search-index-case7-implicit");
315+
let index1 = wait_for_index(&coll0, &name).await;
316+
assert_eq!(index1.get_str("type"), Ok("search"));
317+
318+
let name = coll0
319+
.create_search_index(
320+
SearchIndexModel::builder()
321+
.name(String::from("test-search-index-case7-explicit"))
322+
.index_type(SearchIndexType::Search)
323+
.definition(doc! { "mappings": { "dynamic": false } })
324+
.build(),
325+
)
326+
.await
327+
.unwrap();
328+
assert_eq!(name, "test-search-index-case7-explicit");
329+
let index2 = wait_for_index(&coll0, &name).await;
330+
assert_eq!(index2.get_str("type"), Ok("search"));
331+
332+
let name = coll0
333+
.create_search_index(
334+
SearchIndexModel::builder()
335+
.name(String::from("test-search-index-case7-vector"))
336+
.index_type(SearchIndexType::VectorSearch)
337+
.definition(doc! {
338+
"fields": [{
339+
"type": "vector",
340+
"path": "plot_embedding",
341+
"numDimensions": 1536,
342+
"similarity": "euclidean",
343+
}]
344+
})
345+
.build(),
346+
)
347+
.await
348+
.unwrap();
349+
assert_eq!(name, "test-search-index-case7-vector");
350+
let index3 = wait_for_index(&coll0, &name).await;
351+
assert_eq!(index3.get_str("type"), Ok("vectorSearch"));
352+
}
353+
354+
// SearchIndex Case 8: Driver requires explicit type to create a vector search index
355+
#[tokio::test]
356+
async fn search_index_requires_explicit_vector() {
357+
if env::var("INDEX_MANAGEMENT_TEST_PROSE").is_err() {
358+
log_uncaptured("Skipping index management prose test: INDEX_MANAGEMENT_TEST_PROSE not set");
359+
return;
360+
}
361+
let client = Client::test_builder().build().await;
362+
let coll_name = ObjectId::new().to_hex();
363+
let db = client.database("search_index_test");
364+
db.create_collection(&coll_name).await.unwrap();
365+
let coll0 = db.collection::<Document>(&coll_name);
366+
367+
let result = coll0
368+
.create_search_index(
369+
SearchIndexModel::builder()
370+
.name(String::from("test-search-index-case8-error"))
371+
.definition(doc! {
372+
"fields": [{
373+
"type": "vector",
374+
"path": "plot_embedding",
375+
"numDimensions": 1536,
376+
"similarity": "euclidean",
377+
}]
378+
})
379+
.build(),
380+
)
381+
.await;
382+
assert!(
383+
result
384+
.as_ref()
385+
.is_err_and(|e| e.to_string().contains("Attribute mappings missing")),
386+
"invalid result: {:?}",
387+
result
388+
);
389+
}

0 commit comments

Comments
 (0)