Skip to content

Commit

Permalink
fix pgsql table_column ddl
Browse files Browse the repository at this point in the history
  • Loading branch information
zgq committed May 10, 2024
1 parent 4d1b91d commit a753f49
Show file tree
Hide file tree
Showing 2 changed files with 123 additions and 19 deletions.
103 changes: 93 additions & 10 deletions ...ins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLMetaData.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import ai.chat2db.spi.util.SortUtils;
import com.google.common.collect.Lists;
import jakarta.validation.constraints.NotEmpty;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.sql.Connection;
Expand All @@ -19,8 +20,6 @@
import java.util.*;
import java.util.stream.Collectors;

import static ai.chat2db.plugin.postgresql.consts.SQLConst.DROP_TYPE_SQL;
import static ai.chat2db.plugin.postgresql.consts.SQLConst.TABLE_DEF_FUNCTION_SQL;
import static ai.chat2db.spi.util.SortUtils.sortDatabase;

public class PostgreSQLMetaData extends DefaultMetaService implements MetaData {
Expand Down Expand Up @@ -102,17 +101,101 @@ public List<Trigger> triggers(Connection connection, String databaseName, String
});
}

public static final String SELECT_TABLE_COLUMN_SQL = """
SELECT
ic.*,
d.description AS column_comment
FROM information_schema.columns ic
LEFT JOIN
pg_namespace n ON n.nspname = ic.table_schema
LEFT JOIN
pg_class c ON c.relname = ic.table_name AND c.relnamespace = n.oid
LEFT JOIN
pg_attribute a ON a.attrelid = c.oid AND a.attname = ic.column_name
LEFT JOIN
pg_description d ON d.objoid = c.oid AND d.objsubid = a.attnum
WHERE ic.table_schema = '%s'
AND ic.table_name = '%s'
ORDER BY ic.ordinal_position;""";

