distsql: fix wrong schema version when snapshot has been set (#15258)

distsql/request_builder.go

@@ -17,6 +17,7 @@ import (
 	"math"
 
 	"github.com/pingcap/parser/mysql"
+	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
@@ -166,7 +167,11 @@ func (builder *RequestBuilder) SetFromSessionVars(sv *variable.SessionVars) *Req
 	builder.Request.NotFillCache = sv.StmtCtx.NotFillCache
 	builder.Request.Priority = builder.getKVPriority(sv)
 	builder.Request.ReplicaRead = sv.ReplicaRead
-	builder.Request.SchemaVar = sv.TxnCtx.SchemaVersion
+	if sv.SnapshotInfoschema != nil {
+		builder.Request.SchemaVar = infoschema.GetInfoSchemaBySessionVars(sv).SchemaMetaVersion()
+	} else {
+		builder.Request.SchemaVar = sv.TxnCtx.SchemaVersion
+	}
 	return builder
 }
 

distsql/request_builder_test.go

@@ -18,6 +18,7 @@ import (
 	"testing"
 
 	. "github.com/pingcap/check"
+	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
@@ -587,3 +588,24 @@ func (s *testSuite) TestRequestBuilder7(c *C) {
 
 	c.Assert(actual, DeepEquals, expect)
 }
+
+func (s *testSuite) TestRequestBuilder8(c *C) {
+	sv := variable.NewSessionVars()
+	sv.SnapshotInfoschema = infoschema.MockInfoSchemaWithSchemaVer(nil, 10000)
+	actual, err := (&RequestBuilder{}).
+		SetFromSessionVars(sv).
+		Build()
+	c.Assert(err, IsNil)
+	expect := &kv.Request{
+		Tp:             0,
+		StartTs:        0x0,
+		Data:           []uint8(nil),
+		Concurrency:    15,
+		IsolationLevel: 0,
+		Priority:       0,
+		MemTracker:     (*memory.Tracker)(nil),
+		ReplicaRead:    0x1,
+		SchemaVar:      10000,
+	}
+	c.Assert(actual, DeepEquals, expect)
+}

infoschema/infoschema.go

@@ -24,6 +24,7 @@ import (
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/meta/autoid"
 	"github.com/pingcap/tidb/sessionctx"
+	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/table"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/logutil"
@@ -164,6 +165,30 @@ func MockInfoSchema(tbList []*model.TableInfo) InfoSchema {
 	return result
 }
 
+// MockInfoSchemaWithSchemaVer only serves for test.
+func MockInfoSchemaWithSchemaVer(tbList []*model.TableInfo, schemaVer int64) InfoSchema {
+	result := &infoSchema{}
+	result.schemaMap = make(map[string]*schemaTables)
+	result.sortedTablesBuckets = make([]sortedTables, bucketCount)
+	dbInfo := &model.DBInfo{ID: 0, Name: model.NewCIStr("test"), Tables: tbList}
+	tableNames := &schemaTables{
+		dbInfo: dbInfo,
+		tables: make(map[string]table.Table),
+	}
+	result.schemaMap["test"] = tableNames
+	for _, tb := range tbList {
+		tbl := table.MockTableFromMeta(tb)
+		tableNames.tables[tb.Name.L] = tbl
+		bucketIdx := tableBucketIdx(tb.ID)
+		result.sortedTablesBuckets[bucketIdx] = append(result.sortedTablesBuckets[bucketIdx], tbl)
+	}
+	for i := range result.sortedTablesBuckets {
+		sort.Sort(result.sortedTablesBuckets[i])
+	}
+	result.schemaMetaVersion = schemaVer
+	return result
+}
+
 var _ InfoSchema = (*infoSchema)(nil)
 
 func (is *infoSchema) SchemaByName(schema model.CIStr) (val *model.DBInfo, ok bool) {
@@ -443,7 +468,12 @@ func HasAutoIncrementColumn(tbInfo *model.TableInfo) (bool, string) {
 // GetInfoSchema gets TxnCtx InfoSchema if snapshot schema is not set,
 // Otherwise, snapshot schema is returned.
 func GetInfoSchema(ctx sessionctx.Context) InfoSchema {
-	sessVar := ctx.GetSessionVars()
+	return GetInfoSchemaBySessionVars(ctx.GetSessionVars())
+}
+
+// GetInfoSchemaBySessionVars gets TxnCtx InfoSchema if snapshot schema is not set,
+// Otherwise, snapshot schema is returned.
+func GetInfoSchemaBySessionVars(sessVar *variable.SessionVars) InfoSchema {
 	var is InfoSchema
 	if snap := sessVar.SnapshotInfoschema; snap != nil {
 		is = snap.(InfoSchema)