|
@@ -76,20 +76,43 @@ func (b *Bluk) Run() (*mongo.BulkWriteResult, error) {
|
|
|
//
|
|
|
type MgoIter struct {
|
|
|
Cursor *mongo.Cursor
|
|
|
+ Ctx context.Context
|
|
|
}
|
|
|
|
|
|
func (mt *MgoIter) Next(result interface{}) bool {
|
|
|
if mt.Cursor != nil {
|
|
|
- if mt.Cursor.Next(nil) {
|
|
|
- err := mt.Cursor.Decode(result)
|
|
|
+ if mt.Cursor.Next(mt.Ctx) {
|
|
|
+ rType := reflect.TypeOf(result)
|
|
|
+ rVal := reflect.ValueOf(result)
|
|
|
+ if rType.Kind() == reflect.Ptr {
|
|
|
+ rType = rType.Elem()
|
|
|
+ rVal = rVal.Elem()
|
|
|
+ }
|
|
|
+ var err error
|
|
|
+ if rType.Kind() == reflect.Map {
|
|
|
+ r := make(map[string]interface{})
|
|
|
+ err = mt.Cursor.Decode(&r)
|
|
|
+ if rVal.CanSet() {
|
|
|
+ rVal.Set(reflect.ValueOf(r))
|
|
|
+ } else {
|
|
|
+ for it := rVal.MapRange(); it.Next(); {
|
|
|
+ rVal.SetMapIndex(it.Key(), reflect.Value{})
|
|
|
+ }
|
|
|
+ for it := reflect.ValueOf(r).MapRange(); it.Next(); {
|
|
|
+ rVal.SetMapIndex(it.Key(), it.Value())
|
|
|
+ }
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ err = mt.Cursor.Decode(&result)
|
|
|
+ }
|
|
|
if err != nil {
|
|
|
log.Println("mgo cur err", err.Error())
|
|
|
- mt.Cursor.Close(nil)
|
|
|
+ mt.Cursor.Close(mt.Ctx)
|
|
|
return false
|
|
|
}
|
|
|
return true
|
|
|
} else {
|
|
|
- mt.Cursor.Close(nil)
|
|
|
+ mt.Cursor.Close(mt.Ctx)
|
|
|
return false
|
|
|
}
|
|
|
} else {
|
|
@@ -226,6 +249,7 @@ func (ms *MgoSess) Iter() *MgoIter {
|
|
|
log.Println("mgo find err", err.Error())
|
|
|
} else {
|
|
|
it.Cursor = cur
|
|
|
+ it.Ctx = ms.M.Ctx
|
|
|
}
|
|
|
return it
|
|
|
}
|