From b7c22aa31d062b77a96493519a6b50ca9b627971 Mon Sep 17 00:00:00 2001
From: disksing <i@disksing.com>
Date: Thu, 13 May 2021 21:29:39 +0800
Subject: [PATCH] kv: remove UnionStore interface (#24625)

---
 ddl/index.go                          |  2 +-
 ddl/index_change_test.go              |  2 +-
 executor/admin.go                     |  2 +-
 executor/admin_test.go                | 44 +++++++++++++--------------
 executor/distsql_test.go              |  2 +-
 executor/executor_test.go             |  4 +--
 kv/interface_mock_test.go             |  4 ---
 kv/kv.go                              |  2 --
 kv/union_store.go                     | 17 +----------
 store/driver/txn/txn_driver.go        |  4 ---
 store/driver/txn/unionstore_driver.go | 36 ----------------------
 table/index.go                        |  6 ++--
 table/tables/index.go                 | 27 ++++++++--------
 table/tables/index_test.go            | 12 ++++----
 table/tables/tables.go                |  4 +--
 util/admin/admin.go                   |  2 +-
 util/mock/context.go                  |  8 -----
 17 files changed, 54 insertions(+), 124 deletions(-)

diff --git a/ddl/index.go b/ddl/index.go
index b1b4303d7a0f1..fdba6c65008f6 100644
--- a/ddl/index.go
+++ b/ddl/index.go
@@ -1343,7 +1343,7 @@ func (w *cleanUpIndexWorker) BackfillDataInTxn(handleRange reorgBackfillTask) (t
 			// we fetch records row by row, so records will belong to
 			// index[0], index[1] ... index[n-1], index[0], index[1] ...
 			// respectively. So indexes[i%n] is the index of idxRecords[i].
-			err := w.indexes[i%n].Delete(w.sessCtx.GetSessionVars().StmtCtx, txn.GetUnionStore(), idxRecord.vals, idxRecord.handle)
+			err := w.indexes[i%n].Delete(w.sessCtx.GetSessionVars().StmtCtx, txn, idxRecord.vals, idxRecord.handle)
 			if err != nil {
 				return errors.Trace(err)
 			}
diff --git a/ddl/index_change_test.go b/ddl/index_change_test.go
index 0a54b6b25e694..dfdfc7111c372 100644
--- a/ddl/index_change_test.go
+++ b/ddl/index_change_test.go
@@ -198,7 +198,7 @@ func checkIndexExists(ctx sessionctx.Context, tbl table.Table, indexValue interf
 	if err != nil {
 		return errors.Trace(err)
 	}
-	doesExist, _, err := idx.Exist(ctx.GetSessionVars().StmtCtx, txn.GetUnionStore(), types.MakeDatums(indexValue), kv.IntHandle(handle))
+	doesExist, _, err := idx.Exist(ctx.GetSessionVars().StmtCtx, txn, types.MakeDatums(indexValue), kv.IntHandle(handle))
 	if err != nil {
 		return errors.Trace(err)
 	}
diff --git a/executor/admin.go b/executor/admin.go
index 7e15a24e667ca..be46e39a1d2a3 100644
--- a/executor/admin.go
+++ b/executor/admin.go
@@ -575,7 +575,7 @@ func (e *CleanupIndexExec) deleteDanglingIdx(txn kv.Transaction, values map[stri
 				return errors.Trace(errors.Errorf("batch keys are inconsistent with handles"))
 			}
 			for _, handleIdxVals := range handleIdxValsGroup.([][]types.Datum) {
-				if err := e.index.Delete(e.ctx.GetSessionVars().StmtCtx, txn.GetUnionStore(), handleIdxVals, handle); err != nil {
+				if err := e.index.Delete(e.ctx.GetSessionVars().StmtCtx, txn, handleIdxVals, handle); err != nil {
 					return err
 				}
 				e.removeCnt++
diff --git a/executor/admin_test.go b/executor/admin_test.go
index 20095eb59a0ba..35e3d08345d63 100644
--- a/executor/admin_test.go
+++ b/executor/admin_test.go
@@ -135,7 +135,7 @@ func (s *testSuite5) TestAdminRecoverIndex(c *C) {
 	sc := s.ctx.GetSessionVars().StmtCtx
 	txn, err := s.store.Begin()
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(1), kv.IntHandle(1))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(1), kv.IntHandle(1))
 	c.Assert(err, IsNil)
 	err = txn.Commit(context.Background())
 	c.Assert(err, IsNil)
@@ -158,7 +158,7 @@ func (s *testSuite5) TestAdminRecoverIndex(c *C) {
 
 	txn, err = s.store.Begin()
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(10), kv.IntHandle(10))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(10), kv.IntHandle(10))
 	c.Assert(err, IsNil)
 	err = txn.Commit(context.Background())
 	c.Assert(err, IsNil)
@@ -172,15 +172,15 @@ func (s *testSuite5) TestAdminRecoverIndex(c *C) {
 
 	txn, err = s.store.Begin()
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(1), kv.IntHandle(1))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(1), kv.IntHandle(1))
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(2), kv.IntHandle(2))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(2), kv.IntHandle(2))
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(3), kv.IntHandle(3))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(3), kv.IntHandle(3))
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(10), kv.IntHandle(10))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(10), kv.IntHandle(10))
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(20), kv.IntHandle(20))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(20), kv.IntHandle(20))
 	c.Assert(err, IsNil)
 	err = txn.Commit(context.Background())
 	c.Assert(err, IsNil)