@Override
public String tableDDL(Connection connection, String databaseName, String schemaName, String tableName) {
SQLExecutor.getInstance().execute(connection, String.format(DROP_TYPE_SQL,schemaName,"tabledefs"), resultSet -> null);
SQLExecutor.getInstance().execute(connection, TABLE_DEF_FUNCTION_SQL, resultSet -> null);
String ddlSql = String.format("select * from pg_get_tabledef('%s','%s',false,'COMMENTS') as ddl;", schemaName, tableName);
return SQLExecutor.getInstance().execute(connection, ddlSql, resultSet -> {
if (resultSet.next()) {
return resultSet.getString("ddl");
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("CREATE TABLE ");
sqlBuilder.append("\"").append(tableName).append("\"").append(" (").append(" ").append("\n");
try (ResultSet columnInfo = connection.createStatement().executeQuery(String.format(SELECT_TABLE_COLUMN_SQL, schemaName, tableName));
) {
List<TableColumn> tableColumns = new ArrayList<>();
while (columnInfo.next()) {
TableColumn tableColumn = new TableColumn();
tableColumn.setDatabaseName(databaseName);
tableColumn.setSchemaName(schemaName);
tableColumn.setTableName(tableName);
tableColumn.setName(columnInfo.getString("column_name"));
tableColumn.setDefaultValue(columnInfo.getString("column_default"));
tableColumn.setNullable(columnInfo.getString("is_nullable").equals("YES") ? 1 : 0);
tableColumn.setComment(columnInfo.getString("column_comment"));
String dataType = columnInfo.getString("data_type").toUpperCase();
if (PostgreSQLColumnTypeEnum.CHARACTER_VARYING.getColumnType().getTypeName().equals(dataType)) {
tableColumn.setColumnType(PostgreSQLColumnTypeEnum.VARCHAR.name());
} else if (PostgreSQLColumnTypeEnum.CHARACTER.getColumnType().getTypeName().equals(dataType)) {
tableColumn.setColumnType(PostgreSQLColumnTypeEnum.CHAR.name());
} else if (PostgreSQLColumnTypeEnum.TIME_WITHOUT_TIME_ZONE.getColumnType().getTypeName().equals(dataType)) {
tableColumn.setColumnType(PostgreSQLColumnTypeEnum.TIME.name());
} else if (PostgreSQLColumnTypeEnum.TIMESTAMP_WITHOUT_TIME_ZONE.getColumnType().getTypeName().equals(dataType)) {
tableColumn.setColumnType(PostgreSQLColumnTypeEnum.TIMESTAMP.name());
} else if (PostgreSQLColumnTypeEnum.SMALLINT.getColumnType().getTypeName().equals(dataType)) {
String defaultValue = tableColumn.getDefaultValue();
if (StringUtils.isNotBlank(defaultValue) && defaultValue.contains("nextval")) {
tableColumn.setColumnType(PostgreSQLColumnTypeEnum.SMALLSERIAL.name());
tableColumn.setDefaultValue(null);
} else {
tableColumn.setColumnType(PostgreSQLColumnTypeEnum.SMALLINT.name());
}
} else if (PostgreSQLColumnTypeEnum.INTEGER.getColumnType().getTypeName().equals(dataType)) {
String defaultValue = tableColumn.getDefaultValue();
if (StringUtils.isNotBlank(defaultValue) && defaultValue.contains("nextval")) {
tableColumn.setColumnType(PostgreSQLColumnTypeEnum.SERIAL.name());
tableColumn.setDefaultValue(null);
} else {
tableColumn.setColumnType(PostgreSQLColumnTypeEnum.INTEGER.name());
}
} else {
tableColumn.setColumnType(dataType);
}
if (Arrays.asList(PostgreSQLColumnTypeEnum.BIT.name(),
PostgreSQLColumnTypeEnum.VARBIT.name(),
PostgreSQLColumnTypeEnum.CHAR.name(),
PostgreSQLColumnTypeEnum.VARCHAR.getColumnType().getTypeName())
.contains(tableColumn.getColumnType())) {
tableColumn.setColumnSize(columnInfo.getInt("character_maximum_length"));
}
if (Arrays.asList(PostgreSQLColumnTypeEnum.NUMERIC.name(),
PostgreSQLColumnTypeEnum.DECIMAL.name())
.contains(tableColumn.getColumnType())) {
tableColumn.setColumnSize(columnInfo.getInt("numeric_precision"));
tableColumn.setDecimalDigits(columnInfo.getInt("numeric_scale"));
}
if (Arrays.asList(PostgreSQLColumnTypeEnum.TIME.name(),
PostgreSQLColumnTypeEnum.TIMESTAMP.name(),
PostgreSQLColumnTypeEnum.TIME_WITH_TIME_ZONE.getColumnType().getTypeName(),
PostgreSQLColumnTypeEnum.TIMESTAMP_WITH_TIME_ZONE.getColumnType().getTypeName())
.contains(tableColumn.getColumnType())) {
tableColumn.setColumnSize(columnInfo.getInt("datetime_precision"));
}
PostgreSQLColumnTypeEnum typeEnum = PostgreSQLColumnTypeEnum.getByType(tableColumn.getColumnType());
sqlBuilder.append("\t").append(typeEnum.buildCreateColumnSql(tableColumn)).append(",\n");
tableColumns.add(tableColumn);
}
return null;
});
if (CollectionUtils.isEmpty(tableColumns)) {
return sqlBuilder.toString();
}
sqlBuilder = new StringBuilder(sqlBuilder.substring(0, sqlBuilder.length() - 2));
sqlBuilder.append("\n)\n");
} catch (SQLException e) {
throw new RuntimeException(e);
}
return sqlBuilder.toString();
}


Expand Down
39 changes: 30 additions & 9 deletions ...-postgresql/src/main/java/ai/chat2db/plugin/postgresql/type/PostgreSQLColumnTypeEnum.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,20 +15,28 @@ public enum PostgreSQLColumnTypeEnum implements ColumnBuilder {

BIGSERIAL("BIGSERIAL", false, false, true, false, false, false, true, true, false, false),
BIT("BIT", true, false, true, false, false, false, true, true, false, false),
BIT_VARYING("BIT VARYING", true, false, true, false, false, false, true, true, false, false),

BOOL("BOOL", false, false, true, false, false, false, true, true, false, false),
BOX("BOX", false, false, true, false, false, false, true, true, false, false),
BYTEA("BYTEA", false, false, true, false, false, false, true, true, false, false),
CHAR("CHAR", true, false, true, false, false, true, true, true, false, false),
CHARACTER("CHARACTER", true, false, true, false, false, true, true, true, false, false),
CHARACTER_VARYING("CHARACTER VARYING", true, false, true, false, false, true, true, true, false, false),
CIDR("CIDR", false, false, true, false, false, false, true, true, false, false),
CIRCLE("CIRCLE", false, false, true, false, false, false, true, true, false, false),
DATE("DATE", false, false, true, false, false, false, true, true, false, false),
DECIMAL("DECIMAL", true, false, true, false, false, false, true, true, false, false),
FLOAT4("FLOAT4", false, false, true, false, false, false, true, true, false, false),
REAL("REAL", false, false, true, false, false, false, true, true, false, false),
FLOAT8("FLOAT8", false, false, true, false, false, false, true, true, false, false),
DOUBLE_PRECISION("DOUBLE PRECISION", false, false, true, false, false, false, true, true, false, false),
INET("INET", false, false, true, false, false, false, true, true, false, false),
INT2("INT2", false, false, true, false, false, false, true, true, false, false),
SMALLINT("SMALLINT", false, false, true, false, false, false, true, true, false, false),
INT4("INT4", false, false, true, false, false, false, true, true, false, false),
INT8("INT8", false, false, true, false, false, false, true, true, false, false),
BIG_INT("BIGINT", false, false, true, false, false, false, true, true, false, false),
INTERVAL("INTERVAL", false, false, true, false, false, false, true, true, false, false),
JSON("JSON", false, false, true, false, false, false, true, true, false, false),
JSONB("JSONB", false, false, true, false, false, false, true, true, false, false),
Expand All @@ -41,13 +49,18 @@ public enum PostgreSQLColumnTypeEnum implements ColumnBuilder {
POINT("POINT", false, false, true, false, false, false, true, true, false, false),
POLYGON("POLYGON", false, false, true, false, false, false, true, true, false, false),
SERIAL("SERIAL", false, false, true, false, false, false, true, true, false, false),
INTEGER("INTEGER", false, false, true, false, false, false, true, true, false, false),
SERIAL2("SERIAL2", false, false, true, false, false, false, true, true, false, false),
SERIAL4("SERIAL4", false, false, true, false, false, false, true, true, false, false),
SERIAL8("SERIAL8", false, false, true, false, false, false, true, true, false, false),
SMALLSERIAL("SMALLSERIAL", false, false, true, false, false, false, true, true, false, false),
TEXT("TEXT", false, false, true, false, false, true, true, true, false, false),
TIME("TIME", true, false, true, false, false, false, true, true, false, false),
TIME_WITHOUT_TIME_ZONE("TIME WITHOUT TIME ZONE", true, false, true, false, false, false, true, true, false, false),
TIME_WITH_TIME_ZONE("TIME WITH TIME ZONE", true, false, true, false, false, false, true, true, false, false),
TIMESTAMP("TIMESTAMP", true, false, true, false, false, false, true, true, false, false),
TIMESTAMP_WITHOUT_TIME_ZONE("TIMESTAMP WITHOUT TIME ZONE", true, false, true, false, false, false, true, true, false, false),
TIMESTAMP_WITH_TIME_ZONE("TIMESTAMP WITH TIME ZONE", true, false, true, false, false, false, true, true, false, false),
TIMESTAMPTZ("TIMESTAMPTZ", true, false, true, false, false, false, true, true, false, false),
TIMETZ("TIMETZ", true, false, true, false, false, false, true, true, false, false),
TSQUERY("TSQUERY", false, false, true, false, false, false, true, true, false, false),
Expand All @@ -57,8 +70,7 @@ public enum PostgreSQLColumnTypeEnum implements ColumnBuilder {
VARBIT("VARBIT", true, false, true, false, false, false, true, true, false, false),
VARCHAR("VARCHAR", true, false, true, false, false, true, true, true, false, false),
XML("XML", false, false, true, false, false, false, true, true, false, false),

;
BOOLEAN("BOOLEAN", false, false, true, false, false, false, true, true, false, false);;

private static Map<String, PostgreSQLColumnTypeEnum> COLUMN_TYPE_MAP = Maps.newHashMap();

Expand Down Expand Up @@ -152,19 +164,19 @@ public String buildComment(TableColumn column, PostgreSQLColumnTypeEnum type) {
return "";
}
return StringUtils.join("COMMENT ON COLUMN", " \"", column.getTableName(),
"\".\"", column.getName(), "\" IS '", column.getComment(), "';");
"\".\"", column.getName(), "\" IS '", column.getComment(), "';");
}

private String buildDefaultValue(TableColumn column, PostgreSQLColumnTypeEnum type) {
if (!type.getColumnType().isSupportDefaultValue() || StringUtils.isEmpty(column.getDefaultValue())) {
return "";
}

if("EMPTY_STRING".equalsIgnoreCase(column.getDefaultValue().trim())){
if ("EMPTY_STRING".equalsIgnoreCase(column.getDefaultValue().trim())) {
return StringUtils.join("DEFAULT ''");
}

if("NULL".equalsIgnoreCase(column.getDefaultValue().trim())){
if ("NULL".equalsIgnoreCase(column.getDefaultValue().trim())) {
return StringUtils.join("DEFAULT NULL");
}

Expand Down Expand Up @@ -196,27 +208,36 @@ private String buildNullable(TableColumn column, PostgreSQLColumnTypeEnum type)
private String buildDataType(TableColumn column, PostgreSQLColumnTypeEnum type) {
String columnType = type.columnType.getTypeName();
if (Arrays.asList(VARCHAR, CHAR).contains(type)) {
if (column.getColumnSize() == null ) {
if (column.getColumnSize() == null) {
return columnType;
}
return StringUtils.join(columnType, "(", column.getColumnSize(), ")");
}

if (Arrays.asList(VARBIT, BIT).contains(type)) {
if (column.getColumnSize() == null ) {
if (Arrays.asList(VARBIT, BIT, BIT_VARYING).contains(type)) {
if (column.getColumnSize() == null) {
return columnType;
}
return StringUtils.join(columnType, "(", column.getColumnSize(), ")");
}

if (Arrays.asList(TIME, TIMETZ, TIMESTAMPTZ, TIMESTAMP).contains(type)) {
if (Arrays.asList(TIME, TIMETZ, TIMESTAMPTZ).contains(type)) {
if (column.getColumnSize() == null || column.getColumnSize() == 0) {
return columnType;
} else {
return StringUtils.join(columnType, "(", column.getColumnSize(), ")");
}
}

if (Arrays.asList(TIMESTAMP_WITH_TIME_ZONE, TIME_WITH_TIME_ZONE).contains(type)) {
if (column.getColumnSize() == null || column.getColumnSize() == 0) {
return columnType;
} else {
String[] strings = columnType.split(" ");
return strings[0] + "(" + column.getColumnSize() + ")" + " " + strings[1] + " " + strings[2] + " " + strings[3];
}
}

if (Arrays.asList(DECIMAL, NUMERIC).contains(type)) {
if (column.getColumnSize() == null && column.getDecimalDigits() == null) {
return columnType;
Expand Down

0 comments on commit a753f49

Please sign in to comment.