From a5cafcce10334080f8748439579d5d13f435e310 Mon Sep 17 00:00:00 2001
From: crazycs520 <crazycs520@gmail.com>
Date: Mon, 25 Feb 2019 15:21:19 +0800
Subject: [PATCH 1/2] *: add rollback for load data statement

---
 server/conn.go | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/server/conn.go b/server/conn.go
index 2ec4bbe523d25..5621bc113d772 100644
--- a/server/conn.go
+++ b/server/conn.go
@@ -954,19 +954,24 @@ func (cc *clientConn) handleLoadData(ctx context.Context, loadDataInfo *executor
 	}
 	loadDataInfo.SetMessage()
 
-	if err = loadDataInfo.Ctx.StmtCommit(); err != nil {
-		return errors.Trace(err)
-	}
-	txn, err := loadDataInfo.Ctx.Txn(true)
 	if err != nil {
+		loadDataInfo.Ctx.StmtRollback()
+	} else {
+		err = loadDataInfo.Ctx.StmtCommit()
+	}
+
+	if txn, err1 := loadDataInfo.Ctx.Txn(true); err1 == nil {
 		if txn != nil && txn.Valid() {
-			if err1 := txn.Rollback(); err1 != nil {
-				log.Errorf("load data rollback failed: %v", err1)
+			if err != nil {
+				if err1 := txn.Rollback(); err1 != nil {
+					log.Errorf("load data rollback failed: %v", err1)
+				}
+				return errors.Trace(err)
 			}
 		}
-		return errors.Trace(err)
+	} else {
+		log.Error(err1)
 	}
-
 	return errors.Trace(cc.ctx.CommitTxn(sessionctx.SetCommitCtx(ctx, loadDataInfo.Ctx)))
 }
 

From 8f7996db35b6f9df6ce2092282d73c80e0bb560f Mon Sep 17 00:00:00 2001
From: crazycs520 <crazycs520@gmail.com>
Date: Wed, 27 Feb 2019 15:46:25 +0800
Subject: [PATCH 2/2] address comment

---
 server/conn.go | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/server/conn.go b/server/conn.go
index 5621bc113d772..761caa591a86a 100644
--- a/server/conn.go
+++ b/server/conn.go
@@ -960,7 +960,10 @@ func (cc *clientConn) handleLoadData(ctx context.Context, loadDataInfo *executor
 		err = loadDataInfo.Ctx.StmtCommit()
 	}
 
-	if txn, err1 := loadDataInfo.Ctx.Txn(true); err1 == nil {
+	var txn kv.Transaction
+	var err1 error
+	txn, err1 = loadDataInfo.Ctx.Txn(true)
+	if err1 == nil {
 		if txn != nil && txn.Valid() {
 			if err != nil {
 				if err1 := txn.Rollback(); err1 != nil {
@@ -968,11 +971,11 @@ func (cc *clientConn) handleLoadData(ctx context.Context, loadDataInfo *executor
 				}
 				return errors.Trace(err)
 			}
+			return errors.Trace(cc.ctx.CommitTxn(sessionctx.SetCommitCtx(ctx, loadDataInfo.Ctx)))
 		}
-	} else {
-		log.Error(err1)
 	}
-	return errors.Trace(cc.ctx.CommitTxn(sessionctx.SetCommitCtx(ctx, loadDataInfo.Ctx)))
+	// Should never reach here.
+	panic(err1)
 }
 
 // handleLoadStats does the additional work after processing the 'load stats' query.