Golang llmgo.Session类代码示例

// Execute performs the DeleteOp on a given session, yielding the reply when
// successful (and an error otherwise).
func (op *DeleteOp) Execute(session *mgo.Session) (Replyable, error) {
	session.SetSocketTimeout(0)
	if err := mgo.ExecOpWithoutReply(session, &op.DeleteOp); err != nil {
		return nil, err
	}
	return nil, nil
}

// Execute performs the GetMoreOp on a given session, yielding the reply when
// successful (and an error otherwise).
func (op *GetMoreOp) Execute(session *mgo.Session) (Replyable, error) {
	session.SetSocketTimeout(0)
	before := time.Now()

	_, _, data, resultReply, err := mgo.ExecOpWithReply(session, &op.GetMoreOp)
	after := time.Now()

	mgoReply, ok := resultReply.(*mgo.ReplyOp)
	if !ok {
		panic("reply from execution was not the correct type")
	}

	reply := &ReplyOp{
		ReplyOp: *mgoReply,
		Docs:    make([]bson.Raw, 0, len(data)),
	}

	for _, d := range data {
		dataDoc := bson.Raw{}
		err = bson.Unmarshal(d, &dataDoc)
		if err != nil {
			return nil, err
		}
		reply.Docs = append(reply.Docs, dataDoc)
	}

	reply.Latency = after.Sub(before)
	return reply, nil
}

func getTestDBPort(session *mgo.Session) (string, error) {
	if port := os.Getenv("DB_PORT"); port != "" {
		return port, nil
	}

	result := struct {
		Members []struct {
			Name     string `bson:"name"`
			StateStr string `bson:"stateStr"`
		} `bson:"members"`
	}{}

	err := session.DB("admin").Run("replSetGetStatus", &result)
	if err != nil && err.Error() != "not running with --replSet" {
		return "", err
	}

	if err != nil && err.Error() == "not running with --replSet" {
		return fmt.Sprintf("%d", defaultTestPort), nil
	}

	for _, member := range result.Members {
		if member.StateStr == "PRIMARY" {
			return strings.Split(member.Name, ":")[1], nil
		}
	}

	return "", fmt.Errorf("unable to determine database port")
}

func getPrimaryPort(session *mgo.Session) (string, error) {

	result := struct {
		Members []struct {
			Name     string `bson:"name"`
			StateStr string `bson:"stateStr"`
		} `bson:"members"`
	}{}

	res := &struct {
		Msg string
	}{}
	session.Run("ismaster", res)
	isMongosTestServer = (res.Msg == "isdbgrid")
	if isMongosTestServer {
		return "", nil
	}

	err := session.DB("admin").Run("replSetGetStatus", &result)

	if err != nil {
		if err.Error() == "not running with --replSet" {
			return "", nil
		}
		return "", err
	}

	for _, member := range result.Members {
		if member.StateStr == "PRIMARY" {
			return strings.Split(member.Name, ":")[1], nil
		}
	}

	return "", fmt.Errorf("replset status has no primary")
}

// Execute performs the CommandOp on a given session, yielding the reply when
// successful (and an error otherwise).
func (op *CommandOp) Execute(session *mgo.Session) (Replyable, error) {
	session.SetSocketTimeout(0)

	before := time.Now()
	metadata, commandReply, replyData, resultReply, err := mgo.ExecOpWithReply(session, &op.CommandOp)
	after := time.Now()
	if err != nil {
		return nil, err
	}
	mgoCommandReplyOp, ok := resultReply.(*mgo.CommandReplyOp)
	if !ok {
		panic("reply from execution was not the correct type")
	}
	commandReplyOp := &CommandReplyOp{
		CommandReplyOp: *mgoCommandReplyOp,
	}

	commandReplyOp.Metadata = &bson.Raw{}
	err = bson.Unmarshal(metadata, commandReplyOp.Metadata)
	if err != nil {
		return nil, err
	}
	commandReplyAsRaw := &bson.Raw{}
	err = bson.Unmarshal(commandReply, commandReplyAsRaw)
	if err != nil {
		return nil, err
	}
	commandReplyOp.CommandReply = commandReplyAsRaw
	doc := &struct {
		Cursor struct {
			FirstBatch []bson.Raw `bson:"firstBatch"`
			NextBatch  []bson.Raw `bson:"nextBatch"`
		} `bson:"cursor"`
	}{}
	err = commandReplyAsRaw.Unmarshal(&doc)
	if err != nil {
		return nil, err
	}

	if doc.Cursor.FirstBatch != nil {
		commandReplyOp.Docs = doc.Cursor.FirstBatch
	} else if doc.Cursor.NextBatch != nil {
		commandReplyOp.Docs = doc.Cursor.NextBatch
	}

	for _, d := range replyData {
		dataDoc := &bson.Raw{}
		err = bson.Unmarshal(d, &dataDoc)
		if err != nil {
			return nil, err
		}
		commandReplyOp.OutputDocs = append(commandReplyOp.OutputDocs, dataDoc)
	}
	commandReplyOp.Latency = after.Sub(before)
	return commandReplyOp, nil

}