From c83cb44875595559c38a8c0b8306692b2ec1afec Mon Sep 17 00:00:00 2001 From: Andres Taylor Date: Thu, 12 Sep 2024 10:53:24 +0200 Subject: [PATCH] add trace primitive Signed-off-by: Andres Taylor --- go/vt/vtgate/engine/fake_vcursor_test.go | 4 ++ go/vt/vtgate/engine/primitive.go | 2 + go/vt/vtgate/engine/trace.go | 72 +++++++++++++++++++ .../planbuilder/operator_transformers.go | 7 -- .../planbuilder/testdata/vexplain_cases.json | 28 ++++++++ go/vt/vtgate/planbuilder/vexplain.go | 19 +++++ go/vt/vtgate/vcursor_impl.go | 4 ++ 7 files changed, 129 insertions(+), 7 deletions(-) create mode 100644 go/vt/vtgate/engine/trace.go diff --git a/go/vt/vtgate/engine/fake_vcursor_test.go b/go/vt/vtgate/engine/fake_vcursor_test.go index e7b38658d68..c9682bd4895 100644 --- a/go/vt/vtgate/engine/fake_vcursor_test.go +++ b/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 { diff --git a/go/vt/vtgate/engine/primitive.go b/go/vt/vtgate/engine/primitive.go index 2f6f1c46db2..8d20e776208 100644 --- a/go/vt/vtgate/engine/primitive.go +++ b/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 diff --git a/go/vt/vtgate/engine/trace.go b/go/vt/vtgate/engine/trace.go new file mode 100644 index 00000000000..e7264068342 --- /dev/null +++ b/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", + } +} diff --git a/go/vt/vtgate/planbuilder/operator_transformers.go b/go/vt/vtgate/planbuilder/operator_transformers.go index 6c04e1fde93..ca6266b6a5f 100644 --- a/go/vt/vtgate/planbuilder/operator_transformers.go +++ b/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 } diff --git a/go/vt/vtgate/planbuilder/testdata/vexplain_cases.json b/go/vt/vtgate/planbuilder/testdata/vexplain_cases.json index 630e59f3526..6be05dff078 100644 --- a/go/vt/vtgate/planbuilder/testdata/vexplain_cases.json +++ b/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", diff --git a/go/vt/vtgate/planbuilder/vexplain.go b/go/vt/vtgate/planbuilder/vexplain.go index 21a35f02967..934a2424706 100644 --- a/go/vt/vtgate/planbuilder/vexplain.go +++ b/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 { diff --git a/go/vt/vtgate/vcursor_impl.go b/go/vt/vtgate/vcursor_impl.go index c3c015e5906..75e87d1e646 100644 --- a/go/vt/vtgate/vcursor_impl.go +++ b/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() +}