add the trace engine primitive

Signed-off-by: Andres Taylor <[email protected]>

go/vt/vtgate/engine/fake_vcursor_test.go

@@ -132,6 +132,10 @@ func (t *noopVCursor) UnresolvedTransactions(ctx context.Context, keyspace strin
 	panic("implement me")
 }
 
+func (t *noopVCursor) StartPrimitiveTrace() func() map[int]RowsReceived {
+	panic("implement me")
+}
+
 func (t *noopVCursor) SetExec(ctx context.Context, name string, value string) error {
 	panic("implement me")
 }
@@ -874,6 +878,10 @@ func (f *loggingVCursor) UnresolvedTransactions(_ context.Context, _ string) ([]
 	return f.transactionStatusOutput, nil
 }
 
+func (f *loggingVCursor) StartPrimitiveTrace() func() map[int]RowsReceived {
+	panic("implement me")
+}
+
 // SQLParser implements VCursor
 func (t *loggingVCursor) SQLParser() *sqlparser.Parser {
 	if t.parser == nil {

go/vt/vtgate/engine/plan_description.go

@@ -48,6 +48,8 @@ type PrimitiveDescription struct {
 	ID        PrimitiveID
 	InputName string
 	Inputs    []PrimitiveDescription
+
+	Stats RowsReceived
 }
 
 // MarshalJSON serializes the PlanDescription into a JSON representation.
@@ -97,6 +99,14 @@ func (pd PrimitiveDescription) MarshalJSON() ([]byte, error) {
 			return nil, err
 		}
 	}
+	if len(pd.Stats) > 0 {
+		if err := marshalAdd(prepend, buf, "NoOfCalls", len(pd.Stats)); err != nil {
+			return nil, err
+		}
+		if err := marshalAdd(prepend, buf, "Rows", pd.Stats); err != nil {
+			return nil, err
+		}
+	}
 	err := addMap(pd.Other, buf)
 	if err != nil {
 		return nil, err

go/vt/vtgate/engine/primitive.go

@@ -140,6 +140,10 @@ type (
 
 		// UnresolvedTransactions reads the state of all the unresolved atomic transactions in the given keyspace.
 		UnresolvedTransactions(ctx context.Context, keyspace string) ([]*querypb.TransactionMetadata, error)
+
+		// StartPrimitiveTrace starts a trace for the given primitive,
+		// and returns a function to get the trace logs after the primitive execution.
+		StartPrimitiveTrace() func() map[int]RowsReceived
 	}
 
 	// SessionActions gives primitives ability to interact with the session state

go/vt/vtgate/engine/trace.go

@@ -0,0 +1,134 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package engine
+
+import (
+	"context"
+	"encoding/json"
+
+	"vitess.io/vitess/go/mysql/collations"
+
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+)
+
+var _ Primitive = (*Trace)(nil)
+
+type Trace struct {
+	identifiablePrimitive
+	Inner Primitive
+}
+
+type RowsReceived []int
+
+func (t *Trace) RouteType() string {
+	return t.Inner.RouteType()
+}
+
+func (t *Trace) GetKeyspaceName() string {
+	return t.Inner.GetKeyspaceName()
+}
+
+func (t *Trace) GetTableName() string {
+	return t.Inner.GetTableName()
+}
+
+func getFields() []*querypb.Field {
+	return []*querypb.Field{{
+		Name:    "Trace",
+		Type:    sqltypes.VarChar,
+		Charset: uint32(collations.SystemCollation.Collation),
+		Flags:   uint32(querypb.MySqlFlag_NOT_NULL_FLAG),
+	}}
+}
+
+func (t *Trace) GetFields(context.Context, VCursor, map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
+	return &sqltypes.Result{Fields: getFields()}, nil
+}
+
+func (t *Trace) NeedsTransaction() bool {
+	return t.Inner.NeedsTransaction()
+}
+
+func preWalk(desc PrimitiveDescription, f func(PrimitiveDescription)) {
+	f(desc)
+	for _, input := range desc.Inputs {
+		preWalk(input, f)
+	}
+}
+
+func (t *Trace) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+	getOpStats := vcursor.StartPrimitiveTrace()
+	_, err := t.Inner.TryExecute(ctx, vcursor, bindVars, wantfields)
+	if err != nil {
+		return nil, err
+	}
+
+	return t.getExplainTraceOutput(getOpStats)
+}
+
+func (t *Trace) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+	getOpsStats := vcursor.StartPrimitiveTrace()
+	noop := func(result *sqltypes.Result) error { return nil }
+	err := t.Inner.TryStreamExecute(ctx, vcursor, bindVars, wantfields, noop)
+	if err != nil {
+		return err
+	}
+
+	res, err := t.getExplainTraceOutput(getOpsStats)
+	if err != nil {
+		return err
+	}
+
+	return callback(res)
+}
+
+func (t *Trace) getExplainTraceOutput(getOpStats func() map[int]RowsReceived) (*sqltypes.Result, error) {
+	description := PrimitiveToPlanDescription(t.Inner)
+	statsMap := getOpStats()
+
+	// let's add the stats to the description
+	preWalk(description, func(desc PrimitiveDescription) {
+		stats, found := statsMap[int(desc.ID)]
+		if !found {
+			return
+		}
+		desc.Stats = stats
+	})
+
+	output, err := json.MarshalIndent(description, "", "\t")
+	if err != nil {
+		return nil, err
+	}
+
+	return &sqltypes.Result{
+		Fields: getFields(),
+		Rows: []sqltypes.Row{{
+			sqltypes.NewVarChar(string(output)),
+		}},
+	}, nil
+}
+
+func (t *Trace) Inputs() ([]Primitive, []map[string]any) {
+	return []Primitive{t.Inner}, nil
+}
+
+func (t *Trace) description() PrimitiveDescription {
+	return PrimitiveDescription{
+		OperatorType: "Trace",
+	}
+}

