Make data-branch keywords context-aware (#23197)

Make data-branch keywords context-aware to avoid conflict with the user's input SQLs.

Approved by: @XuPeng-SH, @iamlinjunhong

Author: gouhongshen

pkg/frontend/data_branch.go

@@ -577,7 +577,9 @@ func tryFlushDeletesOrReplace(
 	ses *Session,
 	bh BackgroundExec,
 	tblStuff tableStuff,
-	newWrapped *batchWithKind,
+	newKind string,
+	newValsLen int,
+	newRowCnt int,
 	deleteCnt *int,
 	deletesBuf *bytes.Buffer,
 	replaceCnt *int,
@@ -611,14 +613,15 @@ func tryFlushDeletesOrReplace(
 	}
 
 	// if wrapped is nil, means force flush
-	if newWrapped != nil {
-		newSize := newWrapped.batch.Size()
-		if newWrapped.kind == diffDelete {
-			if deletesBuf.Len()+newSize >= maxSqlBatchSize || *deleteCnt >= maxSqlBatchCnt {
+	if newKind != "" {
+		if newKind == diffDelete {
+			if deletesBuf.Len()+newValsLen >= maxSqlBatchSize ||
+				*deleteCnt+newRowCnt >= maxSqlBatchCnt {
 				return flushDeletes()
 			}
 		} else {
-			if replaceBuf.Len()+newSize >= maxSqlBatchSize || *replaceCnt >= maxSqlBatchCnt {
+			if replaceBuf.Len()+newValsLen >= maxSqlBatchSize ||
+				*replaceCnt+newRowCnt >= maxSqlBatchCnt {
 				return flushReplace()
 			}
 		}
@@ -658,11 +661,13 @@ func mergeDiffs(
 		replaceIntoVals = acquireBuffer()
 		deleteFromVals  = acquireBuffer()
 		firstErr        error
+		tmpValsBuffer   = acquireBuffer()
 	)
 
 	defer func() {
 		releaseBuffer(replaceIntoVals)
 		releaseBuffer(deleteFromVals)
+		releaseBuffer(tmpValsBuffer)
 	}()
 
 	defer func() {
@@ -682,18 +687,36 @@ func mergeDiffs(
 			continue
 		}
 
-		if err = tryFlushDeletesOrReplace(
-			ctx, ses, bh, tblStuff, &wrapped,
-			&deleteCnt, deleteFromVals, &replaceCnt, replaceIntoVals, nil,
+		tmpValsBuffer.Reset()
+
+		if err = constructValsFromBatch(
+			ctx, ses, tblStuff, wrapped, tmpValsBuffer, tmpValsBuffer,
 		); err != nil {
 			firstErr = err
 			cancel()
 			tblStuff.retPool.releaseRetBatch(wrapped.batch, false)
 			continue
 		}
 
-		if err = constructValsFromBatch(
-			ctx, ses, tblStuff, wrapped, deleteFromVals, replaceIntoVals,
+		if tmpValsBuffer.Len() == 0 {
+			tblStuff.retPool.releaseRetBatch(wrapped.batch, false)
+			continue
+		}
+
+		newRowCnt := wrapped.batch.RowCount()
+		newValsLen := tmpValsBuffer.Len()
+		if wrapped.kind == diffDelete {
+			if deleteFromVals.Len() > 0 {
+				newValsLen++
+			}
+		} else {
+			if replaceIntoVals.Len() > 0 {
+				newValsLen++
+			}
+		}
+		if err = tryFlushDeletesOrReplace(
+			ctx, ses, bh, tblStuff, wrapped.kind, newValsLen, newRowCnt,
+			&deleteCnt, deleteFromVals, &replaceCnt, replaceIntoVals, nil,
 		); err != nil {
 			firstErr = err
 			cancel()
@@ -702,17 +725,25 @@ func mergeDiffs(
 		}
 
 		if wrapped.kind == diffDelete {
-			deleteCnt += wrapped.batch.RowCount()
+			if deleteFromVals.Len() > 0 {
+				deleteFromVals.WriteString(",")
+			}
+			deleteFromVals.Write(tmpValsBuffer.Bytes())
+			deleteCnt += newRowCnt
 		} else {
-			replaceCnt += wrapped.batch.RowCount()
+			if replaceIntoVals.Len() > 0 {
+				replaceIntoVals.WriteString(",")
+			}
+			replaceIntoVals.Write(tmpValsBuffer.Bytes())
+			replaceCnt += newRowCnt
 		}
 
 		tblStuff.retPool.releaseRetBatch(wrapped.batch, false)
 	}
 
 	if err = tryFlushDeletesOrReplace(
-		ctx, ses, bh, tblStuff, nil,
-		&deleteCnt, deleteFromVals, &replaceCnt, replaceIntoVals, nil,
+		ctx, ses, bh, tblStuff, "",
+		0, 0, &deleteCnt, deleteFromVals, &replaceCnt, replaceIntoVals, nil,
 	); err != nil {
 		if firstErr == nil {
 			firstErr = err
@@ -836,6 +867,7 @@ func satisfyDiffOutputOpt(
 
 			deleteFromValsBuffer  = acquireBuffer()
 			replaceIntoValsBuffer = acquireBuffer()
+			tmpValsBuffer         = acquireBuffer()
 
 			fileHint     string
 			fullFilePath string
@@ -854,6 +886,7 @@ func satisfyDiffOutputOpt(
 			}
 			releaseBuffer(deleteFromValsBuffer)
 			releaseBuffer(replaceIntoValsBuffer)
+			releaseBuffer(tmpValsBuffer)
 		}()
 
 		if fullFilePath, fileHint, writeFile, release, cleanup, err = prepareFSForDiffAsFile(
@@ -881,29 +914,56 @@ func satisfyDiffOutputOpt(
 			}
 
 			if wrapped.name == tblStuff.tarRel.GetTableName() {
-				if err = tryFlushDeletesOrReplace(
-					ctx, ses, bh, tblStuff, &wrapped,
-					&deleteCnt, deleteFromValsBuffer, &replaceCnt, replaceIntoValsBuffer,
-					writeFile,
+				tmpValsBuffer.Reset()
+
+				if err = constructValsFromBatch(
+					ctx, ses, tblStuff, wrapped, tmpValsBuffer, tmpValsBuffer,
 				); err != nil {
 					first = err
 					cancel()
 					tblStuff.retPool.releaseRetBatch(wrapped.batch, false)
 					continue
 				}
 
-				if err = constructValsFromBatch(
-					ctx, ses, tblStuff, wrapped, deleteFromValsBuffer, replaceIntoValsBuffer,
+				if tmpValsBuffer.Len() == 0 {
+					tblStuff.retPool.releaseRetBatch(wrapped.batch, false)
+					continue
+				}
+
+				newRowCnt := wrapped.batch.RowCount()
+				newValsLen := tmpValsBuffer.Len()
+				if wrapped.kind == diffDelete {
+					if deleteFromValsBuffer.Len() > 0 {
+						newValsLen++
+					}
+				} else {
+					if replaceIntoValsBuffer.Len() > 0 {
+						newValsLen++
+					}
+				}
+				if err = tryFlushDeletesOrReplace(
+					ctx, ses, bh, tblStuff, wrapped.kind, newValsLen, newRowCnt,
+					&deleteCnt, deleteFromValsBuffer, &replaceCnt, replaceIntoValsBuffer,
+					writeFile,
 				); err != nil {
 					first = err
 					cancel()
 					tblStuff.retPool.releaseRetBatch(wrapped.batch, false)
 					continue
 				}
+
 				if wrapped.kind == diffDelete {
-					deleteCnt += int(wrapped.batch.RowCount())
+					if deleteFromValsBuffer.Len() > 0 {
+						deleteFromValsBuffer.WriteString(",")
+					}
+					deleteFromValsBuffer.Write(tmpValsBuffer.Bytes())
+					deleteCnt += newRowCnt
 				} else {
-					replaceCnt += int(wrapped.batch.RowCount())
+					if replaceIntoValsBuffer.Len() > 0 {
+						replaceIntoValsBuffer.WriteString(",")
+					}
+					replaceIntoValsBuffer.Write(tmpValsBuffer.Bytes())
+					replaceCnt += newRowCnt
 				}
 			}
 
@@ -915,8 +975,8 @@ func satisfyDiffOutputOpt(
 		}
 
 		if err = tryFlushDeletesOrReplace(
-			ctx, ses, bh, tblStuff, nil,
-			&deleteCnt, deleteFromValsBuffer, &replaceCnt, replaceIntoValsBuffer,
+			ctx, ses, bh, tblStuff, "",
+			0, 0, &deleteCnt, deleteFromValsBuffer, &replaceCnt, replaceIntoValsBuffer,
 			writeFile,
 		); err != nil {
 			first = err