001package cn.mybatis.mp.generator.config;
002
003
004import db.sql.api.DbType;
005import lombok.Getter;
006
007import javax.sql.DataSource;
008import java.sql.Connection;
009import java.sql.DriverManager;
010import java.sql.SQLException;
011import java.util.Properties;
012
013@Getter
014public class DataBaseConfig {
015
016    private final DbType dbType;
017    private final DataSource dataSource;
018    private final String url;
019    private final String username;
020    private final String password;
021    private String schema;
022    private String databaseName;
023
024    public DataBaseConfig(DbType dbType, DataSource dataSource) {
025        this.dataSource = dataSource;
026        this.dbType = dbType;
027
028        this.url = null;
029        this.username = null;
030        this.password = null;
031    }
032
033    public DataBaseConfig(String url, String username, String password) {
034
035        this.dbType = getDbType(url);
036        this.url = url;
037        this.username = username;
038        this.password = password;
039
040        this.dataSource = null;
041    }
042
043    public static Connection getConnection(String url, String username, String password) {
044        Properties properties = new Properties();
045        properties.put("user", username);
046        properties.put("password", password);
047        DbType dbType = getDbType(url);
048        addAdditionalJdbcProperties(properties, dbType);
049        try {
050            return DriverManager.getConnection(url, properties);
051        } catch (SQLException e) {
052            throw new RuntimeException(e);
053        }
054    }
055
056    private static DbType getDbType(String jdbcUrl) {
057        if (jdbcUrl.contains(":mysql:") || jdbcUrl.contains(":cobar:")) {
058            return DbType.MYSQL;
059        } else if (jdbcUrl.contains(":oracle:")) {
060            return DbType.ORACLE;
061        } else if (jdbcUrl.contains(":postgresql:")) {
062            return DbType.PGSQL;
063        } else if (jdbcUrl.contains(":sqlserver:")) {
064            return DbType.SQL_SERVER;
065        } else {
066            throw new RuntimeException("Unrecognized database type");
067        }
068    }
069
070    private static void addAdditionalJdbcProperties(Properties properties, DbType dbType) {
071        switch (dbType) {
072            case MYSQL:
073                properties.put("remarks", "true");
074                properties.put("useInformationSchema", "true");
075                break;
076            case ORACLE:
077                properties.put("remarks", "true");
078                properties.put("remarksReporting", "true");
079                break;
080            default: {
081
082            }
083        }
084    }
085
086    public DataBaseConfig schema(String schema) {
087        this.schema = schema;
088        return this;
089    }
090
091    public DataBaseConfig databaseName(String databaseName) {
092        this.databaseName = databaseName;
093        return this;
094    }
095
096    public Connection getConnection() {
097        try {
098            if (this.dataSource != null) {
099                return this.getDataSource().getConnection();
100            } else {
101                return getConnection(this.url, this.username, this.password);
102            }
103        } catch (SQLException e) {
104            throw new RuntimeException(e);
105        }
106    }
107}