go/vt/vtgate/planbuilder/vexplain.go

@@ -38,6 +38,8 @@ func buildVExplainPlan(ctx context.Context, vexplainStmt *sqlparser.VExplainStmt
 		return buildVExplainLoggingPlan(ctx, vexplainStmt, reservedVars, vschema, enableOnlineDDL, enableDirectDDL)
 	case sqlparser.PlanVExplainType:
 		return buildVExplainVtgatePlan(ctx, vexplainStmt.Statement, reservedVars, vschema, enableOnlineDDL, enableDirectDDL)
+	case sqlparser.TraceVExplainType:
+		return buildVExplainTracePlan(ctx, vexplainStmt, reservedVars, vschema, enableOnlineDDL, enableDirectDDL)
 	}
 	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] unexpected vtexplain type: %s", vexplainStmt.Type.ToString())
 }
@@ -166,3 +168,21 @@ func explainPlan(explain *sqlparser.ExplainStmt, reservedVars *sqlparser.Reserve
 		SingleShardOnly:   true,
 	}, tables...), nil
 }
+
+func buildVExplainTracePlan(ctx context.Context, explainStatement sqlparser.Statement, reservedVars *sqlparser.ReservedVars, vschema plancontext.VSchema, enableOnlineDDL, enableDirectDDL bool) (*planResult, error) {
+	innerInstruction, err := createInstructionFor(ctx, sqlparser.String(explainStatement), explainStatement, reservedVars, vschema, enableOnlineDDL, enableDirectDDL)
+	if err != nil {
+		return nil, err
+	}
+
+	// We'll go over the primitive tree and assign unique IDs
+	id := 1
+	engine.PreOrderTraverse(innerInstruction.primitive, func(primitive engine.Primitive) {
+		primitive.SetID(engine.PrimitiveID(id))
+		id++
+	})
+
+	// We'll set the trace engine as the root primitive
+	innerInstruction.primitive = &engine.Trace{Inner: innerInstruction.primitive}
+	return innerInstruction, nil
+}