PMM-12451 MongoDB Explain - V3

agent/runner/actions/mongodb_explain_action.go

@@ -18,9 +18,9 @@ import (
 	"context"
 	"fmt"
 	"path/filepath"
+	"strings"
 	"time"
 
-	"github.com/percona/percona-toolkit/src/go/mongolib/proto"
 	"github.com/pkg/errors"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
@@ -39,6 +39,15 @@ type mongodbExplainAction struct {
 	dsn     string
 }
 
+type explain struct {
+	Ns                 string `json:"ns"`
+	Op                 string `json:"op"`
+	Query              bson.D `json:"query,omitempty"`
+	Command            bson.D `json:"command,omitempty"`
+	OriginatingCommand bson.D `json:"originatingCommand,omitempty"`
+	UpdateObj          bson.D `json:"updateobj,omitempty"`
+}
+
 var errCannotExplain = fmt.Errorf("cannot explain this type of query")
 
 // NewMongoDBExplainAction creates a MongoDB EXPLAIN query Action.
@@ -89,28 +98,165 @@ func (a *mongodbExplainAction) Run(ctx context.Context) ([]byte, error) {
 	}
 	defer client.Disconnect(ctx) //nolint:errcheck
 
-	var eq proto.ExampleQuery
+	return explainForQuery(ctx, client, a.params.Query)
+}
+
+func (a *mongodbExplainAction) sealed() {}
+
+func (e explain) prepareCommand() (bson.D, error) {
+	command := e.Command
+
+	switch e.Op {
+	case "query":
+		if len(command) == 0 {
+			command = e.Query
+		}
+
+		if len(command) == 0 || command[0].Key != "find" {
+			return bson.D{
+				{Key: "find", Value: e.getCollection()},
+				{Key: "filter", Value: command},
+			}, nil
+		}
+
+		if len(command) != 0 && command[0].Key == "query" {
+			return bson.D{
+				{Key: "find", Value: e.getCollection()},
+				{Key: "filter", Value: command[0].Value},
+			}, nil
+		}
+
+		return dropDBField(command), nil
+	case "update":
+		if len(command) == 0 {
+			command = bson.D{
+				{Key: "q", Value: e.Query},
+				{Key: "u", Value: e.UpdateObj},
+			}
+		}
+
+		return bson.D{
+			{Key: "update", Value: e.getCollection()},
+			{Key: "updates", Value: []any{command}},
+		}, nil
+	case "remove":
+		if len(command) == 0 {
+			command = bson.D{{Key: "q", Value: e.Query}}
+		}
+
+		return bson.D{
+			{Key: "delete", Value: e.getCollection()},
+			{Key: "deletes", Value: []any{command}},
+		}, nil
+	case "getmore":
+		if len(e.OriginatingCommand) == 0 {
+			return bson.D{{Key: "getmore", Value: ""}}, nil
+		}
+
+		command = e.OriginatingCommand
+
+		return dropDBField(command), nil
+	case "command":
+		command = dropDBField(command)
+
+		if len(command) == 0 {
+			return command, nil
+		}
+
+		switch command[0].Key {
+		// Not supported commands.
+		case "dbStats":
+			return nil, errors.Errorf("command %s is not supported for explain", command[0].Key)
+		case "group":
+		default:
+			return command, nil
+		}
+
+		return fixReduceField(command), nil
+	// Not supported operations.
+	case "insert", "drop":
+		return nil, errors.Errorf("operation %s is not supported for explain", e.Op)
+	}
+
+	return command, nil
+}
+
+func (e explain) getDB() string {
+	s := strings.SplitN(e.Ns, ".", 2)
+	if len(s) == 2 {
+		return s[0]
+	}
+
+	return ""
+}
+
+func (e explain) getCollection() string {
+	s := strings.SplitN(e.Ns, ".", 2)
+	if len(s) == 2 {
+		return s[1]
+	}
+
+	return ""
+}
+
+// dropDBField remove DB field to be able run explain on all supported types.
+// Otherwise it could end up with BSON field 'xxx.$db' is a duplicate field.
+func dropDBField(command bson.D) bson.D {
+	for i := range command {
+		if command[i].Key != "$db" {
+			continue
+		}
+
+		if len(command)-1 == i {
+			return command[:i]
+		}
+
+		return append(command[:i], command[i+1:]...)
+	}
+
+	return command
+}
+
+// fixReduceField fixing nil/empty values after unmarshalling funcs.
+func fixReduceField(command bson.D) bson.D {
+	var group bson.D
+	var ok bool
+	if group, ok = command[0].Value.(bson.D); !ok {
+		return command
+	}
 
-	err = bson.UnmarshalExtJSON([]byte(a.params.Query), true, &eq)
+	for i := range group {
+		if group[i].Key == "$reduce" {
+			group[i].Value = "{}"
+			command[0].Value = group
+			break
+		}
+	}
+
+	return command
+}
+
+func explainForQuery(ctx context.Context, client *mongo.Client, query string) ([]byte, error) {
+	var e explain
+	err := bson.UnmarshalExtJSON([]byte(query), false, &e)
 	if err != nil {
-		return nil, errors.Wrapf(err, "Query: %s", a.params.Query)
+		return nil, errors.Wrapf(err, "Query: %s", query)
 	}
 
-	database := "admin"
-	if eq.Db() != "" {
-		database = eq.Db()
+	preparedCommand, err := e.prepareCommand()
+	if err != nil {
+		return nil, errors.Wrap(errCannotExplain, err.Error())
 	}
-	res := client.Database(database).RunCommand(ctx, eq.ExplainCmd())
+	command := bson.D{{Key: "explain", Value: preparedCommand}}
+	res := client.Database(e.getDB()).RunCommand(ctx, command)
 	if res.Err() != nil {
 		return nil, errors.Wrap(errCannotExplain, res.Err().Error())
 	}
 
-	result, err := res.DecodeBytes()
+	result, err := res.Raw()
 	if err != nil {
 		return nil, err
 	}
-	// We need it because result
+
 	return []byte(result.String()), nil
 }
-
-func (a *mongodbExplainAction) sealed() {}