@@ -236,7 +236,7 @@ func (s *testSuite5) TestClusteredIndexAdminRecoverIndex(c *C) {
 	txn, err := s.store.Begin()
 	c.Assert(err, IsNil)
 	cHandle := testutil.MustNewCommonHandle(c, "1", "3")
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(2), cHandle)
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(2), cHandle)
 	c.Assert(err, IsNil)
 	err = txn.Commit(context.Background())
 	c.Assert(err, IsNil)
@@ -269,7 +269,7 @@ func (s *testSuite5) TestAdminRecoverPartitionTableIndex(c *C) {
 		sc := s.ctx.GetSessionVars().StmtCtx
 		txn, err := s.store.Begin()
 		c.Assert(err, IsNil)
-		err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(idxValue), kv.IntHandle(idxValue))
+		err = indexOpr.Delete(sc, txn, types.MakeDatums(idxValue), kv.IntHandle(idxValue))
 		c.Assert(err, IsNil)
 		err = txn.Commit(context.Background())
 		c.Assert(err, IsNil)
@@ -345,13 +345,13 @@ func (s *testSuite5) TestAdminRecoverIndex1(c *C) {
 
 	txn, err := s.store.Begin()
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums("1"), kv.IntHandle(1))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums("1"), kv.IntHandle(1))
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums("2"), kv.IntHandle(2))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums("2"), kv.IntHandle(2))
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums("3"), kv.IntHandle(3))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums("3"), kv.IntHandle(3))
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums("10"), kv.IntHandle(4))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums("10"), kv.IntHandle(4))
 	c.Assert(err, IsNil)
 	err = txn.Commit(context.Background())
 	c.Assert(err, IsNil)
@@ -745,7 +745,7 @@ func (s *testSuite3) TestAdminCheckPartitionTableFailed(c *C) {
 		indexOpr := tables.NewIndex(tblInfo.GetPartitionInfo().Definitions[partitionIdx].ID, tblInfo, idxInfo)
 		txn, err := s.store.Begin()
 		c.Assert(err, IsNil)
-		err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(i), kv.IntHandle(i))
+		err = indexOpr.Delete(sc, txn, types.MakeDatums(i), kv.IntHandle(i))
 		c.Assert(err, IsNil)
 		err = txn.Commit(context.Background())
 		c.Assert(err, IsNil)
@@ -784,7 +784,7 @@ func (s *testSuite3) TestAdminCheckPartitionTableFailed(c *C) {
 		// TODO: fix admin recover for partition table.
 		txn, err = s.store.Begin()
 		c.Assert(err, IsNil)
-		err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(i+8), kv.IntHandle(i+8))
+		err = indexOpr.Delete(sc, txn, types.MakeDatums(i+8), kv.IntHandle(i+8))
 		c.Assert(err, IsNil)
 		err = txn.Commit(context.Background())
 		c.Assert(err, IsNil)
@@ -807,7 +807,7 @@ func (s *testSuite3) TestAdminCheckPartitionTableFailed(c *C) {
 		// TODO: fix admin recover for partition table.
 		txn, err = s.store.Begin()
 		c.Assert(err, IsNil)
-		err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(i+8), kv.IntHandle(i))
+		err = indexOpr.Delete(sc, txn, types.MakeDatums(i+8), kv.IntHandle(i))
 		c.Assert(err, IsNil)
 		err = txn.Commit(context.Background())
 		c.Assert(err, IsNil)
