!14 完成导出数据库文档

Merge pull request !14 from timfruit/feature/screw-export-dbdoc

!14 完成导出数据库文档
Merge pull request !14 from timfruit/feature/screw-export-dbdoc
芋道源码 · Gitee
2 parents 4dc484f5 562f4cb9
Showing 1 changed file with 68 additions and 16 deletions
src/main/java/cn/iocoder/dashboard/modules/infra/controller/doc/InfDbDocController.java
 package cn.iocoder.dashboard.modules.infra.controller.doc;
+import cn.hutool.core.io.FileUtil;
+import cn.hutool.core.io.IoUtil;
+import cn.hutool.core.lang.UUID;
+import cn.hutool.core.util.IdUtil;
 import cn.hutool.extra.servlet.ServletUtil;
+import cn.iocoder.dashboard.util.servlet.ServletUtils;
 import cn.smallbun.screw.core.Configuration;
 import cn.smallbun.screw.core.engine.EngineConfig;
 import cn.smallbun.screw.core.engine.EngineFileType;
@@ -10,18 +15,20 @@ import cn.smallbun.screw.core.process.ProcessConfig;
 import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
 import io.swagger.annotations.Api;
+import io.swagger.annotations.ApiImplicitParam;
+import io.swagger.annotations.ApiImplicitParams;
+import io.swagger.annotations.ApiOperation;
 import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
 import org.springframework.http.MediaType;
+import org.springframework.util.StreamUtils;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RequestMapping;
+import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.RestController;
 import javax.annotation.Resource;
 import javax.servlet.http.HttpServletResponse;
-import javax.sql.DataSource;
-import java.io.File;
-import java.io.FileInputStream;
-import java.io.FileNotFoundException;
+import java.io.*;
 import java.util.Collections;
 @Api(tags = "数据库文档")
@@ -34,34 +41,79 @@ public class InfDbDocController {
     private static final String FILE_OUTPUT_DIR = System.getProperty("java.io.tmpdir") + File.separator
             + "db-doc";
-    private static final EngineFileType FILE_OUTPUT_TYPE = EngineFileType.HTML; // 可以设置 Word 或者 Markdown 格式
     private static final String DOC_FILE_NAME = "数据库文档";
     private static final String DOC_VERSION = "1.0.0";
     private static final String DOC_DESCRIPTION = "文档描述";
-    @Resource
-    private DataSource dataSource;
     @GetMapping("/export-html")
-    public synchronized void exportHtml(HttpServletResponse response) throws FileNotFoundException {
+    @ApiOperation("导出html格式的数据文档")
+    @ApiImplicitParam(name = "deleteFile", value = "是否删除在服务器本地生成的数据库文档", example = "true", dataTypeClass = Boolean.class)
+    public void exportHtml(@RequestParam(defaultValue = "true") Boolean deleteFile,
+                           HttpServletResponse response) throws IOException {
+        doExportFile(EngineFileType.HTML, deleteFile, response);
+    }
+
+    @GetMapping("/export-word")
+    @ApiOperation("导出word格式的数据文档")
+    @ApiImplicitParam(name = "deleteFile", value = "是否删除在服务器本地生成的数据库文档", example = "true", dataTypeClass = Boolean.class)
+    public void exportWord(@RequestParam(defaultValue = "true") Boolean deleteFile,
+                           HttpServletResponse response) throws IOException {
+        doExportFile(EngineFileType.WORD, deleteFile, response);
+    }
+
+    @GetMapping("/export-markdown")
+    @ApiOperation("导出markdown格式的数据文档")
+    @ApiImplicitParam(name = "deleteFile", value = "是否删除在服务器本地生成的数据库文档", example = "true", dataTypeClass = Boolean.class)
+    public void exportMarkdown(@RequestParam(defaultValue = "true") Boolean deleteFile,
+                               HttpServletResponse response) throws IOException {
+        doExportFile(EngineFileType.MD, deleteFile, response);
+    }
+
+    private void doExportFile(EngineFileType fileOutputType, Boolean deleteFile,
+                              HttpServletResponse response) throws IOException {
+        String docFileName = DOC_FILE_NAME + "_" + IdUtil.fastSimpleUUID();
+        String filePath = doExportFile(fileOutputType, docFileName);
+        String downloadFileName = DOC_FILE_NAME + fileOutputType.getFileSuffix(); //下载后的文件名
+        try {
+            // 读取，返回
+            //这里不用hutool工具类，它的中文文件名编码有问题,导致在浏览器下载时有问题
+            ServletUtils.writeAttachment(response, downloadFileName, FileUtil.readBytes(filePath));
+        }finally {
+            handleDeleteFile(deleteFile, filePath);
+        }
+    }
+
+    /**
+     * 输出文件，返回文件路径
+     *
+     * @param fileOutputType 文件类型
+     * @param fileName       文件名, 无需 ".docx" 等文件后缀
+     * @return 生成的文件所在路径
+     */
+    private String doExportFile(EngineFileType fileOutputType, String fileName) {
         try (HikariDataSource dataSource = buildDataSource()) {
             // 创建 screw 的配置
             Configuration config = Configuration.builder()
                     .version(DOC_VERSION)  // 版本
                     .description(DOC_DESCRIPTION) // 描述
                     .dataSource(dataSource) // 数据源
-                    .engineConfig(buildEngineConfig()) // 引擎配置
+                    .engineConfig(buildEngineConfig(fileOutputType, fileName)) // 引擎配置
                     .produceConfig(buildProcessConfig()) // 处理配置
                     .build();
             // 执行 screw，生成数据库文档
             new DocumentationExecute(config).execute();
-            // 读取，返回
-            ServletUtil.write(response,
-                    new FileInputStream(FILE_OUTPUT_DIR + File.separator + DOC_FILE_NAME + FILE_OUTPUT_TYPE.getFileSuffix()),
-                    MediaType.TEXT_HTML_VALUE);
+            return FILE_OUTPUT_DIR + File.separator + fileName + fileOutputType.getFileSuffix();
+        }
+    }
+
+    private void handleDeleteFile(Boolean deleteFile, String filePath) {
+        if (!deleteFile) {
+            return;
         }
+        FileUtil.del(filePath);
     }
     /**
@@ -83,13 +135,13 @@ public class InfDbDocController {
     /**
      * 创建 screw 的引擎配置
      */
-    private static EngineConfig buildEngineConfig() {
+    private static EngineConfig buildEngineConfig(EngineFileType fileOutputType, String docFileName) {
         return EngineConfig.builder()
                 .fileOutputDir(FILE_OUTPUT_DIR) // 生成文件路径
                 .openOutputDir(false) // 打开目录
-                .fileType(FILE_OUTPUT_TYPE) // 文件类型
+                .fileType(fileOutputType) // 文件类型
                 .produceType(EngineTemplateType.freemarker) // 文件类型
-                .fileName(DOC_FILE_NAME) // 自定义文件名称
+                .fileName(docFileName) // 自定义文件名称
                 .build();
     }