add: code generator for H2 database github#737

mybatis-plus-generator/src/main/java/com/baomidou/mybatisplus/generator/config/DataSourceConfig.java

@@ -15,17 +15,28 @@
  */
 package com.baomidou.mybatisplus.generator.config;
 
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+
 import com.baomidou.mybatisplus.annotation.DbType;
 import com.baomidou.mybatisplus.core.toolkit.ExceptionUtils;
-import com.baomidou.mybatisplus.generator.config.converts.*;
-import com.baomidou.mybatisplus.generator.config.querys.*;
+import com.baomidou.mybatisplus.generator.config.converts.DB2TypeConvert;
+import com.baomidou.mybatisplus.generator.config.converts.MySqlTypeConvert;
+import com.baomidou.mybatisplus.generator.config.converts.OracleTypeConvert;
+import com.baomidou.mybatisplus.generator.config.converts.PostgreSqlTypeConvert;
+import com.baomidou.mybatisplus.generator.config.converts.SqlServerTypeConvert;
+import com.baomidou.mybatisplus.generator.config.querys.DB2Query;
+import com.baomidou.mybatisplus.generator.config.querys.H2Query;
+import com.baomidou.mybatisplus.generator.config.querys.MariadbQuery;
+import com.baomidou.mybatisplus.generator.config.querys.MySqlQuery;
+import com.baomidou.mybatisplus.generator.config.querys.OracleQuery;
+import com.baomidou.mybatisplus.generator.config.querys.PostgreSqlQuery;
+import com.baomidou.mybatisplus.generator.config.querys.SqlServerQuery;
+
 import lombok.Data;
 import lombok.experimental.Accessors;
 
