diff --git a/collection.go b/collection.go index 76053ca..3c4646d 100644 --- a/collection.go +++ b/collection.go @@ -540,6 +540,64 @@ func generateDroppedIndex(index []string) string { return res } +// ListIndex returns all indexes in collection +func (c *Collection) ListIndexes(ctx context.Context) ([]opts.IndexModel, error) { + cursor, err := c.collection.Indexes().List(ctx) + if err != nil { + return nil, err + } + defer cursor.Close(ctx) + + var indexInfos []bson.M + err = cursor.All(ctx, &indexInfos) + if err != nil { + return nil, err + } + + res := make([]opts.IndexModel, 0, len(indexInfos)) + for _, indexInfo := range indexInfos { + model, err := transIndexInfoToIndexModel(indexInfo) + if err != nil { + return nil, err + } + res = append(res, model) + } + return res, err +} + +// transIndexInfoToIndexModel transfer bson.M to qmgo opts.IndexModel +func transIndexInfoToIndexModel(indexInfo bson.M) (opts.IndexModel, error) { + bytes, err := bson.Marshal(indexInfo) + res := opts.IndexModel{} + if err != nil { + return res, err + } + + indexOptions := options.IndexOptions{} + err = bson.Unmarshal(bytes, &indexOptions) + if err != nil { + return res, err + } + + indexName := indexInfo["name"].(string) + keySlice := make([]string, 0) + if indexName == "_id_" { + keySlice = append(keySlice, "_id") + } else { + sortFields := strings.Split(indexName, "_") + for i := 0; i < len(sortFields); i += 2 { + field := sortFields[i] + if sortFields[i+1] == "-1" { + field = fmt.Sprintf("-%s", field) + } + keySlice = append(keySlice, field) + } + } + res.Key = keySlice + res.IndexOptions = &indexOptions + return res, nil +} + // DropCollection drops collection // it's safe even collection is not exists func (c *Collection) DropCollection(ctx context.Context) error { diff --git a/collection_test.go b/collection_test.go index 005f84b..6c33b39 100644 --- a/collection_test.go +++ b/collection_test.go @@ -174,6 +174,32 @@ func TestCollection_DropIndex(t *testing.T) { ast.Error(err) } +func TestCollection_ListIndexes(t *testing.T) { + ast := require.New(t) + + cli := initClient("test") + defer cli.DropCollection(context.Background()) + + unique := []string{"id1"} + common := []string{"id2,id3", "id4,-id5"} + cli.EnsureIndexes(context.Background(), unique, common) + + indexes, err := cli.ListIndexes(context.Background()) + indexKeys := make([][]string, 0, len(indexes)) + for _, index := range indexes { + indexKeys = append(indexKeys, index.Key) + } + ast.NoError(err) + ast.Equal(4, len(indexes)) + ast.Contains(indexKeys, []string{"_id"}) + ast.Contains(indexKeys, []string{"id1"}) + ast.Contains(indexKeys, []string{"id2", "id3"}) + ast.Contains(indexKeys, []string{"id4", "-id5"}) + + // same index,error + ast.Error(cli.EnsureIndexes(context.Background(), nil, unique)) +} + func TestCollection_Insert(t *testing.T) { ast := require.New(t) diff --git a/database.go b/database.go index b03a457..829b28a 100644 --- a/database.go +++ b/database.go @@ -17,6 +17,7 @@ import ( "context" opts "github.com/qiniu/qmgo/options" + "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson/bsoncodec" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" @@ -50,6 +51,11 @@ func (d *Database) DropDatabase(ctx context.Context) error { return d.database.Drop(ctx) } +// ListCollections lists all collections in the database. +func (d *Database) ListCollectionNames(ctx context.Context) ([]string, error) { + return d.database.ListCollectionNames(ctx, bson.D{}, nil) +} + // RunCommand executes the given command against the database. // // The runCommand parameter must be a document for the command to be executed. It cannot be nil. diff --git a/database_test.go b/database_test.go index 1dcd8af..fa7de32 100644 --- a/database_test.go +++ b/database_test.go @@ -51,9 +51,14 @@ func TestDatabase(t *testing.T) { ast.Equal(dbName, cli.GetDatabaseName()) coll := cli.Collection(collName) ast.Equal(collName, coll.GetCollectionName()) + res, err := coll.InsertOne(context.Background(), bson.D{{Key: "x", Value: 1}}) + ast.NoError(err) + ast.NotNil(res) + collNames, err := cli.ListCollectionNames(context.Background()) + ast.NoError(err) + ast.Contains(collNames, collName) cli.Collection(collName).DropCollection(context.Background()) cli.DropDatabase(context.Background()) - } func TestRunCommand(t *testing.T) {