Merge pull request kkangert#17 from zhaozhaonihao/main

添加产物下载
kenuoseclab · Jul 28, 2024 · 336a70e · 336a70e
2 parents b3ecc80 + d005c2a
commit 336a70e
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 69 deletions.
diff --git a/kspider-core/src/main/java/top/kangert/kspider/executor/node/event/OutputEventHandler.java b/kspider-core/src/main/java/top/kangert/kspider/executor/node/event/OutputEventHandler.java
@@ -15,7 +15,6 @@
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.stereotype.Component;
 
-import cn.hutool.core.date.DateUtil;
 import cn.hutool.core.exceptions.ExceptionUtil;
 import cn.hutool.core.util.RandomUtil;
 import cn.hutool.core.util.StrUtil;
@@ -141,7 +140,14 @@ private void outputCSV(OutputEventBean eventBean) {
                     printer = cachePrinter.get(key);
                     if (printer == null) {
                         CSVFormat format = CSVFormat.DEFAULT.withHeader(headers.toArray(new String[headers.size()]));
-                        String fileName = spiderConfig.getWorkspace() + File.separator + "files" + File.separator + node.getNodeId() + "_" + DateUtil.format(DateUtil.date(), "yyyyMMddHHmmss") + File.separator + csvName + ".csv";
+                        String fileName = spiderConfig.getWorkspace() + File.separator + "files";
+                        if (null != context.getTaskId() && context.getTaskId() > 0) {
+                            fileName += File.separator + "prod_flowid_taskid" + File.separator + context.getFlowId() + "_" + context.getTaskId();
+                        } else {
+                            fileName += File.separator + "test_flowid" + File.separator + context.getFlowId();
+                        }
+
+                        fileName += File.separator + csvName + ".csv";
                         File csvFile = new File(fileName);
                         if (csvFile.getParentFile() != null) {
                             csvFile.getParentFile().mkdirs();

diff --git a/kspider-core/src/main/java/top/kangert/kspider/service/impl/SpiderTaskServiceImpl.java b/kspider-core/src/main/java/top/kangert/kspider/service/impl/SpiderTaskServiceImpl.java
@@ -27,13 +27,13 @@
 import cn.hutool.core.convert.Convert;
 import cn.hutool.core.date.DateUtil;
 import cn.hutool.core.io.FileUtil;
+import cn.hutool.core.io.IoUtil;
 import cn.hutool.core.lang.TypeReference;
 import cn.hutool.core.util.StrUtil;
 import cn.hutool.core.util.ZipUtil;
 import cn.hutool.json.JSONUtil;
 import lombok.extern.slf4j.Slf4j;
 
-import java.io.BufferedInputStream;
 import java.io.BufferedOutputStream;
 import java.io.File;
 import java.io.IOException;
@@ -42,6 +42,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.io.FileInputStream;
 
 import javax.annotation.PostConstruct;
 import javax.annotation.Resource;
@@ -316,80 +317,35 @@ private void scheduleStart(Long taskId) {
 
     @Override
     public void download(Map<String, Object> params, HttpServletResponse response) {
-        checkParams(params, new String[] { "taskId" });
-        // 获取任务
-        Long taskId = Convert.toLong(params.get("taskId"));
-        SpiderTask spiderTask = queryItem(taskId);
-        String matedata = spiderTask.getMatedata();
-        Map<String, Object> variables = JSONUtil.toBean(matedata, new TypeReference<Map<String, Object>>() {
-        }, false);
-        String fileName = null;
-        String filePath = null;
-        BufferedInputStream inputStream = null;
-        BufferedOutputStream outputStream = null;
-
-        // 判断平台类型
-        if (variables.containsKey("classNo")) {
-            fileName = variables.get("classNo") + ".json";
-            // json文件路径
-            String jsonPath = spiderConfig.getWorkspace() + File.separator + "fileData" + File.separator + "json"
-                    + File.separator + fileName;
-            if (!FileUtil.exist(jsonPath)) {
-                throw new BaseException(ExceptionCodes.FILE_NOT_EXIST);
-            }
-            inputStream = FileUtil.getInputStream(jsonPath);
+        checkParams(params, new String[] { "flowId" });
+        Long flowId = Convert.toLong(params.get("flowId"));
+        String filePath = spiderConfig.getWorkspace() + File.separator + "files";
+        if (params.containsKey("taskId")) {
+            Long taskId = Convert.toLong(params.get("taskId"));
+            filePath += File.separator + "prod_flowid_taskid" + File.separator + flowId + "_" + taskId;
         } else {
-            // zip压缩路径
-            String enterpriseName = Convert.toStr(variables.get("enterpriseName"));
-            String className = Convert.toStr(variables.get("className"));
-            long currentTime = System.currentTimeMillis();
-            // 上级目录
-            String upCatalog = taskId + "_" + enterpriseName;
-            String srcPath = spiderConfig.getWorkspace() + File.separator + "fileData" + File.separator + upCatalog;
-            filePath = spiderConfig.getWorkspace() + File.separator + "fileData" + File.separator + "zip"
-                    + File.separator + enterpriseName + currentTime + ".zip";
-            if (!FileUtil.exist(srcPath)) {
-                throw new BaseException(ExceptionCodes.FILE_NOT_EXIST);
-            }
-            // 线上文件压缩
-            ZipUtil.zip(srcPath, filePath);
-            fileName = taskId + "-" + enterpriseName + "-" + className + ".zip";
-            inputStream = FileUtil.getInputStream(filePath);
+            filePath += File.separator + "test_flowid" + File.separator + flowId;
+        }
+
+        if (!FileUtil.exist(filePath)) {
+            throw new BaseException("文件不存在,请先启动!");
         }
 
-        try {
-            // 重置响应头信息
+        String zipPath = filePath + ".zip";
+        File zipFile = ZipUtil.zip(filePath, zipPath);
+        try (BufferedOutputStream outputStream = new BufferedOutputStream(response.getOutputStream());) {
             response.reset();
-            response.addHeader("Content-Disposition", "attachment;filename=" + URLEncoder.encode(fileName, "UTF-8"));
+            response.setHeader("Content-Disposition", "attachment;filename=" + URLEncoder.encode(zipFile.getName(), "UTF-8"));
             response.setContentType("application/octet-stream;charset=UTF-8");
-            outputStream = new BufferedOutputStream(response.getOutputStream());
-            byte[] readBytes = new byte[inputStream.available()];
-            inputStream.read(readBytes);
-            response.setHeader("Content-Length", "" + readBytes.length);
-            outputStream.write(readBytes);
+            FileInputStream inputStream = new FileInputStream(zipFile);
+            byte[] bytes = IoUtil.readBytes(inputStream);
+            outputStream.write(bytes);
             outputStream.flush();
-        } catch (Exception e) {
-            log.error(e.getMessage());
+        } catch (IOException e) {
+            log.error("文件导出失败", e.getCause());
             throw new BaseException(ExceptionCodes.FILE_EXPORT_FAILED);
         } finally {
-            if (outputStream != null) {
-                try {
-                    outputStream.close();
-                } catch (IOException e) {
-                    log.error(e.getMessage(), e);
-                }
-            }
-            if (inputStream != null) {
-                try {
-                    inputStream.close();
-                } catch (IOException e) {
-                    log.error(e.getMessage(), e);
-                }
-            }
-            // 删除线上文件压缩包
-            if (fileName.lastIndexOf("zip") > 0) {
-                FileUtil.del(filePath);
-            }
+            zipFile.delete();
         }
     }
 }
diff --git a/kspider-web/src/main/java/top/kangert/kspider/websocket/WebSocketEditorHandler.java b/kspider-web/src/main/java/top/kangert/kspider/websocket/WebSocketEditorHandler.java
@@ -35,13 +35,16 @@ protected void handleTextMessage(WebSocketSession session, TextMessage message)
         JSONObject event = JSONUtil.parseObj(message.getPayload());
         // 获取事件类型
         String eventType = event.getStr("eventType");
+        // 获取flowId
+        Long flowId = event.getLong("flowId");
         // 是 debug 类型的事件吗
         boolean isDebug = WebSocketEvent.DEBUG_EVENT_TYPE.equalsIgnoreCase(eventType);
         // test 类型事件或 debug 类型事件
         if (WebSocketEvent.TEST_EVENT_TYPE.equalsIgnoreCase(eventType) || isDebug) {
             // 创建 WebSocket 通信时使用的爬虫上下文
             context = new SpiderWebSocketContext(session);
             context.setDebug(isDebug);
+            context.setFlowId(flowId);
             context.setRunning(true);
             // 异步执行
             CompletableFuture.runAsync(() -> {