-import java.sql.Connection;
-import java.sql.DriverManager;
-import java.sql.SQLException;
-
 /**
  * <p>
  * 数据库配置
@@ -89,6 +100,8 @@ public class DataSourceConfig {
                 case MARIADB:
                     dbQuery = new MariadbQuery();
                     break;
+                case H2:
+                    dbQuery = new H2Query();
                 default:
                     // 默认 MYSQL
                     dbQuery = new MySqlQuery();
@@ -111,13 +124,13 @@ public class DataSourceConfig {
                 dbType = DbType.ORACLE;
             } else if (driverName.contains("postgresql")) {
                 dbType = DbType.POSTGRE_SQL;
-            } else if (driverName.contains("sqlserver")) {
-                dbType = DbType.SQL_SERVER;
             } else if (driverName.contains("db2")) {
                 dbType = DbType.DB2;
             } else if (driverName.contains("mariadb")) {
                 dbType = DbType.MARIADB;
-            } else {
+            } else if(driverName.contains("h2")){
+                dbType = DbType.H2;
+            }else {
                 throw ExceptionUtils.mpe("Unknown type of database!");
             }
         }

mybatis-plus-generator/src/main/java/com/baomidou/mybatisplus/generator/config/builder/ConfigBuilder.java

@@ -42,6 +42,7 @@ import com.baomidou.mybatisplus.generator.config.TemplateConfig;
 import com.baomidou.mybatisplus.generator.config.po.TableField;
 import com.baomidou.mybatisplus.generator.config.po.TableFill;
 import com.baomidou.mybatisplus.generator.config.po.TableInfo;
+import com.baomidou.mybatisplus.generator.config.querys.H2Query;
 import com.baomidou.mybatisplus.generator.config.rules.NamingStrategy;
 
 /**
@@ -571,31 +572,54 @@ public class ConfigBuilder {
         boolean haveId = false;
         List<TableField> fieldList = new ArrayList<>();
         List<TableField> commonFieldList = new ArrayList<>();
+        DbType dbType = dbQuery.dbType();
+        String tableName = tableInfo.getName();
         try {
             String tableFieldsSql = dbQuery.tableFieldsSql();
-            if (DbType.POSTGRE_SQL == dbQuery.dbType()) {
-                tableFieldsSql = String.format(tableFieldsSql, dataSourceConfig.getSchemaName(), tableInfo.getName());
-            } else if (DbType.ORACLE == dbQuery.dbType()) {
-                tableFieldsSql = String.format(tableFieldsSql.replace("#schema", dataSourceConfig.getSchemaName()), tableInfo.getName());
+            Set<String> h2PkColumns = new HashSet<>();
+            if (DbType.POSTGRE_SQL == dbType) {
+                tableFieldsSql = String.format(tableFieldsSql, dataSourceConfig.getSchemaName(), tableName);
+            } else if (DbType.ORACLE == dbType) {
+                tableName = tableName.toUpperCase();
+                tableFieldsSql = String.format(tableFieldsSql.replace("#schema", dataSourceConfig.getSchemaName()), tableName);
+            } else if (DbType.H2 == dbType) {
+                tableName = tableName.toUpperCase();
+                PreparedStatement pkQueryStmt = connection.prepareStatement(String.format(H2Query.PK_QUERY_SQL, tableName));
+                ResultSet pkResults = pkQueryStmt.executeQuery();
+                while (pkResults.next()) {
+                    String primaryKey = pkResults.getString(dbQuery.fieldKey());
+                    if ("TRUE".equalsIgnoreCase(primaryKey)) {
+                        h2PkColumns.add(pkResults.getString(dbQuery.fieldName()));
+                    }
+                }
+                pkResults.close();
+                pkQueryStmt.close();
+                tableFieldsSql = String.format(tableFieldsSql, tableName);
             } else {
-                tableFieldsSql = String.format(tableFieldsSql, tableInfo.getName());
+                tableFieldsSql = String.format(tableFieldsSql, tableName);
             }
             PreparedStatement preparedStatement = connection.prepareStatement(tableFieldsSql);
             ResultSet results = preparedStatement.executeQuery();
             while (results.next()) {
                 TableField field = new TableField();
-                String key = results.getString(dbQuery.fieldKey());
+                String columnName = results.getString(dbQuery.fieldName());
                 // 避免多重主键设置，目前只取第一个找到ID，并放到list中的索引为0的位置
                 boolean isId;
-                if (DbType.DB2 == dbQuery.dbType()) {
-                    isId = StringUtils.isNotEmpty(key) && "1".equals(key);
-                } else {
-                    isId = StringUtils.isNotEmpty(key) && "PRI".equals(key.toUpperCase());
+                if(DbType.H2 == dbType){
+                    isId = h2PkColumns.contains(columnName);
+                }else{
+                    String key = results.getString(dbQuery.fieldKey());
+                    if (DbType.DB2 == dbType) {
+                        isId = StringUtils.isNotEmpty(key) && "1".equals(key);
+                    } else {
+                        isId = StringUtils.isNotEmpty(key) && "PRI".equals(key.toUpperCase());
+                    }
                 }
+
                 // 处理ID
                 if (isId && !haveId) {
                     field.setKeyFlag(true);
-                    if (dbQuery.isKeyIdentity(results)) {
+                    if (DbType.H2 == dbType || dbQuery.isKeyIdentity(results)) {
                         field.setKeyIdentityFlag(true);
                     }
                     haveId = true;
@@ -612,7 +636,7 @@ public class ConfigBuilder {
                     field.setCustomMap(customMap);
                 }
                 // 处理其它信息
-                field.setName(results.getString(dbQuery.fieldName()));
+                field.setName(columnName);
                 field.setType(results.getString(dbQuery.fieldType()));
                 field.setPropertyName(strategyConfig, processName(field.getName(), strategy));
                 field.setColumnType(dataSourceConfig.getTypeConvert().processTypeConvert(globalConfig, field.getType()));
@@ -631,6 +655,8 @@ public class ConfigBuilder {
                 }
                 fieldList.add(field);
             }
+            results.close();
+            preparedStatement.close();
         } catch (SQLException e) {
             System.err.println("SQL Exception：" + e.getMessage());
         }

mybatis-plus-generator/src/main/java/com/baomidou/mybatisplus/generator/config/querys/H2Query.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright (c) 2011-2020, hubin (jobob@qq.com).
+ * <p>
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ * <p>
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+package com.baomidou.mybatisplus.generator.config.querys;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+
+import com.baomidou.mybatisplus.annotation.DbType;
+
+/**
+ * <p>
+ * MySql 表数据查询
+ * </p>
+ *
+ * @author yuxiaobin
+ * @since 2019-01-8
+ */
+public class H2Query extends AbstractDbQuery {
+
+    public static final String PK_QUERY_SQL = "select * from INFORMATION_SCHEMA.INDEXES WHERE TABLE_NAME = '%s'";
+
+
+    @Override
+    public DbType dbType() {
+        return DbType.H2;
+    }
+
+
+    @Override
+    public String tablesSql() {
+        return "SELECT * FROM INFORMATION_SCHEMA.TABLES";
+    }
+
+
+    @Override
+    public String tableFieldsSql() {
+        return "SELECT * FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME= '%s' ";
+    }
+
+
+    @Override
+    public String tableName() {
+        return "TABLE_NAME";
+    }
+
+
+    @Override
+    public String tableComment() {
+        return "REMARKS";
+    }
+
+
+    @Override
+    public String fieldName() {
+        return "COLUMN_NAME";
+    }
+
+
+    @Override
+    public String fieldType() {
+        return "TYPE_NAME";
+    }
+
+
+    @Override
+    public String fieldComment() {
+        return "REMARKS";
+    }
+
+
+    @Override
+    public String fieldKey() {
+        return "PRIMARY_KEY";
+    }
+
+
+    @Override
+    public boolean isKeyIdentity(ResultSet results) throws SQLException {
+        return "auto_increment".equals(results.getString("Extra"));
+    }
+}

