Skip to content

Commit 640ab10

Browse files
tmlx1990tmlx1990
authored
Merge branch 'chat2db:main' into main
2 parents 8625bc2 + 9b1bfad commit 640ab10

File tree

27 files changed

+2091
-201
lines changed

27 files changed

+2091
-201
lines changed

‎chat2db-server/chat2db-plugins/chat2db-dm/src/main/java/ai/chat2db/plugin/dm/DMMetaData.java‎

Lines changed: 62 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,13 @@
88
import ai.chat2db.spi.SqlBuilder;
99
import ai.chat2db.spi.jdbc.DefaultMetaService;
1010
import ai.chat2db.spi.model.*;
11+
import ai.chat2db.spi.sql.Chat2DBContext;
1112
import ai.chat2db.spi.sql.SQLExecutor;
1213
import ai.chat2db.spi.util.SortUtils;
1314
import com.google.common.collect.Lists;
1415
import jakarta.validation.constraints.NotEmpty;
16+
import lombok.extern.slf4j.Slf4j;
17+
import org.apache.commons.collections4.CollectionUtils;
1518
import org.apache.commons.lang3.StringUtils;
1619

1720
import java.sql.Connection;
@@ -20,62 +23,90 @@
2023
import java.util.*;
2124
import java.util.stream.Collectors;
2225

