From e92a92f7f4c842af350338080c52a85e30eb41f8 Mon Sep 17 00:00:00 2001
From: ekexium <eke@fastmail.com>
Date: Thu, 16 Jan 2025 20:19:10 +0800
Subject: [PATCH] *: Update client-go and verify all read ts (#58925)

ref pingcap/tidb#57786
---
 DEPS.bzl                                    | 37 +++++++++-----
 br/pkg/lightning/config/bytesize_test.go    |  4 +-
 go.mod                                      |  4 +-
 go.sum                                      |  8 +--
 pkg/executor/set.go                         |  3 +-
 pkg/executor/test/executor/executor_test.go |  2 +-
 pkg/planner/core/planbuilder.go             |  4 +-
 pkg/sessionctx/context.go                   |  9 ++--
 pkg/sessiontxn/staleread/processor.go       |  2 +-
 pkg/sessiontxn/staleread/util.go            |  2 +-
 pkg/store/copr/BUILD.bazel                  |  1 +
 pkg/store/copr/batch_coprocessor.go         |  2 +-
 pkg/store/copr/batch_request_sender.go      |  5 +-
 pkg/store/copr/mpp.go                       |  2 +-
 pkg/util/mock/BUILD.bazel                   |  1 +
 pkg/util/mock/context.go                    | 56 +++++++++++++++++----
 16 files changed, 100 insertions(+), 42 deletions(-)

diff --git a/DEPS.bzl b/DEPS.bzl
index 0ddba0a8564b3..809709b1cfcc3 100644
--- a/DEPS.bzl
+++ b/DEPS.bzl
@@ -1594,13 +1594,13 @@ def go_deps():
         name = "com_github_docker_go_units",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/docker/go-units",
-        sha256 = "0f2be7dce7b1a0ba6a4a786eb144a3398e9a61afc0eec5799a1520d9906fc58c",
-        strip_prefix = "github.com/docker/go-units@v0.4.0",
+        sha256 = "039d53ebe64af1aefa0be94ce42c621a17a3052c58ad15e5b3f357529beeaff6",
+        strip_prefix = "github.com/docker/go-units@v0.5.0",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/docker/go-units/com_github_docker_go_units-v0.4.0.zip",
-            "http://ats.apps.svc/gomod/github.com/docker/go-units/com_github_docker_go_units-v0.4.0.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/docker/go-units/com_github_docker_go_units-v0.4.0.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/docker/go-units/com_github_docker_go_units-v0.4.0.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/docker/go-units/com_github_docker_go_units-v0.5.0.zip",
+            "http://ats.apps.svc/gomod/github.com/docker/go-units/com_github_docker_go_units-v0.5.0.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/docker/go-units/com_github_docker_go_units-v0.5.0.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/docker/go-units/com_github_docker_go_units-v0.5.0.zip",
         ],
     )
     go_repository(
@@ -5399,6 +5399,19 @@ def go_deps():
             "https://storage.googleapis.com/pingcapmirror/gomod/github.com/niemeyer/pretty/com_github_niemeyer_pretty-v0.0.0-20200227124842-a10e7caefd8e.zip",
         ],
     )
