代码生成模块优化

src/main/java/com/aizuda/boot/modules/gen/controller/GenDatabaseController.java

@@ -49,7 +49,7 @@ public class GenDatabaseController extends ApiController {
     @Permission("gen:database:get")
     @GetMapping("/get")
     public GenDatabase get(@RequestParam Long id) {
-        return genDatabaseService.getById(id);
+        return genDatabaseService.getNoPasswordById(id);
     }
 
     @Operation(summary = "根据 id 修改信息")

src/main/java/com/aizuda/boot/modules/gen/service/IGenDatabaseService.java

@@ -17,5 +17,7 @@ public interface IGenDatabaseService extends IBaseService<GenDatabase> {
 
     Page<GenDatabase> page(Page<GenDatabase> page, GenDatabase genDatabase);
 
+    GenDatabase getNoPasswordById(Long id);
+
     List<SelectOptionVO> listSelectOptions();
 }

src/main/java/com/aizuda/boot/modules/gen/service/IGenTemplateService.java

@@ -1,8 +1,10 @@
 package com.aizuda.boot.modules.gen.service;
 
-import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
-import com.aizuda.service.service.IBaseService;
 import com.aizuda.boot.modules.gen.entity.GenTemplate;
+import com.aizuda.service.service.IBaseService;
+import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
+
+import java.util.List;
 
 /**
  * 代码生成模板表 服务类
@@ -14,4 +16,5 @@ public interface IGenTemplateService extends IBaseService<GenTemplate> {
 
     Page<GenTemplate> page(Page<GenTemplate> page, GenTemplate genTemplate);
 
+    List<GenTemplate> listCheckByIds(List<Long> ids);
 }

src/main/java/com/aizuda/boot/modules/gen/service/impl/GenDatabaseServiceImpl.java

@@ -13,7 +13,6 @@ import com.baomidou.mybatisplus.core.toolkit.Wrappers;
 import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
 import org.springframework.stereotype.Service;
 
-import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -39,7 +38,7 @@ public class GenDatabaseServiceImpl extends BaseServiceImpl<GenDatabaseMapper, G
 
     @Override
     public boolean save(GenDatabase genDatabase) {
-//        this.checkDatabase(genDatabase);
+        this.checkDatabase(genDatabase);
         return super.save(genDatabase);
     }
 
@@ -51,12 +50,12 @@ public class GenDatabaseServiceImpl extends BaseServiceImpl<GenDatabaseMapper, G
     @Override
     public boolean updateById(GenDatabase genDatabase) {
         ApiAssert.fail(null == genDatabase.getId(), "主键不存在无法更新");
-//        this.checkDatabase(genDatabase);
+        this.checkDatabase(genDatabase);
         return super.updateById(genDatabase);
     }
 
     @Override
-    public GenDatabase getById(Serializable id) {
+    public GenDatabase getNoPasswordById(Long id) {
         GenDatabase genDatabase = super.getById(id);
         if (null != genDatabase) {
             // 屏蔽数据源密码

src/main/java/com/aizuda/boot/modules/gen/service/impl/GenTableServiceImpl.java

@@ -10,6 +10,7 @@ import com.aizuda.boot.modules.gen.service.IGenTemplateService;
 import com.aizuda.common.toolkit.CollectionUtils;
 import com.aizuda.core.api.ApiAssert;
 import com.aizuda.core.bean.BaseEntity;
+import com.aizuda.core.exception.ApiException;
 import com.aizuda.service.web.UserSession;
 import com.baomidou.mybatisplus.annotation.FieldFill;
 import com.baomidou.mybatisplus.core.toolkit.StringUtils;
@@ -50,14 +51,18 @@ public class GenTableServiceImpl implements IGenTableService {
 
     @Override
     public List<GenVO> preview(GenDTO dto) {
-        List<GenTemplate> genTemplates = genTemplateService.listByIds(dto.getTemplateIds());
-        ApiAssert.fail(CollectionUtils.isEmpty(genTemplates), "请选择生成模板");
+        List<GenTemplate> genTemplates = genTemplateService.listCheckByIds(dto.getTemplateIds());
 
-        // 生成预览内容
+        // 初始化数据连接配置
         ConfigBuilder configBuilder = this.buildConfigBuilder(dto);
-        TableInfo tableInfo = this.buildConfigBuilder(dto).getTableInfoList().get(0);
-        VelocityTemplateEngine templateEngine = new VelocityTemplateEngine();
-        templateEngine.init(configBuilder);
+
+        // 查询指定表信息
+        List<TableInfo> tableInfos = configBuilder.getTableInfoList();
+        ApiAssert.fail(CollectionUtils.isEmpty(tableInfos), "未找到指定表信息");
+        TableInfo tableInfo = tableInfos.get(0);
+
+        // 初始化模板引擎
+        VelocityTemplateEngine templateEngine = new VelocityTemplateEngine().init(configBuilder);
         return genTemplates.stream().map(t -> {
             GenVO vo = new GenVO();
             try {
@@ -75,35 +80,44 @@ public class GenTableServiceImpl implements IGenTableService {
 
     @Override
     public void download(HttpServletResponse response, GenDTO dto) {
-        List<GenTemplate> genTemplates = genTemplateService.listByIds(dto.getTemplateIds());
-        ApiAssert.fail(CollectionUtils.isEmpty(genTemplates), "请选择生成模板");
+        List<GenTemplate> genTemplates = genTemplateService.listCheckByIds(dto.getTemplateIds());
 
-        // 生成预览内容
+        // 初始化数据连接配置
         ConfigBuilder configBuilder = this.buildConfigBuilder(dto);
-        TableInfo tableInfo = this.buildConfigBuilder(dto).getTableInfoList().get(0);
-        VelocityTemplateEngine templateEngine = new VelocityTemplateEngine();
-        templateEngine.init(configBuilder);
+
+        // 查询指定表信息
+        List<TableInfo> tableInfos = configBuilder.getTableInfoList();
+
+        // 初始化模板引擎
+        VelocityTemplateEngine templateEngine = new VelocityTemplateEngine().init(configBuilder);
 
         // 设置响应内容类型
+        String fileName = "genCode";
+        if (tableInfos.size() < 2) {
+            fileName = dto.getTableName().replaceAll(",", "_");
+        }
+        response.setHeader("Content-Disposition", "attachment; filename=\"" + fileName + ".zip\"");
         response.setContentType("application/zip");
-        response.setHeader("Content-Disposition", "attachment; filename=\"generatedCode.zip\"");
+
         // 创建ZIP输出流
         try (ZipOutputStream zos = new ZipOutputStream(response.getOutputStream())) {
-            // 渲染模板
-            for (GenTemplate gt : genTemplates) {
-                // 创建输出流
-                ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
-                Map<String, Object> objectMap = this.getObjectMap(configBuilder, tableInfo);
-                String context = templateEngine.writer(objectMap, tableInfo.getEntityName(), gt.getTplContent());
-                byteArrayOutputStream.write(context.getBytes());
-
-                // 创建文件内容
-                String packageName = configBuilder.getPackageConfig().getParent();
-                ZipEntry zipEntry = new ZipEntry(packageName.replace('.', '/') + "/" +
-                        String.format(gt.getOutFile(), tableInfo.getEntityName()));
-                zos.putNextEntry(zipEntry);
-                zos.write(byteArrayOutputStream.toByteArray());
-                zos.closeEntry();
+            for (TableInfo ti : tableInfos) {
+                // 渲染模板
+                for (GenTemplate gt : genTemplates) {
+                    // 创建输出流
+                    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+                    Map<String, Object> objectMap = this.getObjectMap(configBuilder, ti);
+                    String context = templateEngine.writer(objectMap, ti.getEntityName(), gt.getTplContent());
+                    byteArrayOutputStream.write(context.getBytes());
+
+                    // 创建文件内容
+                    String packageName = configBuilder.getPackageConfig().getParent();
+                    ZipEntry zipEntry = new ZipEntry(packageName.replace('.', '/') + "/" +
+                            String.format(gt.getOutFile(), ti.getEntityName()));
+                    zos.putNextEntry(zipEntry);
+                    zos.write(byteArrayOutputStream.toByteArray());
+                    zos.closeEntry();
+                }
             }
         } catch (Exception e) {
             ApiAssert.fail("生成文件异常");
@@ -113,8 +127,7 @@ public class GenTableServiceImpl implements IGenTableService {
     private ConfigBuilder buildConfigBuilder(GenDTO dto) {
         DataSourceConfig.Builder dataSource;
         if (null != dto.getDatabaseId()) {
-            GenDatabase gb = genDatabaseService.getById(dto.getDatabaseId());
-            ApiAssert.fail(null == gb, "指定数据源不存在");
+            GenDatabase gb = genDatabaseService.checkById(dto.getDatabaseId());
             dataSource = new DataSourceConfig.Builder(gb.url(), gb.getUsername(), gb.getPassword());
         } else {
             dataSource = new DataSourceConfig.Builder(dsp.getUrl(), dsp.getUsername(), dsp.getPassword());
@@ -146,11 +159,15 @@ public class GenTableServiceImpl implements IGenTableService {
                         "update_by", "update_time", "deleted")
                 .logicDeleteColumnName("deleted")
                 .superClass(BaseEntity.class);
-        // 注入配置
-        InjectionConfig.Builder injection = new InjectionConfig.Builder();
-        // 模板渲染引擎 new CodeTemplateEngine()
-        return new ConfigBuilder(packageInfo.build(), dataSource.build(), strategy.build(), null,
-                globalConfig.build(), injection.build());
+        try {
+            // 注入配置
+            InjectionConfig.Builder injection = new InjectionConfig.Builder();
+            // 模板渲染引擎 new CodeTemplateEngine()
+            return new ConfigBuilder(packageInfo.build(), dataSource.build(), strategy.build(), null,
+                    globalConfig.build(), injection.build());
+        } catch (Throwable t) {
+            throw new ApiException("数据库连接初始化失败");
+        }
     }
 
     private Map<String, Object> getObjectMap(ConfigBuilder config, TableInfo tableInfo) {
@@ -235,7 +252,7 @@ public class GenTableServiceImpl implements IGenTableService {
                 customMap = new HashMap<>();
             }
             String va = annotations.toString();
-            if (org.apache.commons.lang3.StringUtils.isNotBlank(va)) {
+            if (StringUtils.isNotBlank(va)) {
                 va += lnt;
                 tableInfo.addImportPackages("com.aizuda.core.validation.Create");
             } else {

src/main/java/com/aizuda/boot/modules/gen/service/impl/GenTemplateServiceImpl.java

@@ -1,15 +1,18 @@
 package com.aizuda.boot.modules.gen.service.impl;
 
-import com.aizuda.core.api.ApiAssert;
 import com.aizuda.boot.modules.gen.entity.GenTemplate;
 import com.aizuda.boot.modules.gen.mapper.GenTemplateMapper;
 import com.aizuda.boot.modules.gen.service.IGenTemplateService;
+import com.aizuda.common.toolkit.CollectionUtils;
+import com.aizuda.core.api.ApiAssert;
 import com.aizuda.service.service.BaseServiceImpl;
 import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
 import com.baomidou.mybatisplus.core.toolkit.Wrappers;
 import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
 import org.springframework.stereotype.Service;
 
+import java.util.List;
+
 /**
  * 代码生成模板表 服务实现类
  *
@@ -25,6 +28,13 @@ public class GenTemplateServiceImpl extends BaseServiceImpl<GenTemplateMapper, G
         return super.page(page, lqw);
     }
 
+    @Override
+    public List<GenTemplate> listCheckByIds(List<Long> ids) {
+        List<GenTemplate> gts = super.listByIds(ids);
+        ApiAssert.fail(CollectionUtils.isEmpty(gts), "请选择生成模板");
+        return gts;
+    }
+
     @Override
     public boolean updateById(GenTemplate genTemplate) {
         ApiAssert.fail(null == genTemplate.getId(), "主键不存在无法更新");