mybatis-plus/src/test/java/com/baomidou/mybatisplus/test/h2/GenerateCode4H2Database.java

@@ -0,0 +1,75 @@
+package com.baomidou.mybatisplus.test.h2;
+
+import org.junit.FixMethodOrder;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.MethodSorters;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
+
+import com.baomidou.mybatisplus.annotation.DbType;
+import com.baomidou.mybatisplus.generator.AutoGenerator;
+import com.baomidou.mybatisplus.generator.config.DataSourceConfig;
+import com.baomidou.mybatisplus.generator.config.GlobalConfig;
+import com.baomidou.mybatisplus.generator.config.PackageConfig;
+import com.baomidou.mybatisplus.generator.config.StrategyConfig;
+import com.baomidou.mybatisplus.generator.config.querys.H2Query;
+import com.baomidou.mybatisplus.generator.config.rules.NamingStrategy;
+
+/**
+ * <p>
+ * </p>
+ *
+ * @author yuxiaobin
+ * @date 2019/1/8
+ */
+@FixMethodOrder(MethodSorters.JVM)
+@RunWith(SpringJUnit4ClassRunner.class)
+@ContextConfiguration(locations = {"classpath:h2/spring-test-h2.xml"})
+public class GenerateCode4H2Database extends BaseTest{
+
+    @Test
+    public void test() {
+        String dbUrl = "jdbc:h2:mem:test;MODE=mysql";
+        DataSourceConfig dataSourceConfig = new DataSourceConfig();
+        dataSourceConfig.setDbType(DbType.H2)
+            .setDbQuery(new H2Query())
+            .setUrl(dbUrl)
+            .setUsername("sa")
+            .setPassword("")
+            .setDriverName("org.h2.Driver");
+        generate("com.baomidou.mp.h2", dataSourceConfig, "h2user");
+    }
+
+    public void generate(String packageName, DataSourceConfig dataSourceConfig, String tableNames) {
+        boolean serviceClassNameStartWithI = false;
+        boolean enableTableFieldAnnotation = true;
+        GlobalConfig config = new GlobalConfig();
+        StrategyConfig strategyConfig = new StrategyConfig();
+        strategyConfig
+            .setCapitalMode(true)
+            .setEntityLombokModel(false)
+            // .setDbColumnUnderline(true) 改为如下 2 个配置
+            .setNaming(NamingStrategy.underline_to_camel)
+            .setColumnNaming(NamingStrategy.underline_to_camel)
+            .entityTableFieldAnnotationEnable(enableTableFieldAnnotation)
+            .setInclude(tableNames);//修改替换成你需要的表名，多个表名传数组
+        config.setActiveRecord(false)
+            .setAuthor("K神带你飞")
+            .setOutputDir("d:\\codeGen")
+            .setFileOverride(true);
+        if (!serviceClassNameStartWithI) {
+            config.setServiceName("%sService");
+        }
+        new AutoGenerator().setGlobalConfig(config)
+            .setDataSource(dataSourceConfig)
+            .setStrategy(strategyConfig)
+            .setPackageInfo(
+                new PackageConfig()
+                    .setParent(packageName)
+                    .setController("controller")
+                    .setEntity("entity")
+            ).execute();
+    }
+
+}

mybatis-plus/src/test/resources/h2/user.ddl.sql

@@ -11,4 +11,7 @@ CREATE TABLE IF NOT EXISTS  h2user (
 	deleted INT(1) NULL DEFAULT 0 ,
 	PRIMARY KEY (test_id)
 );
+COMMENT ON COLUMN h2user.test_id IS 'PK';
+COMMENT ON COLUMN h2user.name IS 'USERNAME';
+