Skip to content

Commit bbabc25

Browse files
committed
Support cursors on listIndexes and listCollections.
1 parent baa44ca commit bbabc25

File tree

1 file changed

+54
-28
lines changed

1 file changed

+54
-28
lines changed

session.go

Lines changed: 54 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1230,32 +1230,37 @@ func (c *Collection) DropIndex(key ...string) error {
12301230
// See the EnsureIndex method for more details on indexes.
12311231
func (c *Collection) Indexes() (indexes []Index, err error) {
12321232
// Try with a command.
1233-
var cmdResult struct {
1234-
Indexes []indexSpec
1235-
}
1236-
err = c.Database.Run(bson.D{{"listIndexes", c.Name}}, &cmdResult)
1237-
if err == nil {
1238-
for _, spec := range cmdResult.Indexes {
1239-
indexes = append(indexes, indexFromSpec(spec))
1233+
var result struct {
1234+
Indexes []bson.Raw
1235+
1236+
Cursor struct {
1237+
FirstBatch []bson.Raw "firstBatch"
1238+
NS string
1239+
Id int64
12401240
}
1241-
sort.Sort(indexSlice(indexes))
1242-
return indexes, nil
12431241
}
1244-
if err != nil && !isNoCmd(err) {
1242+
var iter *Iter
1243+
err = c.Database.Run(bson.D{{"listIndexes", c.Name}, {"cursor", bson.D{}}}, &result)
1244+
if err == nil {
1245+
firstBatch := result.Indexes
1246+
if firstBatch == nil {
1247+
firstBatch = result.Cursor.FirstBatch
1248+
}
1249+
iter = c.Database.C(result.Cursor.NS).NewIter(nil, firstBatch, result.Cursor.Id, nil)
1250+
} else if isNoCmd(err) {
1251+
// Command not yet supported. Query the database instead.
1252+
iter = c.Database.C("system.indexes").Find(bson.M{"ns": c.FullName}).Iter()
1253+
} else {
12451254
return nil, err
12461255
}
12471256

1248-
// Command not yet supported. Query the database instead.
1249-
query := c.Database.C("system.indexes").Find(bson.M{"ns": c.FullName})
1250-
iter := query.Sort("name").Iter()
1251-
for {
1252-
var spec indexSpec
1253-
if !iter.Next(&spec) {
1254-
break
1255-
}
1257+
var spec indexSpec
1258+
for iter.Next(&spec) {
12561259
indexes = append(indexes, indexFromSpec(spec))
12571260
}
1258-
err = iter.Close()
1261+
if err = iter.Close(); err != nil {
1262+
return nil, err
1263+
}
12591264
sort.Sort(indexSlice(indexes))
12601265
return indexes, nil
12611266
}
@@ -2755,17 +2760,38 @@ func (s *Session) FindRef(ref *DBRef) *Query {
27552760

27562761
// CollectionNames returns the collection names present in the db database.
27572762
func (db *Database) CollectionNames() (names []string, err error) {
2763+
// Clone session and set it to strong mode so that the server
2764+
// used for the query may be safely obtained afterwards, if
2765+
// necessary for iteration when a cursor is received.
2766+
session := db.Session
2767+
cloned := session.Clone()
2768+
cloned.SetMode(Strong, false)
2769+
defer cloned.Close()
2770+
27582771
// Try with a command.
2759-
var cmdResult struct {
2760-
Collections []struct {
2761-
Name string
2772+
var result struct {
2773+
Collections []bson.Raw
2774+
2775+
Cursor struct {
2776+
FirstBatch []bson.Raw "firstBatch"
2777+
NS string
2778+
Id int64
27622779
}
27632780
}
2764-
err = db.Run(bson.D{{"listCollections", 1}}, &cmdResult)
2781+
err = db.Run(bson.D{{"listCollections", 1}, {"cursor", bson.D{}}}, &result)
27652782
if err == nil {
2766-
for _, coll := range cmdResult.Collections {
2783+
firstBatch := result.Collections
2784+
if firstBatch == nil {
2785+
firstBatch = result.Cursor.FirstBatch
2786+
}
2787+
iter := db.C(result.Cursor.NS).NewIter(nil, firstBatch, result.Cursor.Id, nil)
2788+
var coll struct{ Name string }
2789+
for iter.Next(&coll) {
27672790
names = append(names, coll.Name)
27682791
}
2792+
if err := iter.Close(); err != nil {
2793+
return nil, err
2794+
}
27692795
sort.Strings(names)
27702796
return names, err
27712797
}
@@ -2776,10 +2802,10 @@ func (db *Database) CollectionNames() (names []string, err error) {
27762802
// Command not yet supported. Query the database instead.
27772803
nameIndex := len(db.Name) + 1
27782804
iter := db.C("system.namespaces").Find(nil).Iter()
2779-
var result *struct{ Name string }
2780-
for iter.Next(&result) {
2781-
if strings.Index(result.Name, "$") < 0 || strings.Index(result.Name, ".oplog.$") >= 0 {
2782-
names = append(names, result.Name[nameIndex:])
2805+
var coll struct{ Name string }
2806+
for iter.Next(&coll) {
2807+
if strings.Index(coll.Name, "$") < 0 || strings.Index(coll.Name, ".oplog.$") >= 0 {
2808+
names = append(names, coll.Name[nameIndex:])
27832809
}
27842810
}
27852811
if err := iter.Close(); err != nil {

0 commit comments

Comments
 (0)