diff --git a/src/main/java/com/southern/power/grid/listener/DataExcelListener.java b/src/main/java/com/southern/power/grid/listener/DataExcelListener.java index 373de2c..71c87bf 100644 --- a/src/main/java/com/southern/power/grid/listener/DataExcelListener.java +++ b/src/main/java/com/southern/power/grid/listener/DataExcelListener.java @@ -1,5 +1,6 @@ package com.southern.power.grid.listener; +import com.alibaba.excel.EasyExcel; import com.alibaba.excel.context.AnalysisContext; import com.alibaba.excel.event.AnalysisEventListener; import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper; @@ -13,9 +14,11 @@ import lombok.extern.slf4j.Slf4j; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Component; +import java.io.InputStream; import java.util.ArrayList; import java.util.List; import java.util.Objects; +import java.util.concurrent.atomic.AtomicInteger; /** * 数据导入逻辑 -- easyexcel事件监听器 @@ -51,7 +54,6 @@ public class DataExcelListener extends AnalysisEventListener { @Override public void invoke(DataExcelEntity data, AnalysisContext context) { if (!Objects.isNull(data)) { - //TODO 待完善 等志明拿到具体数据再导入到数据库中 cacheList.add(data); // 达到批次数量 → 插入数据库 if (cacheList.size() >= BATCH_SIZE) { @@ -107,7 +109,25 @@ public class DataExcelListener extends AnalysisEventListener { importTaskMapper.update(null, new LambdaUpdateWrapper() .set(ImportTask::getSuccessCount, success) // 成功数 .set(ImportTask::getFailCount, fail) // 失败数 - .set(ImportTask::getTotal, success + fail) // 总数据量 .eq(ImportTask::getTaskNo, task.getTaskNo())); } + + /** + * 获取 Excel 总条数(只读行数,不处理数据) + */ + public int getTotalCount(InputStream inputStream) { + AtomicInteger totalCount = new AtomicInteger(0); + + EasyExcel.read(inputStream, DataExcelEntity.class, new AnalysisEventListener() { + @Override + public void invoke(DataExcelEntity data, AnalysisContext context) { + totalCount.incrementAndGet(); // 每读一行 +1 + } + + @Override + public void doAfterAllAnalysed(AnalysisContext context) {} + }).sheet().doRead(); + + return totalCount.get(); + } } diff --git a/src/main/java/com/southern/power/grid/service/impl/AsyncImportServiceImpl.java b/src/main/java/com/southern/power/grid/service/impl/AsyncImportServiceImpl.java index da3e841..fca7ee6 100644 --- a/src/main/java/com/southern/power/grid/service/impl/AsyncImportServiceImpl.java +++ b/src/main/java/com/southern/power/grid/service/impl/AsyncImportServiceImpl.java @@ -33,13 +33,17 @@ public class AsyncImportServiceImpl { @Async("asyncExecutor") public void doAsyncImport(InputStream inputStream, ImportTask task) { try { + // 获取总条数,先读一遍获取总条数 + int total = dataExcelListener.getTotalCount(inputStream); // 更新为导入中 importTaskMapper.update(null, new LambdaUpdateWrapper() .set(ImportTask::getStatus, ImportTaskStatusEnum.PROCESSING.getCode()) + .set(ImportTask::getTotal, total) .eq(ImportTask::getTaskNo, task.getTaskNo())); // 注入任务 dataExcelListener.setTask(task); + inputStream.reset(); // 流被读过一次了,重置一下 // 采用easyExcel流式读取,一行一行读,不加载整个excel文件,防止OOM EasyExcel.read(inputStream, DataExcelEntity.class, dataExcelListener) .sheet()