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 默认对含主键的表生成两个类:
*
* - XXXPO extends XXXPOKey(业务字段)
* - XXXPOKey(仅主键字段)
*
- * 本插件通过将 primaryKeyColumns 合并到 baseColumns 中并清空 PK 列表,
- * 使 MBG 只生成一个包含所有字段的 PO 类。
+ * 本插件实现:
+ *
+ * - 将 primaryKeyColumns 合并到 baseColumns(PO 类含所有字段)
+ * - 抑制 *Key 类文件生成(返回 false)
+ * - 不清空 primaryKeyColumns(保持 MBG 对主键感知 → 生成 selectByPrimaryKey)
+ * - 重写 selectByPrimaryKey 的参数类型:*Key → *PO
+ *
*/
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 @@
+
+