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.