From bf3eeed34d344cd0a65a178c5fc748621cfe223d Mon Sep 17 00:00:00 2001
From: Harshit Gangal <harshit@planetscale.com>
Date: Fri, 13 Oct 2023 20:30:02 +0530
Subject: [PATCH] analyze plan to account for routing rules

Signed-off-by: Harshit Gangal <harshit@planetscale.com>
---
 go/vt/vtgate/executor_test.go       | 51 +++++++++++++++++++++++++++--
 go/vt/vtgate/planbuilder/builder.go | 21 +++++++-----
 2 files changed, 61 insertions(+), 11 deletions(-)

diff --git a/go/vt/vtgate/executor_test.go b/go/vt/vtgate/executor_test.go
index 2ab45f1ef42..c0c52fa7377 100644
--- a/go/vt/vtgate/executor_test.go
+++ b/go/vt/vtgate/executor_test.go
@@ -1214,7 +1214,6 @@ func TestExecutorOther(t *testing.T) {
 	}
 
 	stmts := []string{
-		"analyze table t1",
 		"describe select * from t1",
 		"explain select * from t1",
 		"repair table t1",
@@ -2190,7 +2189,6 @@ func TestExecutorOtherRead(t *testing.T) {
 	}
 
 	stmts := []string{
-		"analyze table t1",
 		"describe select * from t1",
 		"explain select * from t1",
 		"do 1",
@@ -2222,6 +2220,55 @@ func TestExecutorOtherRead(t *testing.T) {
 	}
 }
 
+func TestExecutorAnalyze(t *testing.T) {
+	executor, sbc1, sbc2, sbclookup, _ := createExecutorEnv(t)
+
+	type cnts struct {
+		Sbc1Cnt      int64
+		Sbc2Cnt      int64
+		SbcLookupCnt int64
+	}
+
+	tcs := []struct {
+		targetStr string
+
+		wantCnts cnts
+	}{{
+		targetStr: "TestExecutor[-]",
+		wantCnts:  cnts{Sbc1Cnt: 1, Sbc2Cnt: 1},
+	}, {
+		targetStr: KsTestUnsharded,
+		wantCnts:  cnts{SbcLookupCnt: 1},
+	}, {
+		targetStr: "TestExecutor",
+		wantCnts:  cnts{Sbc1Cnt: 1, Sbc2Cnt: 1},
+	}, {
+		targetStr: "TestExecutor/-20",
+		wantCnts:  cnts{Sbc1Cnt: 1},
+	}, {
+		targetStr: "TestExecutor[00]",
+		wantCnts:  cnts{Sbc1Cnt: 1},
+	}}
+
+	stmt := "analyze table t1"
+	for _, tc := range tcs {
+		t.Run(tc.targetStr, func(t *testing.T) {
+			sbc1.ExecCount.Store(0)
+			sbc2.ExecCount.Store(0)
+			sbclookup.ExecCount.Store(0)
+
+			_, err := executor.Execute(context.Background(), nil, "TestExecute", NewSafeSession(&vtgatepb.Session{TargetString: tc.targetStr}), stmt, nil)
+			require.NoError(t, err)
+
+			utils.MustMatch(t, tc.wantCnts, cnts{
+				Sbc1Cnt:      sbc1.ExecCount.Load(),
+				Sbc2Cnt:      sbc2.ExecCount.Load(),
+				SbcLookupCnt: sbclookup.ExecCount.Load(),
+			}, "count did not match")
+		})
+	}
+}
+
 func TestExecutorVExplain(t *testing.T) {
 	executor, _, _, _, ctx := createExecutorEnv(t)
 
diff --git a/go/vt/vtgate/planbuilder/builder.go b/go/vt/vtgate/planbuilder/builder.go
index 52b3da90441..a67878d7119 100644
--- a/go/vt/vtgate/planbuilder/builder.go
+++ b/go/vt/vtgate/planbuilder/builder.go
@@ -198,7 +198,7 @@ func createInstructionFor(ctx context.Context, query string, stmt sqlparser.Stat
 	case *sqlparser.OtherAdmin:
 		return buildOtherReadAndAdmin(query, vschema)
 	case *sqlparser.Analyze:
-		return buildAnalyzePlan(stmt, vschema)
+		return buildRoutePlan(stmt, reservedVars, vschema, buildAnalyzePlan)
 	case *sqlparser.Set:
 		return buildSetPlan(stmt, vschema)
 	case *sqlparser.Load:
@@ -239,38 +239,41 @@ func createInstructionFor(ctx context.Context, query string, stmt sqlparser.Stat
 	return nil, vterrors.VT13001(fmt.Sprintf("unexpected statement type: %T", stmt))
 }
 
-func buildAnalyzePlan(stmt *sqlparser.Analyze, vschema plancontext.VSchema) (*planResult, error) {
+func buildAnalyzePlan(stmt sqlparser.Statement, _ *sqlparser.ReservedVars, vschema plancontext.VSchema) (*planResult, error) {
+	analyzeStmt := stmt.(*sqlparser.Analyze)
+
 	var ks *vindexes.Keyspace
 	var err error
 	dest := key.Destination(key.DestinationAllShards{})
-	if !stmt.Table.Qualifier.IsEmpty() && sqlparser.SystemSchema(stmt.Table.Qualifier.String()) {
+
+	if !analyzeStmt.Table.Qualifier.IsEmpty() && sqlparser.SystemSchema(analyzeStmt.Table.Qualifier.String()) {
 		ks, err = vschema.AnyKeyspace()
 		if err != nil {
 			return nil, err
 		}
 	} else {
-		tbl, _, _, _, destKs, err := vschema.FindTableOrVindex(stmt.Table)
+		tbl, _, _, _, destKs, err := vschema.FindTableOrVindex(analyzeStmt.Table)
 		if err != nil {
 			return nil, err
 		}
 		if tbl == nil {
-			return nil, vterrors.VT05004(sqlparser.String(stmt.Table))
+			return nil, vterrors.VT05004(sqlparser.String(analyzeStmt.Table))
 		}
 
 		ks = tbl.Keyspace
 		if destKs != nil {
 			dest = destKs
 		}
-		stmt.Table.Name = tbl.Name
+		analyzeStmt.Table.Name = tbl.Name
 	}
-	stmt.Table.Qualifier = sqlparser.NewIdentifierCS("")
+	analyzeStmt.Table.Qualifier = sqlparser.NewIdentifierCS("")
 
 	prim := &engine.Send{
 		Keyspace:          ks,
 		TargetDestination: dest,
-		Query:             sqlparser.String(stmt),
+		Query:             sqlparser.String(analyzeStmt),
 	}
-	return newPlanResult(prim, sqlparser.String(stmt.Table)), nil
+	return newPlanResult(prim, sqlparser.String(analyzeStmt.Table)), nil
 }
 
 func buildDBDDLPlan(stmt sqlparser.Statement, _ *sqlparser.ReservedVars, vschema plancontext.VSchema) (*planResult, error) {