add the trace engine primitive

Signed-off-by: Andres Taylor <andres@planetscale.com>

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
+}