*: make tidb_replica_read work with set_var hint (#48156)

executor/executor_test.go

@@ -6320,3 +6320,16 @@ func TestProcessInfoOfSubQuery(t *testing.T) {
 	tk2.MustQuery("select 1 from information_schema.processlist where TxnStart != '' and info like 'select%sleep% from t%'").Check(testkit.Rows("1"))
 	wg.Wait()
 }
+
+func TestSetVarHint(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustQuery("select /*+ SET_VAR(tidb_replica_read = 'follower')*/ @@tidb_replica_read;").Check(testkit.Rows("follower"))
+	tk.MustQuery("select @@tidb_replica_read;").Check(testkit.Rows("leader"))
+
+	tk.MustExec("use test;")
+	tk.MustExec("create table t (a int key, b int);")
+	tk.MustExec("insert into t values (1,1);")
+	time.Sleep(time.Second * 2)
+	tk.MustQuery("select /*+ SET_VAR(tidb_replica_read = 'closest-replicas') */  * from t as of timestamp NOW() - INTERVAL 1 SECOND where a=1;").Check(testkit.Rows("1 1"))
+}

sessionctx/variable/session.go

@@ -1837,21 +1837,29 @@ func (s *SessionVars) GetReplicaRead() kv.ReplicaReadType {
 	if s.StmtCtx.HasReplicaReadHint {
 		return kv.ReplicaReadType(s.StmtCtx.ReplicaRead)
 	}
+	replicaRead := s.getReplicaRead()
 	// if closest-adaptive is unavailable, fallback to leader read
-	if s.replicaRead == kv.ReplicaReadClosestAdaptive && !IsAdaptiveReplicaReadEnabled() {
+	if replicaRead == kv.ReplicaReadClosestAdaptive && !IsAdaptiveReplicaReadEnabled() {
 		return kv.ReplicaReadLeader
 	}
-	return s.replicaRead
+	return replicaRead
 }
 
 // SetReplicaRead set SessionVars.replicaRead.
 func (s *SessionVars) SetReplicaRead(val kv.ReplicaReadType) {
 	s.replicaRead = val
 }
 
+func (s *SessionVars) getReplicaRead() kv.ReplicaReadType {
+	if replicaRead, ok := s.getReplicaReadFromStmtVar(); ok {
+		return replicaRead
+	}
+	return s.replicaRead
+}
+
 // IsReplicaReadClosestAdaptive returns whether adaptive closest replica can be enabled.
 func (s *SessionVars) IsReplicaReadClosestAdaptive() bool {
-	return s.replicaRead == kv.ReplicaReadClosestAdaptive && IsAdaptiveReplicaReadEnabled()
+	return s.getReplicaRead() == kv.ReplicaReadClosestAdaptive && IsAdaptiveReplicaReadEnabled()
 }
 
 // GetWriteStmtBufs get pointer of SessionVars.writeStmtBufs.
@@ -2134,6 +2142,12 @@ func (s *SessionVars) setStmtVar(name string, val string) error {
 	return nil
 }
 
+// GetStmtVar gets the value of the statement's system variable.
+func (s *SessionVars) GetStmtVar(name string) (string, bool) {
+	val, ok := s.stmtVars[name]
+	return val, ok
+}
+
 // ClearStmtVars clear temporarily system variables.
 func (s *SessionVars) ClearStmtVars() {
 	s.stmtVars = make(map[string]string)
@@ -3195,3 +3209,21 @@ func (s *SessionVars) GetTiKVClientReadTimeout() uint64 {
 	}
 	return s.TiKVClientReadTimeout
 }
+
+func (s *SessionVars) getReplicaReadFromStmtVar() (kv.ReplicaReadType, bool) {
+	val, ok := s.stmtVars[TiDBReplicaRead]
+	if ok {
+		if strings.EqualFold(val, "follower") {
+			return kv.ReplicaReadFollower, true
+		} else if strings.EqualFold(val, "leader-and-follower") {
+			return kv.ReplicaReadMixed, true
+		} else if strings.EqualFold(val, "leader") || len(val) == 0 {
+			return kv.ReplicaReadLeader, true
+		} else if strings.EqualFold(val, "closest-replicas") {
+			return kv.ReplicaReadClosest, true
+		} else if strings.EqualFold(val, "closest-adaptive") {
+			return kv.ReplicaReadClosestAdaptive, true
+		}
+	}
+	return kv.ReplicaReadLeader, false
+}

sessionctx/variable/sysvar.go

@@ -1763,7 +1763,7 @@ var defaultSysVars = []*SysVar{
 		s.NoopFuncsMode = TiDBOptOnOffWarn(val)
 		return nil
 	}},
-	{Scope: ScopeGlobal | ScopeSession, Name: TiDBReplicaRead, Value: "leader", Type: TypeEnum, PossibleValues: []string{"leader", "follower", "leader-and-follower", "closest-replicas", "closest-adaptive"}, SetSession: func(s *SessionVars, val string) error {
+	{Scope: ScopeGlobal | ScopeSession, Name: TiDBReplicaRead, IsHintUpdatable: true, Value: "leader", Type: TypeEnum, PossibleValues: []string{"leader", "follower", "leader-and-follower", "closest-replicas", "closest-adaptive"}, SetSession: func(s *SessionVars, val string) error {
 		if strings.EqualFold(val, "follower") {
 			s.SetReplicaRead(kv.ReplicaReadFollower)
 		} else if strings.EqualFold(val, "leader-and-follower") {