Skip to content

Commit

Permalink
GODRIVER-3125 Allow to set search index type (#1649)
Browse files Browse the repository at this point in the history
Co-authored-by: Jeroen Vervaeke <[email protected]>
Co-authored-by: Preston Vasquez <[email protected]>
  • Loading branch information
3 people authored Jun 3, 2024
1 parent 9e50281 commit 0f10f5e
Show file tree
Hide file tree
Showing 8 changed files with 351 additions and 16 deletions.
147 changes: 147 additions & 0 deletions mongo/integration/search_index_prose_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -311,4 +311,151 @@ func TestSearchIndexProse(t *testing.T) {
actual := doc.Lookup("latestDefinition").Value
assert.Equal(mt, expected, actual, "unmatched definition")
})

case7CollName, err := uuid.New()
assert.NoError(mt, err, "failed to create random collection name for case #7")

mt.RunOpts("case 7: Driver can successfully handle search index types when creating indexes",
mtest.NewOptions().CollectionName(case7CollName.String()),
func(mt *mtest.T) {
ctx := context.Background()

_, err := mt.Coll.InsertOne(ctx, bson.D{})
require.NoError(mt, err, "failed to insert")

view := mt.Coll.SearchIndexes()

definition := bson.D{{"mappings", bson.D{{"dynamic", false}}}}
indexName := "test-search-index-case7-implicit"
opts := options.SearchIndexes().SetName(indexName)
index, err := view.CreateOne(ctx, mongo.SearchIndexModel{
Definition: definition,
Options: opts,
})
require.NoError(mt, err, "failed to create index")
require.Equal(mt, indexName, index, "unmatched name")
var doc bson.Raw
for doc == nil {
cursor, err := view.List(ctx, opts)
require.NoError(mt, err, "failed to list")

if !cursor.Next(ctx) {
break
}
name := cursor.Current.Lookup("name").StringValue()
queryable := cursor.Current.Lookup("queryable").Boolean()
indexType := cursor.Current.Lookup("type").StringValue()
if name == indexName && queryable {
doc = cursor.Current
assert.Equal(mt, indexType, "search")
} else {
t.Logf("cursor: %s, sleep 5 seconds...", cursor.Current.String())
time.Sleep(5 * time.Second)
}
}

indexName = "test-search-index-case7-explicit"
opts = options.SearchIndexes().SetName(indexName).SetType("search")
index, err = view.CreateOne(ctx, mongo.SearchIndexModel{
Definition: definition,
Options: opts,
})
require.NoError(mt, err, "failed to create index")
require.Equal(mt, indexName, index, "unmatched name")
doc = nil
for doc == nil {
cursor, err := view.List(ctx, opts)
require.NoError(mt, err, "failed to list")

if !cursor.Next(ctx) {
break
}
name := cursor.Current.Lookup("name").StringValue()
queryable := cursor.Current.Lookup("queryable").Boolean()
indexType := cursor.Current.Lookup("type").StringValue()
if name == indexName && queryable {
doc = cursor.Current
assert.Equal(mt, indexType, "search")
} else {
t.Logf("cursor: %s, sleep 5 seconds...", cursor.Current.String())
time.Sleep(5 * time.Second)
}
}

indexName = "test-search-index-case7-vector"
type vectorDefinitionField struct {
Type string `bson:"type"`
Path string `bson:"path"`
NumDimensions int `bson:"numDimensions"`
Similarity string `bson:"similarity"`
}

type vectorDefinition struct {
Fields []vectorDefinitionField `bson:"fields"`
}

opts = options.SearchIndexes().SetName(indexName).SetType("vectorSearch")
index, err = view.CreateOne(ctx, mongo.SearchIndexModel{
Definition: vectorDefinition{
Fields: []vectorDefinitionField{{"vector", "path", 1536, "euclidean"}},
},
Options: opts,
})
require.NoError(mt, err, "failed to create index")
require.Equal(mt, indexName, index, "unmatched name")
doc = nil
for doc == nil {
cursor, err := view.List(ctx, opts)
require.NoError(mt, err, "failed to list")

if !cursor.Next(ctx) {
break
}
name := cursor.Current.Lookup("name").StringValue()
queryable := cursor.Current.Lookup("queryable").Boolean()
indexType := cursor.Current.Lookup("type").StringValue()
if name == indexName && queryable {
doc = cursor.Current
assert.Equal(mt, indexType, "vectorSearch")
} else {
t.Logf("cursor: %s, sleep 5 seconds...", cursor.Current.String())
time.Sleep(5 * time.Second)
}
}
})

