wangchuanjin há 4 anos atrás
pai
commit
c48d92ef32
2 ficheiros alterados com 41 adições e 10 exclusões
  1. 28 4
      src/mongodb/mongodb.go
  2. 13 6
      src/mongodb/mongodb_test.go

+ 28 - 4
src/mongodb/mongodb.go

@@ -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
 }

+ 13 - 6
src/mongodb/mongodb_test.go

@@ -121,16 +121,23 @@ func Test_del(t *testing.T) {
 }
 func Test_itor(t *testing.T) {
 	m := &MongodbSim{
-		MongodbAddr: "192.168.3.128:27090",
+		MongodbAddr: "192.168.3.128:27080",
 		Size:        5,
 		DbName:      "wcj",
-		UserName:    "admin",
-		Password:    "123456",
 	}
 	m.InitPool()
 	sess := m.GetMgoConn()
-	it := sess.DB("wcj").C("test").Find(nil).Iter()
-	for m := make(map[string]interface{}); it.Next(&m); {
-		log.Println(m)
+	it := sess.DB("wcj").C("jyopenplatform_user").Find(map[string]interface{}{
+		"i_status": 1,
+		"i_freeze": 1,
+	}).Select(map[string]interface{}{
+		"a_power":          1,
+		"s_m_openid":       1,
+		"o_rulepro":        1,
+		"s_appid":          1,
+		"s_enterprisename": 1,
+	}).Sort("_id").Iter()
+	for user := make(map[string]interface{}); it.Next(&user); {
+		log.Println(user["_id"], user["s_appid"])
 	}
 }