[release-17.0] VReplication: send unique key name to rowstreamer, which can then use with FORCE INDEX (#14916)

go/vt/vttablet/onlineddl/vrepl.go

@@ -28,6 +28,7 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"net/url"
 	"strconv"
 	"strings"
 
@@ -535,6 +536,7 @@ func (v *VRepl) analyzeBinlogSource(ctx context.Context) {
 		SourceUniqueKeyColumns:       encodeColumns(&v.chosenSourceUniqueKey.Columns),
 		TargetUniqueKeyColumns:       encodeColumns(&v.chosenTargetUniqueKey.Columns),
 		SourceUniqueKeyTargetColumns: encodeColumns(v.chosenSourceUniqueKey.Columns.MappedNamesColumnList(v.sharedColumnsMap)),
+		ForceUniqueKey:               url.QueryEscape(v.chosenSourceUniqueKey.Name),
 	}
 	if len(v.convertCharset) > 0 {
 		rule.ConvertCharset = v.convertCharset

go/vt/vttablet/tabletmanager/vreplication/table_plan_builder.go

@@ -306,6 +306,9 @@ func buildTablePlan(tableName string, rule *binlogdatapb.Rule, colInfos []*Colum
 	if rule.SourceUniqueKeyColumns != "" {
 		commentsList = append(commentsList, fmt.Sprintf(`ukColumns="%s"`, rule.SourceUniqueKeyColumns))
 	}
+	if rule.ForceUniqueKey != "" {
+		commentsList = append(commentsList, fmt.Sprintf(`ukForce="%s"`, rule.ForceUniqueKey))
+	}
 	if len(commentsList) > 0 {
 		comments := sqlparser.Comments{
 			fmt.Sprintf(`/*vt+ %s */`, strings.Join(commentsList, " ")),

go/vt/vttablet/tabletserver/vstreamer/rowstreamer.go

@@ -19,11 +19,13 @@ package vstreamer
 import (
 	"context"
 	"fmt"
+	"net/url"
 	"sync"
 	"time"
 
 	"vitess.io/vitess/go/mysql"
 	"vitess.io/vitess/go/mysql/collations"
+	"vitess.io/vitess/go/sqlescape"
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/textutil"
 	"vitess.io/vitess/go/timer"
@@ -175,12 +177,17 @@ func (rs *rowStreamer) buildPlan() error {
 			return err
 		}
 	}
-
+	if s, found := directives.GetString("ukForce", ""); found {
+		st.PKIndexName, err = url.QueryUnescape(s)
+		if err != nil {
+			return err
+		}
+	}
 	rs.pkColumns, err = rs.buildPKColumns(st)
 	if err != nil {
 		return err
 	}
-	rs.sendQuery, err = rs.buildSelect()
+	rs.sendQuery, err = rs.buildSelect(st)
 	if err != nil {
 		return err
 	}
@@ -230,7 +237,7 @@ func (rs *rowStreamer) buildPKColumns(st *binlogdatapb.MinimalTable) ([]int, err
 	return pkColumns, nil
 }
 
-func (rs *rowStreamer) buildSelect() (string, error) {
+func (rs *rowStreamer) buildSelect(st *binlogdatapb.MinimalTable) (string, error) {
 	buf := sqlparser.NewTrackedBuffer(nil)
 	// We could have used select *, but being explicit is more predictable.
 	buf.Myprintf("select ")
@@ -245,7 +252,17 @@ func (rs *rowStreamer) buildSelect() (string, error) {
 		}
 		prefix = ", "
 	}
-	buf.Myprintf(" from %v", sqlparser.NewIdentifierCS(rs.plan.Table.Name))
+	// If we know the index name that we should be using then tell MySQL
+	// to use it if possible. This helps to ensure that we are able to
+	// leverage the ordering from the index itself and avoid having to
+	// do a FILESORT of all the results. This index should contain all
+	// of the PK columns which are used in the ORDER BY clause below.
+	var indexHint string
+	if st.PKIndexName != "" {
+		indexHint = fmt.Sprintf(" force index (%s)",
+			sqlescape.EscapeID(sqlescape.UnescapeID(st.PKIndexName)))
+	}
+	buf.Myprintf(" from %v%s", sqlparser.NewIdentifierCS(rs.plan.Table.Name), indexHint)
 	if len(rs.lastpk) != 0 {
 		if len(rs.lastpk) != len(rs.pkColumns) {
 			return "", fmt.Errorf("primary key values don't match length: %v vs %v", rs.lastpk, rs.pkColumns)

go/vt/vttablet/tabletserver/vstreamer/rowstreamer_test.go

@@ -33,6 +33,57 @@ import (
 	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
 )
 
+// TestRowStreamerQuery validates that the correct force index hint and order by is added to the rowstreamer query.
+func TestRowStreamerQuery(t *testing.T) {
+	execStatements(t, []string{
+		"create table t1(id int, uk1 int, val varbinary(128), primary key(id), unique key uk2 (uk1))",
+	})
+	defer execStatements(t, []string{
+		"drop table t1",
+	})
+	engine.se.Reload(context.Background())
+	// We need to StreamRows, to get an initialized RowStreamer.
+	// Note that the query passed into StreamRows is overwritten while running the test.
+	err := engine.StreamRows(context.Background(), "select * from t1", nil, func(rows *binlogdatapb.VStreamRowsResponse) error {
+		type testCase struct {
+			directives      string
+			sendQuerySuffix string
+		}
+		queryTemplate := "select %s id, uk1, val from t1"
+		getQuery := func(directives string) string {
+			return fmt.Sprintf(queryTemplate, directives)
+		}
+		sendQueryPrefix := "select id, uk1, val from t1"
+		testCases := []testCase{
+			{"/*vt+ ukColumns=\"uk1\" ukForce=\"uk2\" */", "force index (`uk2`) order by uk1"},
+			{"/*vt+ ukForce=\"uk2\" */", "force index (`uk2`) order by uk1"},
+			{"/*vt+ ukColumns=\"uk1\" */", "order by uk1"},
+		}
+
+		for _, tc := range testCases {
+			t.Run(tc.directives, func(t *testing.T) {
+				var err error
+				var rs *rowStreamer
+				// Depending on the order of the test cases, the index of the engine.rowStreamers slice may change.
+				for _, rs2 := range engine.rowStreamers {
+					if rs2 != nil {
+						rs = rs2
+						break
+					}
+				}
+				require.NotNil(t, rs)
+				rs.query = getQuery(tc.directives)
+				err = rs.buildPlan()
+				require.NoError(t, err)
+				want := fmt.Sprintf("%s %s", sendQueryPrefix, tc.sendQuerySuffix)
+				require.Equal(t, want, rs.sendQuery)
+			})
+		}
+		return nil
+	})
+	require.NoError(t, err)
+}
+
 func TestStreamRowsScan(t *testing.T) {
 	if testing.Short() {
 		t.Skip()

proto/binlogdata.proto

@@ -176,6 +176,9 @@ message Rule {
   // such columns need to have special transofrmation of the data, from an integral format into a
   // string format. e.g. the value 0 needs to be converted to '0'.
   map<string, bool> convert_int_to_enum = 8;
+
+  // ForceUniqueKey gives vtreamer a hint for `FORCE INDEX (...)` usage.
+  string force_unique_key = 9;
 }
 
 // Filter represents a list of ordered rules. The first
@@ -440,6 +443,10 @@ message MinimalTable {
   string name = 1;
   repeated query.Field fields = 2;
   repeated int64 p_k_columns = 3;
+  // This will be PRIMARY when the actual primary key is used and it
+  // will be the name of the Primary Key equivalent if one is used
+  // instead. Otherwise it will be empty.
+  string p_k_index_name = 4;
 }
 
 message MinimalSchema {