case8CollName, err := uuid.New()
assert.NoError(mt, err, "failed to create random collection name for case #8")

mt.RunOpts("case 8: Driver requires explicit type to create a vector search index",
mtest.NewOptions().CollectionName(case8CollName.String()),
func(mt *mtest.T) {
ctx := context.Background()

_, err := mt.Coll.InsertOne(ctx, bson.D{})
require.NoError(mt, err, "failed to insert")

view := mt.Coll.SearchIndexes()

type vectorDefinitionField struct {
Type string `bson:"type"`
Path string `bson:"path"`
NumDimensions int `bson:"numDimensions"`
Similarity string `bson:"similarity"`
}

type vectorDefinition struct {
Fields []vectorDefinitionField `bson:"fields"`
}

const indexName = "test-search-index-case7-vector"
opts := options.SearchIndexes().SetName(indexName)
_, err = view.CreateOne(ctx, mongo.SearchIndexModel{
Definition: vectorDefinition{
Fields: []vectorDefinitionField{{"vector", "plot_embedding", 1536, "euclidean"}},
},
Options: opts,
})
assert.ErrorContains(mt, err, "Attribute mappings missing")
})
}
4 changes: 4 additions & 0 deletions mongo/integration/unified/collection_operation_execution.go
Original file line number Diff line number Diff line change
Expand Up @@ -326,6 +326,7 @@ func executeCreateSearchIndex(ctx context.Context, operation *operation) (*opera
var m struct {
Definition interface{}
Name *string
Type *string
}
err = bson.Unmarshal(val.Document(), &m)
if err != nil {
Expand All @@ -334,6 +335,7 @@ func executeCreateSearchIndex(ctx context.Context, operation *operation) (*opera
model.Definition = m.Definition
model.Options = options.SearchIndexes()
model.Options.Name = m.Name
model.Options.Type = m.Type
default:
return nil, fmt.Errorf("unrecognized createSearchIndex option %q", key)
}
Expand Down Expand Up @@ -369,6 +371,7 @@ func executeCreateSearchIndexes(ctx context.Context, operation *operation) (*ope
var m struct {
Definition interface{}
Name *string
Type *string
}
err = bson.Unmarshal(val.Value, &m)
if err != nil {
Expand All @@ -379,6 +382,7 @@ func executeCreateSearchIndexes(ctx context.Context, operation *operation) (*ope
Options: options.SearchIndexes(),
}
model.Options.Name = m.Name
model.Options.Type = m.Type
models = append(models, model)
}
default:
Expand Down
7 changes: 7 additions & 0 deletions mongo/options/searchindexoptions.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ package options
// SearchIndexesOptions represents options that can be used to configure a SearchIndexView.
type SearchIndexesOptions struct {
Name *string
Type *string
}

// SearchIndexes creates a new SearchIndexesOptions instance.
Expand All @@ -22,6 +23,12 @@ func (sio *SearchIndexesOptions) SetName(name string) *SearchIndexesOptions {
return sio
}

// SetType sets the value for the Type field.
func (sio *SearchIndexesOptions) SetType(typ string) *SearchIndexesOptions {
sio.Type = &typ
return sio
}

// CreateSearchIndexesOptions represents options that can be used to configure a SearchIndexView.CreateOne or
// SearchIndexView.CreateMany operation.
type CreateSearchIndexesOptions struct {
Expand Down
3 changes: 3 additions & 0 deletions mongo/search_index_view.go
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,9 @@ func (siv SearchIndexView) CreateMany(
if model.Options != nil && model.Options.Name != nil {
indexes = bsoncore.AppendStringElement(indexes, "name", *model.Options.Name)
}
if model.Options != nil && model.Options.Type != nil {
indexes = bsoncore.AppendStringElement(indexes, "type", *model.Options.Type)
}
indexes = bsoncore.AppendDocumentElement(indexes, "definition", definition)

indexes, err = bsoncore.AppendDocumentEnd(indexes, iidx)
Expand Down
72 changes: 68 additions & 4 deletions testdata/index-management/createSearchIndex.json
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,8 @@
"mappings": {
"dynamic": true
}
}
},
"type": "search"
}
},
"expectError": {
Expand All @@ -73,7 +74,8 @@
"mappings": {
"dynamic": true
}
}
},
"type": "search"
}
],
"$db": "database0"
Expand All @@ -97,7 +99,8 @@
"dynamic": true
}
},
"name": "test index"
"name": "test index",
"type": "search"
}
},
"expectError": {
Expand All @@ -121,7 +124,68 @@
"dynamic": true
}
},
"name": "test index"
"name": "test index",
"type": "search"
}
],
"$db": "database0"
}
}
}
]
}
]
},
{
"description": "create a vector search index",
"operations": [
{
"name": "createSearchIndex",
"object": "collection0",
"arguments": {
"model": {
"definition": {
"fields": [
{
"type": "vector",
"path": "plot_embedding",
"numDimensions": 1536,
"similarity": "euclidean"
}
]
},
"name": "test index",
"type": "vectorSearch"
}
},
"expectError": {
"isError": true,
"errorContains": "Atlas"
}
}
],
"expectEvents": [
{
"client": "client0",
"events": [
{
"commandStartedEvent": {
"command": {
"createSearchIndexes": "collection0",
"indexes": [
{
"definition": {
"fields": [
{
"type": "vector",
"path": "plot_embedding",
"numDimensions": 1536,
"similarity": "euclidean"
}
]
},
"name": "test index",
"type": "vectorSearch"
}
],
"$db": "database0"
Expand Down
30 changes: 26 additions & 4 deletions testdata/index-management/createSearchIndex.yml
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ tests:
- name: createSearchIndex
object: *collection0
arguments:
model: { definition: &definition { mappings: { dynamic: true } } }
model: { definition: &definition { mappings: { dynamic: true } } , type: 'search' }
expectError:
# This test always errors in a non-Atlas environment. The test functions as a unit test by asserting
# that the driver constructs and sends the correct command.
Expand All @@ -39,15 +39,15 @@ tests:
- commandStartedEvent:
command:
createSearchIndexes: *collection0
indexes: [ { definition: *definition } ]
indexes: [ { definition: *definition, type: 'search'} ]
$db: *database0

- description: "name provided for an index definition"
operations:
- name: createSearchIndex
object: *collection0
arguments:
model: { definition: &definition { mappings: { dynamic: true } } , name: 'test index' }
model: { definition: &definition { mappings: { dynamic: true } } , name: 'test index', type: 'search' }
expectError:
# This test always errors in a non-Atlas environment. The test functions as a unit test by asserting
# that the driver constructs and sends the correct command.
Expand All @@ -60,5 +60,27 @@ tests:
- commandStartedEvent:
command:
createSearchIndexes: *collection0
indexes: [ { definition: *definition, name: 'test index' } ]
indexes: [ { definition: *definition, name: 'test index', type: 'search' } ]
$db: *database0

- description: "create a vector search index"
operations:
- name: createSearchIndex
object: *collection0
arguments:
model: { definition: &definition { fields: [ {"type": "vector", "path": "plot_embedding", "numDimensions": 1536, "similarity": "euclidean"} ] }
, name: 'test index', type: 'vectorSearch' }
expectError:
# This test always errors in a non-Atlas environment. The test functions as a unit test by asserting
# that the driver constructs and sends the correct command.
# The expected error message was changed in SERVER-83003. Check for the substring "Atlas" shared by both error messages.
isError: true
errorContains: Atlas
expectEvents:
- client: *client0
events:
- commandStartedEvent:
command:
createSearchIndexes: *collection0
indexes: [ { definition: *definition, name: 'test index', type: 'vectorSearch' } ]
$db: *database0
Loading

0 comments on commit 0f10f5e

Please sign in to comment.