@@ -842,7 +842,7 @@ func (s *testSuite5) TestAdminCheckTableFailed(c *C) {
 	// Index c2 is missing 11.
 	txn, err := s.store.Begin()
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(-10), kv.IntHandle(-1))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(-10), kv.IntHandle(-1))
 	c.Assert(err, IsNil)
 	err = txn.Commit(context.Background())
 	c.Assert(err, IsNil)
@@ -873,7 +873,7 @@ func (s *testSuite5) TestAdminCheckTableFailed(c *C) {
 	// Index c2 has two more values than table data: 10, 13, and these handles have correlative record.
 	txn, err = s.store.Begin()
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(0), kv.IntHandle(0))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(0), kv.IntHandle(0))
 	c.Assert(err, IsNil)
 	// Make sure the index value "19" is smaller "21". Then we scan to "19" before "21".
 	_, err = indexOpr.Create(s.ctx, txn, types.MakeDatums(19), kv.IntHandle(10), nil)
@@ -890,9 +890,9 @@ func (s *testSuite5) TestAdminCheckTableFailed(c *C) {
 	// Two indices have the same handle.
 	txn, err = s.store.Begin()
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(13), kv.IntHandle(2))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(13), kv.IntHandle(2))
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(12), kv.IntHandle(2))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(12), kv.IntHandle(2))
 	c.Assert(err, IsNil)
 	err = txn.Commit(context.Background())
 	c.Assert(err, IsNil)
@@ -906,7 +906,7 @@ func (s *testSuite5) TestAdminCheckTableFailed(c *C) {
 	c.Assert(err, IsNil)
 	_, err = indexOpr.Create(s.ctx, txn, types.MakeDatums(12), kv.IntHandle(2), nil)
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(20), kv.IntHandle(10))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(20), kv.IntHandle(10))
 	c.Assert(err, IsNil)
 	err = txn.Commit(context.Background())
 	c.Assert(err, IsNil)
