增加referenceBean功能

gradle/wrapper/gradle-wrapper.properties

@@ -1,5 +1,5 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://downloads.gradle-dn.com/distributions/gradle-6.9-bin.zip
+distributionUrl=./gradle-6.9-bin.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists

src/main/java/com/sjhy/plugin/dto/TableInfoDTO.java

@@ -47,6 +47,7 @@ private TableInfoDTO(PsiClass psiClass) {
         this.savePackageName = "";
         this.savePath = "";
         this.saveModelName = "";
+        this.referenceBean = "";
         this.fullColumn = new ArrayList<>();
         for (PsiField field : psiClass.getAllFields()) {
             this.fullColumn.add(new ColumnInfoDTO(field));
@@ -61,6 +62,7 @@ private TableInfoDTO(DbTable dbTable) {
         this.savePackageName = "";
         this.savePath = "";
         this.saveModelName = "";
+        this.referenceBean = "";
         this.fullColumn = new ArrayList<>();
         // 处理所有列
         JBIterable<? extends DasColumn> columns = DasUtil.getColumns(dbTable);
@@ -73,6 +75,9 @@ private static void merge(TableInfoDTO oldData, TableInfoDTO newData) {
         if (oldData == null || CollectionUtil.isEmpty(oldData.getFullColumn())) {
             return;
         }
+        if (!StringUtils.isEmpty(oldData.getReferenceBean())) {
+            newData.referenceBean = oldData.getReferenceBean();
+        }
         if (!StringUtils.isEmpty(oldData.getPreName())) {
             newData.preName = oldData.getPreName();
         }
@@ -187,6 +192,10 @@ private static void merge(TableInfoDTO oldData, TableInfoDTO newData) {
      * 保存的model名称
      */
     private String saveModelName;
+    /**
+     * 参考的bean名
+     */
+    private String referenceBean;
 
     public TableInfo toTableInfo(PsiClass psiClass) {
         TableInfo tableInfo = new TableInfo();
@@ -235,6 +244,7 @@ public TableInfo toTableInfo(DbTable dbTable) {
         tableInfo.setFullColumn(new ArrayList<>());
         tableInfo.setPkColumn(new ArrayList<>());
         tableInfo.setOtherColumn(new ArrayList<>());
+        tableInfo.setReferenceBean(this.getReferenceBean());
         // 列
         JBIterable<? extends DasColumn> columns = DasUtil.getColumns(dbTable);
         Map<String, DasColumn> nameToObj = new HashMap<>(columns.size());
@@ -271,6 +281,7 @@ public static TableInfoDTO valueOf(TableInfo tableInfo) {
         dto.setComment(tableInfo.getComment());
         dto.setSavePackageName(tableInfo.getSavePackageName());
         dto.setSaveModelName(tableInfo.getSaveModelName());
+        dto.setReferenceBean(tableInfo.getReferenceBean());
         dto.setFullColumn(new ArrayList<>());
         // 处理列
         for (ColumnInfo columnInfo : tableInfo.getFullColumn()) {

src/main/java/com/sjhy/plugin/entity/TableInfo.java

@@ -69,4 +69,13 @@ public class TableInfo {
      * 保存的model名称
      */
     private String saveModelName;
+    /**
+     * 参考的bean名
+     */
+    private String referenceBean;
+    /**
+     * 保存的包名称
+     */
+    private String saveFullPackageName;
+
 }

src/main/java/com/sjhy/plugin/service/impl/CodeGenerateServiceImpl.java

@@ -21,6 +21,9 @@
 import com.sjhy.plugin.service.TableInfoSettingsService;
 import com.sjhy.plugin.tool.*;
 
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
 import java.util.*;
 import java.util.stream.Collectors;
 
@@ -46,6 +49,10 @@ public class CodeGenerateServiceImpl implements CodeGenerateService {
      * 缓存数据工具
      */
     private CacheDataUtils cacheDataUtils;
+    /**
+     * 名称工具
+     */
+    private NameUtils nameUtils;
     /**
      * 导入包时过滤的包前缀
      */
@@ -54,6 +61,7 @@ public class CodeGenerateServiceImpl implements CodeGenerateService {
     public CodeGenerateServiceImpl(Project project) {
         this.project = project;
         this.moduleManager = ModuleManager.getInstance(project);
+        this.nameUtils = NameUtils.getInstance();
         this.tableInfoService = TableInfoSettingsService.getInstance();
         this.cacheDataUtils = CacheDataUtils.getInstance();
     }
@@ -76,6 +84,15 @@ public void generate(Collection<Template> templates, GenerateOptions generateOpt
             selectedTableInfo = tableInfoService.getTableInfo(cacheDataUtils.getSelectDbTable());
             tableInfoList = cacheDataUtils.getDbTableList().stream().map(item -> tableInfoService.getTableInfo(item)).collect(Collectors.toList());
         }
+        FindRelevancyBeanFileVisitor fileVisitor = null;
+        if (!StringUtils.isEmpty(selectedTableInfo.getReferenceBean())) {
+            fileVisitor = new FindRelevancyBeanFileVisitor(selectedTableInfo.getReferenceBean());
+            try {
+                Files.walkFileTree(Paths.get(project.getBasePath()), fileVisitor);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
         // 校验选中表的保存路径是否正确
         if (StringUtils.isEmpty(selectedTableInfo.getSavePath())) {
             if (selectedTableInfo.getObj() != null) {
@@ -110,21 +127,23 @@ public void generate(Collection<Template> templates, GenerateOptions generateOpt
         }
 
         // 生成代码
-        generate(templates, tableInfoList, generateOptions, null);
+        generate(templates, tableInfoList, fileVisitor, generateOptions, null);
     }
 
     /**
      * 生成代码，并自动保存到对应位置
      *
      * @param templates       模板
      * @param tableInfoList   表信息对象
+     * @param fileVisitor
      * @param generateOptions 生成配置
      * @param otherParam      其他参数
      */
-    public void generate(Collection<Template> templates, Collection<TableInfo> tableInfoList, GenerateOptions generateOptions, Map<String, Object> otherParam) {
+    public void generate(Collection<Template> templates, Collection<TableInfo> tableInfoList, FindRelevancyBeanFileVisitor fileVisitor, GenerateOptions generateOptions, Map<String, Object> otherParam) {
         if (CollectionUtil.isEmpty(templates) || CollectionUtil.isEmpty(tableInfoList)) {
             return;
         }
+        boolean isReferenceBean = fileVisitor != null;
         // 处理模板，注入全局变量（克隆一份，防止篡改）
         templates = CloneUtils.cloneByJson(templates, new TypeReference<ArrayList<Template>>() {
         });
@@ -151,25 +170,45 @@ public void generate(Collection<Template> templates, Collection<TableInfo> table
             // 设置额外代码生成服务
             param.put("generateService", new ExtraCodeGenerateUtils(this, tableInfo, generateOptions));
             for (Template template : templates) {
+                String templateName = template.getName().replace(".vm", "");
+                boolean hasReferenceBean = isReferenceBean && fileVisitor.isExist(templateName);
+
                 Callback callback = new Callback();
                 callback.setWriteFile(true);
                 callback.setReformat(generateOptions.getReFormat());
-                // 默认名称
-                callback.setFileName(tableInfo.getName() + "Default.java");
-                // 默认路径
-                callback.setSavePath(tableInfo.getSavePath());
-                // 设置回调对象
+                // 默认路径和默认名称
+                if (hasReferenceBean) {
+                    callback.setSavePath(fileVisitor.getPath(templateName));
+                    callback.setFileName(fileVisitor.getFileName(templateName).replace(fileVisitor.getReferenceBeanName(), tableInfo.getName()));
+                    if (!templateName.contains("xml")) {
+                        tableInfo.setSavePackageName(fileVisitor.getCutOffPackage(templateName));
+                        tableInfo.setSaveFullPackageName(fileVisitor.getPackage(templateName));
+                    }
+                }else {
+                    callback.setSavePath(tableInfo.getSavePath());
+                    callback.setFileName(tableInfo.getName() + "Default.java");
+                }
                 param.put("callback", callback);
+
+
                 // 开始生成
                 String code = VelocityUtils.generate(template.getCode(), param);
-                // 设置一个默认保存路径与默认文件名
-                String path = callback.getSavePath();
-                path = path.replace("\\", "/");
-                // 针对相对路径进行处理
-                if (path.startsWith(".")) {
-                    path = project.getBasePath() + path.substring(1);
+
+
+                // 重新填充callback. 有可能在生成代码的过程中，被模板方法修改
+                if (hasReferenceBean) {
+                    callback.setSavePath(fileVisitor.getPath(templateName));
+                }else {
+                    // 设置一个默认保存路径与默认文件名
+                    String path = callback.getSavePath();
+                    path = path.replace("\\", "/");
+                    // 针对相对路径进行处理
+                    if (path.startsWith(".")) {
+                        path = project.getBasePath() + path.substring(1);
+                    }
+                    callback.setSavePath(path);
                 }
-                callback.setSavePath(path);
+
                 new SaveFile(project, code, callback, generateOptions).write();
             }
         }

src/main/java/com/sjhy/plugin/tool/ExtraCodeGenerateUtils.java

@@ -54,6 +54,6 @@ public void run(String templateName, Map<String, Object> param) {
             return;
         }
         // 生成代码
-        codeGenerateService.generate(Collections.singletonList(currTemplate), Collections.singletonList(this.tableInfo), this.generateOptions, param);
+        codeGenerateService.generate(Collections.singletonList(currTemplate), Collections.singletonList(this.tableInfo), null, this.generateOptions, param);
     }
 }

src/main/java/com/sjhy/plugin/tool/FindRelevancyBeanFileVisitor.java

@@ -0,0 +1,141 @@
+package com.sjhy.plugin.tool;
+
+import com.intellij.util.ExceptionUtil;
+import org.apache.commons.lang3.StringUtils;
+
+import java.io.IOException;
+import java.nio.file.FileVisitResult;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Stream;
+
+/**
+ * @author Created by cgb
+ */
+public class FindRelevancyBeanFileVisitor extends SimpleFileVisitor<Path> {
+    /**
+     * 文件路径
+     */
+    Map<String, String> rbNameAndPathMap = new HashMap<>();
+    /**
+     * 文件截断的packet 兼容旧模板
+     * 保存的是:com.sjhy.plugin
+     */
+    Map<String, String> rbNameAndCutOffPackageMap = new HashMap<>();
+    /**
+     * 文件全packet
+     * 保存的是:package com.sjhy.plugin.tool;
+     */
+    Map<String, String> rbNameAndPackageMap = new HashMap<>();
+    /**
+     * 文件名称
+     */
+    Map<String, String> rbNameAndFileNameMap = new HashMap<>();
+    String referenceBeanName;
+
+    NameUtils nameUtils = NameUtils.getInstance();
+
+    public FindRelevancyBeanFileVisitor(String referenceBeanName) {
+        // 转换为类名, 兼容下划线和首字母小写
+        if (referenceBeanName.contains("_")) {
+            this.referenceBeanName = nameUtils.getClassName(referenceBeanName);
+        } else {
+            this.referenceBeanName = referenceBeanName;
+        }
+
+        this.referenceBeanName = nameUtils.firstUpperCase(this.referenceBeanName);
+    }
+
+    @Override
+    public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) {
+        String fileSourceName = file.getFileName().toString();
+
+        // TODO 后续可以配置忽略的文件和文件夹
+        if (fileSourceName.contains(".class")) {
+            return FileVisitResult.SKIP_SUBTREE;
+        }
+
+        if (fileSourceName.contains(referenceBeanName)) {
+            // 去掉类名中的bean名称. 获取模板名称并小写 DefaultBeanDto.java => DefaultDto.java => defaultDto.java
+            String templateName = nameUtils.firstLowerCase(fileSourceName.replace(referenceBeanName, ""));
+
+            rbNameAndPathMap.put(templateName, file.getParent().toString());
+            rbNameAndPackageMap.put(templateName, readPackageLine(file));
+            rbNameAndCutOffPackageMap.put(templateName, cutOffPackage(rbNameAndPackageMap.get(templateName)));
+            rbNameAndFileNameMap.put(templateName, fileSourceName);
+        }
+        return FileVisitResult.SKIP_SUBTREE;
+    }
+
+    /**
+     * 读取文件的第一行
+     *
+     * @param file 文件
+     * @return 第一行
+     */
+    private String readPackageLine(Path file) {
+        try (Stream<String> lines = Files.lines(file)) {
+            Optional<String> first = lines
+                    .filter(StringUtils::isNotBlank)
+                    .findFirst();
+
+            return first.orElse(StringUtils.EMPTY);
+
+        } catch (IOException e) {
+            e.printStackTrace();
+            ExceptionUtil.rethrow(e);
+        }
+
+        return StringUtils.EMPTY;
+    }
+
+    /**
+     * 截断package
+     * 兼容以前的模板
+     *
+     * @param packageStr 完整package(package com.sjhy.plugin.tool;)
+     * @return 截断的package(com.sjhy.plugin)
+     */
+    private String cutOffPackage(String packageStr) {
+        // impl的要多截取一次
+        if (packageStr.endsWith(".impl;")) {
+            packageStr = packageStr.replace(".impl;", "").trim();
+        }
+
+        int lastIndex = packageStr.lastIndexOf(".");
+        if (lastIndex != -1) {
+            packageStr = packageStr.substring(0, lastIndex);
+        }
+        return packageStr.replace("package", "").trim();
+    }
+
+    public boolean isExist(String name) {
+        return rbNameAndPathMap.containsKey(name);
+    }
+
+
+    public String getPath(String name) {
+        return rbNameAndPathMap.get(name);
+    }
+
+    public String getCutOffPackage(String name) {
+        return rbNameAndCutOffPackageMap.get(name);
+    }
+
+    public String getPackage(String name) {
+        return rbNameAndPackageMap.get(name);
+    }
+
+    public String getFileName(String name) {
+        return rbNameAndFileNameMap.get(name);
+    }
+
+    public String getReferenceBeanName() {
+        return referenceBeanName;
+    }
+}