26+
@Slf4j
2327
public class DMMetaData extends DefaultMetaService implements MetaData {
2428

25-
private List<String> systemSchemas = Arrays.asList("CTISYS", "SYS","SYSDBA","SYSSSO","SYSAUDITOR");
29+
private List<String> systemSchemas = Arrays.asList("CTISYS", "SYS", "SYSDBA", "SYSSSO", "SYSAUDITOR");
2630

2731
@Override
2832
public List<Schema> schemas(Connection connection, String databaseName) {
2933
List<Schema> schemas = SQLExecutor.getInstance().schemas(connection, databaseName, null);
3034
return SortUtils.sortSchema(schemas, systemSchemas);
3135
}
32-
private String format(String tableName){
36+
37+
private String format(String tableName) {
3338
return "\"" + tableName + "\"";
34-
}
39+
}
3540

3641
private static String tableDDL = "SELECT dbms_metadata.get_ddl('TABLE', '%s','%s') as ddl FROM dual ;";
37-
private static String tableComment = "select COMMENTS from dba_tab_comments where OWNER='%s' and TABLE_TYPE='TABLE' and TABLE_NAME='%s';";
38-
private static String columnComment = "SELECT COLNAME,COMMENT$ FROM SYS.SYSCOLUMNCOMMENTS where SCHNAME = '%s' and TVNAME = '%s' and TABLE_TYPE = 'TABLE';";
3942

40-
public String tableDDL(Connection connection, String databaseName, String schemaName, String tableName) {
43+
public String tableDDL(Connection connection, String databaseName, String schemaName, String tableName) {
4144
String tableDDLSql = String.format(tableDDL, tableName, schemaName);
42-
String tableCommentSql = String.format(tableComment, schemaName, tableName);
43-
String columnCommentSql = String.format(columnComment, schemaName, tableName);
4445
StringBuilder ddlBuilder = new StringBuilder();
4546
SQLExecutor.getInstance().execute(connection, tableDDLSql, resultSet -> {
4647
if (resultSet.next()) {
4748
String ddl = resultSet.getString("ddl");
4849
ddlBuilder.append(ddl).append("\n");
4950
}
5051
});
51-
SQLExecutor.getInstance().execute(connection, tableCommentSql, resultSet -> {
52-
if (resultSet.next()) {
53-
String comments = resultSet.getString("COMMENTS");
54-
if (Objects.nonNull(comments)) {
55-
ddlBuilder.append("COMMENT ON TABLE ").append(format(schemaName)).append(".").append(format(tableName))
56-
.append(" IS ").append(comments).append(";").append("\n");
52+
MetaData metaData = Chat2DBContext.getMetaData();
53+
List<Table> tables = metaData.tables(connection, databaseName, schemaName, tableName);
54+
if (CollectionUtils.isNotEmpty(tables)) {
55+
String tableComment = tables.get(0).getComment();
56+
if (StringUtils.isNotBlank(tableComment)) {
57+
ddlBuilder.append("COMMENT ON TABLE ").append(format(schemaName)).append(".").append(format(tableName))
58+
.append(" IS '").append(tableComment.replace("'", "''")).append("'").append(";").append("\n");
59+
}
60+
}
61+
List<TableColumn> columns = metaData.columns(connection, databaseName, schemaName, tableName);
62+
if (CollectionUtils.isNotEmpty(columns)) {
63+
for (TableColumn column : columns) {
64+
String columnName = column.getName();
65+
String comment = column.getComment();
66+
if (StringUtils.isNotBlank(comment)) {
67+
ddlBuilder.append("COMMENT ON COLUMN ").append(format(schemaName)).append(".").append(format(tableName))
68+
.append(".").append(format(columnName)).append(" IS ")
69+
.append("'").append(comment.replace("'", "''"))
70+
.append("';").append("\n");
5771
}
5872
}
59-
});
60-
SQLExecutor.getInstance().execute(connection, columnCommentSql, resultSet -> {
61-
while (resultSet.next()) {
62-
String columnName = resultSet.getString("COLNAME");
63-
String comment = resultSet.getString("COMMENT$");
64-
ddlBuilder.append("COMMENT ON COLUMN ").append(format(schemaName)).append(".").append(format(tableName))
65-
.append(".").append(format(columnName)).append(" IS ").append("'").append(comment).append("';").append("\n");
73+
}
74+
if (tableName.startsWith("V$")){
75+
return ddlBuilder.toString();
76+
}
77+
List<TableIndex> indexes = metaData.indexes(connection, databaseName, schemaName, tableName);
78+
if (CollectionUtils.isNotEmpty(indexes)) {
79+
for (TableIndex index : indexes) {
80+
String indexName = index.getName();
81+
if (StringUtils.isNotBlank(indexName)) {
82+
String sql = "select DBMS_METADATA.GET_DDL('INDEX','%s') as INDEX_DDL";
83+
try {
84+
SQLExecutor.getInstance().execute(connection, String.format(sql,indexName), resultSet -> {
85+
if (resultSet.next()) {
86+
ddlBuilder.append(resultSet.getString("INDEX_DDL")).append("\n");
87+
}
88+
});
89+
} catch (Exception e) {
90+
log.warn("Failed to get the DDL of the index.");
91+
for (TableIndex tableIndex : indexes) {
92+
DMIndexTypeEnum indexTypeEnum = DMIndexTypeEnum.getByType(tableIndex.getType());
93+
ddlBuilder.append("\n").append(indexTypeEnum.buildIndexScript(tableIndex)).append(";");
94+
}
95+
}
96+
}
6697
}
67-
});
98+
}
6899
return ddlBuilder.toString();
69100
}
70101

71102
private static String ROUTINES_SQL
72-
= "SELECT OWNER, NAME, TEXT FROM ALL_SOURCE WHERE TYPE = '%s' AND OWNER = '%s' AND NAME = '%s' ORDER BY LINE";
103+
= "SELECT OWNER, NAME, TEXT FROM ALL_SOURCE WHERE TYPE = '%s' AND OWNER = '%s' AND NAME = '%s' ORDER BY LINE";
73104

74105
@Override
75106
public Function function(Connection connection, @NotEmpty String databaseName, String schemaName,
76-
String functionName) {
107+
String functionName) {
77108

78-
String sql = String.format(ROUTINES_SQL, "PROC",schemaName, functionName);
109+
String sql = String.format(ROUTINES_SQL, "PROC", schemaName, functionName);
79110
return SQLExecutor.getInstance().execute(connection, sql, resultSet -> {
80111
StringBuilder sb = new StringBuilder();
81112
while (resultSet.next()) {
@@ -94,8 +125,8 @@ public Function function(Connection connection, @NotEmpty String databaseName, S
94125

95126
@Override
96127
public Procedure procedure(Connection connection, @NotEmpty String databaseName, String schemaName,
97-
String procedureName) {
98-
String sql = String.format(ROUTINES_SQL, "PROC", schemaName,procedureName);
128+
String procedureName) {
129+
String sql = String.format(ROUTINES_SQL, "PROC", schemaName, procedureName);
99130
return SQLExecutor.getInstance().execute(connection, sql, resultSet -> {
100131
StringBuilder sb = new StringBuilder();
101132
while (resultSet.next()) {
@@ -111,8 +142,8 @@ public Procedure procedure(Connection connection, @NotEmpty String databaseName,
111142
}
112143

113144
private static String TRIGGER_SQL
114-
= "SELECT OWNER, TRIGGER_NAME, TABLE_OWNER, TABLE_NAME, TRIGGERING_TYPE, TRIGGERING_EVENT, STATUS, TRIGGER_BODY "
115-
+ "FROM ALL_TRIGGERS WHERE OWNER = '%s' AND TRIGGER_NAME = '%s'";
145+
= "SELECT OWNER, TRIGGER_NAME, TABLE_OWNER, TABLE_NAME, TRIGGERING_TYPE, TRIGGERING_EVENT, STATUS, TRIGGER_BODY "
146+
+ "FROM ALL_TRIGGERS WHERE OWNER = '%s' AND TRIGGER_NAME = '%s'";
116147

117148
private static String TRIGGER_SQL_LIST = "SELECT OWNER, TRIGGER_NAME FROM ALL_TRIGGERS WHERE OWNER = '%s'";
118149

@@ -134,7 +165,7 @@ public List<Trigger> triggers(Connection connection, String databaseName, String
134165

135166
@Override
136167
public Trigger trigger(Connection connection, @NotEmpty String databaseName, String schemaName,
137-
String triggerName) {
168+
String triggerName) {
138169

139170
String sql = String.format(TRIGGER_SQL, schemaName, triggerName);
140171
return SQLExecutor.getInstance().execute(connection, sql, resultSet -> {
@@ -150,7 +181,7 @@ public Trigger trigger(Connection connection, @NotEmpty String databaseName, Str
150181
}
151182

152183
private static String VIEW_SQL
153-
= "SELECT OWNER, VIEW_NAME, TEXT FROM ALL_VIEWS WHERE OWNER = '%s' AND VIEW_NAME = '%s'";
184+
= "SELECT OWNER, VIEW_NAME, TEXT FROM ALL_VIEWS WHERE OWNER = '%s' AND VIEW_NAME = '%s'";
154185

155186
@Override
156187
public Table view(Connection connection, String databaseName, String schemaName, String viewName) {

‎chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlDBManage.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,8 @@ public void updateProcedure(Connection connection, String databaseName, String s
134134
} catch (Exception e) {
135135
connection.rollback();
136136
throw new RuntimeException(e);
137+
}finally {
138+
connection.setAutoCommit(true);
137139
}
138140

139141
}

‎chat2db-server/chat2db-plugins/chat2db-oracle/src/main/java/ai/chat2db/plugin/oracle/OracleMetaData.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,11 @@ public List<TableColumn> columns(Connection connection, String databaseName, Str
118118
log.error("getDefaultValue error",e);
119119
}
120120
tableColumn.setName(resultSet.getString("COLUMN_NAME"));
121-
tableColumn.setColumnType(resultSet.getString("DATA_TYPE"));
121+
String dataType = resultSet.getString("DATA_TYPE");
122+
if(dataType.contains("(")){
123+
dataType = dataType.substring(0,dataType.indexOf("(")).trim();
124+
}
125+
tableColumn.setColumnType(dataType);
122126
Integer dataPrecision = resultSet.getInt("DATA_PRECISION");
123127
if(resultSet.getString("DATA_PRECISION") != null) {
124128
tableColumn.setColumnSize(dataPrecision);

‎chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerMetaData.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ public String tableDDL(Connection connection, String databaseName, String schema
9191
});
9292
}
9393

94-
private static String SELECT_TABLES_SQL = "SELECT t.name AS TableName, mm.value as comment FROM sys.tables t LEFT JOIN(SELECT * from sys.extended_properties ep where ep.minor_id = 0 AND ep.name = 'MS_Description') mm ON t.object_id = mm.major_id WHERE t.schema_id= SCHEMA_ID('%S')";
94+
private static String SELECT_TABLES_SQL = "SELECT t.name AS TableName, mm.value as comment FROM sys.tables t LEFT JOIN(SELECT * from sys.extended_properties ep where ep.minor_id = 0 AND ep.name = 'MS_Description') mm ON t.object_id = mm.major_id WHERE t.schema_id= SCHEMA_ID('%s')";
9595

9696
@Override
9797
public List<Table> tables(Connection connection, String databaseName, String schemaName, String tableName) {
@@ -412,4 +412,5 @@ public List<String> getSystemDatabases() {
412412
public List<String> getSystemSchemas() {
413413
return systemSchemas;
414414
}
415+
415416
}

‎chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/type/SqlServerColumnTypeEnum.java‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,11 +98,16 @@ public enum SqlServerColumnTypeEnum implements ColumnBuilder {
9898
XML("XML", false, false, true, false, false, false, true, true),
9999

100100

101+
OTHER("OTHER", false, false, true, false, false, false, true, true),
101102
;
102103
private ColumnType columnType;
103104

104105
public static SqlServerColumnTypeEnum getByType(String dataType) {
105-
return COLUMN_TYPE_MAP.get(dataType.toUpperCase());
106+
SqlServerColumnTypeEnum typeEnum = COLUMN_TYPE_MAP.get(dataType.toUpperCase());
107+
if (typeEnum == null) {
108+
return OTHER;
109+
}
110+
return typeEnum;
106111
}
107112

108113
private static Map<String, SqlServerColumnTypeEnum> COLUMN_TYPE_MAP = Maps.newHashMap();
@@ -255,7 +260,9 @@ private String buildDataType(TableColumn column, SqlServerColumnTypeEnum type) {
255260
}
256261
return script.toString();
257262
}
258-
263+
if(OTHER.equals(columnType)){
264+
return column.getColumnType();
265+
}
259266

260267
return columnType;
261268
}

‎chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,13 +100,13 @@ public DataResult<ExecuteResult> executeUpdate(DlExecuteParam param) {
100100
List<String> sqlList = SqlUtils.parse(param.getSql(), dbType);
101101
Connection connection = Chat2DBContext.getConnection();
102102
try {
103-
connection.setAutoCommit(false);
103+
// connection.setAutoCommit(false);
104104
for (String originalSql : sqlList) {
105105
ExecuteResult executeResult = executor.executeUpdate(originalSql, connection, 1);
106106
dataResult.setData(executeResult);
107107
addOperationLog(executeResult);
108108
}
109-
connection.commit();
109+
// connection.commit();
110110
} catch (Exception e) {
111111
log.error("executeUpdate error", e);
112112
dataResult.setSuccess(false);

‎chat2db-server/chat2db-spi/pom.xml‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,10 +80,15 @@
8080
<version>1.19.0</version> <!-- Make sure to use the latest version -->
8181
</dependency>
8282

83-
<!-- https://mvnrepository.com/artifact/org.mongodb/bson -->
8483
<dependency>
85-
<groupId>org.mongodb</groupId>
86-
<artifactId>bson</artifactId>
84+
<groupId>org.antlr</groupId>
85+
<artifactId>antlr4</artifactId>
86+
<version>4.9.1</version>
87+
</dependency>
88+
<dependency>
89+
<groupId>com.oceanbase</groupId>
90+
<artifactId>ob-sql-parser</artifactId>
91+
<version>1.2.1</version>
8792
</dependency>
8893
</dependencies>
8994
<build>

‎chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/enums/DataTypeEnum.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,7 @@ public static String getStringValue(String value) {
163163
}
164164
value = value.replace("\\", "\\\\");
165165
value = value.replace("'", "\\'");
166+
value = value.replace("\"", "\\\"");
166167
return "'" + value + "'";
167168
}
168169

‎chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultDBManage.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,13 +67,13 @@ private void close(Connection connection,Session session,SSHInfo ssh){
6767
if (connection != null) {
6868
try {
6969
connection.close();
70-
} catch (SQLException e) {
70+
} catch (Exception e) {
7171
}
7272
}
7373
if (session != null) {
7474
try {
7575
session.delPortForwardingL(Integer.parseInt(ssh.getLocalPort()));
76-
} catch (JSchException e) {
76+
} catch (Exception e) {
7777
}
7878
try {
7979
session.disconnect();

‎chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultSqlBuilder.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ public String buildSqlByQuery(QueryResult queryResult) {
9696
String tableName = queryResult.getTableName();
9797
StringBuilder stringBuilder = new StringBuilder();
9898
MetaData metaSchema = Chat2DBContext.getMetaData();
99+
String dbType = Chat2DBContext.getDBConfig().getDbType();
99100
List<String> keyColumns = getPrimaryColumns(headerList);
100101
for (int i = 0; i < operations.size(); i++) {
101102
ResultOperation operation = operations.get(i);
@@ -104,14 +105,19 @@ public String buildSqlByQuery(QueryResult queryResult) {
104105
String sql = "";
105106
if ("UPDATE".equalsIgnoreCase(operation.getType())) {
106107
sql = getUpdateSql(tableName, headerList, row, odlRow, metaSchema, keyColumns, false);
108+
if("MYSQL".equalsIgnoreCase(dbType)){
109+
sql = sql + " LIMIT 1";
110+
}
107111
} else if ("CREATE".equalsIgnoreCase(operation.getType())) {
108112
sql = getInsertSql(tableName, headerList, row, metaSchema);
109113
} else if ("DELETE".equalsIgnoreCase(operation.getType())) {
110114
sql = getDeleteSql(tableName, headerList, odlRow, metaSchema, keyColumns);
115+
if("MYSQL".equalsIgnoreCase(dbType)){
116+
sql = sql + " LIMIT 1";
117+
}
111118
} else if ("UPDATE_COPY".equalsIgnoreCase(operation.getType())) {
112119
sql = getUpdateSql(tableName, headerList, row, row, metaSchema, keyColumns, true);
113120
}
114-
115121
stringBuilder.append(sql + ";\n");
116122
}
117123
return stringBuilder.toString();

0 commit comments

Comments
 (0)