@@ -917,7 +917,7 @@ func (s *testSuite5) TestAdminCheckTableFailed(c *C) {
 	// Recover records.
 	txn, err = s.store.Begin()
 	c.Assert(err, IsNil)
-	err = indexOpr.Delete(sc, txn.GetUnionStore(), types.MakeDatums(19), kv.IntHandle(10))
+	err = indexOpr.Delete(sc, txn, types.MakeDatums(19), kv.IntHandle(10))
 	c.Assert(err, IsNil)
 	_, err = indexOpr.Create(s.ctx, txn, types.MakeDatums(20), kv.IntHandle(10), nil)
 	c.Assert(err, IsNil)
diff --git a/executor/distsql_test.go b/executor/distsql_test.go
index eca6e1d016b40..6f06fd550f0b4 100644
--- a/executor/distsql_test.go
+++ b/executor/distsql_test.go
@@ -241,7 +241,7 @@ func (s *testSuite3) TestInconsistentIndex(c *C) {
 	for i := 0; i < 10; i++ {
 		txn, err := s.store.Begin()
 		c.Assert(err, IsNil)
-		err = idxOp.Delete(ctx.GetSessionVars().StmtCtx, txn.GetUnionStore(), types.MakeDatums(i+10), kv.IntHandle(100+i))
+		err = idxOp.Delete(ctx.GetSessionVars().StmtCtx, txn, types.MakeDatums(i+10), kv.IntHandle(100+i))
 		c.Assert(err, IsNil)
 		err = txn.Commit(context.Background())
 		c.Assert(err, IsNil)
diff --git a/executor/executor_test.go b/executor/executor_test.go
index 3b168636606ed..7b433e71e60ac 100644
--- a/executor/executor_test.go
+++ b/executor/executor_test.go
@@ -3852,9 +3852,9 @@ func (s *testSuite) TestCheckIndex(c *C) {
 	// table     data (handle, data): (1, 10), (2, 20), (4, 40)
 	txn, err = s.store.Begin()
 	c.Assert(err, IsNil)
-	err = idx.Delete(sc, txn.GetUnionStore(), types.MakeDatums(int64(30)), kv.IntHandle(3))
+	err = idx.Delete(sc, txn, types.MakeDatums(int64(30)), kv.IntHandle(3))
 	c.Assert(err, IsNil)
-	err = idx.Delete(sc, txn.GetUnionStore(), types.MakeDatums(int64(20)), kv.IntHandle(2))
+	err = idx.Delete(sc, txn, types.MakeDatums(int64(20)), kv.IntHandle(2))
 	c.Assert(err, IsNil)
 	err = txn.Commit(context.Background())
 	c.Assert(err, IsNil)
diff --git a/kv/interface_mock_test.go b/kv/interface_mock_test.go
index 2388c4f48b9f3..e1d41f1693088 100644
--- a/kv/interface_mock_test.go
+++ b/kv/interface_mock_test.go
@@ -106,10 +106,6 @@ func (t *mockTxn) GetSnapshot() Snapshot {
 	return nil
 }
 
-func (t *mockTxn) GetUnionStore() UnionStore {
-	return nil
-}
-
 func (t *mockTxn) NewStagingBuffer() MemBuffer {
 	return nil
 }
diff --git a/kv/kv.go b/kv/kv.go
index 1fad79d641009..e5ab4eed6f812 100644
--- a/kv/kv.go
+++ b/kv/kv.go
@@ -174,8 +174,6 @@ type Transaction interface {
 	GetMemBuffer() MemBuffer
 	// GetSnapshot returns the Snapshot binding to this transaction.
 	GetSnapshot() Snapshot
-	// GetUnionStore returns the UnionStore binding to this transaction.
-	GetUnionStore() UnionStore
 	// SetVars sets variables to the transaction.
 	SetVars(vars interface{})
 	// GetVars gets variables from the transaction.
diff --git a/kv/union_store.go b/kv/union_store.go
index 0e9a6768c5ebc..1d12e0a7fd37b 100644
--- a/kv/union_store.go
+++ b/kv/union_store.go
@@ -15,21 +15,6 @@ package kv
 
 // UnionStore is a store that wraps a snapshot for read and a MemBuffer for buffered write.
 // Also, it provides some transaction related utilities.
+// TODO: Remove after upgrading BR.
 type UnionStore interface {
-	Retriever
-
-	// HasPresumeKeyNotExists returns whether the key presumed key not exists error for the lazy check.
-	HasPresumeKeyNotExists(k Key) bool
-	// UnmarkPresumeKeyNotExists deletes the key presume key not exists error flag for the lazy check.
-	UnmarkPresumeKeyNotExists(k Key)
-
-	// SetOption sets an option with a value, when val is nil, uses the default
-	// value of this option.
-	SetOption(opt int, val interface{})
-	// DelOption deletes an option.
-	DelOption(opt int)
-	// GetOption gets an option.
-	GetOption(opt int) interface{}
-	// GetMemBuffer return the MemBuffer binding to this unionStore.
-	GetMemBuffer() MemBuffer
 }
diff --git a/store/driver/txn/txn_driver.go b/store/driver/txn/txn_driver.go
index 4d5ce77034312..72c1aac8c8e71 100644
--- a/store/driver/txn/txn_driver.go
+++ b/store/driver/txn/txn_driver.go
@@ -120,10 +120,6 @@ func (txn *tikvTxn) GetMemBuffer() kv.MemBuffer {
 	return newMemBuffer(txn.KVTxn.GetMemBuffer())
 }
 
-func (txn *tikvTxn) GetUnionStore() kv.UnionStore {
-	return &tikvUnionStore{txn.KVTxn.GetUnionStore()}
-}
-
 func (txn *tikvTxn) SetOption(opt int, val interface{}) {
 	switch opt {
 	case kv.BinlogInfo:
diff --git a/store/driver/txn/unionstore_driver.go b/store/driver/txn/unionstore_driver.go
index 5a2f56bfe4233..b54bffa139700 100644
--- a/store/driver/txn/unionstore_driver.go
+++ b/store/driver/txn/unionstore_driver.go
@@ -111,42 +111,6 @@ func (m *memBuffer) SnapshotGetter() kv.Getter {
 	return newKVGetter(m.MemDB.SnapshotGetter())
 }
 
-//tikvUnionStore implements kv.UnionStore
-type tikvUnionStore struct {
-	*unionstore.KVUnionStore
-}
-
-func (u *tikvUnionStore) GetMemBuffer() kv.MemBuffer {
-	return newMemBuffer(u.KVUnionStore.GetMemBuffer())
-}
-
-func (u *tikvUnionStore) Get(ctx context.Context, k kv.Key) ([]byte, error) {
-	data, err := u.KVUnionStore.Get(ctx, k)
-	return data, derr.ToTiDBErr(err)
-}
-
-func (u *tikvUnionStore) HasPresumeKeyNotExists(k kv.Key) bool {
-	return u.KVUnionStore.HasPresumeKeyNotExists(k)
-}
-
-func (u *tikvUnionStore) UnmarkPresumeKeyNotExists(k kv.Key) {
-	u.KVUnionStore.UnmarkPresumeKeyNotExists(k)
-}
-
-func (u *tikvUnionStore) Iter(k kv.Key, upperBound kv.Key) (kv.Iterator, error) {
-	it, err := u.KVUnionStore.Iter(k, upperBound)
-	return newKVIterator(it), derr.ToTiDBErr(err)
-}
-
-// IterReverse creates a reversed Iterator positioned on the first entry which key is less than k.
-// The returned iterator will iterate from greater key to smaller key.
-// If k is nil, the returned iterator will be positioned at the last key.
-// TODO: Add lower bound limit
-func (u *tikvUnionStore) IterReverse(k kv.Key) (kv.Iterator, error) {
-	it, err := u.KVUnionStore.IterReverse(k)
-	return newKVIterator(it), derr.ToTiDBErr(err)
-}
-
 type tikvGetter struct {
 	unionstore.Getter
 }
diff --git a/table/index.go b/table/index.go
index 5a9f32fbbfd3f..336efb7f574c2 100644
--- a/table/index.go
+++ b/table/index.go
@@ -66,11 +66,11 @@ type Index interface {
 	// Create supports insert into statement.
 	Create(ctx sessionctx.Context, txn kv.Transaction, indexedValues []types.Datum, h kv.Handle, handleRestoreData []types.Datum, opts ...CreateIdxOptFunc) (kv.Handle, error)
 	// Delete supports delete from statement.
-	Delete(sc *stmtctx.StatementContext, us kv.UnionStore, indexedValues []types.Datum, h kv.Handle) error
+	Delete(sc *stmtctx.StatementContext, txn kv.Transaction, indexedValues []types.Datum, h kv.Handle) error
 	// Drop supports drop table, drop index statements.
-	Drop(us kv.UnionStore) error
+	Drop(txn kv.Transaction) error
 	// Exist supports check index exists or not.
-	Exist(sc *stmtctx.StatementContext, us kv.UnionStore, indexedValues []types.Datum, h kv.Handle) (bool, kv.Handle, error)
+	Exist(sc *stmtctx.StatementContext, txn kv.Transaction, indexedValues []types.Datum, h kv.Handle) (bool, kv.Handle, error)
 	// GenIndexKey generates an index key.
 	GenIndexKey(sc *stmtctx.StatementContext, indexedValues []types.Datum, h kv.Handle, buf []byte) (key []byte, distinct bool, err error)
 	// Seek supports where clause.
diff --git a/table/tables/index.go b/table/tables/index.go
index 8b4630d47f70d..ab6296390fbfd 100644
--- a/table/tables/index.go
+++ b/table/tables/index.go
@@ -184,9 +184,8 @@ func (c *index) Create(sctx sessionctx.Context, txn kv.Transaction, indexedValue
 		return nil, err
 	}
 
-	us := txn.GetUnionStore()
 	if !distinct || skipCheck || opt.Untouched {
-		err = us.GetMemBuffer().Set(key, idxVal)
+		err = txn.GetMemBuffer().Set(key, idxVal)
 		return nil, err
 	}
 
@@ -202,18 +201,18 @@ func (c *index) Create(sctx sessionctx.Context, txn kv.Transaction, indexedValue
 
 	var value []byte
 	if sctx.GetSessionVars().LazyCheckKeyNotExists() {
-		value, err = us.GetMemBuffer().Get(ctx, key)
+		value, err = txn.GetMemBuffer().Get(ctx, key)
 	} else {
-		value, err = us.Get(ctx, key)
+		value, err = txn.Get(ctx, key)
 	}
 	if err != nil && !kv.IsErrNotFound(err) {
 		return nil, err
 	}
 	if err != nil || len(value) == 0 {
 		if sctx.GetSessionVars().LazyCheckKeyNotExists() && err != nil {
-			err = us.GetMemBuffer().SetWithFlags(key, idxVal, tikvstore.SetPresumeKeyNotExists)
+			err = txn.GetMemBuffer().SetWithFlags(key, idxVal, tikvstore.SetPresumeKeyNotExists)
 		} else {
-			err = us.GetMemBuffer().Set(key, idxVal)
+			err = txn.GetMemBuffer().Set(key, idxVal)
 		}
 		return nil, err
 	}
@@ -226,22 +225,22 @@ func (c *index) Create(sctx sessionctx.Context, txn kv.Transaction, indexedValue
 }
 
 // Delete removes the entry for handle h and indexedValues from KV index.
-func (c *index) Delete(sc *stmtctx.StatementContext, us kv.UnionStore, indexedValues []types.Datum, h kv.Handle) error {
+func (c *index) Delete(sc *stmtctx.StatementContext, txn kv.Transaction, indexedValues []types.Datum, h kv.Handle) error {
 	key, distinct, err := c.GenIndexKey(sc, indexedValues, h, nil)
 	if err != nil {
 		return err
 	}
 	if distinct {
-		err = us.GetMemBuffer().DeleteWithFlags(key, tikvstore.SetNeedLocked)
+		err = txn.GetMemBuffer().DeleteWithFlags(key, tikvstore.SetNeedLocked)
 	} else {
-		err = us.GetMemBuffer().Delete(key)
+		err = txn.GetMemBuffer().Delete(key)
 	}
 	return err
 }
 
 // Drop removes the KV index from store.
-func (c *index) Drop(us kv.UnionStore) error {
-	it, err := us.Iter(c.prefix, c.prefix.PrefixNext())
+func (c *index) Drop(txn kv.Transaction) error {
+	it, err := txn.Iter(c.prefix, c.prefix.PrefixNext())
 	if err != nil {
 		return err
 	}
@@ -252,7 +251,7 @@ func (c *index) Drop(us kv.UnionStore) error {
 		if !it.Key().HasPrefix(c.prefix) {
 			break
 		}
-		err := us.GetMemBuffer().Delete(it.Key())
+		err := txn.GetMemBuffer().Delete(it.Key())
 		if err != nil {
 			return err
 		}
@@ -298,13 +297,13 @@ func (c *index) SeekFirst(r kv.Retriever) (iter table.IndexIterator, err error)
 	return &indexIter{it: it, idx: c, prefix: c.prefix, colInfos: colInfos, tps: tps}, nil
 }
 
-func (c *index) Exist(sc *stmtctx.StatementContext, us kv.UnionStore, indexedValues []types.Datum, h kv.Handle) (bool, kv.Handle, error) {
+func (c *index) Exist(sc *stmtctx.StatementContext, txn kv.Transaction, indexedValues []types.Datum, h kv.Handle) (bool, kv.Handle, error) {
 	key, distinct, err := c.GenIndexKey(sc, indexedValues, h, nil)
 	if err != nil {
 		return false, nil, err
 	}
 
-	value, err := us.Get(context.TODO(), key)
+	value, err := txn.Get(context.TODO(), key)
 	if kv.IsErrNotFound(err) {
 		return false, nil, nil
 	}
diff --git a/table/tables/index_test.go b/table/tables/index_test.go
index 9345e86bab185..2c0a417746d42 100644
--- a/table/tables/index_test.go
+++ b/table/tables/index_test.go
@@ -104,15 +104,15 @@ func (s *testIndexSuite) TestIndex(c *C) {
 	c.Assert(h.IntValue(), Equals, int64(1))
 	it.Close()
 	sc := &stmtctx.StatementContext{TimeZone: time.Local}
-	exist, _, err := index.Exist(sc, txn.GetUnionStore(), values, kv.IntHandle(100))
+	exist, _, err := index.Exist(sc, txn, values, kv.IntHandle(100))
 	c.Assert(err, IsNil)
 	c.Assert(exist, IsFalse)
 
-	exist, _, err = index.Exist(sc, txn.GetUnionStore(), values, kv.IntHandle(1))
+	exist, _, err = index.Exist(sc, txn, values, kv.IntHandle(1))
 	c.Assert(err, IsNil)
 	c.Assert(exist, IsTrue)
 
-	err = index.Delete(sc, txn.GetUnionStore(), values, kv.IntHandle(1))
+	err = index.Delete(sc, txn, values, kv.IntHandle(1))
 	c.Assert(err, IsNil)
 
 	it, err = index.SeekFirst(txn)
@@ -132,7 +132,7 @@ func (s *testIndexSuite) TestIndex(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(hit, IsFalse)
 
-	err = index.Drop(txn.GetUnionStore())
+	err = index.Drop(txn)
 	c.Assert(err, IsNil)
 
 	it, hit, err = index.Seek(sc, txn, values)
@@ -194,12 +194,12 @@ func (s *testIndexSuite) TestIndex(c *C) {
 	c.Assert(h.IntValue(), Equals, int64(1))
 	it.Close()
 
-	exist, h, err = index.Exist(sc, txn.GetUnionStore(), values, kv.IntHandle(1))
+	exist, h, err = index.Exist(sc, txn, values, kv.IntHandle(1))
 	c.Assert(err, IsNil)
 	c.Assert(h.IntValue(), Equals, int64(1))
 	c.Assert(exist, IsTrue)
 
-	exist, h, err = index.Exist(sc, txn.GetUnionStore(), values, kv.IntHandle(2))
+	exist, h, err = index.Exist(sc, txn, values, kv.IntHandle(2))
 	c.Assert(err, NotNil)
 	c.Assert(h.IntValue(), Equals, int64(1))
 	c.Assert(exist, IsTrue)
diff --git a/table/tables/tables.go b/table/tables/tables.go
index a6a4180401d6a..37e6ab1103a89 100644
--- a/table/tables/tables.go
+++ b/table/tables/tables.go
@@ -1152,7 +1152,7 @@ func (t *TableCommon) removeRowIndices(ctx sessionctx.Context, h kv.Handle, rec
 			logutil.BgLogger().Info("remove row index failed", zap.Any("index", v.Meta()), zap.Uint64("txnStartTS", txn.StartTS()), zap.String("handle", h.String()), zap.Any("record", rec), zap.Error(err))
 			return err
 		}
-		if err = v.Delete(ctx.GetSessionVars().StmtCtx, txn.GetUnionStore(), vals, h); err != nil {
+		if err = v.Delete(ctx.GetSessionVars().StmtCtx, txn, vals, h); err != nil {
 			if v.Meta().State != model.StatePublic && kv.ErrNotExist.Equal(err) {
 				// If the index is not in public state, we may have not created the index,
 				// or already deleted the index, so skip ErrNotExist error.
@@ -1167,7 +1167,7 @@ func (t *TableCommon) removeRowIndices(ctx sessionctx.Context, h kv.Handle, rec
 
 // removeRowIndex implements table.Table RemoveRowIndex interface.
 func (t *TableCommon) removeRowIndex(sc *stmtctx.StatementContext, h kv.Handle, vals []types.Datum, idx table.Index, txn kv.Transaction) error {
-	return idx.Delete(sc, txn.GetUnionStore(), vals, h)
+	return idx.Delete(sc, txn, vals, h)
 }
 
 // buildIndexForRow implements table.Table BuildIndexForRow interface.
diff --git a/util/admin/admin.go b/util/admin/admin.go
index 20217a53c1b6d..608040a89d6b0 100644
--- a/util/admin/admin.go
+++ b/util/admin/admin.go
@@ -389,7 +389,7 @@ func CheckRecordAndIndex(sessCtx sessionctx.Context, txn kv.Transaction, t table
 				vals1[i] = colDefVal
 			}
 		}
-		isExist, h2, err := idx.Exist(sc, txn.GetUnionStore(), vals1, h1)
+		isExist, h2, err := idx.Exist(sc, txn, vals1, h1)
 		if kv.ErrKeyExists.Equal(err) {
 			record1 := &RecordData{Handle: h1, Values: vals1}
 			record2 := &RecordData{Handle: h2, Values: vals1}
diff --git a/util/mock/context.go b/util/mock/context.go
index 4b329e0ff1f55..d6a5f1d913902 100644
--- a/util/mock/context.go
+++ b/util/mock/context.go
@@ -59,14 +59,6 @@ func (txn *wrapTxn) Valid() bool {
 	return txn.Transaction != nil && txn.Transaction.Valid()
 }
 
-// GetUnionStore implements GetUnionStore
-func (txn *wrapTxn) GetUnionStore() kv.UnionStore {
-	if txn.Transaction == nil {
-		return nil
-	}
-	return txn.Transaction.GetUnionStore()
-}
-
 func (txn *wrapTxn) CacheTableInfo(id int64, info *model.TableInfo) {
 	if txn.Transaction == nil {
 		return