diff --git a/kv/memdb/memory_mutation.go b/kv/memdb/memory_mutation.go
index 154f825b3..7fb221c72 100644
--- a/kv/memdb/memory_mutation.go
+++ b/kv/memdb/memory_mutation.go
@@ -56,6 +56,10 @@ func NewMemoryBatch(tx kv.Tx) *MemoryMutation {
 	}
 }
 
+func (m *MemoryMutation) UpdateTxn(tx kv.Tx) {
+	m.db = tx
+}
+
 func (m *MemoryMutation) isTableCleared(table string) bool {
 	_, ok := m.clearedTables[table]
 	return ok