Merge #153111

153111: import: fix bulk summary override bug when import with multiple files r=ZhouXing19 a=ZhouXing19

When the import source is more than 1 csv file, there is a risk that the bulk summary is incorrectly reset after the first csv file progress is processed, leading to the actual progress being overriden with a nil value. This commit is to fix it.

More internal discussion can be found [here](https://cockroachlabs.slack.com/archives/C02DSDS9TM1/p1756991625692929).

Epic: None
Release note (bug fix): fix bulk summary override bug with IMPORT with multiple files.


Co-authored-by: ZhouXing19 <[email protected]>

Author: craig[bot]

pkg/sql/importer/import_processor.go

@@ -431,10 +431,10 @@ func ingestKvs(
 	pkIndexAdder.SetOnFlush(func(summary kvpb.BulkOpSummary) {
 		for i, emitted := range writtenRow {
 			atomic.StoreInt64(&pkFlushedRow[i], emitted)
-			bulkSummaryMu.Lock()
-			bulkSummaryMu.summary.Add(summary)
-			bulkSummaryMu.Unlock()
 		}
+		bulkSummaryMu.Lock()
+		bulkSummaryMu.summary.Add(summary)
+		bulkSummaryMu.Unlock()
 		if indexAdder.IsEmpty() {
 			for i, emitted := range writtenRow {
 				atomic.StoreInt64(&idxFlushedRow[i], emitted)
@@ -444,10 +444,10 @@ func ingestKvs(
 	indexAdder.SetOnFlush(func(summary kvpb.BulkOpSummary) {
 		for i, emitted := range writtenRow {
 			atomic.StoreInt64(&idxFlushedRow[i], emitted)
-			bulkSummaryMu.Lock()
-			bulkSummaryMu.summary.Add(summary)
-			bulkSummaryMu.Unlock()
 		}
+		bulkSummaryMu.Lock()
+		bulkSummaryMu.summary.Add(summary)
+		bulkSummaryMu.Unlock()
 	})
 
 	// offsets maps input file ID to a slot in our progress tracking slices.
@@ -473,12 +473,12 @@ func ingestKvs(
 				prog.ResumePos[file] = idx
 			}
 			prog.CompletedFraction[file] = math.Float32frombits(atomic.LoadUint32(&writtenFraction[offset]))
-			// Write down the summary of how much we've ingested since the last update.
-			bulkSummaryMu.Lock()
-			prog.BulkSummary = bulkSummaryMu.summary
-			bulkSummaryMu.summary.Reset()
-			bulkSummaryMu.Unlock()
 		}
+		// Write down the summary of how much we've ingested since the last update.
+		bulkSummaryMu.Lock()
+		prog.BulkSummary = bulkSummaryMu.summary
+		bulkSummaryMu.summary.Reset()
+		bulkSummaryMu.Unlock()
 		select {
 		case progCh <- prog:
 		case <-ctx.Done():

pkg/sql/importer/import_processor_test.go

@@ -774,6 +774,11 @@ func TestCSVImportCanBeResumed(t *testing.T) {
 	sqlDB.CheckQueryResults(t, `SELECT id FROM t ORDER BY id`,
 		sqlDB.QueryStr(t, `SELECT generate_series(0, $1)`, csv1.numRows-1),
 	)
+
+	js = queryJobUntil(t, sqlDB.DB, jobID, func(js jobState) bool { return js.prog.Summary.EntryCounts != nil })
+	for _, e := range js.prog.Summary.EntryCounts {
+		require.Equal(t, int64(csv1.numRows), e)
+	}
 }
 
 func TestCSVImportMarksFilesFullyProcessed(t *testing.T) {
@@ -875,6 +880,11 @@ func TestCSVImportMarksFilesFullyProcessed(t *testing.T) {
 
 	// Verify that after resume we have not processed any additional rows.
 	assert.Zero(t, importSummary.Rows)
+
+	js = queryJobUntil(t, sqlDB.DB, jobID, func(js jobState) bool { return js.prog.Summary.EntryCounts != nil })
+	for _, e := range js.prog.Summary.EntryCounts {
+		require.Equal(t, int64(csv1.numRows+csv2.numRows+csv3.numRows), e)
+	}
 }
 
 func (ses *generatedStorage) externalStorageFactory() cloud.ExternalStorageFactory {