fix(vdiff): preserve PK definition order in getSourcePKCols

go/vt/vttablet/tabletmanager/vdiff/framework_test.go

@@ -23,6 +23,7 @@ import (
 	"io"
 	"os"
 	"regexp"
+	"slices"
 	"strings"
 	"sync"
 	"testing"
@@ -112,6 +113,16 @@ var (
 				Name:    "nopkwithpke",
 				Columns: []string{"c1", "c2", "c3"},
 				Fields:  sqltypes.MakeTestFields("c1|c2|c3", "int64|int64|int64"),
+			}, {
+				// t2 is a cross-table MoveTables source: its physical PK is c0,
+				// which is renamed to the target's c1 by the filter
+				// "select c0 as c1, c2 from t2". getSourcePKCols resolves the
+				// source schema from the FROM table (t2) and matches the
+				// physical PK c0 via its underlying ColName.
+				Name:              "t2",
+				Columns:           []string{"c0", "c2"},
+				PrimaryKeyColumns: []string{"c0"},
+				Fields:            sqltypes.MakeTestFields("c0|c2", "int64|int64"),
 			},
 		},
 	}
@@ -124,6 +135,7 @@ var (
 		"datze":       5,
 		"nopk":        6,
 		"nopkwithpke": 7,
+		"t2":          8,
 	}
 )
 
@@ -435,27 +447,50 @@ func (dbc *realDBClient) SupportsCapability(capability capabilities.FlavorCapabi
 
 type fakeTMClient struct {
 	tmclient.TabletManagerClient
-	schema    *tabletmanagerdatapb.SchemaDefinition
-	vrQueries map[int]map[string]*querypb.QueryResult
-	waitpos   map[int]string
-	vrpos     map[int]string
-	pos       map[int]string
+	schema     *tabletmanagerdatapb.SchemaDefinition
+	vrQueries  map[int]map[string]*querypb.QueryResult
+	waitpos    map[int]string
+	vrpos      map[int]string
+	pos        map[int]string
+	pkeResults map[string]*sqltypes.Result
 }
 
 func newFakeTMClient() *fakeTMClient {
 	return &fakeTMClient{
-		vrQueries: make(map[int]map[string]*querypb.QueryResult),
-		waitpos:   make(map[int]string),
-		vrpos:     make(map[int]string),
-		pos:       make(map[int]string),
+		vrQueries:  make(map[int]map[string]*querypb.QueryResult),
+		waitpos:    make(map[int]string),
+		vrpos:      make(map[int]string),
+		pos:        make(map[int]string),
+		pkeResults: make(map[string]*sqltypes.Result),
 	}
 }
 
 // Close satisfies the TabletManagerClient interface.
 func (tmc *fakeTMClient) Close() {}
 
 func (tmc *fakeTMClient) GetSchema(ctx context.Context, tablet *topodatapb.Tablet, request *tabletmanagerdatapb.GetSchemaRequest) (*tabletmanagerdatapb.SchemaDefinition, error) {
-	return tmc.schema, nil
+	if len(request.Tables) == 0 {
+		return tmc.schema, nil
+	}
+	filtered := &tabletmanagerdatapb.SchemaDefinition{
+		TableDefinitions: make([]*tabletmanagerdatapb.TableDefinition, 0),
+	}
+	for _, td := range tmc.schema.TableDefinitions {
+		if slices.Contains(request.Tables, td.Name) {
+			filtered.TableDefinitions = append(filtered.TableDefinitions, td)
+		}
+	}
+	return filtered, nil
+}
+
+func (tmc *fakeTMClient) ExecuteFetchAsApp(ctx context.Context, tablet *topodatapb.Tablet, usePool bool, req *tabletmanagerdatapb.ExecuteFetchAsAppRequest) (*querypb.QueryResult, error) {
+	query := string(req.Query)
+	for tableName, result := range tmc.pkeResults {
+		if strings.Contains(query, tableName) {
+			return sqltypes.ResultToProto3(result), nil
+		}
+	}
+	return sqltypes.ResultToProto3(noResults), nil
 }
 
 // setVRResults allows you to specify VReplicationExec queries and their results. You can specify
@@ -564,6 +599,10 @@ func newTestVDiffEnv(t *testing.T) *testVDiffEnv {
 	vdiffenv.vre.Open(t.Context())
 
 	vdiffenv.tmc.schema = testSchema
+	vdiffenv.tmc.pkeResults["nopkwithpke"] = sqltypes.MakeTestResult(
+		sqltypes.MakeTestFields("column_name|index_name", "varchar|varchar"),
+		"c3|c3",
+	)
 	// We need to add t1, which we use for a full VDiff in TestVDiff, to
 	// the schema engine with the PK val.
 	st := &schema.Table{

go/vt/vttablet/tabletmanager/vdiff/schema.go

@@ -64,6 +64,7 @@ const (
 	sqlUpdateTableState          = "update _vt.vdiff_table set state = %a where vdiff_id = %a and table_name = %a"
 	sqlUpdateTableStateAndReport = "update _vt.vdiff_table set state = %a, rows_compared = %a, report = %a where vdiff_id = %a and table_name = %a"
 	sqlUpdateTableMismatch       = "update _vt.vdiff_table set mismatch = true where vdiff_id = %a and table_name = %a"
+	sqlClearTableMismatch        = "update _vt.vdiff_table set mismatch = false where vdiff_id = %a and table_name = %a"
 
 	sqlGetIncompleteTables = "select table_name as table_name from _vt.vdiff_table where vdiff_id = %a and state != 'completed' order by table_name"
 )