@@ -3,16 +3,15 @@ package qmgo
33import (
44 "context"
55 "fmt"
6- "reflect"
7- "strings"
8-
96 "github.com/qiniu/qmgo/field"
107 "github.com/qiniu/qmgo/hook"
118 opts "github.com/qiniu/qmgo/options"
129 "go.mongodb.org/mongo-driver/bson"
1310 "go.mongodb.org/mongo-driver/mongo"
1411 "go.mongodb.org/mongo-driver/mongo/options"
1512 "go.mongodb.org/mongo-driver/x/bsonx"
13+ "reflect"
14+ "strings"
1615)
1716
1817// Collection is a handle to a MongoDB collection
@@ -294,77 +293,104 @@ func (c *Collection) Aggregate(ctx context.Context, pipeline interface{}) Aggreg
294293// Example:indexes = []string{"idx1", "-idx2", "idx3,idx4"}
295294// Three indexes will be created, index idx1 with ascending order, index idx2 with descending order, idex3 and idex4 are Compound ascending sort index
296295// Reference: https://docs.mongodb.com/manual/reference/command/createIndexes/
297- func (c * Collection ) ensureIndex (ctx context.Context , indexes []string , isUnique bool ) {
296+ func (c * Collection ) ensureIndex (ctx context.Context , indexes []opts. IndexModel ) error {
298297 var indexModels []mongo.IndexModel
299298
300299 // 组建[]mongo.IndexModel
301300 for _ , idx := range indexes {
302301 var model mongo.IndexModel
303302 var keysDoc bsonx.Doc
304303
305- colIndexArr := strings .Split (idx , "," )
306- for _ , field := range colIndexArr {
304+ for _ , field := range idx .Key {
307305 key , n := SplitSortField (field )
308306
309307 keysDoc = keysDoc .Append (key , bsonx .Int32 (n ))
310308 }
311-
309+ iOptions := options .Index ().SetUnique (idx .Unique ).SetBackground (idx .Background ).SetSparse (idx .Sparse )
310+ if idx .ExpireAfterSeconds != nil {
311+ iOptions .SetExpireAfterSeconds (* idx .ExpireAfterSeconds )
312+ }
312313 model = mongo.IndexModel {
313314 Keys : keysDoc ,
314- Options : options . Index (). SetUnique ( isUnique ) ,
315+ Options : iOptions ,
315316 }
316317
317318 indexModels = append (indexModels , model )
318319 }
319320
320321 if len (indexModels ) == 0 {
321- return
322+ return nil
322323 }
323324
324- var err error
325- var res []string
326- res , err = c .collection .Indexes ().CreateMany (ctx , indexModels )
327-
325+ res , err := c .collection .Indexes ().CreateMany (ctx , indexModels )
328326 if err != nil || len (res ) == 0 {
329- s := fmt .Sprint ("<MongoDB.C>: " , c .collection .Name (), " Index: " , indexes , " error: " , err , "res: " , res )
330- panic ( s )
327+ fmt .Println ("<MongoDB.C>: " , c .collection .Name (), " Index: " , indexes , " error: " , err , "res: " , res )
328+ return err
331329 }
332- return
330+ return nil
333331}
334332
333+ // EnsureIndexes Deprecated
334+ // Recommend to use CreateIndexes / CreateOneIndex for more function)
335335// EnsureIndexes creates unique and non-unique indexes in collection
336- func (c * Collection ) EnsureIndexes (ctx context.Context , uniques []string , indexes []string ) {
337- // 创建唯一索引
338- if len (uniques ) != 0 {
339- c .ensureIndex (ctx , uniques , true )
336+ // the combination of indexes is different from CreateIndexes:
337+ // if uniques/indexes is []string{"name"}, means create index "name"
338+ // if uniques/indexes is []string{"name,-age","uid"} means create Compound indexes: name and -age, then create one index: uid
339+ func (c * Collection ) EnsureIndexes (ctx context.Context , uniques []string , indexes []string ) (err error ) {
340+ uniqueModel , indexesModel := []opts.IndexModel {}, []opts.IndexModel {}
341+ for _ , v := range uniques {
342+ vv := strings .Split (v , "," )
343+ model := opts.IndexModel {Key : vv , Unique : true }
344+ uniqueModel = append (uniqueModel , model )
345+ }
346+ if err = c .CreateIndexes (ctx , uniqueModel ); err != nil {
347+ return
340348 }
341349
342- // 创建普通索引
343- if len (indexes ) != 0 {
344- c .ensureIndex (ctx , indexes , false )
350+ for _ , v := range indexes {
351+ vv := strings .Split (v , "," )
352+ model := opts.IndexModel {Key : vv }
353+ indexesModel = append (uniqueModel , model )
354+ }
355+ if err = c .CreateIndexes (ctx , indexesModel ); err != nil {
356+ return
345357 }
358+ return
359+ }
346360
361+ // CreateIndexes creates multiple indexes in collection
362+ // If the Key in opts.IndexModel is []string{"name"}, means create index: name
363+ // If the Key in opts.IndexModel is []string{"name","-age"} means create Compound indexes: name and -age
364+ func (c * Collection ) CreateIndexes (ctx context.Context , indexes []opts.IndexModel ) (err error ) {
365+ // 创建普通索引
366+ err = c .ensureIndex (ctx , indexes )
347367 return
348368}
349369
350- // DropIndexes drop indexes in collection, indexes that be dropped should be in line with inputting indexes
351- func (c * Collection ) DropIndexes (ctx context.Context , indexes []string ) error {
370+ // CreateIndex creates one index
371+ // If the Key in opts.IndexModel is []string{"name"}, means create index name
372+ // If the Key in opts.IndexModel is []string{"name","-age"} means drop Compound indexes: name and -age
373+ func (c * Collection ) CreateOneIndex (ctx context.Context , indexes opts.IndexModel ) error {
374+ // 创建普通索引
375+ return c .ensureIndex (ctx , []opts.IndexModel {indexes })
352376
353- var err error
354- for _ , index := range indexes {
355- _ , err = c .collection .Indexes ().DropOne (ctx , generateDroppedIndex (index ))
356- if err != nil {
357- return err
358- }
377+ }
378+
379+ // DropIndex drop indexes in collection, indexes that be dropped should be in line with inputting indexes
380+ // The indexes is []string{"name"} means drop index: name
381+ // The indexes is []string{"name","-age"} means drop Compound indexes: name and -age
382+ func (c * Collection ) DropIndex (ctx context.Context , indexes []string ) error {
383+ _ , err := c .collection .Indexes ().DropOne (ctx , generateDroppedIndex (indexes ))
384+ if err != nil {
385+ return err
359386 }
360387 return err
361388}
362389
363- // generate indexes that store in mongo which may consist more than one index(like "index1, index2" is stored as "index1_1_index2_1")
364- func generateDroppedIndex (index string ) string {
390+ // generate indexes that store in mongo which may consist more than one index(like []string{ "index1"," index2"} is stored as "index1_1_index2_1")
391+ func generateDroppedIndex (index [] string ) string {
365392 var res string
366- s := strings .Split (index , "," )
367- for _ , e := range s {
393+ for _ , e := range index {
368394 key , sort := SplitSortField (e )
369395 n := key + "_" + fmt .Sprint (sort )
370396 if len (res ) == 0 {
0 commit comments