From d545085d9c92bfa05e950613051ae9469adb0b38 Mon Sep 17 00:00:00 2001 From: SwallowGG <1558143046@qq.com> Date: Thu, 2 Nov 2023 22:58:15 +0800 Subject: [PATCH 1/7] =?UTF-8?q?Support=20=E4=BA=BA=E5=A4=A7=E9=87=91?= =?UTF-8?q?=E4=BB=93=20TABLE=20create=20and=20edit.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../plugin/kingbase/KingBaseDBManage.java | 48 ++++ .../plugin/kingbase/KingBaseMetaData.java | 181 +++++++++++++ .../kingbase/builder/KingBaseSqlBuilder.java | 209 +++++++++++++++ .../kingbase/type/KingBaseColumnTypeEnum.java | 244 ++++++++++++++++++ .../kingbase/type/KingBaseIndexTypeEnum.java | 180 +++++++++++++ .../plugin/postgresql/PostgreSQLMetaData.java | 5 +- .../api/controller/rdb/TriggerController.java | 6 +- .../chat2db/spi/jdbc/DefaultMetaService.java | 14 +- .../java/ai/chat2db/spi/model/Database.java | 2 + .../main/java/ai/chat2db/spi/model/Table.java | 5 + 10 files changed, 887 insertions(+), 7 deletions(-) create mode 100644 chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/builder/KingBaseSqlBuilder.java create mode 100644 chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/type/KingBaseColumnTypeEnum.java create mode 100644 chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/type/KingBaseIndexTypeEnum.java diff --git a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseDBManage.java b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseDBManage.java index 4ec37558..4a790d9c 100644 --- a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseDBManage.java @@ -4,10 +4,58 @@ import java.sql.Connection; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; +import ai.chat2db.spi.sql.Chat2DBContext; +import ai.chat2db.spi.sql.ConnectInfo; import ai.chat2db.spi.sql.SQLExecutor; +import org.apache.commons.lang3.StringUtils; public class KingBaseDBManage extends DefaultDBManage implements DBManage { + @Override + public void connectDatabase(Connection connection, String database) { + try { + ConnectInfo connectInfo = Chat2DBContext.getConnectInfo(); + if (!StringUtils.isEmpty(connectInfo.getSchemaName())) { + SQLExecutor.getInstance().execute(connection, "SET search_path TO \"" + connectInfo.getSchemaName() + "\""); + } + } catch (Exception e) { + } + } + + @Override + public Connection getConnection(ConnectInfo connectInfo) { + String url = connectInfo.getUrl(); + String database = connectInfo.getDatabaseName(); + if (database != null && !database.isEmpty()) { + url = replaceDatabaseInJdbcUrl(url, database); + } + connectInfo.setUrl(url); + + return super.getConnection(connectInfo); + } + + + public String replaceDatabaseInJdbcUrl(String url, String newDatabase) { + // 先在"?"字符处分割字符串，处理查询参数 + String[] urlAndParams = url.split("\\?"); + String urlWithoutParams = urlAndParams[0]; + + // 在URL中的"/"字符处分割字符串 + String[] parts = urlWithoutParams.split("/"); + + // 取最后一部分，即数据库名，并替换为新的数据库名 + parts[parts.length - 1] = newDatabase; + + // 将修改后的部分重新组合成URL + String newUrlWithoutParams = String.join("/", parts); + + // 如果存在查询参数，重新添加 + String newUrl = urlAndParams.length > 1 ? newUrlWithoutParams + "?" + urlAndParams[1] : newUrlWithoutParams; + + return newUrl; + } + + @Override public void dropTable(Connection connection, String databaseName, String schemaName, String tableName) { String sql = "drop table if exists " +tableName; diff --git a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseMetaData.java b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseMetaData.java index 1f3404e5..e1805d8f 100644 --- a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseMetaData.java +++ b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseMetaData.java @@ -1,7 +1,188 @@ package ai.chat2db.plugin.kingbase; +import ai.chat2db.plugin.kingbase.builder.KingBaseSqlBuilder; +import ai.chat2db.plugin.kingbase.type.KingBaseColumnTypeEnum; +import ai.chat2db.plugin.kingbase.type.KingBaseIndexTypeEnum; import ai.chat2db.spi.MetaData; +import ai.chat2db.spi.SqlBuilder; import ai.chat2db.spi.jdbc.DefaultMetaService; +import ai.chat2db.spi.model.*; +import ai.chat2db.spi.sql.SQLExecutor; +import com.google.common.collect.Lists; +import jakarta.validation.constraints.NotEmpty; +import org.apache.commons.lang3.StringUtils; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.*; +import java.util.stream.Collectors; + +import static ai.chat2db.spi.util.SortUtils.sortDatabase; public class KingBaseMetaData extends DefaultMetaService implements MetaData { + + + private static final String SELECT_KEY_INDEX = "SELECT ccu.table_schema AS Foreign_schema_name, ccu.table_name AS Foreign_table_name, ccu.column_name AS Foreign_column_name, constraint_type AS Constraint_type, tc.CONSTRAINT_NAME AS Key_name, tc.TABLE_NAME, kcu.Column_name, tc.is_deferrable, tc.initially_deferred FROM information_schema.table_constraints AS tc JOIN information_schema.key_column_usage AS kcu ON tc.CONSTRAINT_NAME = kcu.CONSTRAINT_NAME JOIN information_schema.constraint_column_usage AS ccu ON ccu.constraint_name = tc.constraint_name WHERE tc.TABLE_SCHEMA = '%s' AND tc.TABLE_NAME = '%s';"; + + + private List systemDatabases = Arrays.asList("SAMPLES", "SECURITY"); + + @Override + public List databases(Connection connection) { + List list = SQLExecutor.getInstance().executeSql(connection, "SELECT datname FROM sys_database", resultSet -> { + List databases = new ArrayList<>(); + try { + while (resultSet.next()) { + String dbName = resultSet.getString("datname"); + if ("template0".equalsIgnoreCase(dbName) || "template1".equalsIgnoreCase(dbName) || + "template2".equalsIgnoreCase(dbName)) { + continue; + } + Database database = new Database(); + database.setName(dbName); + databases.add(database); + } + } catch (SQLException e) { + throw new RuntimeException(e); + } + return databases; + }); + return sortDatabase(list, systemDatabases, connection); + } + + private static final String SELECT_TABLE_INDEX = "SELECT tmp.INDISPRIMARY AS Index_primary, tmp.TABLE_SCHEM, tmp.TABLE_NAME, tmp.NON_UNIQUE, tmp.INDEX_QUALIFIER, tmp.INDEX_NAME AS Key_name, tmp.indisclustered, tmp.ORDINAL_POSITION AS Seq_in_index, trim(BOTH '\"' FROM sys_get_indexdef( tmp.CI_OID, tmp.ORDINAL_POSITION, FALSE) ) AS Column_name, CASE tmp.AM_NAME WHEN 'btree' THEN CASE tmp.I_INDOPTION [ tmp.ORDINAL_POSITION - 1 ] & 1 :: SMALLINT WHEN 1 THEN 'D' ELSE'A' END ELSE NULL END AS Collation, tmp.CARDINALITY, tmp.PAGES, tmp.FILTER_CONDITION , tmp.AM_NAME AS Index_method, tmp.DESCRIPTION AS Index_comment FROM ( SELECT n.nspname AS TABLE_SCHEM, ct.relname AS TABLE_NAME, NOT i.indisunique AS NON_UNIQUE, NULL AS INDEX_QUALIFIER, ci.relname AS INDEX_NAME, i.INDISPRIMARY , i.indisclustered , ( information_schema._sys_expandarray ( i.indkey ) ).n AS ORDINAL_POSITION, ci.reltuples AS CARDINALITY, ci.relpages AS PAGES, sys_get_expr ( i.indpred, i.indrelid ) AS FILTER_CONDITION, ci.OID AS CI_OID, i.indoption AS I_INDOPTION, am.amname AS AM_NAME , d.description FROM sys_class ct JOIN sys_namespace n ON ( ct.relnamespace = n.OID ) JOIN sys_index i ON ( ct.OID = i.indrelid ) JOIN sys_class ci ON ( ci.OID = i.indexrelid ) JOIN sys_am am ON ( ci.relam = am.OID ) left outer join sys_description d on i.indexrelid = d.objoid WHERE n.nspname = '%s' AND ct.relname = '%s' ) AS tmp"; + + @Override + public List indexes(Connection connection, String databaseName, String schemaName, String tableName) { + + String constraintSql = String.format(SELECT_KEY_INDEX, schemaName, tableName); + Map constraintMap = new HashMap(); + LinkedHashMap foreignMap = new LinkedHashMap(); + SQLExecutor.getInstance().execute(connection, constraintSql, resultSet -> { + while (resultSet.next()) { + String keyName = resultSet.getString("Key_name"); + String constraintType = resultSet.getString("Constraint_type"); + constraintMap.put(keyName, constraintType); + if (StringUtils.equalsIgnoreCase(constraintType, KingBaseIndexTypeEnum.FOREIGN.getKeyword())) { + TableIndex tableIndex = foreignMap.get(keyName); + String columnName = resultSet.getString("Column_name"); + if (tableIndex == null) { + tableIndex = new TableIndex(); + tableIndex.setDatabaseName(databaseName); + tableIndex.setSchemaName(schemaName); + tableIndex.setTableName(tableName); + tableIndex.setName(keyName); + tableIndex.setForeignSchemaName(resultSet.getString("Foreign_schema_name")); + tableIndex.setForeignTableName(resultSet.getString("Foreign_table_name")); + tableIndex.setForeignColumnNamelist(Lists.newArrayList(columnName)); + tableIndex.setType(KingBaseIndexTypeEnum.FOREIGN.getName()); + foreignMap.put(keyName, tableIndex); + } else { + tableIndex.getForeignColumnNamelist().add(columnName); + } + } + } + return null; + }); + + String sql = String.format(SELECT_TABLE_INDEX, schemaName, tableName); + return SQLExecutor.getInstance().execute(connection, sql, resultSet -> { + LinkedHashMap map = new LinkedHashMap(foreignMap); + + while (resultSet.next()) { + String keyName = resultSet.getString("Key_name"); + TableIndex tableIndex = map.get(keyName); + if (tableIndex != null) { + List columnList = tableIndex.getColumnList(); + columnList.add(getTableIndexColumn(resultSet)); + columnList = columnList.stream().sorted(Comparator.comparing(TableIndexColumn::getOrdinalPosition)) + .collect(Collectors.toList()); + tableIndex.setColumnList(columnList); + } else { + TableIndex index = new TableIndex(); + index.setDatabaseName(databaseName); + index.setSchemaName(schemaName); + index.setTableName(tableName); + index.setName(keyName); + index.setUnique(!StringUtils.equals("t", resultSet.getString("NON_UNIQUE"))); + index.setMethod(resultSet.getString("Index_method")); + index.setComment(resultSet.getString("Index_comment")); + List tableIndexColumns = new ArrayList<>(); + tableIndexColumns.add(getTableIndexColumn(resultSet)); + index.setColumnList(tableIndexColumns); + String constraintType = constraintMap.get(keyName); + if (StringUtils.equals("t", resultSet.getString("Index_primary"))) { + index.setType(KingBaseIndexTypeEnum.PRIMARY.getName()); + } else if (StringUtils.equalsIgnoreCase(constraintType, KingBaseIndexTypeEnum.UNIQUE.getName())) { + index.setType(KingBaseIndexTypeEnum.UNIQUE.getName()); + } else { + index.setType(KingBaseIndexTypeEnum.NORMAL.getName()); + } + map.put(keyName, index); + } + } + return map.values().stream().collect(Collectors.toList()); + }); + + } + + private TableIndexColumn getTableIndexColumn(ResultSet resultSet) throws SQLException { + TableIndexColumn tableIndexColumn = new TableIndexColumn(); + tableIndexColumn.setColumnName(resultSet.getString("Column_name")); + tableIndexColumn.setOrdinalPosition(resultSet.getShort("Seq_in_index")); + tableIndexColumn.setCollation(resultSet.getString("Collation")); + tableIndexColumn.setAscOrDesc(resultSet.getString("Collation")); + return tableIndexColumn; + } + + private static String ROUTINES_SQL = " SELECT p.proname, p.prokind, sys_catalog.sys_get_functiondef(p.oid) as \"code\" FROM sys_catalog.sys_proc p where p.proname='%s'"; + + @Override + public Function function(Connection connection, @NotEmpty String databaseName, String schemaName, + String functionName) { + + String sql = String.format(ROUTINES_SQL, "f", functionName); + return SQLExecutor.getInstance().execute(connection, sql, resultSet -> { + Function function = new Function(); + function.setDatabaseName(databaseName); + function.setSchemaName(schemaName); + function.setFunctionName(functionName); + if (resultSet.next()) { + function.setFunctionBody(resultSet.getString("code")); + } + return function; + }); + } + + @Override + public Procedure procedure(Connection connection, @NotEmpty String databaseName, String schemaName, + String procedureName) { + String sql = String.format(ROUTINES_SQL, procedureName); + return SQLExecutor.getInstance().execute(connection, sql, resultSet -> { + Procedure procedure = new Procedure(); + procedure.setDatabaseName(databaseName); + procedure.setSchemaName(schemaName); + procedure.setProcedureName(procedureName); + if (resultSet.next()) { + procedure.setProcedureBody(resultSet.getString("code")); + } + return procedure; + }); + } + + @Override + public SqlBuilder getSqlBuilder() { + return new KingBaseSqlBuilder(); + } + + @Override + public TableMeta getTableMeta(String databaseName, String schemaName, String tableName) { + return TableMeta.builder() + .columnTypes(KingBaseColumnTypeEnum.getTypes()) + //.charsets(PostgreSQLCharsetEnum.getCharsets()) + //.collations(PostgreSQLCollationEnum.getCollations()) + .indexTypes(KingBaseIndexTypeEnum.getIndexTypes()) + .build(); + } } diff --git a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/builder/KingBaseSqlBuilder.java b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/builder/KingBaseSqlBuilder.java new file mode 100644 index 00000000..bccae895 --- /dev/null +++ b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/builder/KingBaseSqlBuilder.java @@ -0,0 +1,209 @@ +package ai.chat2db.plugin.kingbase.builder; + +import ai.chat2db.plugin.kingbase.type.KingBaseColumnTypeEnum; +import ai.chat2db.plugin.kingbase.type.KingBaseIndexTypeEnum; +import ai.chat2db.spi.SqlBuilder; +import ai.chat2db.spi.jdbc.DefaultSqlBuilder; +import ai.chat2db.spi.model.*; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.commons.lang3.BooleanUtils; +import org.apache.commons.lang3.StringUtils; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + + +public class KingBaseSqlBuilder extends DefaultSqlBuilder implements SqlBuilder { + @Override + public String buildCreateTableSql(Table table) { + StringBuilder script = new StringBuilder(); + script.append("CREATE TABLE "); + script.append("\"").append(table.getName()).append("\"").append(" (").append(" ").append("\n"); + // append column + for (TableColumn column : table.getColumnList()) { + if (StringUtils.isBlank(column.getName()) || StringUtils.isBlank(column.getColumnType())) { + continue; + } + KingBaseColumnTypeEnum typeEnum = KingBaseColumnTypeEnum.getByType(column.getColumnType()); + script.append("\t").append(typeEnum.buildCreateColumnSql(column)).append(",\n"); + } + Map> tableIndexMap = table.getIndexList().stream() + .collect(Collectors.partitioningBy(v -> KingBaseIndexTypeEnum.NORMAL.getName().equals(v.getType()))); + // append constraint key + List constraintList = tableIndexMap.get(Boolean.FALSE); + if (CollectionUtils.isNotEmpty(constraintList)) { + for (TableIndex index : constraintList) { + if (StringUtils.isBlank(index.getName()) || StringUtils.isBlank(index.getType())) { + continue; + } + KingBaseIndexTypeEnum indexTypeEnum = KingBaseIndexTypeEnum.getByType(index.getType()); + script.append("\t").append("").append(indexTypeEnum.buildIndexScript(index)); + script.append(",\n"); + } + + } + script = new StringBuilder(script.substring(0, script.length() - 2)); + script.append("\n)"); + if(StringUtils.isNotBlank(table.getTablespace())){ + script.append(" TABLESPACE \"").append(table.getTablespace()).append("\";"); + }else { + script.append(" TABLESPACE \"SYS_DEFAULT\";"); + } + // append index + List tableIndexList = tableIndexMap.get(Boolean.TRUE); + for (TableIndex tableIndex : tableIndexList) { + if (StringUtils.isBlank(tableIndex.getName()) || StringUtils.isBlank(tableIndex.getType())) { + continue; + } + script.append("\n"); + KingBaseIndexTypeEnum indexTypeEnum = KingBaseIndexTypeEnum.getByType(tableIndex.getType()); + script.append("").append(indexTypeEnum.buildIndexScript(tableIndex)).append(";"); + } + + // append comment + if (StringUtils.isNotBlank(table.getComment())) { + script.append("\n"); + script.append("COMMENT ON TABLE").append(" ").append("\"").append(table.getName()).append("\" IS '") + .append(table.getComment()).append("';\n"); + } + List tableColumnList = table.getColumnList().stream().filter(v -> StringUtils.isNotBlank(v.getComment())).toList(); + for (TableColumn tableColumn : tableColumnList) { + KingBaseColumnTypeEnum typeEnum = KingBaseColumnTypeEnum.getByType(tableColumn.getColumnType()); + script.append(typeEnum.buildComment(tableColumn, typeEnum)).append("\n"); + ; + } + List indexList = table.getIndexList().stream().filter(v -> StringUtils.isNotBlank(v.getComment())).toList(); + for (TableIndex index : indexList) { + KingBaseIndexTypeEnum indexEnum = KingBaseIndexTypeEnum.getByType(index.getType()); + script.append(indexEnum.buildIndexComment(index)).append("\n"); + } + + return script.toString(); + } + + @Override + public String buildModifyTaleSql(Table oldTable, Table newTable) { + StringBuilder script = new StringBuilder(); + if (!StringUtils.equalsIgnoreCase(oldTable.getName(), newTable.getName())) { + script.append("ALTER TABLE ").append("\"").append(oldTable.getName()).append("\""); + script.append("\t").append("RENAME TO ").append("\"").append(newTable.getName()).append("\"").append(";\n"); + + } + newTable.setColumnList(newTable.getColumnList().stream().filter(v -> StringUtils.isNotBlank(v.getEditStatus())).toList()); + newTable.setIndexList(newTable.getIndexList().stream().filter(v -> StringUtils.isNotBlank(v.getEditStatus())).toList()); + + //update name + List columnNameList = newTable.getColumnList().stream().filter(v -> + v.getOldName() != null && !StringUtils.equals(v.getOldName(), v.getName())).toList(); + for (TableColumn tableColumn : columnNameList) { + script.append("ALTER TABLE ").append("\"").append(newTable.getName()).append("\" ").append("RENAME COLUMN \"") + .append(tableColumn.getOldName()).append("\" TO \"").append(tableColumn.getName()).append("\";\n"); + } + + Map> tableIndexMap = newTable.getIndexList().stream() + .collect(Collectors.partitioningBy(v -> KingBaseIndexTypeEnum.NORMAL.getName().equals(v.getType()))); + StringBuilder scriptModify = new StringBuilder(); + Boolean modify = false; + scriptModify.append("ALTER TABLE ").append("\"").append(newTable.getName()).append("\" \n"); + // append modify column + for (TableColumn tableColumn : newTable.getColumnList()) { + KingBaseColumnTypeEnum typeEnum = KingBaseColumnTypeEnum.getByType(tableColumn.getColumnType()); + scriptModify.append("\t").append(typeEnum.buildModifyColumn(tableColumn)).append(",\n"); + modify = true; + + } + + // append modify constraint + for (TableIndex tableIndex : tableIndexMap.get(Boolean.FALSE)) { + if (StringUtils.isNotBlank(tableIndex.getType())) { + KingBaseIndexTypeEnum indexTypeEnum = KingBaseIndexTypeEnum.getByType(tableIndex.getType()); + scriptModify.append("\t").append(indexTypeEnum.buildModifyIndex(tableIndex)).append(",\n"); + modify = true; + } + } + + if (BooleanUtils.isTrue(modify)) { + script.append(scriptModify); + script = new StringBuilder(script.substring(0, script.length() - 2)); + script.append(";\n"); + } + + // append modify index + for (TableIndex tableIndex : tableIndexMap.get(Boolean.TRUE)) { + if (StringUtils.isNotBlank(tableIndex.getEditStatus()) && StringUtils.isNotBlank(tableIndex.getType())) { + KingBaseIndexTypeEnum indexTypeEnum = KingBaseIndexTypeEnum.getByType(tableIndex.getType()); + script.append(indexTypeEnum.buildModifyIndex(tableIndex)).append(";\n"); + } + } + + // append comment + if (!StringUtils.equals(oldTable.getComment(), newTable.getComment())) { + script.append("\n"); + script.append("COMMENT ON TABLE").append(" ").append("\"").append(newTable.getName()).append("\" IS '") + .append(newTable.getComment()).append("';\n"); + } + for (TableColumn tableColumn : newTable.getColumnList()) { + KingBaseColumnTypeEnum typeEnum = KingBaseColumnTypeEnum.getByType(tableColumn.getColumnType()); + script.append(typeEnum.buildComment(tableColumn, typeEnum)).append("\n"); + ; + } + List indexList = newTable.getIndexList().stream().filter(v -> StringUtils.isNotBlank(v.getComment())).toList(); + for (TableIndex index : indexList) { + KingBaseIndexTypeEnum indexEnum = KingBaseIndexTypeEnum.getByType(index.getType()); + script.append(indexEnum.buildIndexComment(index)).append("\n"); + } + + return script.toString(); + } + + @Override + public String pageLimit(String sql, int offset, int pageNo, int pageSize) { + StringBuilder sqlStr = new StringBuilder(sql.length() + 17); + sqlStr.append(sql); + if (offset == 0) { + sqlStr.append(" LIMIT "); + sqlStr.append(pageSize); + } else { + sqlStr.append(" LIMIT "); + sqlStr.append(pageSize); + sqlStr.append(" OFFSET "); + sqlStr.append(offset); + } + return sqlStr.toString(); + } + + @Override + public String buildCreateDatabaseSql(Database database) { + StringBuilder sqlBuilder = new StringBuilder(); + sqlBuilder.append("CREATE DATABASE "+database.getName()); + String owner = database.getOwner(); + if (StringUtils.isBlank(owner)) { + owner = "SYSTEM"; + } + sqlBuilder.append(" WITH OWNER = \"").append(owner).append("\""); + if (StringUtils.isNotBlank(database.getCharset())) { + sqlBuilder.append(" ENCODING ").append(database.getCharset()).append(""); + } + sqlBuilder.append(";\n"); + + if (StringUtils.isNotBlank(database.getComment())) { + sqlBuilder.append("COMMENT ON DATABASE ").append(database.getName()).append(" IS '").append(database.getComment()).append("';"); + } + return sqlBuilder.toString(); + } + + + @Override + public String buildCreateSchemaSql(Schema schema){ + StringBuilder sqlBuilder = new StringBuilder(); + sqlBuilder.append("CREATE SCHEMA "+schema.getName()+""); + String owner = schema.getOwner(); + if(StringUtils.isBlank(schema.getOwner())){ + owner = "SYSTEM"; + } + sqlBuilder.append(" AUTHORIZATION \"").append(owner).append("\""); + return sqlBuilder.toString(); + } + +} diff --git a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/type/KingBaseColumnTypeEnum.java b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/type/KingBaseColumnTypeEnum.java new file mode 100644 index 00000000..1eec1a77 --- /dev/null +++ b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/type/KingBaseColumnTypeEnum.java @@ -0,0 +1,244 @@ +package ai.chat2db.plugin.kingbase.type; + +import ai.chat2db.spi.ColumnBuilder; +import ai.chat2db.spi.enums.EditStatus; +import ai.chat2db.spi.model.ColumnType; +import ai.chat2db.spi.model.TableColumn; +import com.google.common.collect.Maps; +import org.apache.commons.lang3.StringUtils; + +import java.util.Arrays; +import java.util.List; +import java.util.Map; + +public enum KingBaseColumnTypeEnum 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), + 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), + + 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), + CHAR("CHAR", true, false, true, false, false, true, true, true, false, false), + + CID("CID", false, false, true, false, false, false, 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), + + CLOB("CLOB", 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), + FLOAT8("FLOAT8", false, false, true, false, false, false, true, true, false, false), + + INTEGER("INTEGER", 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), + INT4("INT4", false, false, true, false, false, false, true, true, false, false), + INT8("INT8", 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), + LINE("LINE", false, false, true, false, false, false, true, true, false, false), + LSEG("LSEG", false, false, true, false, false, false, true, true, false, false), + MACADDR("MACADDR", false, false, true, false, false, false, true, true, false, false), + MONEY("MONEY", false, false, true, false, false, false, true, true, false, false), + NUMERIC("NUMERIC", true, false, true, false, false, false, true, true, false, false), + PATH("PATH", false, false, true, false, false, false, true, true, false, false), + 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), + 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), + TIMESTAMP("TIMESTAMP", 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), + TSVECTOR("TSVECTOR", false, false, true, false, false, false, true, true, false, false), + TXID_SNAPSHOT("TXID_SNAPSHOT", false, false, true, false, false, false, true, true, false, false), + UUID("UUID", false, false, true, false, false, false, true, true, false, false), + 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), + + ; + + private static Map COLUMN_TYPE_MAP = Maps.newHashMap(); + + static { + for (KingBaseColumnTypeEnum value : KingBaseColumnTypeEnum.values()) { + COLUMN_TYPE_MAP.put(value.getColumnType().getTypeName(), value); + } + } + + private ColumnType columnType; + + + KingBaseColumnTypeEnum(String dataTypeName, boolean supportLength, boolean supportScale, boolean supportNullable, boolean supportAutoIncrement, boolean supportCharset, boolean supportCollation, boolean supportComments, boolean supportDefaultValue, boolean supportExtent, boolean supportValue) { + this.columnType = new ColumnType(dataTypeName, supportLength, supportScale, supportNullable, supportAutoIncrement, supportCharset, supportCollation, supportComments, supportDefaultValue, supportExtent, supportValue, false); + } + + public static KingBaseColumnTypeEnum getByType(String dataType) { + return COLUMN_TYPE_MAP.get(dataType.toUpperCase()); + } + + public static List getTypes() { + return Arrays.stream(KingBaseColumnTypeEnum.values()).map(columnTypeEnum -> + columnTypeEnum.getColumnType() + ).toList(); + } + + public ColumnType getColumnType() { + return columnType; + } + + @Override + public String buildCreateColumnSql(TableColumn column) { + KingBaseColumnTypeEnum type = COLUMN_TYPE_MAP.get(column.getColumnType().toUpperCase()); + if (type == null) { + return ""; + } + StringBuilder script = new StringBuilder(); + + script.append("\"").append(column.getName()).append("\"").append(" "); + + script.append(buildDataType(column, type)).append(" "); + + + script.append(buildCollation(column, type)).append(" "); + + script.append(buildNullable(column, type)).append(" "); + + script.append(buildDefaultValue(column, type)).append(" "); + + return script.toString(); + } + + private String buildCollation(TableColumn column, KingBaseColumnTypeEnum type) { + if (!type.getColumnType().isSupportCollation() || StringUtils.isEmpty(column.getCollationName())) { + return ""; + } + return StringUtils.join("\"", column.getCollationName(), "\""); + } + + @Override + public String buildModifyColumn(TableColumn column) { + + if (EditStatus.DELETE.name().equals(column.getEditStatus())) { + return StringUtils.join("DROP COLUMN `", column.getName() + "`"); + } + if (EditStatus.ADD.name().equals(column.getEditStatus())) { + return StringUtils.join("ADD COLUMN ", buildCreateColumnSql(column)); + } + if (EditStatus.MODIFY.name().equals(column.getEditStatus())) { + StringBuilder script = new StringBuilder(); + script.append("ALTER COLUMN \"").append(column.getName()).append("\" TYPE ").append(buildDataType(column, this)).append(",\n"); + if (column.getNullable() != null && 1 == column.getNullable()) { + script.append("\t").append("ALTER COLUMN \"").append(column.getName()).append("\" DROP NOT NULL ,\n"); + } else { + script.append("\t").append("ALTER COLUMN \"").append(column.getName()).append("\" SET NOT NULL ,\n"); + + } + String defaultValue = buildDefaultValue(column, this); + if (StringUtils.isNotBlank(defaultValue)) { + script.append("ALTER COLUMN \"").append(column.getName()).append("\" SET ").append(defaultValue).append(",\n"); + } + script = new StringBuilder(script.substring(0, script.length() - 2)); + return script.toString(); + } + return ""; + } + + public String buildComment(TableColumn column, KingBaseColumnTypeEnum type) { + if (!this.columnType.isSupportComments() || column.getComment() == null + || EditStatus.DELETE.name().equals(column.getEditStatus())) { + return ""; + } + return StringUtils.join("COMMENT ON COLUMN", " \"", column.getTableName(), + "\".\"", column.getName(), "\" IS '", column.getComment(), "';"); + } + + private String buildDefaultValue(TableColumn column, KingBaseColumnTypeEnum type) { + if (!type.getColumnType().isSupportDefaultValue() || StringUtils.isEmpty(column.getDefaultValue())) { + return ""; + } + + if("EMPTY_STRING".equalsIgnoreCase(column.getDefaultValue().trim())){ + return StringUtils.join("DEFAULT ''"); + } + + if("NULL".equalsIgnoreCase(column.getDefaultValue().trim())){ + return StringUtils.join("DEFAULT NULL"); + } + + if (Arrays.asList(CHAR, VARCHAR).contains(type)) { + return StringUtils.join("DEFAULT '", column.getDefaultValue(), "'"); + } + + if (Arrays.asList(TIMESTAMP, TIME, TIMETZ, TIMESTAMPTZ, DATE).contains(type)) { + if ("CURRENT_TIMESTAMP".equalsIgnoreCase(column.getDefaultValue().trim())) { + return StringUtils.join("DEFAULT ", column.getDefaultValue()); + } + return StringUtils.join("DEFAULT '", column.getDefaultValue(), "'"); + } + + return StringUtils.join("DEFAULT ", column.getDefaultValue()); + } + + private String buildNullable(TableColumn column, KingBaseColumnTypeEnum type) { + if (!type.getColumnType().isSupportNullable()) { + return ""; + } + if (column.getNullable() != null && 1 == column.getNullable()) { + return "NULL"; + } else { + return "NOT NULL"; + } + } + + private String buildDataType(TableColumn column, KingBaseColumnTypeEnum type) { + String columnType = type.columnType.getTypeName(); + if (Arrays.asList(VARCHAR, CHAR,CHARACTER).contains(type)) { + if (column.getColumnSize() == null ) { + return columnType; + } + return StringUtils.join(columnType, "(", column.getColumnSize(), ")"); + } + + if (Arrays.asList(VARBIT, BIT).contains(type)) { + if (column.getColumnSize() == null ) { + return columnType; + } + return StringUtils.join(columnType, "(", column.getColumnSize(), ")"); + } + + if (Arrays.asList(TIME, TIMETZ, TIMESTAMPTZ, TIMESTAMP).contains(type)) { + if (column.getColumnSize() == null || column.getColumnSize() == 0) { + return columnType; + } else { + return StringUtils.join(columnType, "(", column.getColumnSize(), ")"); + } + } + + if (Arrays.asList(DECIMAL, NUMERIC).contains(type)) { + if (column.getColumnSize() == null && column.getDecimalDigits() == null) { + return columnType; + } + if (column.getColumnSize() != null && column.getDecimalDigits() == null) { + return StringUtils.join(columnType, "(", column.getColumnSize() + ")"); + } else { + return StringUtils.join(columnType, "(", column.getColumnSize() + "," + column.getDecimalDigits() + ")"); + } + } + return columnType; + } + +} diff --git a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/type/KingBaseIndexTypeEnum.java b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/type/KingBaseIndexTypeEnum.java new file mode 100644 index 00000000..a92b0511 --- /dev/null +++ b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/type/KingBaseIndexTypeEnum.java @@ -0,0 +1,180 @@ +package ai.chat2db.plugin.kingbase.type; + +import ai.chat2db.spi.enums.EditStatus; +import ai.chat2db.spi.model.IndexType; +import ai.chat2db.spi.model.TableIndex; +import ai.chat2db.spi.model.TableIndexColumn; +import org.apache.commons.collections4.CollectionUtils; +import org.apache.commons.lang3.BooleanUtils; +import org.apache.commons.lang3.StringUtils; + +import java.util.Arrays; +import java.util.List; + +public enum KingBaseIndexTypeEnum { + + PRIMARY("Primary", "PRIMARY KEY"), + + FOREIGN("Foreign", "FOREIGN KEY"), + + NORMAL("Normal", "INDEX"), + + UNIQUE("Unique", "UNIQUE"), + ; + + private String name; + private String keyword; + + private IndexType indexType; + + + KingBaseIndexTypeEnum(String name, String keyword) { + this.name = name; + this.keyword = keyword; + this.indexType =new IndexType(name); + } + + public static KingBaseIndexTypeEnum getByType(String type) { + for (KingBaseIndexTypeEnum value : KingBaseIndexTypeEnum.values()) { + if (value.name.equalsIgnoreCase(type)) { + return value; + } + } + return null; + } + + public static List getIndexTypes() { + return Arrays.asList(KingBaseIndexTypeEnum.values()).stream().map(KingBaseIndexTypeEnum::getIndexType).collect(java.util.stream.Collectors.toList()); + } + + public IndexType getIndexType() { + return indexType; + } + + public String getName() { + return name; + } + + public String getKeyword() { + return keyword; + } + + public String buildIndexScript(TableIndex tableIndex) { + StringBuilder script = new StringBuilder(); + if (NORMAL.equals(this)) { + script.append("CREATE").append(" "); + script.append(buildIndexUnique(tableIndex)).append(" "); + script.append(buildIndexConcurrently(tableIndex)).append(" "); + script.append(buildIndexName(tableIndex)).append(" "); + script.append("ON ").append("\"").append(tableIndex.getTableName()).append("\"").append(" "); + script.append(buildIndexMethod(tableIndex)).append(" "); + script.append(buildIndexColumn(tableIndex)); + } else { + script.append("CONSTRAINT").append(" "); + script.append(buildIndexName(tableIndex)).append(" "); + script.append(keyword).append(" "); + script.append(buildIndexColumn(tableIndex)); + script.append(buildForeignColum(tableIndex)); + } + return script.toString(); + } + + private String buildForeignColum(TableIndex tableIndex) { + if (FOREIGN.equals(this)) { + StringBuilder script = new StringBuilder(); + script.append(" REFERENCES "); + if (StringUtils.isNotBlank(tableIndex.getForeignSchemaName())) { + script.append(tableIndex.getForeignSchemaName()).append("."); + } + if (StringUtils.isNotBlank(tableIndex.getForeignTableName())) { + script.append(tableIndex.getForeignTableName()).append(" "); + } + if (CollectionUtils.isNotEmpty(tableIndex.getForeignColumnNamelist())) { + script.append("("); + for (String column : tableIndex.getForeignColumnNamelist()) { + if (StringUtils.isNotBlank(column)) { + script.append("\"").append(column).append("\"").append(","); + } + } + script.deleteCharAt(script.length() - 1); + script.append(")"); + } + return script.toString(); + } + return ""; + } + + private String buildIndexMethod(TableIndex tableIndex) { + if (StringUtils.isNotBlank(tableIndex.getMethod())) { + return "USING " + tableIndex.getMethod(); + } else { + return ""; + } + } + + private String buildIndexConcurrently(TableIndex tableIndex) { + if (BooleanUtils.isTrue(tableIndex.getConcurrently())) { + return "CONCURRENTLY"; + } else { + return ""; + } + } + + private String buildIndexUnique(TableIndex tableIndex) { + if (BooleanUtils.isTrue(tableIndex.getUnique())) { + return "UNIQUE " + keyword; + } else { + return keyword; + } + } + + public String buildIndexComment(TableIndex tableIndex) { + if (StringUtils.isBlank(tableIndex.getComment()) || EditStatus.DELETE.name().equals(tableIndex.getEditStatus())) { + return ""; + } else if (NORMAL.equals(this)) { + return StringUtils.join("COMMENT ON INDEX", " ", + "\"", tableIndex.getName(), "\" IS '", tableIndex.getComment(), "';"); + } else { + return StringUtils.join("COMMENT ON CONSTRAINT", " \"", tableIndex.getName(), "\" ON \"", tableIndex.getSchemaName(), + "\".\"", tableIndex.getTableName(), "\" IS '", tableIndex.getComment(), "';"); + } + } + + private String buildIndexColumn(TableIndex tableIndex) { + StringBuilder script = new StringBuilder(); + script.append("("); + for (TableIndexColumn column : tableIndex.getColumnList()) { + if (StringUtils.isNotBlank(column.getColumnName())) { + script.append("\"").append(column.getColumnName()).append("\"").append(","); + } + } + script.deleteCharAt(script.length() - 1); + script.append(")"); + return script.toString(); + } + + private String buildIndexName(TableIndex tableIndex) { + return "\"" + tableIndex.getName() + "\""; + } + + public String buildModifyIndex(TableIndex tableIndex) { + boolean isNormal = NORMAL.equals(this); + if (EditStatus.DELETE.name().equals(tableIndex.getEditStatus())) { + return buildDropIndex(tableIndex); + } + if (EditStatus.MODIFY.name().equals(tableIndex.getEditStatus())) { + return StringUtils.join(buildDropIndex(tableIndex), isNormal ? ";\n" : ",\n\tADD ", buildIndexScript(tableIndex)); + } + if (EditStatus.ADD.name().equals(tableIndex.getEditStatus())) { + return StringUtils.join(isNormal ? "" : "ADD ", buildIndexScript(tableIndex)); + } + return ""; + } + + private String buildDropIndex(TableIndex tableIndex) { + if (NORMAL.equals(this)) { + return StringUtils.join("DROP INDEX \"", tableIndex.getOldName(), "\""); + } + return StringUtils.join("DROP CONSTRAINT \"", tableIndex.getOldName(), "\""); + } +} diff --git a/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLMetaData.java b/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLMetaData.java index 78e98656..ef1413fa 100644 --- a/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLMetaData.java +++ b/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLMetaData.java @@ -75,9 +75,7 @@ public class PostgreSQLMetaData extends DefaultMetaService implements MetaData { private static final String SELECT_TABLE_INDEX = "SELECT tmp.INDISPRIMARY AS Index_primary, tmp.TABLE_SCHEM, tmp.TABLE_NAME, tmp.NON_UNIQUE, tmp.INDEX_QUALIFIER, tmp.INDEX_NAME AS Key_name, tmp.indisclustered, tmp.ORDINAL_POSITION AS Seq_in_index, TRIM ( BOTH '\"' FROM pg_get_indexdef ( tmp.CI_OID, tmp.ORDINAL_POSITION, FALSE ) ) AS Column_name,CASE tmp.AM_NAME WHEN 'btree' THEN CASE tmp.I_INDOPTION [ tmp.ORDINAL_POSITION - 1 ] & 1 :: SMALLINT WHEN 1 THEN 'D' ELSE'A' END ELSE NULL END AS Collation, tmp.CARDINALITY, tmp.PAGES, tmp.FILTER_CONDITION , tmp.AM_NAME AS Index_method, tmp.DESCRIPTION AS Index_comment FROM ( SELECT n.nspname AS TABLE_SCHEM, ct.relname AS TABLE_NAME, NOT i.indisunique AS NON_UNIQUE, NULL AS INDEX_QUALIFIER, ci.relname AS INDEX_NAME,i.INDISPRIMARY , i.indisclustered , ( information_schema._pg_expandarray ( i.indkey ) ).n AS ORDINAL_POSITION, ci.reltuples AS CARDINALITY, ci.relpages AS PAGES, pg_get_expr ( i.indpred, i.indrelid ) AS FILTER_CONDITION, ci.OID AS CI_OID, i.indoption AS I_INDOPTION, am.amname AS AM_NAME , d.description FROM pg_class ct JOIN pg_namespace n ON ( ct.relnamespace = n.OID ) JOIN pg_index i ON ( ct.OID = i.indrelid ) JOIN pg_class ci ON ( ci.OID = i.indexrelid ) JOIN pg_am am ON ( ci.relam = am.OID ) left outer join pg_description d on i.indexrelid = d.objoid WHERE n.nspname = '%s' AND ct.relname = '%s' ) AS tmp ;"; - private static String ROUTINES_SQL - = " SELECT p.proname, p.prokind, pg_catalog.pg_get_functiondef(p.oid) as \"code\" FROM pg_catalog.pg_proc p " - + "where p.prokind = '%s' and p.proname='%s';"; + private static String ROUTINES_SQL = "SELECT p.proname, p.prokind, pg_catalog.pg_get_functiondef(p.oid) as \"code\" FROM pg_catalog.pg_proc p where p.prokind = '%s' and p.proname='%s'"; private static String TRIGGER_SQL = "SELECT n.nspname AS \"schema\", c.relname AS \"table_name\", t.tgname AS \"trigger_name\", t.tgenabled AS " + "\"enabled\", pg_get_triggerdef(t.oid) AS \"trigger_body\" FROM pg_trigger t JOIN pg_class c ON c.oid = t" @@ -301,4 +299,5 @@ public class PostgreSQLMetaData extends DefaultMetaService implements MetaData { .indexTypes(PostgreSQLIndexTypeEnum.getIndexTypes()) .build(); } + } diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/TriggerController.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/TriggerController.java index 3ee67f6e..03524568 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/TriggerController.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/TriggerController.java @@ -9,6 +9,7 @@ import ai.chat2db.server.web.api.controller.rdb.request.TriggerDetailRequest; import ai.chat2db.server.web.api.controller.rdb.request.TriggerPageRequest; import ai.chat2db.spi.model.Trigger; import jakarta.validation.Valid; +import org.apache.commons.collections4.CollectionUtils; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.web.bind.annotation.GetMapping; import org.springframework.web.bind.annotation.RequestMapping; @@ -25,8 +26,9 @@ public class TriggerController { @GetMapping("/list") public WebPageResult list(@Valid TriggerPageRequest request) { ListResult listResult = triggerService.triggers(request.getDatabaseName(), request.getSchemaName()); - return WebPageResult.of(listResult.getData(), Long.valueOf(listResult.getData().size()), 1, - listResult.getData().size()); + Long total = CollectionUtils.isNotEmpty(listResult.getData()) ? Long.valueOf(listResult.getData().size()) : 0L; + return WebPageResult.of(listResult.getData(), total, 1, + listResult.getData().size()); } @GetMapping("/detail") diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultMetaService.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultMetaService.java index 8f39e920..29a2b94e 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultMetaService.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultMetaService.java @@ -10,6 +10,7 @@ import ai.chat2db.spi.MetaData; import ai.chat2db.spi.SqlBuilder; import ai.chat2db.spi.model.*; import ai.chat2db.spi.sql.SQLExecutor; +import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.lang3.StringUtils; /** @@ -49,7 +50,11 @@ public class DefaultMetaService implements MetaData { @Override public List functions(Connection connection, String databaseName, String schemaName) { - return SQLExecutor.getInstance().functions(connection, StringUtils.isEmpty(databaseName) ? null : databaseName, StringUtils.isEmpty(schemaName) ? null : schemaName); + List functions = SQLExecutor.getInstance().functions(connection, StringUtils.isEmpty(databaseName) ? null : databaseName, StringUtils.isEmpty(schemaName) ? null : schemaName); + if(CollectionUtils.isEmpty(functions)){ + return functions; + } + return functions.stream().filter(function -> StringUtils.isNotBlank(function.getFunctionName())).collect(Collectors.toList()); } @Override @@ -59,7 +64,12 @@ public class DefaultMetaService implements MetaData { @Override public List procedures(Connection connection, String databaseName, String schemaName) { - return SQLExecutor.getInstance().procedures(connection, StringUtils.isEmpty(databaseName) ? null : databaseName, StringUtils.isEmpty(schemaName) ? null : schemaName); + List procedures = SQLExecutor.getInstance().procedures(connection, StringUtils.isEmpty(databaseName) ? null : databaseName, StringUtils.isEmpty(schemaName) ? null : schemaName); + + if(CollectionUtils.isEmpty(procedures)){ + return procedures; + } + return procedures.stream().filter(function -> StringUtils.isNotBlank(function.getProcedureName())).collect(Collectors.toList()); } @Override diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Database.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Database.java index 299b2cef..71630626 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Database.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Database.java @@ -38,4 +38,6 @@ public class Database implements Serializable { private String charset; private String collation; + + private String owner; } diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Table.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Table.java index 4ee22fec..640f2180 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Table.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Table.java @@ -86,8 +86,13 @@ public class Table { private String collate; + private Long incrementValue; + private String partition; + + + private String tablespace; } From 14461ed59d4ee436684336af147da44a4eaef557 Mon Sep 17 00:00:00 2001 From: SwallowGG <1558143046@qq.com> Date: Thu, 2 Nov 2023 23:41:30 +0800 Subject: [PATCH 2/7] Simplify the update and delete statements for query results when the table contains primary keys --- .../plugin/kingbase/KingBaseMetaData.java | 4 + .../core/impl/DlTemplateServiceImpl.java | 110 +++++++++++++----- 2 files changed, 85 insertions(+), 29 deletions(-) diff --git a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseMetaData.java b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseMetaData.java index e1805d8f..288e9add 100644 --- a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseMetaData.java +++ b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseMetaData.java @@ -185,4 +185,8 @@ public class KingBaseMetaData extends DefaultMetaService implements MetaData { .indexTypes(KingBaseIndexTypeEnum.getIndexTypes()) .build(); } + @Override + public String getMetaDataName(String... names) { + return Arrays.stream(names).filter(name -> StringUtils.isNotBlank(name)).map(name -> "\"" + name + "\"").collect(Collectors.joining(".")); + } } diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java index 0892e1cd..42af6474 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java @@ -6,11 +6,14 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Optional; +import java.util.stream.Collectors; import ai.chat2db.server.domain.api.param.*; import ai.chat2db.server.domain.api.param.operation.OperationLogCreateParam; import ai.chat2db.server.domain.api.service.OperationLogService; +import ai.chat2db.server.domain.api.service.TableService; import ai.chat2db.spi.MetaData; +import ai.chat2db.spi.model.TableIndex; import ai.chat2db.spi.sql.ConnectInfo; import com.alibaba.druid.DbType; import com.alibaba.druid.sql.PagerUtils; @@ -37,6 +40,7 @@ import ai.chat2db.spi.util.SqlUtils; import com.google.common.collect.Lists; import lombok.extern.slf4j.Slf4j; import org.apache.commons.collections4.CollectionUtils; +import org.apache.commons.lang3.ObjectUtils; import org.apache.commons.lang3.StringUtils; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.stereotype.Service; @@ -53,6 +57,9 @@ public class DlTemplateServiceImpl implements DlTemplateService { @Autowired private OperationLogService operationLogService; + @Autowired + private TableService tableService; + @Override public ListResult execute(DlExecuteParam param) { if (StringUtils.isBlank(param.getSql())) { @@ -246,19 +253,20 @@ public class DlTemplateServiceImpl implements DlTemplateService { public DataResult updateSelectResult(UpdateSelectResultParam param) { StringBuilder stringBuilder = new StringBuilder(); MetaData metaSchema = Chat2DBContext.getMetaData(); + List keyColumns = getPrimaryColumns(param); for (int i = 0; i < param.getOperations().size(); i++) { SelectResultOperation operation = param.getOperations().get(i); - List row = operation.getDataList(); List odlRow = operation.getOldDataList(); String sql = ""; if ("UPDATE".equalsIgnoreCase(operation.getType())) { - sql = getUpdateSql(param, row, odlRow, metaSchema); + sql = getUpdateSql(param, row, odlRow, metaSchema,keyColumns,false); } else if ("CREATE".equalsIgnoreCase(operation.getType())) { sql = getInsertSql(param, row, metaSchema); - } else if ("DELETE".equalsIgnoreCase(operation.getType())) { - sql = getDeleteSql(param, odlRow, metaSchema); + sql = getDeleteSql(param, odlRow, metaSchema,keyColumns); + } else if ("UPDATE_COPY".equalsIgnoreCase(operation.getType())) { + sql = getUpdateSql(param, row, row, metaSchema,keyColumns,true); } stringBuilder.append(sql + ";\n"); @@ -266,38 +274,76 @@ public class DlTemplateServiceImpl implements DlTemplateService { return DataResult.of(stringBuilder.toString()); } - private String getDeleteSql(UpdateSelectResultParam param, List row, MetaData metaSchema) { + private List getPrimaryColumns(UpdateSelectResultParam param){ + TableQueryParam tableQueryParam = new TableQueryParam(); + tableQueryParam.setTableName(param.getTableName()); + tableQueryParam.setSchemaName(param.getSchemaName()); + tableQueryParam.setDatabaseName(param.getDatabaseName()); + List tableIndices = tableService.queryIndexes(tableQueryParam); + if(!CollectionUtils.isEmpty(tableIndices)) { + return tableIndices.stream() + .filter(tableIndex -> "PRIMARY".equalsIgnoreCase(tableIndex.getType())) + .map(tableIndex -> { + if(!CollectionUtils.isEmpty(tableIndex.getColumnList())) { + return tableIndex.getColumnList().get(0).getColumnName(); + } + return null; + }).collect(Collectors.toList()); + } + return Lists.newArrayList(); + } + + private String getDeleteSql(UpdateSelectResultParam param, List row, MetaData metaSchema,List keyColumns) { StringBuilder script = new StringBuilder(); script.append("DELETE FROM ").append(param.getTableName()).append(""); - script.append(buildWhere(param.getHeaderList(), row, metaSchema)); + script.append(buildWhere(param.getHeaderList(), row, metaSchema,keyColumns)); return script.toString(); } - private String buildWhere(List
headerList, List row, MetaData metaSchema) { + private String buildWhere(List
headerList, List row, MetaData metaSchema,List keyColumns) { StringBuilder script = new StringBuilder(); script.append(" where "); - for (int i = 1; i < row.size(); i++) { - String oldValue = row.get(i); - Header header = headerList.get(i); - String value = SqlUtils.getSqlValue(oldValue, header.getDataType()); - if (value == null) { - script.append(metaSchema.getMetaDataName(header.getName())) - .append(" is null and "); - } else { - script.append(metaSchema.getMetaDataName(header.getName())) - .append(" = ") - .append(value) - .append(" and "); + if(CollectionUtils.isEmpty(keyColumns)) { + for (int i = 1; i < row.size(); i++) { + String oldValue = row.get(i); + Header header = headerList.get(i); + String value = SqlUtils.getSqlValue(oldValue, header.getDataType()); + if (value == null) { + script.append(metaSchema.getMetaDataName(header.getName())) + .append(" is null and "); + } else { + script.append(metaSchema.getMetaDataName(header.getName())) + .append(" = ") + .append(value) + .append(" and "); + } + } + }else{ + for (int i = 1; i < row.size(); i++) { + String oldValue = row.get(i); + Header header = headerList.get(i); + String columnName = header.getName(); + if(keyColumns.contains(columnName)) { + String value = SqlUtils.getSqlValue(oldValue, header.getDataType()); + if (value == null) { + script.append(metaSchema.getMetaDataName(columnName)) + .append(" is null and "); + } else { + script.append(metaSchema.getMetaDataName(columnName)) + .append(" = ") + .append(value) + .append(" and "); + } + } } } - script.delete(script.length() - 4, script.length()); return script.toString(); } private String getInsertSql(UpdateSelectResultParam param, List row, MetaData metaSchema) { - if (CollectionUtils.isEmpty(row)) { + if (CollectionUtils.isEmpty(row) || ObjectUtils.allNull(row.toArray())) { return ""; } StringBuilder script = new StringBuilder(); @@ -305,16 +351,21 @@ public class DlTemplateServiceImpl implements DlTemplateService { .append(" ("); for (int i = 1; i < row.size(); i++) { Header header = param.getHeaderList().get(i); - script.append(metaSchema.getMetaDataName(header.getName())) - .append(","); + String newValue = row.get(i); + if(newValue!=null) { + script.append(metaSchema.getMetaDataName(header.getName())) + .append(","); + } } script.deleteCharAt(script.length() - 1); script.append(") VALUES ("); for (int i = 1; i < row.size(); i++) { String newValue = row.get(i); - Header header = param.getHeaderList().get(i); - script.append(SqlUtils.getSqlValue(newValue, header.getDataType())) - .append(","); + if(newValue!=null) { + Header header = param.getHeaderList().get(i); + script.append(SqlUtils.getSqlValue(newValue, header.getDataType())) + .append(","); + } } script.deleteCharAt(script.length() - 1); script.append(")"); @@ -323,7 +374,8 @@ public class DlTemplateServiceImpl implements DlTemplateService { } - private String getUpdateSql(UpdateSelectResultParam param, List row, List odlRow, MetaData metaSchema) { + private String getUpdateSql(UpdateSelectResultParam param, List row, List odlRow, MetaData metaSchema, + List keyColumns,boolean copy){ StringBuilder script = new StringBuilder(); if (CollectionUtils.isEmpty(row) || CollectionUtils.isEmpty(odlRow)) { return ""; @@ -332,7 +384,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { for (int i = 1; i < row.size(); i++) { String newValue = row.get(i); String oldValue = odlRow.get(i); - if (StringUtils.equals(newValue, oldValue)) { + if (StringUtils.equals(newValue, oldValue) && !copy) { continue; } Header header = param.getHeaderList().get(i); @@ -343,7 +395,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { .append(","); } script.deleteCharAt(script.length() - 1); - script.append(buildWhere(param.getHeaderList(), odlRow, metaSchema)); + script.append(buildWhere(param.getHeaderList(), odlRow, metaSchema,keyColumns)); return script.toString(); } From ef917b8143fd59850e4a609bd42975fc8adc1911 Mon Sep 17 00:00:00 2001 From: SwallowGG <1558143046@qq.com> Date: Fri, 3 Nov 2023 00:26:54 +0800 Subject: [PATCH 3/7] Simplify the update and delete statements for query results when the table contains primary keys --- .../domain/api/param/DlExecuteParam.java | 6 + .../core/impl/DlTemplateServiceImpl.java | 122 +++++++++++------- .../request/SelectResultUpdateRequest.java | 4 +- .../controller/rdb/vo/ExecuteResultVO.java | 3 +- .../web/api/controller/rdb/vo/HeaderVO.java | 31 ----- .../java/ai/chat2db/spi/model/Header.java | 19 ++- 6 files changed, 105 insertions(+), 80 deletions(-) delete mode 100644 chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/vo/HeaderVO.java diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/DlExecuteParam.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/DlExecuteParam.java index ccffac01..cae98b5c 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/DlExecuteParam.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/DlExecuteParam.java @@ -36,6 +36,12 @@ public class DlExecuteParam { @NotNull private String databaseName; + + /** + * schema名称 + */ + private String schemaName; + /** * 分页编码 * 只有select语句才有 diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java index 42af6474..82c43435 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java @@ -2,10 +2,7 @@ package ai.chat2db.server.domain.core.impl; import java.sql.Connection; import java.sql.SQLException; -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; -import java.util.Optional; +import java.util.*; import java.util.stream.Collectors; import ai.chat2db.server.domain.api.param.*; @@ -13,7 +10,7 @@ import ai.chat2db.server.domain.api.param.operation.OperationLogCreateParam; import ai.chat2db.server.domain.api.service.OperationLogService; import ai.chat2db.server.domain.api.service.TableService; import ai.chat2db.spi.MetaData; -import ai.chat2db.spi.model.TableIndex; +import ai.chat2db.spi.model.*; import ai.chat2db.spi.sql.ConnectInfo; import com.alibaba.druid.DbType; import com.alibaba.druid.sql.PagerUtils; @@ -31,8 +28,6 @@ import ai.chat2db.server.tools.common.util.EasyCollectionUtils; import ai.chat2db.server.tools.common.util.I18nUtils; import ai.chat2db.spi.enums.DataTypeEnum; import ai.chat2db.spi.enums.SqlTypeEnum; -import ai.chat2db.spi.model.ExecuteResult; -import ai.chat2db.spi.model.Header; import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.SQLExecutor; import ai.chat2db.spi.util.JdbcUtils; @@ -40,6 +35,7 @@ import ai.chat2db.spi.util.SqlUtils; import com.google.common.collect.Lists; import lombok.extern.slf4j.Slf4j; import org.apache.commons.collections4.CollectionUtils; +import org.apache.commons.collections4.ListUtils; import org.apache.commons.lang3.ObjectUtils; import org.apache.commons.lang3.StringUtils; import org.springframework.beans.factory.annotation.Autowired; @@ -154,9 +150,9 @@ public class DlTemplateServiceImpl implements DlTemplateService { log.warn("解析sql失败:{}", originalSql, e); } ExecuteResult executeResult = null; - if (SqlTypeEnum.SELECT.getCode().equals(sqlType) && !SqlUtils.hasPageLimit(originalSql,dbType)) { + if (SqlTypeEnum.SELECT.getCode().equals(sqlType) && !SqlUtils.hasPageLimit(originalSql, dbType)) { String pageLimit = Chat2DBContext.getSqlBuilder().pageLimit(originalSql, offset, pageNo, pageSize); - if(StringUtils.isNotBlank(pageLimit)) { + if (StringUtils.isNotBlank(pageLimit)) { executeResult = execute(pageLimit, 0, count); } } @@ -183,16 +179,16 @@ public class DlTemplateServiceImpl implements DlTemplateService { executeResult.setHasNextPage(Boolean.FALSE); } - // Splice row numbers - List
newHeaderList = new ArrayList<>(); - newHeaderList.add(Header.builder() + List
headers = executeResult.getHeaderList(); + if (executeResult.isCanEdit()) { + headers = setColumnInfo(headers, executeResult.getTableName(), param.getSchemaName(), param.getDatabaseName()); + } + Header rowNumberHeader = Header.builder() .name(I18nUtils.getMessage("sqlResult.rowNumber")) .dataType(DataTypeEnum.CHAT2DB_ROW_NUMBER - .getCode()).build()); - if (executeResult.getHeaderList() != null) { - newHeaderList.addAll(executeResult.getHeaderList()); - } - executeResult.setHeaderList(newHeaderList); + .getCode()).build(); + + executeResult.setHeaderList(ListUtils.union(Arrays.asList(rowNumberHeader), headers)); if (executeResult.getDataList() != null) { int rowNumberIncrement = 1 + Math.max(pageNo - 1, 0) * pageSize; for (int i = 0; i < executeResult.getDataList().size(); i++) { @@ -260,13 +256,13 @@ public class DlTemplateServiceImpl implements DlTemplateService { List odlRow = operation.getOldDataList(); String sql = ""; if ("UPDATE".equalsIgnoreCase(operation.getType())) { - sql = getUpdateSql(param, row, odlRow, metaSchema,keyColumns,false); + sql = getUpdateSql(param, row, odlRow, metaSchema, keyColumns, false); } else if ("CREATE".equalsIgnoreCase(operation.getType())) { sql = getInsertSql(param, row, metaSchema); } else if ("DELETE".equalsIgnoreCase(operation.getType())) { - sql = getDeleteSql(param, odlRow, metaSchema,keyColumns); + sql = getDeleteSql(param, odlRow, metaSchema, keyColumns); } else if ("UPDATE_COPY".equalsIgnoreCase(operation.getType())) { - sql = getUpdateSql(param, row, row, metaSchema,keyColumns,true); + sql = getUpdateSql(param, row, row, metaSchema, keyColumns, true); } stringBuilder.append(sql + ";\n"); @@ -274,37 +270,32 @@ public class DlTemplateServiceImpl implements DlTemplateService { return DataResult.of(stringBuilder.toString()); } - private List getPrimaryColumns(UpdateSelectResultParam param){ - TableQueryParam tableQueryParam = new TableQueryParam(); - tableQueryParam.setTableName(param.getTableName()); - tableQueryParam.setSchemaName(param.getSchemaName()); - tableQueryParam.setDatabaseName(param.getDatabaseName()); - List tableIndices = tableService.queryIndexes(tableQueryParam); - if(!CollectionUtils.isEmpty(tableIndices)) { - return tableIndices.stream() - .filter(tableIndex -> "PRIMARY".equalsIgnoreCase(tableIndex.getType())) - .map(tableIndex -> { - if(!CollectionUtils.isEmpty(tableIndex.getColumnList())) { - return tableIndex.getColumnList().get(0).getColumnName(); - } - return null; - }).collect(Collectors.toList()); + private List getPrimaryColumns(UpdateSelectResultParam param) { + List
headerList = param.getHeaderList(); + if (CollectionUtils.isEmpty(headerList)) { + return Lists.newArrayList(); + } + List keyColumns = Lists.newArrayList(); + for (Header header : headerList) { + if (header.getPrimaryKey() != null && header.getPrimaryKey()) { + keyColumns.add(header.getName()); + } } return Lists.newArrayList(); } - private String getDeleteSql(UpdateSelectResultParam param, List row, MetaData metaSchema,List keyColumns) { + private String getDeleteSql(UpdateSelectResultParam param, List row, MetaData metaSchema, List keyColumns) { StringBuilder script = new StringBuilder(); script.append("DELETE FROM ").append(param.getTableName()).append(""); - script.append(buildWhere(param.getHeaderList(), row, metaSchema,keyColumns)); + script.append(buildWhere(param.getHeaderList(), row, metaSchema, keyColumns)); return script.toString(); } - private String buildWhere(List
headerList, List row, MetaData metaSchema,List keyColumns) { + private String buildWhere(List
headerList, List row, MetaData metaSchema, List keyColumns) { StringBuilder script = new StringBuilder(); script.append(" where "); - if(CollectionUtils.isEmpty(keyColumns)) { + if (CollectionUtils.isEmpty(keyColumns)) { for (int i = 1; i < row.size(); i++) { String oldValue = row.get(i); Header header = headerList.get(i); @@ -319,12 +310,12 @@ public class DlTemplateServiceImpl implements DlTemplateService { .append(" and "); } } - }else{ + } else { for (int i = 1; i < row.size(); i++) { String oldValue = row.get(i); Header header = headerList.get(i); String columnName = header.getName(); - if(keyColumns.contains(columnName)) { + if (keyColumns.contains(columnName)) { String value = SqlUtils.getSqlValue(oldValue, header.getDataType()); if (value == null) { script.append(metaSchema.getMetaDataName(columnName)) @@ -352,7 +343,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { for (int i = 1; i < row.size(); i++) { Header header = param.getHeaderList().get(i); String newValue = row.get(i); - if(newValue!=null) { + if (newValue != null) { script.append(metaSchema.getMetaDataName(header.getName())) .append(","); } @@ -361,7 +352,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { script.append(") VALUES ("); for (int i = 1; i < row.size(); i++) { String newValue = row.get(i); - if(newValue!=null) { + if (newValue != null) { Header header = param.getHeaderList().get(i); script.append(SqlUtils.getSqlValue(newValue, header.getDataType())) .append(","); @@ -375,7 +366,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { private String getUpdateSql(UpdateSelectResultParam param, List row, List odlRow, MetaData metaSchema, - List keyColumns,boolean copy){ + List keyColumns, boolean copy) { StringBuilder script = new StringBuilder(); if (CollectionUtils.isEmpty(row) || CollectionUtils.isEmpty(odlRow)) { return ""; @@ -395,10 +386,51 @@ public class DlTemplateServiceImpl implements DlTemplateService { .append(","); } script.deleteCharAt(script.length() - 1); - script.append(buildWhere(param.getHeaderList(), odlRow, metaSchema,keyColumns)); + script.append(buildWhere(param.getHeaderList(), odlRow, metaSchema, keyColumns)); return script.toString(); } + private List
setColumnInfo(List
headers, String tableName, String schemaName, String databaseName) { + TableQueryParam tableQueryParam = new TableQueryParam(); + tableQueryParam.setTableName(tableName); + tableQueryParam.setSchemaName(schemaName); + tableQueryParam.setDatabaseName(databaseName); + List columns = tableService.queryColumns(tableQueryParam); + if (CollectionUtils.isEmpty(columns)) { + return headers; + } + Map columnMap = columns.stream().collect(Collectors.toMap(TableColumn::getName, tableColumn -> tableColumn)); + + List tableIndices = tableService.queryIndexes(tableQueryParam); + if (!CollectionUtils.isEmpty(tableIndices)) { + for (TableIndex tableIndex : tableIndices) { + if ("PRIMARY".equalsIgnoreCase(tableIndex.getType())) { + List columnList = tableIndex.getColumnList(); + if (!CollectionUtils.isEmpty(columnList)) { + for (TableIndexColumn tableIndexColumn : columnList) { + TableColumn tableColumn = columnMap.get(tableIndexColumn.getColumnName()); + if (tableColumn != null) { + tableColumn.setPrimaryKey(true); + } + } + } + } + } + } + for (Header header : headers) { + TableColumn tableColumn = columnMap.get(header.getName()); + if (tableColumn != null) { + header.setPrimaryKey(tableColumn.getPrimaryKey()); + header.setComment(tableColumn.getComment()); + header.setDefaultValue(tableColumn.getDefaultValue()); + header.setNullable(tableColumn.getNullable()); + header.setColumnSize(tableColumn.getColumnSize()); + header.setDecimalDigits(tableColumn.getDecimalDigits()); + } + } + return headers; + } + private ExecuteResult execute(String sql, Integer offset, Integer count) { ExecuteResult executeResult; try { diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/request/SelectResultUpdateRequest.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/request/SelectResultUpdateRequest.java index 0fb02315..b87b173e 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/request/SelectResultUpdateRequest.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/request/SelectResultUpdateRequest.java @@ -3,7 +3,7 @@ package ai.chat2db.server.web.api.controller.rdb.request; import ai.chat2db.server.domain.api.param.SelectResultOperation; import ai.chat2db.server.web.api.controller.data.source.request.DataSourceBaseRequest; import ai.chat2db.server.web.api.controller.data.source.request.DataSourceConsoleRequestInfo; -import ai.chat2db.server.web.api.controller.rdb.vo.HeaderVO; +import ai.chat2db.spi.model.Header; import jakarta.validation.constraints.NotEmpty; import jakarta.validation.constraints.NotNull; import lombok.Data; @@ -16,7 +16,7 @@ public class SelectResultUpdateRequest extends DataSourceBaseRequest implements /** * 展示头的列表 */ - private List headerList; + private List
headerList; /** * 修改后数据的列表 diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/vo/ExecuteResultVO.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/vo/ExecuteResultVO.java index 1ea6aadb..69147ade 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/vo/ExecuteResultVO.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/vo/ExecuteResultVO.java @@ -3,6 +3,7 @@ package ai.chat2db.server.web.api.controller.rdb.vo; import java.util.List; +import ai.chat2db.spi.model.Header; import lombok.Data; /** @@ -46,7 +47,7 @@ public class ExecuteResultVO { /** * 展示头的列表 */ - private List headerList; + private List
headerList; /** * 数据的列表 diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/vo/HeaderVO.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/vo/HeaderVO.java deleted file mode 100644 index af5596d2..00000000 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/vo/HeaderVO.java +++ /dev/null @@ -1,31 +0,0 @@ -package ai.chat2db.server.web.api.controller.rdb.vo; - -import ai.chat2db.spi.enums.DataTypeEnum; - -import lombok.AllArgsConstructor; -import lombok.Data; -import lombok.NoArgsConstructor; -import lombok.experimental.SuperBuilder; - -/** - * 单元格头 - * - * @author Jiaju Zhuang - */ -@Data -@SuperBuilder -@NoArgsConstructor -@AllArgsConstructor -public class HeaderVO { - /** - * 单元格类型 - * - * @see DataTypeEnum - */ - private String dataType; - - /** - * 展示的名字 - */ - private String name; -} diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Header.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Header.java index 9fe67797..8935a240 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Header.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/model/Header.java @@ -15,7 +15,7 @@ import lombok.experimental.SuperBuilder; @SuperBuilder @NoArgsConstructor @AllArgsConstructor -public class Header { +public class Header{ /** * 单元格类型 * @@ -27,4 +27,21 @@ public class Header { * 展示的名字 */ private String name; + + + private Boolean primaryKey; + + + private String comment; + + private String defaultValue; + + private Integer autoIncrement; + + private Integer nullable; + + private Integer columnSize; + + private Integer decimalDigits; + } From a7dc7402218b4cc8b5c5d2512286e51256ec1a8c Mon Sep 17 00:00:00 2001 From: SwallowGG <1558143046@qq.com> Date: Fri, 3 Nov 2023 12:05:32 +0800 Subject: [PATCH 4/7] Simplify the update and delete statements for query results when the table contains primary keys --- .../plugin/kingbase/KingBaseDBManage.java | 1 + .../core/impl/DlTemplateServiceImpl.java | 7 ++--- .../domain/core/util/MetaNameUtils.java | 26 +++++++++++++++++++ .../common/util/EasyCollectionUtils.java | 17 ++++++++---- 4 files changed, 43 insertions(+), 8 deletions(-) create mode 100644 chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/util/MetaNameUtils.java diff --git a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseDBManage.java b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseDBManage.java index 4a790d9c..4f967a60 100644 --- a/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-kingbase/src/main/java/ai/chat2db/plugin/kingbase/KingBaseDBManage.java @@ -19,6 +19,7 @@ public class KingBaseDBManage extends DefaultDBManage implements DBManage { SQLExecutor.getInstance().execute(connection, "SET search_path TO \"" + connectInfo.getSchemaName() + "\""); } } catch (Exception e) { + e.printStackTrace(); } } diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java index 82c43435..b5fe79ab 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java @@ -9,6 +9,7 @@ import ai.chat2db.server.domain.api.param.*; import ai.chat2db.server.domain.api.param.operation.OperationLogCreateParam; import ai.chat2db.server.domain.api.service.OperationLogService; import ai.chat2db.server.domain.api.service.TableService; +import ai.chat2db.server.domain.core.util.MetaNameUtils; import ai.chat2db.spi.MetaData; import ai.chat2db.spi.model.*; import ai.chat2db.spi.sql.ConnectInfo; @@ -180,7 +181,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { } List
headers = executeResult.getHeaderList(); - if (executeResult.isCanEdit()) { + if (executeResult.getSuccess() && executeResult.isCanEdit() && CollectionUtils.isNotEmpty(headers)){ headers = setColumnInfo(headers, executeResult.getTableName(), param.getSchemaName(), param.getDatabaseName()); } Header rowNumberHeader = Header.builder() @@ -188,7 +189,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { .dataType(DataTypeEnum.CHAT2DB_ROW_NUMBER .getCode()).build(); - executeResult.setHeaderList(ListUtils.union(Arrays.asList(rowNumberHeader), headers)); + executeResult.setHeaderList(EasyCollectionUtils.union(Arrays.asList(rowNumberHeader), headers)); if (executeResult.getDataList() != null) { int rowNumberIncrement = 1 + Math.max(pageNo - 1, 0) * pageSize; for (int i = 0; i < executeResult.getDataList().size(); i++) { @@ -392,7 +393,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { private List
setColumnInfo(List
headers, String tableName, String schemaName, String databaseName) { TableQueryParam tableQueryParam = new TableQueryParam(); - tableQueryParam.setTableName(tableName); + tableQueryParam.setTableName(MetaNameUtils.getMetaName(tableName)); tableQueryParam.setSchemaName(schemaName); tableQueryParam.setDatabaseName(databaseName); List columns = tableService.queryColumns(tableQueryParam); diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/util/MetaNameUtils.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/util/MetaNameUtils.java new file mode 100644 index 00000000..eab4d644 --- /dev/null +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/util/MetaNameUtils.java @@ -0,0 +1,26 @@ +package ai.chat2db.server.domain.core.util; + + +import org.apache.commons.lang3.StringUtils; + +public class MetaNameUtils { + + public static String getMetaName(String tableName) { + if(StringUtils.isBlank(tableName)){ + return tableName; + } + if(tableName.startsWith("`") && tableName.endsWith("`")){ + return tableName.substring(1,tableName.length()-1); + } + if(tableName.startsWith("\"") && tableName.endsWith("\"")){ + return tableName.substring(1,tableName.length()-1); + } + if(tableName.startsWith("'") && tableName.endsWith("'")){ + return tableName.substring(1,tableName.length()-1); + } + if(tableName.startsWith("[") && tableName.endsWith("]")){ + return tableName.substring(1,tableName.length()-1); + } + return tableName; + } +} diff --git a/chat2db-server/chat2db-server-tools/chat2db-server-tools-common/src/main/java/ai/chat2db/server/tools/common/util/EasyCollectionUtils.java b/chat2db-server/chat2db-server-tools/chat2db-server-tools-common/src/main/java/ai/chat2db/server/tools/common/util/EasyCollectionUtils.java index 2576da48..9f1ad43a 100644 --- a/chat2db-server/chat2db-server-tools/chat2db-server-tools-common/src/main/java/ai/chat2db/server/tools/common/util/EasyCollectionUtils.java +++ b/chat2db-server/chat2db-server-tools/chat2db-server-tools-common/src/main/java/ai/chat2db/server/tools/common/util/EasyCollectionUtils.java @@ -1,10 +1,6 @@ package ai.chat2db.server.tools.common.util; -import java.util.Collection; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import java.util.Set; +import java.util.*; import java.util.concurrent.ConcurrentHashMap; import java.util.function.Function; import java.util.function.Predicate; @@ -189,4 +185,15 @@ public class EasyCollectionUtils { return t -> seen.putIfAbsent(keyExtractor.apply(t), Boolean.TRUE) == null; } + public static List union(List list1, List list2) { + ArrayList result = new ArrayList(); + if(list1 != null && list1.size()>0) { + result.addAll(list1); + } + if(list2!= null && list2.size()>0) { + result.addAll(list2); + } + return result; + } + } From 52263ff97e7489eaa1607cac5f95c947c7fc0a99 Mon Sep 17 00:00:00 2001 From: SwallowGG <1558143046@qq.com> Date: Fri, 3 Nov 2023 12:06:36 +0800 Subject: [PATCH 5/7] Simplify the update and delete statements for query results when the table contains primary keys --- .../chat2db/server/domain/core/impl/DlTemplateServiceImpl.java | 1 + 1 file changed, 1 insertion(+) diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java index b5fe79ab..88455ec0 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java @@ -396,6 +396,7 @@ public class DlTemplateServiceImpl implements DlTemplateService { tableQueryParam.setTableName(MetaNameUtils.getMetaName(tableName)); tableQueryParam.setSchemaName(schemaName); tableQueryParam.setDatabaseName(databaseName); + tableQueryParam.setRefresh(true); List columns = tableService.queryColumns(tableQueryParam); if (CollectionUtils.isEmpty(columns)) { return headers; From 2284c32b8fd12505c582e59f736290ee6def920e Mon Sep 17 00:00:00 2001 From: robin <850379744@qq.com> Date: Fri, 3 Nov 2023 13:28:44 +0800 Subject: [PATCH 6/7] =?UTF-8?q?=E6=9C=AC=E5=9C=B0=E8=BF=9E=E6=8E=A5mysql?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../chat2db/client/Chat2DBAIStreamClient.java | 1 + .../listener/OpenAIEventSourceListener.java | 55 +++++++++++++++++-- 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/chat2db/client/Chat2DBAIStreamClient.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/chat2db/client/Chat2DBAIStreamClient.java index 1f63d4ab..34436ca3 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/chat2db/client/Chat2DBAIStreamClient.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/chat2db/client/Chat2DBAIStreamClient.java @@ -196,6 +196,7 @@ public class Chat2DBAIStreamClient { try { ChatCompletion chatCompletion = ChatCompletion.builder() .messages(chatMessages) + .model(this.model) .stream(true) .build(); diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/openai/listener/OpenAIEventSourceListener.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/openai/listener/OpenAIEventSourceListener.java index 2f0bcef0..2ae383eb 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/openai/listener/OpenAIEventSourceListener.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/openai/listener/OpenAIEventSourceListener.java @@ -2,8 +2,19 @@ package ai.chat2db.server.web.api.controller.ai.openai.listener; import java.util.Objects; +import ai.chat2db.server.domain.api.enums.AiSqlSourceEnum; +import ai.chat2db.server.domain.api.model.Config; +import ai.chat2db.server.domain.api.service.ConfigService; +import ai.chat2db.server.tools.base.wrapper.result.DataResult; +import ai.chat2db.server.web.api.controller.ai.baichuan.model.BaichuanChatCompletions; +import ai.chat2db.server.web.api.controller.ai.baichuan.model.BaichuanChatMessage; +import ai.chat2db.server.web.api.controller.ai.chat2db.client.Chat2dbAIClient; +import ai.chat2db.server.web.api.controller.ai.fastchat.model.FastChatMessage; import ai.chat2db.server.web.api.controller.ai.response.ChatCompletionResponse; +import ai.chat2db.server.web.api.controller.ai.zhipu.model.ZhipuChatCompletions; +import ai.chat2db.server.web.api.util.ApplicationContextUtil; +import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.ObjectMapper; import com.unfbx.chatgpt.entity.chat.Message; import lombok.SneakyThrows; @@ -54,16 +65,50 @@ public class OpenAIEventSourceListener extends EventSourceListener { return; } ObjectMapper mapper = new ObjectMapper(); + mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false); + ConfigService configService = ApplicationContextUtil.getBean(ConfigService.class); + DataResult chat2dbModel = configService.find(Chat2dbAIClient.CHAT2DB_OPENAI_MODEL); + String model = Objects.nonNull(chat2dbModel.getData()) ? chat2dbModel.getData().getContent() : AiSqlSourceEnum.OPENAI.getCode(); + AiSqlSourceEnum aiSqlSourceEnum = AiSqlSourceEnum.getByName(model); + String text = ""; + String completionId = null; // 读取Json - ChatCompletionResponse completionResponse = mapper.readValue(data, ChatCompletionResponse.class); - String text = completionResponse.getChoices().get(0).getDelta() == null - ? completionResponse.getChoices().get(0).getText() - : completionResponse.getChoices().get(0).getDelta().getContent(); + switch (aiSqlSourceEnum) { + case BAICHUANAI: + BaichuanChatCompletions chatCompletions = mapper.readValue(data, BaichuanChatCompletions.class); + for (BaichuanChatMessage message : chatCompletions.getData().getMessages()) { + if (message != null) { + if (message.getContent() != null) { + text = message.getContent(); + } + } + } + break; + case ZHIPUAI: + ZhipuChatCompletions zhipuChatCompletions = mapper.readValue(data, ZhipuChatCompletions.class); + text = zhipuChatCompletions.getData(); + if (Objects.isNull(text)) { + for (FastChatMessage message : zhipuChatCompletions.getBody().getChoices()) { + if (message != null && message.getContent() != null) { + text = message.getContent(); + } + } + } + break; + default: + ChatCompletionResponse completionResponse = mapper.readValue(data, ChatCompletionResponse.class); + text = completionResponse.getChoices().get(0).getDelta() == null + ? completionResponse.getChoices().get(0).getText() + : completionResponse.getChoices().get(0).getDelta().getContent(); + completionId = completionResponse.getId(); + break; + } + Message message = new Message(); if (text != null) { message.setContent(text); sseEmitter.send(SseEmitter.event() - .id(completionResponse.getId()) + .id(completionId) .data(message) .reconnectTime(3000)); } From b13ffbfc8b355301a076bd8ce208bd09e6f53470 Mon Sep 17 00:00:00 2001 From: robin <850379744@qq.com> Date: Fri, 3 Nov 2023 13:34:22 +0800 Subject: [PATCH 7/7] =?UTF-8?q?Revert=20"=E6=9C=AC=E5=9C=B0=E8=BF=9E?= =?UTF-8?q?=E6=8E=A5mysql"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 2284c32b8fd12505c582e59f736290ee6def920e. --- .../chat2db/client/Chat2DBAIStreamClient.java | 1 - .../listener/OpenAIEventSourceListener.java | 55 ++----------------- 2 files changed, 5 insertions(+), 51 deletions(-) diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/chat2db/client/Chat2DBAIStreamClient.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/chat2db/client/Chat2DBAIStreamClient.java index 34436ca3..1f63d4ab 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/chat2db/client/Chat2DBAIStreamClient.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/chat2db/client/Chat2DBAIStreamClient.java @@ -196,7 +196,6 @@ public class Chat2DBAIStreamClient { try { ChatCompletion chatCompletion = ChatCompletion.builder() .messages(chatMessages) - .model(this.model) .stream(true) .build(); diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/openai/listener/OpenAIEventSourceListener.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/openai/listener/OpenAIEventSourceListener.java index 2ae383eb..2f0bcef0 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/openai/listener/OpenAIEventSourceListener.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/openai/listener/OpenAIEventSourceListener.java @@ -2,19 +2,8 @@ package ai.chat2db.server.web.api.controller.ai.openai.listener; import java.util.Objects; -import ai.chat2db.server.domain.api.enums.AiSqlSourceEnum; -import ai.chat2db.server.domain.api.model.Config; -import ai.chat2db.server.domain.api.service.ConfigService; -import ai.chat2db.server.tools.base.wrapper.result.DataResult; -import ai.chat2db.server.web.api.controller.ai.baichuan.model.BaichuanChatCompletions; -import ai.chat2db.server.web.api.controller.ai.baichuan.model.BaichuanChatMessage; -import ai.chat2db.server.web.api.controller.ai.chat2db.client.Chat2dbAIClient; -import ai.chat2db.server.web.api.controller.ai.fastchat.model.FastChatMessage; import ai.chat2db.server.web.api.controller.ai.response.ChatCompletionResponse; -import ai.chat2db.server.web.api.controller.ai.zhipu.model.ZhipuChatCompletions; -import ai.chat2db.server.web.api.util.ApplicationContextUtil; -import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.ObjectMapper; import com.unfbx.chatgpt.entity.chat.Message; import lombok.SneakyThrows; @@ -65,50 +54,16 @@ public class OpenAIEventSourceListener extends EventSourceListener { return; } ObjectMapper mapper = new ObjectMapper(); - mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false); - ConfigService configService = ApplicationContextUtil.getBean(ConfigService.class); - DataResult chat2dbModel = configService.find(Chat2dbAIClient.CHAT2DB_OPENAI_MODEL); - String model = Objects.nonNull(chat2dbModel.getData()) ? chat2dbModel.getData().getContent() : AiSqlSourceEnum.OPENAI.getCode(); - AiSqlSourceEnum aiSqlSourceEnum = AiSqlSourceEnum.getByName(model); - String text = ""; - String completionId = null; // 读取Json - switch (aiSqlSourceEnum) { - case BAICHUANAI: - BaichuanChatCompletions chatCompletions = mapper.readValue(data, BaichuanChatCompletions.class); - for (BaichuanChatMessage message : chatCompletions.getData().getMessages()) { - if (message != null) { - if (message.getContent() != null) { - text = message.getContent(); - } - } - } - break; - case ZHIPUAI: - ZhipuChatCompletions zhipuChatCompletions = mapper.readValue(data, ZhipuChatCompletions.class); - text = zhipuChatCompletions.getData(); - if (Objects.isNull(text)) { - for (FastChatMessage message : zhipuChatCompletions.getBody().getChoices()) { - if (message != null && message.getContent() != null) { - text = message.getContent(); - } - } - } - break; - default: - ChatCompletionResponse completionResponse = mapper.readValue(data, ChatCompletionResponse.class); - text = completionResponse.getChoices().get(0).getDelta() == null - ? completionResponse.getChoices().get(0).getText() - : completionResponse.getChoices().get(0).getDelta().getContent(); - completionId = completionResponse.getId(); - break; - } - + ChatCompletionResponse completionResponse = mapper.readValue(data, ChatCompletionResponse.class); + String text = completionResponse.getChoices().get(0).getDelta() == null + ? completionResponse.getChoices().get(0).getText() + : completionResponse.getChoices().get(0).getDelta().getContent(); Message message = new Message(); if (text != null) { message.setContent(text); sseEmitter.send(SseEmitter.event() - .id(completionId) + .id(completionResponse.getId()) .data(message) .reconnectTime(3000)); }