add the plan building and json handling parts of vexplain keys

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

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

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

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

go/vt/vtgate/planbuilder/vexplain.go

@@ -32,14 +32,22 @@ 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)
 	case sqlparser.PlanVExplainType:
 		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 {