diff --git a/src/main/java/com/example/workutils/biz/mybatisgenerator/service/generator/SuppressKeyPlugin.java b/src/main/java/com/example/workutils/biz/mybatisgenerator/service/generator/SuppressKeyPlugin.java index ea15bb3..a98adbf 100644 --- a/src/main/java/com/example/workutils/biz/mybatisgenerator/service/generator/SuppressKeyPlugin.java +++ b/src/main/java/com/example/workutils/biz/mybatisgenerator/service/generator/SuppressKeyPlugin.java @@ -3,6 +3,13 @@ package com.example.workutils.biz.mybatisgenerator.service.generator; import org.mybatis.generator.api.IntrospectedColumn; import org.mybatis.generator.api.IntrospectedTable; import org.mybatis.generator.api.PluginAdapter; +import org.mybatis.generator.api.dom.java.FullyQualifiedJavaType; +import org.mybatis.generator.api.dom.java.Interface; +import org.mybatis.generator.api.dom.java.Method; +import org.mybatis.generator.api.dom.java.Parameter; +import org.mybatis.generator.api.dom.java.TopLevelClass; +import org.mybatis.generator.api.dom.xml.Attribute; +import org.mybatis.generator.api.dom.xml.XmlElement; import java.lang.reflect.Field; import java.util.ArrayList; @@ -10,15 +17,20 @@ import java.util.List; import java.util.Properties; /** - * 自定义插件:抑制 *Key 类的生成,将主键字段直接合并到 PO 类中。 + * 自定义插件:抑制 *Key 类的生成,将主键字段合并到 PO 类中。 *

* MBG 默认对含主键的表生成两个类: *

- * 本插件通过将 primaryKeyColumns 合并到 baseColumns 中并清空 PK 列表, - * 使 MBG 只生成一个包含所有字段的 PO 类。 + * 本插件实现: + *
    + *
  1. 将 primaryKeyColumns 合并到 baseColumns(PO 类含所有字段)
  2. + *
  3. 抑制 *Key 类文件生成(返回 false)
  4. + *
  5. 不清空 primaryKeyColumns(保持 MBG 对主键感知 → 生成 selectByPrimaryKey)
  6. + *
  7. 重写 selectByPrimaryKey 的参数类型:*Key → *PO
  8. + *
*/ public class SuppressKeyPlugin extends PluginAdapter { @@ -35,18 +47,14 @@ public class SuppressKeyPlugin extends PluginAdapter { return true; } - /** - * 将主键字段合并到基础字段中,阻止 MBG 生成独立的 *Key 类。 - */ + // ==================== 初始化:合并主键列到 PO ==================== + @Override public void initialized(IntrospectedTable introspectedTable) { if (!suppressKey) return; - try { - // 通过反射获取 primaryKeyColumns 和 baseColumns Field pkField = getField(introspectedTable.getClass(), "primaryKeyColumns"); Field baseField = getField(introspectedTable.getClass(), "baseColumns"); - if (pkField == null || baseField == null) return; pkField.setAccessible(true); @@ -59,30 +67,82 @@ public class SuppressKeyPlugin extends PluginAdapter { if (pkColumns == null || pkColumns.isEmpty()) return; - // 将主键列全部移到 baseColumns 开头 + // PK 列合并到 baseColumns 开头(去重) List merged = new ArrayList<>(pkColumns); if (baseColumns != null) { - merged.addAll(baseColumns); + for (IntrospectedColumn bc : baseColumns) { + boolean isPk = false; + for (IntrospectedColumn pk : pkColumns) { + if (bc.getActualColumnName().equals(pk.getActualColumnName())) { + isPk = true; + break; + } + } + if (!isPk) merged.add(bc); + } } baseField.set(introspectedTable, merged); - pkField.set(introspectedTable, new ArrayList<>()); - - } catch (Exception e) { - // 反射失败时静默忽略,回退到默认行为 + // ★ 不清空 primaryKeyColumns + } catch (Exception ignored) { } } + // ==================== 抑制 *Key 类文件 ==================== + + @Override + public boolean modelPrimaryKeyClassGenerated(TopLevelClass topLevelClass, + IntrospectedTable introspectedTable) { + return !suppressKey; + } + + // ==================== 重写 selectByPrimaryKey 参数 ==================== + + /** + * MBG 构造的 selectByPrimaryKey 参数是 XxxPOKey key。 + * 由于 Parameter 不可变,清空参数列表后重建为 XxxPO record。 + */ + @Override + public boolean clientSelectByPrimaryKeyMethodGenerated(Method method, + Interface interfaze, IntrospectedTable introspectedTable) { + if (!suppressKey) return true; + + FullyQualifiedJavaType recordType = new FullyQualifiedJavaType( + introspectedTable.getBaseRecordType()); + method.getParameters().clear(); + method.getParameters().add(new Parameter(recordType, "record")); + interfaze.addImportedType(recordType); + return true; + } + + /** + * MBG 构造的 XML parameterType 是 XxxPOKey,替换为 XxxPO。 + * 由于 Attribute 不可变,用 List.set() 替换旧条目。 + */ + @Override + public boolean sqlMapSelectByPrimaryKeyElementGenerated(XmlElement element, + IntrospectedTable introspectedTable) { + if (!suppressKey) return true; + + String baseRecordType = introspectedTable.getBaseRecordType(); + List attrs = element.getAttributes(); + for (int i = 0; i < attrs.size(); i++) { + if ("parameterType".equals(attrs.get(i).getName())) { + attrs.set(i, new Attribute("parameterType", baseRecordType)); + } + } + return true; + } + + // ================================================================= + private static Field getField(Class clazz, String name) { - Field field; try { - field = clazz.getDeclaredField(name); + return clazz.getDeclaredField(name); } catch (NoSuchFieldException e) { - // 可能字段在父类中 if (clazz.getSuperclass() != null) { return getField(clazz.getSuperclass(), name); } return null; } - return field; } } diff --git a/src/main/resources/mybatis-generator.xml b/src/main/resources/mybatis-generator.xml index ccf1fae..9ccebc2 100644 --- a/src/main/resources/mybatis-generator.xml +++ b/src/main/resources/mybatis-generator.xml @@ -11,6 +11,8 @@ + +