add trace primitive

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

go/vt/vtgate/engine/fake_vcursor_test.go

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

go/vt/vtgate/engine/primitive.go

@@ -140,6 +140,8 @@ type (
 
 		// UnresolvedTransactions reads the state of all the unresolved atomic transactions in the given keyspace.
 		UnresolvedTransactions(ctx context.Context, keyspace string) ([]*querypb.TransactionMetadata, error)
+
+		EnableOperatorTracing()
 	}
 
 	// SessionActions gives primitives ability to interact with the session state

go/vt/vtgate/engine/trace.go

@@ -0,0 +1,72 @@
+/*
+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"
+
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+)
+
+var _ Primitive = (*Trace)(nil)
+
+type Trace struct {
+	identifiablePrimitive
+	Inner Primitive
+}
+
+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 (t *Trace) GetFields(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
+	vcursor.EnableOperatorTracing()
+	return t.Inner.GetFields(ctx, vcursor, bindVars)
+}
+
+func (t *Trace) NeedsTransaction() bool {
+	return t.Inner.NeedsTransaction()
+}
+
+func (t *Trace) TryExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+	vcursor.EnableOperatorTracing()
+	return t.Inner.TryExecute(ctx, vcursor, bindVars, wantfields)
+}
+
+func (t *Trace) TryStreamExecute(ctx context.Context, vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+	vcursor.EnableOperatorTracing()
+	return t.Inner.TryStreamExecute(ctx, vcursor, bindVars, wantfields, callback)
+}
+
+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/operator_transformers.go

@@ -44,13 +44,6 @@ func transformToPrimitive(ctx *plancontext.PlanningContext, op operators.Operato
 		return nil, err
 	}
 
-	// We'll go over the primitive tree and assign unique IDs
-	id := 1
-	engine.PreOrderTraverse(primitive, func(primitive engine.Primitive) {
-		primitive.SetID(engine.PrimitiveID(id))
-		id++
-	})
-
 	return primitive, nil
 }
 

go/vt/vtgate/planbuilder/testdata/vexplain_cases.json

@@ -42,6 +42,34 @@
       ]
     }
   },
+  {
+    "comment": "vexplain trace",
+    "query": "vexplain trace select * from user",
+    "plan": {
+      "QueryType": "EXPLAIN",
+      "Original": "vexplain trace select * from user",
+      "Instructions": {
+        "OperatorType": "Trace",
+        "Inputs": [
+          {
+            "ID": 1,
+            "OperatorType": "Route",
+            "Variant": "Scatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select * from `user` where 1 != 1",
+            "Query": "select * from `user`",
+            "Table": "`user`"
+          }
+        ]
+      },
+      "TablesUsed": [
+        "user.user"
+      ]
+    }
+  },
   {
     "comment": "vexplain table",
     "query": "vexplain ALL select * from user",

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.Statement, reservedVars, vschema, enableOnlineDDL, enableDirectDDL)
 	}
 	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] unexpected vtexplain type: %s", vexplainStmt.Type.ToString())
 }
@@ -103,6 +105,23 @@ func buildVExplainVtgatePlan(ctx context.Context, explainStatement sqlparser.Sta
 	return newPlanResult(engine.NewRowsPrimitive(rows, fields)), 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++
+	})
+
+	innerInstruction.primitive = &engine.Trace{Inner: innerInstruction.primitive}
+	return innerInstruction, nil
+}
+
 func buildVExplainLoggingPlan(ctx context.Context, explain *sqlparser.VExplainStmt, reservedVars *sqlparser.ReservedVars, vschema plancontext.VSchema, enableOnlineDDL, enableDirectDDL bool) (*planResult, error) {
 	input, err := createInstructionFor(ctx, sqlparser.String(explain.Statement), explain.Statement, reservedVars, vschema, enableOnlineDDL, enableDirectDDL)
 	if err != nil {

go/vt/vtgate/vcursor_impl.go

@@ -1510,3 +1510,7 @@ func (vc *vcursorImpl) UpdateForeignKeyChecksState(fkStateFromQuery *bool) {
 func (vc *vcursorImpl) GetForeignKeyChecksState() *bool {
 	return vc.fkChecksState
 }
+
+func (vc *vcursorImpl) EnableOperatorTracing() {
+	vc.logStats.EnableOpStats()
+}