+    go_repository(
+        name = "com_github_ninedraft_israce",
+        build_file_proto_mode = "disable_global",
+        importpath = "github.com/ninedraft/israce",
+        sha256 = "bbecd2498bb29bede456e197d22ef5888626577890b194940fc0f6c724c4ba57",
+        strip_prefix = "github.com/ninedraft/israce@v0.0.3",
+        urls = [
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/ninedraft/israce/com_github_ninedraft_israce-v0.0.3.zip",
+            "http://ats.apps.svc/gomod/github.com/ninedraft/israce/com_github_ninedraft_israce-v0.0.3.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/ninedraft/israce/com_github_ninedraft_israce-v0.0.3.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/ninedraft/israce/com_github_ninedraft_israce-v0.0.3.zip",
+        ],
+    )
     go_repository(
         name = "com_github_nishanths_exhaustive",
         build_file_proto_mode = "disable_global",
@@ -7041,13 +7054,13 @@ def go_deps():
         name = "com_github_tikv_client_go_v2",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/tikv/client-go/v2",
-        sha256 = "587d22d21daa1f44b18b0c2325fcb4233af71a985f397a36aa9db8796777b8f2",
-        strip_prefix = "github.com/tikv/client-go/v2@v2.0.8-0.20241212025239-0dc41295f929",
+        sha256 = "06e6f353ff0ed627eaca82ebf7d7899e2651da92697522465aa1e1bbb189a8ec",
+        strip_prefix = "github.com/tikv/client-go/v2@v2.0.8-0.20250115080352-c1b98e6cebab",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20241212025239-0dc41295f929.zip",
-            "http://ats.apps.svc/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20241212025239-0dc41295f929.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20241212025239-0dc41295f929.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20241212025239-0dc41295f929.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20250115080352-c1b98e6cebab.zip",
+            "http://ats.apps.svc/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20250115080352-c1b98e6cebab.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20250115080352-c1b98e6cebab.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/tikv/client-go/v2/com_github_tikv_client_go_v2-v2.0.8-0.20250115080352-c1b98e6cebab.zip",
         ],
     )
     go_repository(
diff --git a/br/pkg/lightning/config/bytesize_test.go b/br/pkg/lightning/config/bytesize_test.go
index 46968777056e5..bb76cd11b25a7 100644
--- a/br/pkg/lightning/config/bytesize_test.go
+++ b/br/pkg/lightning/config/bytesize_test.go
@@ -61,7 +61,7 @@ func TestByteSizeTOMLDecode(t *testing.T) {
 		},
 		{
 			input: "x = 'invalid value'",
-			err:   "invalid size: 'invalid value'",
+			err:   "strconv.ParseFloat: parsing \"invalid\": invalid syntax",
 		},
 		{
 			input: "x = true",
@@ -85,7 +85,7 @@ func TestByteSizeTOMLDecode(t *testing.T) {
 		},
 		{
 			input: "x = 2020-01-01T00:00:00Z",
-			err:   "invalid size: '2020-01-01T00:00:00Z'",
+			err:   "strconv.ParseFloat: parsing \"2020-01-01T00:00:00\": invalid syntax",
 		},
 		{
 			input: "x = ['100000']",
diff --git a/go.mod b/go.mod
index 7e213519ed566..6149febd4d68a 100644
--- a/go.mod
+++ b/go.mod
@@ -34,7 +34,7 @@ require (
 	github.com/danjacques/gofslock v0.0.0-20191023191349-0a45f885bc37
 	github.com/dgraph-io/ristretto v0.1.1
 	github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13
-	github.com/docker/go-units v0.4.0
+	github.com/docker/go-units v0.5.0
 	github.com/dolthub/swiss v0.2.1
 	github.com/emirpasic/gods v1.18.1
 	github.com/fatanugraha/noloopclosure v0.1.1
@@ -104,7 +104,7 @@ require (
 	github.com/stretchr/testify v1.9.0
 	github.com/tdakkota/asciicheck v0.2.0
 	github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2
-	github.com/tikv/client-go/v2 v2.0.8-0.20241212025239-0dc41295f929
+	github.com/tikv/client-go/v2 v2.0.8-0.20250115080352-c1b98e6cebab
 	github.com/tikv/pd/client v0.0.0-20240724132535-fcb34c90790c
 	github.com/timakin/bodyclose v0.0.0-20230421092635-574207250966
 	github.com/twmb/murmur3 v1.1.6
diff --git a/go.sum b/go.sum
index b394d17476f08..ec9c6375f96fd 100644
--- a/go.sum
+++ b/go.sum
@@ -248,8 +248,8 @@ github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13/go.mod h1:SqUrOPUn
 github.com/dgryski/go-sip13 v0.0.0-20181026042036-e10d5fee7954/go.mod h1:vAd38F8PWV+bWy6jNmig1y/TA+kYO4g3RSRF0IAv0no=
 github.com/dnaeon/go-vcr v1.1.0 h1:ReYa/UBrRyQdant9B4fNHGoCNKw6qh6P0fsdGmZpR7c=
 github.com/dnaeon/go-vcr v1.1.0/go.mod h1:M7tiix8f0r6mKKJ3Yq/kqU1OYf3MnfmBWVbPx/yU9ko=
-github.com/docker/go-units v0.4.0 h1:3uh0PgVws3nIA0Q+MwDC8yjEPf9zjRfZZWXZYDct3Tw=
-github.com/docker/go-units v0.4.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
+github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4=
+github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
 github.com/dolthub/maphash v0.1.0 h1:bsQ7JsF4FkkWyrP3oCnFJgrCUAFbFf3kOl4L/QxPDyQ=
 github.com/dolthub/maphash v0.1.0/go.mod h1:gkg4Ch4CdCDu5h6PMriVLawB7koZ+5ijb9puGMV50a4=
 github.com/dolthub/swiss v0.2.1 h1:gs2osYs5SJkAaH5/ggVJqXQxRXtWshF6uE0lgR/Y3Gw=
@@ -996,8 +996,8 @@ github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 h1:mbAskLJ0oJf
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2/go.mod h1:2PfKggNGDuadAa0LElHrByyrz4JPZ9fFx6Gs7nx7ZZU=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a h1:J/YdBZ46WKpXsxsW93SG+q0F8KI+yFrcIDT4c/RNoc4=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a/go.mod h1:h4xBhSNtOeEosLJ4P7JyKXX7Cabg7AVkWCK5gV2vOrM=
-github.com/tikv/client-go/v2 v2.0.8-0.20241212025239-0dc41295f929 h1:dvY5kl35L+aroDl3HQzOx4J7N6sdd8TiXEV+GhNemtU=
-github.com/tikv/client-go/v2 v2.0.8-0.20241212025239-0dc41295f929/go.mod h1:37p0ryKaieJbBpVDWnaPi2ZS6UFqkgpsemBLkGX2FvM=
+github.com/tikv/client-go/v2 v2.0.8-0.20250115080352-c1b98e6cebab h1:Hk2E8LtNDILE9bXpk1mVxPaTuJ0xvo6dnY75oKvgXzY=
+github.com/tikv/client-go/v2 v2.0.8-0.20250115080352-c1b98e6cebab/go.mod h1:zXX9NBGF4U3joK/GEldf3Fggi7Q6JA2u1ozh7Phhv8E=
 github.com/tikv/pd/client v0.0.0-20240724132535-fcb34c90790c h1:oZygf/SCdTUhjoHuZRE85EBgK0oA6LjikpWuJqqjM8U=
 github.com/tikv/pd/client v0.0.0-20240724132535-fcb34c90790c/go.mod h1:NW6Af689Jw1FDxjq+WL0nqOdmQ1XT0ly2R1SIKfQuUw=
 github.com/timakin/bodyclose v0.0.0-20230421092635-574207250966 h1:quvGphlmUVU+nhpFa4gg4yJyTRJ13reZMDHrKwYw53M=
diff --git a/pkg/executor/set.go b/pkg/executor/set.go
index 1fcce5863c6f6..0bc3d936fce58 100644
--- a/pkg/executor/set.go
+++ b/pkg/executor/set.go
@@ -213,7 +213,8 @@ func (e *SetExecutor) setSysVariable(ctx context.Context, name string, v *expres
 	newSnapshotTS := getSnapshotTSByName()
 	newSnapshotIsSet := newSnapshotTS > 0 && newSnapshotTS != oldSnapshotTS
 	if newSnapshotIsSet {
-		err = sessionctx.ValidateSnapshotReadTS(ctx, e.Ctx().GetStore(), newSnapshotTS)
+		isStaleRead := name == variable.TiDBTxnReadTS
+		err = sessionctx.ValidateSnapshotReadTS(ctx, e.Ctx().GetStore(), newSnapshotTS, isStaleRead)
 		if name != variable.TiDBTxnReadTS {
 			// Also check gc safe point for snapshot read.
 			// We don't check snapshot with gc safe point for read_ts
diff --git a/pkg/executor/test/executor/executor_test.go b/pkg/executor/test/executor/executor_test.go
index 251218e90207f..deadf3b439ed0 100644
--- a/pkg/executor/test/executor/executor_test.go
+++ b/pkg/executor/test/executor/executor_test.go
@@ -3062,7 +3062,7 @@ func TestStaleReadAtFutureTime(t *testing.T) {
 
 	tk := testkit.NewTestKit(t, store)
 	// Setting tx_read_ts to a time in the future will fail. (One day before the 2038 problem)
-	tk.MustGetErrMsg("set @@tx_read_ts = '2038-01-18 03:14:07'", "cannot set read timestamp to a future time")
+	tk.MustContainErrMsg("set @@tx_read_ts = '2038-01-18 03:14:07'", "cannot set read timestamp to a future time")
 	// TxnReadTS Is not updated if check failed.
 	require.Zero(t, tk.Session().GetSessionVars().TxnReadTS.PeakTxnReadTS())
 }
diff --git a/pkg/planner/core/planbuilder.go b/pkg/planner/core/planbuilder.go
index 83bc087ad90f6..7fc58a36301c3 100644
--- a/pkg/planner/core/planbuilder.go
+++ b/pkg/planner/core/planbuilder.go
@@ -3696,7 +3696,7 @@ func (b *PlanBuilder) buildSimple(ctx context.Context, node ast.StmtNode) (Plan,
 			if err != nil {
 				return nil, err
 			}
-			if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS); err != nil {
+			if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS, true); err != nil {
 				return nil, err
 			}
 			p.StaleTxnStartTS = startTS
@@ -3710,7 +3710,7 @@ func (b *PlanBuilder) buildSimple(ctx context.Context, node ast.StmtNode) (Plan,
 			if err != nil {
 				return nil, err
 			}
-			if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS); err != nil {
+			if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS, true); err != nil {
 				return nil, err
 			}
 			p.StaleTxnStartTS = startTS
diff --git a/pkg/sessionctx/context.go b/pkg/sessionctx/context.go
index 5ebb1052cd384..7b9aa6abf3fb4 100644
--- a/pkg/sessionctx/context.go
+++ b/pkg/sessionctx/context.go
@@ -213,9 +213,12 @@ const (
 	LastExecuteDDL basicCtxType = 3
 )
 
-// ValidateSnapshotReadTS strictly validates that readTS does not exceed the PD timestamp
-func ValidateSnapshotReadTS(ctx context.Context, store kv.Storage, readTS uint64) error {
-	return store.GetOracle().ValidateSnapshotReadTS(ctx, readTS, &oracle.Option{TxnScope: oracle.GlobalTxnScope})
+// ValidateSnapshotReadTS strictly validates that readTS does not exceed the PD timestamp.
+// For read requests to the storage, the check can be implicitly performed when sending the RPC request. So this
+// function is only needed when it's not proper to delay the check to when RPC requests are being sent (e.g., `BEGIN`
+// statements that don't make reading operation immediately).
+func ValidateSnapshotReadTS(ctx context.Context, store kv.Storage, readTS uint64, isStaleRead bool) error {
+	return store.GetOracle().ValidateReadTS(ctx, readTS, isStaleRead, &oracle.Option{TxnScope: oracle.GlobalTxnScope})
 }
 
 // SysProcTracker is used to track background sys processes
diff --git a/pkg/sessiontxn/staleread/processor.go b/pkg/sessiontxn/staleread/processor.go
index 393c3e7c378bb..e1a3f547fd11b 100644
--- a/pkg/sessiontxn/staleread/processor.go
+++ b/pkg/sessiontxn/staleread/processor.go
@@ -285,7 +285,7 @@ func parseAndValidateAsOf(ctx context.Context, sctx sessionctx.Context, asOf *as
 		return 0, err
 	}
 
-	if err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), ts); err != nil {
+	if err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), ts, true); err != nil {
 		return 0, err
 	}
 
diff --git a/pkg/sessiontxn/staleread/util.go b/pkg/sessiontxn/staleread/util.go
index 09f3edc2dbe0b..01791c6437900 100644
--- a/pkg/sessiontxn/staleread/util.go
+++ b/pkg/sessiontxn/staleread/util.go
@@ -77,7 +77,7 @@ func CalculateTsWithReadStaleness(ctx context.Context, sctx sessionctx.Context,
 		// If the final calculated exceeds the min safe ts, we are not sure whether the ts is safe to read (note that
 		// reading with a ts larger than PD's max allocated ts + 1 is unsafe and may break linearizability).
 		// So in this case, do an extra check on it.
-		err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), readTS)
+		err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), readTS, true)
 		if err != nil {
 			return 0, err
 		}
diff --git a/pkg/store/copr/BUILD.bazel b/pkg/store/copr/BUILD.bazel
index 9bebc01364fbb..e8c17b8094c8d 100644
--- a/pkg/store/copr/BUILD.bazel
+++ b/pkg/store/copr/BUILD.bazel
@@ -55,6 +55,7 @@ go_library(
         "@com_github_tikv_client_go_v2//config",
         "@com_github_tikv_client_go_v2//error",
         "@com_github_tikv_client_go_v2//metrics",
+        "@com_github_tikv_client_go_v2//oracle",
         "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_client_go_v2//tikvrpc",
         "@com_github_tikv_client_go_v2//tikvrpc/interceptor",
diff --git a/pkg/store/copr/batch_coprocessor.go b/pkg/store/copr/batch_coprocessor.go
index c33e8f9f6e112..19c43a567e9c8 100644
--- a/pkg/store/copr/batch_coprocessor.go
+++ b/pkg/store/copr/batch_coprocessor.go
@@ -1292,7 +1292,7 @@ func (b *batchCopIterator) retryBatchCopTask(ctx context.Context, bo *backoff.Ba
 const TiFlashReadTimeoutUltraLong = 3600 * time.Second
 
 func (b *batchCopIterator) handleTaskOnce(ctx context.Context, bo *backoff.Backoffer, task *batchCopTask) ([]*batchCopTask, error) {
-	sender := NewRegionBatchRequestSender(b.store.GetRegionCache(), b.store.GetTiKVClient(), b.enableCollectExecutionInfo)
+	sender := NewRegionBatchRequestSender(b.store.GetRegionCache(), b.store.GetTiKVClient(), b.store.store.GetOracle(), b.enableCollectExecutionInfo)
 	var regionInfos = make([]*coprocessor.RegionInfo, 0, len(task.regionInfos))
 	for _, ri := range task.regionInfos {
 		regionInfos = append(regionInfos, ri.toCoprocessorRegionInfo())
diff --git a/pkg/store/copr/batch_request_sender.go b/pkg/store/copr/batch_request_sender.go
index ccb138f7753c3..5c6d9a6cbe192 100644
--- a/pkg/store/copr/batch_request_sender.go
+++ b/pkg/store/copr/batch_request_sender.go
@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/tidb/pkg/config"
 	tikverr "github.com/tikv/client-go/v2/error"
+	"github.com/tikv/client-go/v2/oracle"
 	"github.com/tikv/client-go/v2/tikv"
 	"github.com/tikv/client-go/v2/tikvrpc"
 	"google.golang.org/grpc/codes"
@@ -56,9 +57,9 @@ type RegionBatchRequestSender struct {
 }
 
 // NewRegionBatchRequestSender creates a RegionBatchRequestSender object.
-func NewRegionBatchRequestSender(cache *RegionCache, client tikv.Client, enableCollectExecutionInfo bool) *RegionBatchRequestSender {
+func NewRegionBatchRequestSender(cache *RegionCache, client tikv.Client, oracle oracle.Oracle, enableCollectExecutionInfo bool) *RegionBatchRequestSender {
 	return &RegionBatchRequestSender{
-		RegionRequestSender:        tikv.NewRegionRequestSender(cache.RegionCache, client),
+		RegionRequestSender:        tikv.NewRegionRequestSender(cache.RegionCache, client, oracle),
 		enableCollectExecutionInfo: enableCollectExecutionInfo,
 	}
 }
diff --git a/pkg/store/copr/mpp.go b/pkg/store/copr/mpp.go
index cd0695a3e0d9d..cc5d361c6d73e 100644
--- a/pkg/store/copr/mpp.go
+++ b/pkg/store/copr/mpp.go
@@ -138,7 +138,7 @@ func (c *MPPClient) DispatchMPPTask(param kv.DispatchMPPTaskParam) (resp *mpp.Di
 	// Or else it's the task without region, which always happens in high layer task without table.
 	// In that case
 	if originalTask != nil {
-		sender := NewRegionBatchRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), param.EnableCollectExecutionInfo)
+		sender := NewRegionBatchRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.store.GetOracle(), param.EnableCollectExecutionInfo)
 		rpcResp, retry, _, err = sender.SendReqToAddr(bo, originalTask.ctx, originalTask.regionInfos, wrappedReq, tikv.ReadTimeoutMedium)
 		// No matter what the rpc error is, we won't retry the mpp dispatch tasks.
 		// TODO: If we want to retry, we must redo the plan fragment cutting and task scheduling.
diff --git a/pkg/util/mock/BUILD.bazel b/pkg/util/mock/BUILD.bazel
index 75ac693889df1..88c26bfd03067 100644
--- a/pkg/util/mock/BUILD.bazel
+++ b/pkg/util/mock/BUILD.bazel
@@ -22,6 +22,7 @@ go_library(
         "//pkg/sessionctx/variable",
         "//pkg/util",
         "//pkg/util/disk",
+        "//pkg/util/logutil",
         "//pkg/util/memory",
         "//pkg/util/sli",
         "//pkg/util/sqlexec",
diff --git a/pkg/util/mock/context.go b/pkg/util/mock/context.go
index 7c6b52c96bce5..18e02c07d40e4 100644
--- a/pkg/util/mock/context.go
+++ b/pkg/util/mock/context.go
@@ -32,6 +32,7 @@ import (
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/disk"
+	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/pingcap/tidb/pkg/util/memory"
 	"github.com/pingcap/tidb/pkg/util/sli"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
@@ -67,7 +68,7 @@ type wrapTxn struct {
 }
 
 func (txn *wrapTxn) validOrPending() bool {
-	return txn.tsFuture != nil || txn.Transaction.Valid()
+	return txn.tsFuture != nil || (txn.Transaction != nil && txn.Transaction.Valid())
 }
 
 func (txn *wrapTxn) pending() bool {
@@ -173,7 +174,15 @@ func (c *Context) GetSessionVars() *variable.SessionVars {
 }
 
 // Txn implements sessionctx.Context Txn interface.
-func (c *Context) Txn(bool) (kv.Transaction, error) {
+func (c *Context) Txn(active bool) (kv.Transaction, error) {
+	if active {
+		if !c.txn.validOrPending() {
+			err := c.newTxn(context.Background())
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
 	return &c.txn, nil
 }
 
@@ -253,10 +262,12 @@ func (c *Context) GetSessionPlanCache() sessionctx.PlanCache {
 	return c.pcache
 }
 
-// NewTxn implements the sessionctx.Context interface.
-func (c *Context) NewTxn(context.Context) error {
+// newTxn Creates new transaction on the session context.
+func (c *Context) newTxn(ctx context.Context) error {
 	if c.Store == nil {
-		return errors.New("store is not set")
+		logutil.Logger(ctx).Warn("mock.Context: No store is specified when trying to create new transaction. A fake transaction will be created. Note that this is unrecommended usage.")
+		c.fakeTxn()
+		return nil
 	}
 	if c.txn.Valid() {
 		err := c.txn.Commit(c.ctx)
@@ -273,14 +284,41 @@ func (c *Context) NewTxn(context.Context) error {
 	return nil
 }
 
-// NewStaleTxnWithStartTS implements the sessionctx.Context interface.
-func (c *Context) NewStaleTxnWithStartTS(ctx context.Context, _ uint64) error {
-	return c.NewTxn(ctx)
+// fakeTxn is used to let some tests pass in the context without an available kv.Storage. Once usages to access
+// transactions without a kv.Storage are removed, this type should also be removed.
+// New code should never use this.
+type fakeTxn struct {
+	// The inner should always be nil.
+	kv.Transaction
+	startTS uint64
+}
+
+func (t *fakeTxn) StartTS() uint64 {
+	return t.startTS
+}
+
+func (*fakeTxn) SetDiskFullOpt(_ kvrpcpb.DiskFullOpt) {}
+
+func (*fakeTxn) SetOption(_ int, _ any) {}
+
+func (*fakeTxn) Get(ctx context.Context, _ kv.Key) ([]byte, error) {
+	// Check your implementation if you meet this error. It's dangerous if some calculation relies on the data but the
+	// read result is faked.
+	logutil.Logger(ctx).Warn("mock.Context: No store is specified but trying to access data from a transaction.")
+	return nil, nil
+}
+
+func (*fakeTxn) Valid() bool { return true }
+
+func (c *Context) fakeTxn() {
+	c.txn.Transaction = &fakeTxn{
+		startTS: 1,
+	}
 }
 
 // RefreshTxnCtx implements the sessionctx.Context interface.
 func (c *Context) RefreshTxnCtx(ctx context.Context) error {
-	return errors.Trace(c.NewTxn(ctx))
+	return errors.Trace(c.newTxn(ctx))
 }
 
 // RollbackTxn indicates an expected call of RollbackTxn.