diff --git a/go/vt/vtgate/executor_vexplain_test.go b/go/vt/vtgate/executor_vexplain_test.go index 9a0668be3cb..a89221b91c2 100644 --- a/go/vt/vtgate/executor_vexplain_test.go +++ b/go/vt/vtgate/executor_vexplain_test.go @@ -17,6 +17,7 @@ limitations under the License. package vtgate import ( + "context" "testing" "github.com/stretchr/testify/require" @@ -109,3 +110,19 @@ func TestSimpleVexplainTrace(t *testing.T) { gotRowString := gotResult.Rows[0][0].ToString() require.Equal(t, expectedRowString, gotRowString) } + +func TestVExplainKeys(t *testing.T) { + executor, _, _, _, _ := createExecutorEnv(t) + + query := "vexplain keys select count(*), col2 from music group by col2" + session := NewSafeSession(&vtgatepb.Session{TargetString: "@primary"}) + gotResult, err := executor.Execute(context.Background(), nil, "Execute", session, query, nil) + require.NoError(t, err) + + expectedRowString := `{ + "StatementType": "SELECT" +}` + + gotRowString := gotResult.Rows[0][0].ToString() + require.Equal(t, expectedRowString, gotRowString) +} diff --git a/go/vt/vtgate/planbuilder/operators/keys.go b/go/vt/vtgate/planbuilder/operators/keys.go new file mode 100644 index 00000000000..e9cba011d0e --- /dev/null +++ b/go/vt/vtgate/planbuilder/operators/keys.go @@ -0,0 +1,35 @@ +/* +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 operators + +import ( + "vitess.io/vitess/go/vt/sqlparser" + "vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext" +) + +type VExplainKeys struct { + GroupingColumns []string `json:"Grouping Columns,omitempty"` + TableName []string `json:"TableName,omitempty"` + JoinColumns []string `json:"JoinColumns,omitempty"` + FilterColumns []string `json:"FilterColumns,omitempty"` + StatementType string `json:"StatementType"` +} + +func GetVExplainKeys(ctx *plancontext.PlanningContext, stmt sqlparser.Statement) (result VExplainKeys) { + result.StatementType = sqlparser.ASTToStatementType(stmt).String() + return +} diff --git a/go/vt/vtgate/planbuilder/testdata/vexplain_cases.json b/go/vt/vtgate/planbuilder/testdata/vexplain_cases.json index 630e59f3526..cafdc2851af 100644 --- a/go/vt/vtgate/planbuilder/testdata/vexplain_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/vexplain_cases.json @@ -69,5 +69,20 @@ "user.user" ] } + }, + { + "comment": "vexplain keys", + "query": "vexplain keys select * from user", + "plan": { + "QueryType": "EXPLAIN", + "Original": "vexplain keys select * from user", + "Instructions": { + "OperatorType": "Rows", + "Fields": { + "ColumnUsage": "VARCHAR" + }, + "RowCount": 1 + } + } } ] diff --git a/go/vt/vtgate/planbuilder/vexplain.go b/go/vt/vtgate/planbuilder/vexplain.go index f31eeee1714..f66af7bfc33 100644 --- a/go/vt/vtgate/planbuilder/vexplain.go +++ b/go/vt/vtgate/planbuilder/vexplain.go @@ -32,7 +32,13 @@ import ( "vitess.io/vitess/go/vt/vtgate/vindexes" ) -func buildVExplainPlan(ctx context.Context, vexplainStmt *sqlparser.VExplainStmt, reservedVars *sqlparser.ReservedVars, vschema plancontext.VSchema, enableOnlineDDL, enableDirectDDL bool) (*planResult, error) { +func buildVExplainPlan( + ctx context.Context, + vexplainStmt *sqlparser.VExplainStmt, + reservedVars *sqlparser.ReservedVars, + vschema plancontext.VSchema, + enableOnlineDDL, enableDirectDDL bool, +) (*planResult, error) { switch vexplainStmt.Type { case sqlparser.QueriesVExplainType, sqlparser.AllVExplainType: return buildVExplainLoggingPlan(ctx, vexplainStmt, reservedVars, vschema, enableOnlineDDL, enableDirectDDL) @@ -40,6 +46,8 @@ func buildVExplainPlan(ctx context.Context, vexplainStmt *sqlparser.VExplainStmt return buildVExplainVtgatePlan(ctx, vexplainStmt.Statement, reservedVars, vschema, enableOnlineDDL, enableDirectDDL) case sqlparser.TraceVExplainType: return buildVExplainTracePlan(ctx, vexplainStmt.Statement, reservedVars, vschema, enableOnlineDDL, enableDirectDDL) + case sqlparser.KeysVExplainType: + return buildVExplainKeysPlan(vexplainStmt.Statement, vschema) } return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "[BUG] unexpected vtexplain type: %s", vexplainStmt.Type.ToString()) } @@ -89,22 +97,33 @@ func buildVExplainVtgatePlan(ctx context.Context, explainStatement sqlparser.Sta if err != nil { return nil, err } - description := engine.PrimitiveToPlanDescription(innerInstruction.primitive, nil) - output, err := json.MarshalIndent(description, "", "\t") + + return getJsonResultPlan( + engine.PrimitiveToPlanDescription(innerInstruction.primitive, nil), + "JSON", + ) +} + +// getJsonResultPlan marshals the given struct into a JSON string and returns it as a planResult. +func getJsonResultPlan(v any, colName string) (*planResult, error) { + output, err := json.MarshalIndent(v, "", "\t") if err != nil { return nil, err } - fields := []*querypb.Field{ - {Name: "JSON", Type: querypb.Type_VARCHAR}, - } - rows := []sqltypes.Row{ - { - sqltypes.NewVarChar(string(output)), - }, - } + fields := []*querypb.Field{{Name: colName, Type: querypb.Type_VARCHAR}} + rows := []sqltypes.Row{{sqltypes.NewVarChar(string(output))}} return newPlanResult(engine.NewRowsPrimitive(rows, fields)), nil } +func buildVExplainKeysPlan(statement sqlparser.Statement, vschema plancontext.VSchema) (*planResult, error) { + ctx, err := plancontext.CreatePlanningContext(statement, sqlparser.NewReservedVars("", sqlparser.BindVars{}), vschema, querypb.ExecuteOptions_Gen4) + if err != nil { + return nil, err + } + result := operators.GetVExplainKeys(ctx, statement) + return getJsonResultPlan(result, "ColumnUsage") +} + 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 {