Merge pull request #6010 from planetscale/at-plandescriptions

Improve query plan output

go/vt/vtgate/engine/delete.go

@@ -20,7 +20,8 @@ import (
 	"fmt"
 	"time"
 
-	"vitess.io/vitess/go/jsonutil"
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/key"
 	"vitess.io/vitess/go/vt/srvtopo"
@@ -41,45 +42,6 @@ type Delete struct {
 	noInputs
 }
 
-// MarshalJSON serializes the Delete into a JSON representation.
-// It's used for testing and diagnostics.
-func (del *Delete) MarshalJSON() ([]byte, error) {
-	var tname, vindexName, ksidVindexName string
-	if del.Table != nil {
-		tname = del.Table.Name.String()
-	}
-	if del.Vindex != nil {
-		vindexName = del.Vindex.String()
-	}
-	if del.KsidVindex != nil {
-		ksidVindexName = del.KsidVindex.String()
-	}
-	marshalDelete := struct {
-		Opcode               string
-		Keyspace             *vindexes.Keyspace   `json:",omitempty"`
-		Query                string               `json:",omitempty"`
-		Vindex               string               `json:",omitempty"`
-		Values               []sqltypes.PlanValue `json:",omitempty"`
-		Table                string               `json:",omitempty"`
-		OwnedVindexQuery     string               `json:",omitempty"`
-		KsidVindex           string               `json:",omitempty"`
-		MultiShardAutocommit bool                 `json:",omitempty"`
-		QueryTimeout         int                  `json:",omitempty"`
-	}{
-		Opcode:               del.RouteType(),
-		Keyspace:             del.Keyspace,
-		Query:                del.Query,
-		Vindex:               vindexName,
-		Values:               del.Values,
-		Table:                tname,
-		OwnedVindexQuery:     del.OwnedVindexQuery,
-		KsidVindex:           ksidVindexName,
-		MultiShardAutocommit: del.MultiShardAutocommit,
-		QueryTimeout:         del.QueryTimeout,
-	}
-	return jsonutil.MarshalNoEscape(marshalDelete)
-}
-
 var delName = map[DMLOpcode]string{
 	Unsharded:     "DeleteUnsharded",
 	Equal:         "DeleteEqual",
@@ -234,3 +196,35 @@ func (del *Delete) execDeleteByDestination(vcursor VCursor, bindVars map[string]
 	res, errs := vcursor.ExecuteMultiShard(rss, queries, true /* rollbackOnError */, autocommit)
 	return res, vterrors.Aggregate(errs)
 }
+
+func (del *Delete) description() PrimitiveDescription {
+	other := map[string]interface{}{
+		"Query":                del.Query,
+		"Table":                del.GetTableName(),
+		"OwnedVindexQuery":     del.OwnedVindexQuery,
+		"MultiShardAutocommit": del.MultiShardAutocommit,
+		"QueryTimeout":         del.QueryTimeout,
+	}
+
+	addFieldsIfNotEmpty(del.DML, other)
+
+	return PrimitiveDescription{
+		OperatorType:     "Delete",
+		Keyspace:         del.Keyspace,
+		Variant:          del.Opcode.String(),
+		TargetTabletType: topodatapb.TabletType_MASTER,
+		Other:            other,
+	}
+}
+
+func addFieldsIfNotEmpty(dml DML, other map[string]interface{}) {
+	if dml.Vindex != nil {
+		other["Vindex"] = dml.Vindex.String()
+	}
+	if dml.KsidVindex != nil {
+		other["KsidVindex"] = dml.KsidVindex.String()
+	}
+	if len(dml.Values) > 0 {
+		other["Values"] = dml.Values
+	}
+}

go/vt/vtgate/engine/dml.go

@@ -82,3 +82,14 @@ const (
 	// e.g: UPDATE `keyspace[-]`.x1 SET foo=1
 	ByDestination
 )
+
+var opcodeName = map[DMLOpcode]string{
+	Unsharded:     "Unsharded",
+	Equal:         "Equal",
+	Scatter:       "Scatter",
+	ByDestination: "ByDestination",
+}
+
+func (op DMLOpcode) String() string {
+	return opcodeName[op]
+}

go/vt/vtgate/engine/fake_primitive_test.go

@@ -137,3 +137,7 @@ func wrapStreamExecute(prim Primitive, vcursor VCursor, bindVars map[string]*que
 	}
 	return result, err
 }
+
+func (f *fakePrimitive) description() PrimitiveDescription {
+	return PrimitiveDescription{OperatorType: "fake"}
+}

go/vt/vtgate/engine/insert.go

@@ -23,7 +23,8 @@ import (
 	"strings"
 	"time"
 
-	"vitess.io/vitess/go/jsonutil"
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/key"
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -115,41 +116,6 @@ func NewInsert(opcode InsertOpcode, keyspace *vindexes.Keyspace, vindexValues []
 	}
 }
 
-// MarshalJSON serializes the Insert into a JSON representation.
-// It's used for testing and diagnostics.
-func (ins *Insert) MarshalJSON() ([]byte, error) {
-	var tname string
-	if ins.Table != nil {
-		tname = ins.Table.Name.String()
-	}
-	marshalInsert := struct {
-		Opcode               InsertOpcode
-		Keyspace             *vindexes.Keyspace   `json:",omitempty"`
-		Query                string               `json:",omitempty"`
-		Values               []sqltypes.PlanValue `json:",omitempty"`
-		Table                string               `json:",omitempty"`
-		Generate             *Generate            `json:",omitempty"`
-		Prefix               string               `json:",omitempty"`
-		Mid                  []string             `json:",omitempty"`
-		Suffix               string               `json:",omitempty"`
-		MultiShardAutocommit bool                 `json:",omitempty"`
-		QueryTimeout         int                  `json:",omitempty"`
-	}{
-		Opcode:               ins.Opcode,
-		Keyspace:             ins.Keyspace,
-		Query:                ins.Query,
-		Values:               ins.VindexValues,
-		Table:                tname,
-		Generate:             ins.Generate,
-		Prefix:               ins.Prefix,
-		Mid:                  ins.Mid,
-		Suffix:               ins.Suffix,
-		MultiShardAutocommit: ins.MultiShardAutocommit,
-		QueryTimeout:         ins.QueryTimeout,
-	}
-	return jsonutil.MarshalNoEscape(marshalInsert)
-}
-
 // Generate represents the instruction to generate
 // a value from a sequence.
 type Generate struct {
@@ -186,6 +152,11 @@ var insName = map[InsertOpcode]string{
 	InsertShardedIgnore: "InsertShardedIgnore",
 }
 
+// String returns the opcode
+func (code InsertOpcode) String() string {
+	return strings.ReplaceAll(insName[code], "Insert", "")
+}
+
 // MarshalJSON serializes the InsertOpcode as a JSON string.
 // It's used for testing and diagnostics.
 func (code InsertOpcode) MarshalJSON() ([]byte, error) {
@@ -607,3 +578,19 @@ func (ins *Insert) processUnowned(vcursor VCursor, vindexColumnsKeys [][]sqltype
 func insertVarName(col sqlparser.ColIdent, rowNum int) string {
 	return "_" + col.CompliantName() + strconv.Itoa(rowNum)
 }
+
+func (ins *Insert) description() PrimitiveDescription {
+	other := map[string]interface{}{
+		"Query":                ins.Query,
+		"TableName":            ins.GetTableName(),
+		"MultiShardAutocommit": ins.MultiShardAutocommit,
+		"QueryTimeout":         ins.QueryTimeout,
+	}
+	return PrimitiveDescription{
+		OperatorType:     "Insert",
+		Keyspace:         ins.Keyspace,
+		Variant:          ins.Opcode.String(),
+		TargetTabletType: topodatapb.TabletType_MASTER,
+		Other:            other,
+	}
+}

go/vt/vtgate/engine/join.go

@@ -18,9 +18,9 @@ package engine
 
 import (
 	"fmt"
+	"strings"
 
 	"vitess.io/vitess/go/sqltypes"
-
 	querypb "vitess.io/vitess/go/vt/proto/query"
 )
 
@@ -254,3 +254,15 @@ func combineVars(bv1, bv2 map[string]*querypb.BindVariable) map[string]*querypb.
 	}
 	return out
 }
+
+func (jn *Join) description() PrimitiveDescription {
+	other := map[string]interface{}{
+		"TableName":         jn.GetTableName(),
+		"JoinColumnIndexes": strings.Trim(strings.Join(strings.Fields(fmt.Sprint(jn.Cols)), ","), "[]"),
+	}
+	return PrimitiveDescription{
+		OperatorType: "Join",
+		Variant:      jn.Opcode.String(),
+		Other:        other,
+	}
+}

go/vt/vtgate/engine/limit.go

@@ -17,7 +17,6 @@ limitations under the License.
 package engine
 
 import (
-	"encoding/json"
 	"fmt"
 	"io"
 
@@ -35,23 +34,6 @@ type Limit struct {
 	Input  Primitive
 }
 
-// MarshalJSON serializes the Limit into a JSON representation.
-// It's used for testing and diagnostics.
-func (l *Limit) MarshalJSON() ([]byte, error) {
-	marshalLimit := struct {
-		Opcode string
-		Count  sqltypes.PlanValue
-		Offset sqltypes.PlanValue
-		Input  Primitive
-	}{
-		Opcode: "Limit",
-		Count:  l.Count,
-		Offset: l.Offset,
-		Input:  l.Input,
-	}
-	return json.Marshal(marshalLimit)
-}
-
 // RouteType returns a description of the query routing type used by the primitive
 func (l *Limit) RouteType() string {
 	return l.Input.RouteType()
@@ -200,3 +182,19 @@ func (l *Limit) fetchOffset(bindVars map[string]*querypb.BindVariable) (int, err
 	}
 	return offset, nil
 }
+
+func (l *Limit) description() PrimitiveDescription {
+	other := map[string]interface{}{}
+
+	if !l.Count.IsNull() {
+		other["Count"] = l.Count.Value
+	}
+	if !l.Offset.IsNull() {
+		other["Offset"] = l.Offset.Value
+	}
+
+	return PrimitiveDescription{
+		OperatorType: "Limit",
+		Other:        other,
+	}
+}

go/vt/vtgate/engine/memory_sort.go

@@ -18,10 +18,11 @@ package engine
 
 import (
 	"container/heap"
-	"encoding/json"
 	"fmt"
 	"math"
+	"reflect"
 	"sort"
+	"strings"
 
 	"vitess.io/vitess/go/sqltypes"
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -41,23 +42,6 @@ type MemorySort struct {
 	TruncateColumnCount int `json:",omitempty"`
 }
 
-// MarshalJSON serializes the MemorySort into a JSON representation.
-// It's used for testing and diagnostics.
-func (ms *MemorySort) MarshalJSON() ([]byte, error) {
-	marshalMemorySort := struct {
-		Opcode  string
-		MaxRows sqltypes.PlanValue
-		OrderBy []OrderbyParams
-		Input   Primitive
-	}{
-		Opcode:  "MemorySort",
-		MaxRows: ms.UpperLimit,
-		OrderBy: ms.OrderBy,
-		Input:   ms.Input,
-	}
-	return json.Marshal(marshalMemorySort)
-}
-
 // RouteType returns a description of the query routing type used by the primitive.
 func (ms *MemorySort) RouteType() string {
 	return ms.Input.RouteType()
@@ -184,6 +168,44 @@ func (ms *MemorySort) fetchCount(bindVars map[string]*querypb.BindVariable) (int
 	return count, nil
 }
 
+func (ms *MemorySort) description() PrimitiveDescription {
+	orderByIndexes := GenericJoin(ms.OrderBy, orderByParamsToString)
+	value := ms.UpperLimit.Value
+	other := map[string]interface{}{"OrderBy": orderByIndexes}
+	if !value.IsNull() {
+		other["UpperLimit"] = value.String()
+	}
+	return PrimitiveDescription{
+		OperatorType: "Sort",
+		Variant:      "Memory",
+		Other:        other,
+	}
+}
+
+func orderByParamsToString(i interface{}) string {
+	return i.(OrderbyParams).String()
+}
+
+//GenericJoin will iterate over arrays, slices or maps, and executes the f function to get a
+//string representation of each element, and then uses strings.Join() join all the strings into a single one
+func GenericJoin(input interface{}, f func(interface{}) string) string {
+	sl := reflect.ValueOf(input)
+	var keys []string
+	switch sl.Kind() {
+	case reflect.Slice:
+		for i := 0; i < sl.Len(); i++ {
+			keys = append(keys, f(sl.Index(i).Interface()))
+		}
+	case reflect.Map:
+		for _, k := range sl.MapKeys() {
+			keys = append(keys, f(k.Interface()))
+		}
+	default:
+		panic("GenericJoin doesn't know how to deal with " + sl.Kind().String())
+	}
+	return strings.Join(keys, ", ")
+}
+
 // sortHeap is sorted based on the orderBy params.
 // Implementation is similar to scatterHeap
 type sortHeap struct {

go/vt/vtgate/engine/merge_sort.go

@@ -155,6 +155,17 @@ func (ms *MergeSort) StreamExecute(vcursor VCursor, bindVars map[string]*querypb
 	return nil
 }
 
+func (ms *MergeSort) description() PrimitiveDescription {
+	other := map[string]interface{}{
+		"OrderBy": ms.OrderBy,
+	}
+	return PrimitiveDescription{
+		OperatorType: "Sort",
+		Variant:      "Merge",
+		Other:        other,
+	}
+}
+
 // streamHandle is the rendez-vous point between each stream and the merge-sorter.
 // The fields channel is used by the stream to transmit the field info, which
 // is the first packet. Following this, the stream sends each row to the row