Update Server::handleMessage()

cybergarage · Aug 22, 2024 · 908a6d9 · 908a6d9
1 parent 13f2849
commit 908a6d9
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/mongo/server.go b/mongo/server.go
@@ -400,7 +400,18 @@ func (server *Server) handleMessage(conn *Conn, reqMsg protocol.Message) (protoc
 
 	switch reqMsg.OpCode() {
 	case protocol.OpMsg:
-		resMsg = protocol.NewMsgWithBody(resDoc)
+		msg, _ := reqMsg.(*OpMsg)
+		q, err := message.NewQueryWithMessage(msg)
+		if err != nil {
+			return nil, err
+		}
+		queryType := q.Type()
+		switch queryType {
+		case message.Insert, message.Delete, message.Update, message.Find, message.KillCursors:
+			resMsg = protocol.NewMsgWithBody(resDoc)
+		default:
+			resMsg = protocol.NewReplyWithDocuments(resDocs)
+		}
 	case protocol.OpQuery:
 		replyMsg := protocol.NewReplyWithDocument(resDoc)
 		replyMsg.SetResponseFlags(protocol.AwaitCapable)