diff --git a/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java b/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java index 5f079170..ca944cd3 100644 --- a/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-db2/src/main/java/ai/chat2db/plugin/db2/DB2DBManage.java @@ -4,6 +4,7 @@ import ai.chat2db.plugin.db2.constant.SQLConstant; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; import ai.chat2db.spi.model.AsyncContext; +import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.ConnectInfo; import ai.chat2db.spi.sql.SQLExecutor; @@ -18,6 +19,10 @@ import java.sql.SQLException; @Slf4j public class DB2DBManage extends DefaultDBManage implements DBManage { + private static String PROCEDURE_SQL = "SELECT COUNT(*) AS procedure_count\n" + + "FROM SYSCAT.PROCEDURES\n" + + "WHERE PROCNAME = '%s';"; + @Override public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { exportTables(connection, databaseName, schemaName, asyncContext); @@ -92,6 +97,44 @@ public class DB2DBManage extends DefaultDBManage implements DBManage { } } + @Override + public void updateProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) throws SQLException { + try { + connection.setAutoCommit(false); + String procedureBody = procedure.getProcedureBody(); + boolean isCreateOrReplace = procedureBody.trim().toUpperCase().startsWith("CREATE OR REPLACE "); + + if (procedureBody == null || !procedureBody.trim().toUpperCase().startsWith("CREATE")) { + throw new IllegalArgumentException("No CREATE statement found."); + } + + String procedureNewName = getSchemaOrProcedureName(procedureBody, schemaName, procedure); + if (!procedureNewName.equals(procedure.getProcedureName())) { + procedureBody = procedureBody.replace(procedure.getProcedureName(), procedureNewName); + } + String checkProcedureSQL = String.format(PROCEDURE_SQL, procedure.getProcedureName().toUpperCase()); + String finalProcedureBody = procedureBody; + SQLExecutor.getInstance().execute(connection, checkProcedureSQL, resultSet -> { + if (resultSet.next()) { + int count = resultSet.getInt(1); + if (count >= 1) { + if (isCreateOrReplace) { + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet2 -> {}); + } else { + throw new SQLException("Procedure with the same name already exists."); + } + } + } + }); + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet -> {}); + } catch (Exception e) { + connection.rollback(); + throw new RuntimeException(e); + } finally { + connection.setAutoCommit(true); + } + } + @Override public void connectDatabase(Connection connection, String database) { ConnectInfo connectInfo = Chat2DBContext.getConnectInfo(); @@ -113,12 +156,20 @@ public class DB2DBManage extends DefaultDBManage implements DBManage { } @Override - public void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName,boolean copyData) throws SQLException { + public void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName, boolean copyData) throws SQLException { String sql = "CREATE TABLE " + newTableName + " LIKE " + tableName + " INCLUDING INDEXES"; SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); - if(copyData){ + if (copyData) { sql = "INSERT INTO " + newTableName + " SELECT * FROM " + tableName; SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } } + + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } } diff --git a/chat2db-server/chat2db-plugins/chat2db-dm/src/main/java/ai/chat2db/plugin/dm/DMDBManage.java b/chat2db-server/chat2db-plugins/chat2db-dm/src/main/java/ai/chat2db/plugin/dm/DMDBManage.java index 81297090..4a1df161 100644 --- a/chat2db-server/chat2db-plugins/chat2db-dm/src/main/java/ai/chat2db/plugin/dm/DMDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-dm/src/main/java/ai/chat2db/plugin/dm/DMDBManage.java @@ -3,6 +3,7 @@ package ai.chat2db.plugin.dm; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; import ai.chat2db.spi.model.AsyncContext; +import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.ConnectInfo; import ai.chat2db.spi.sql.SQLExecutor; @@ -29,6 +30,12 @@ public class DMDBManage extends DefaultDBManage implements DBManage { = "SELECT OWNER, TRIGGER_NAME, TABLE_OWNER, TABLE_NAME, TRIGGERING_TYPE, TRIGGERING_EVENT, STATUS, TRIGGER_BODY " + "FROM ALL_TRIGGERS WHERE OWNER = '%s' AND TRIGGER_NAME = '%s'"; + private static String PROCEDURE_SQL = "SELECT COUNT(*)\n" + + "FROM DBA_OBJECTS\n" + + "WHERE OBJECT_TYPE = 'PROCEDURE' \n" + + "AND OWNER = '%s' \n" + + "AND OBJECT_NAME = '%s'"; + @Override public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { exportTables(connection, databaseName, schemaName, asyncContext); @@ -50,11 +57,11 @@ public class DMDBManage extends DefaultDBManage implements DBManage { public void exportTable(Connection connection, String databaseName, String tableName, String schemaName, AsyncContext asyncContext) throws SQLException { String sql = """ - SELECT - (SELECT comments FROM user_tab_comments WHERE table_name = '%s') AS comments, - (SELECT dbms_metadata.get_ddl('TABLE', '%s', '%s') FROM dual) AS ddl - FROM dual; - """; + SELECT + (SELECT comments FROM user_tab_comments WHERE table_name = '%s') AS comments, + (SELECT dbms_metadata.get_ddl('TABLE', '%s', '%s') FROM dual) AS ddl + FROM dual; + """; try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(String.format(sql, tableName, tableName, schemaName))) { String formatSchemaName = format(schemaName); String formatTableName = format(tableName); @@ -79,7 +86,7 @@ public class DMDBManage extends DefaultDBManage implements DBManage { private void exportTableColumnComment(Connection connection, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { String sql = String.format("select COLNAME,COMMENT$ from SYS.SYSCOLUMNCOMMENTS\n" + - "where SCHNAME = '%s' and TVNAME = '%s'and TABLE_TYPE = 'TABLE';", schemaName, tableName); + "where SCHNAME = '%s' and TVNAME = '%s'and TABLE_TYPE = 'TABLE';", schemaName, tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { while (resultSet.next()) { StringBuilder sqlBuilder = new StringBuilder(); @@ -154,6 +161,44 @@ public class DMDBManage extends DefaultDBManage implements DBManage { } } + @Override + public void updateProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) throws SQLException { + try { + connection.setAutoCommit(false); + String procedureBody = procedure.getProcedureBody(); + boolean isCreateOrReplace = procedureBody.trim().toUpperCase().startsWith("CREATE OR REPLACE "); + + if (procedureBody == null || !procedureBody.trim().toUpperCase().startsWith("CREATE")) { + throw new IllegalArgumentException("No CREATE statement found."); + } + + String procedureNewName = getSchemaOrProcedureName(procedureBody, schemaName, procedure); + if (!procedureNewName.equals(procedure.getProcedureName())) { + procedureBody = procedureBody.replace(procedure.getProcedureName(), procedureNewName); + } + String checkProcedureSQL = String.format(PROCEDURE_SQL, schemaName.toUpperCase(),procedure.getProcedureName().toUpperCase()); + String finalProcedureBody = procedureBody; + SQLExecutor.getInstance().execute(connection, checkProcedureSQL, resultSet -> { + if (resultSet.next()) { + int count = resultSet.getInt(1); + if (count >= 1) { + if (isCreateOrReplace) { + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet2 -> {}); + } else { + throw new SQLException("Procedure with the same name already exists."); + } + } + } + }); + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet -> {}); + } catch (Exception e) { + connection.rollback(); + throw new RuntimeException(e); + } finally { + connection.setAutoCommit(true); + } + } + @Override public void connectDatabase(Connection connection, String database) { ConnectInfo connectInfo = Chat2DBContext.getConnectInfo(); @@ -173,4 +218,12 @@ public class DMDBManage extends DefaultDBManage implements DBManage { String sql = "DROP TABLE IF EXISTS " + tableName; SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } + + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } } 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 42db2114..8a546f61 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 @@ -5,6 +5,7 @@ import java.sql.SQLException; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; +import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.ConnectInfo; import ai.chat2db.spi.sql.SQLExecutor; @@ -14,6 +15,32 @@ import org.apache.commons.lang3.StringUtils; @Slf4j public class KingBaseDBManage extends DefaultDBManage implements DBManage { + @Override + public void updateProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) throws SQLException { + try { + connection.setAutoCommit(false); + String procedureBody = procedure.getProcedureBody(); + String parameterSignature = extractParameterSignature(procedureBody); + + if (procedureBody == null || !procedureBody.trim().toUpperCase().startsWith("CREATE")) { + throw new IllegalArgumentException("No CREATE statement found."); + } + + String procedureNewName = getSchemaOrProcedureName(procedureBody, schemaName, procedure); + if (!procedureNewName.equals(procedure.getProcedureName())) { + procedureBody = procedureBody.replace(procedure.getProcedureName(), procedureNewName); + } + String dropSql = "DROP PROCEDURE IF EXISTS " + procedureNewName + parameterSignature; + SQLExecutor.getInstance().execute(connection, dropSql, resultSet -> {}); + SQLExecutor.getInstance().execute(connection, procedureBody, resultSet -> {}); + } catch (Exception e) { + connection.rollback(); + throw new RuntimeException(e); + } finally { + connection.setAutoCommit(true); + } + } + @Override public void connectDatabase(Connection connection, String database) { try { @@ -66,6 +93,15 @@ public class KingBaseDBManage extends DefaultDBManage implements DBManage { SQLExecutor.getInstance().execute(connection,sql, resultSet -> null); } + @Override + public void deleteProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) { + String procedureBody = procedure.getProcedureBody(); + String parameterSignature = extractParameterSignature(procedureBody); + String procedureNewName = getSchemaOrProcedureName(procedure.getProcedureBody(), schemaName, procedure); + String sql = "DROP PROCEDURE " + procedureNewName + parameterSignature; + SQLExecutor.getInstance().execute(connection, sql, resultSet -> {}); + } + @Override public void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName,boolean copyData) throws SQLException { String sql = ""; @@ -76,4 +112,28 @@ public class KingBaseDBManage extends DefaultDBManage implements DBManage { } SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } + + private String extractParameterSignature(String input) { + int depth = 0, start = -1; + for (int i = 0; i < input.length(); i++) { + char c = input.charAt(i); + if (c == '(') { + if (depth++ == 0) start = i; + } else if (c == ')' && --depth == 0 && start != -1) { + return "(" + input.substring(start + 1, i) + ")"; + } + } + if (depth == 0) { + return ""; + } + return null; + } + + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } } diff --git a/chat2db-server/chat2db-plugins/chat2db-mariadb/src/main/java/ai/chat2db/plugin/mariadb/MariaDBManage.java b/chat2db-server/chat2db-plugins/chat2db-mariadb/src/main/java/ai/chat2db/plugin/mariadb/MariaDBManage.java index 3bcf1571..83f968cc 100644 --- a/chat2db-server/chat2db-plugins/chat2db-mariadb/src/main/java/ai/chat2db/plugin/mariadb/MariaDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-mariadb/src/main/java/ai/chat2db/plugin/mariadb/MariaDBManage.java @@ -1,12 +1,74 @@ package ai.chat2db.plugin.mariadb; import java.sql.Connection; +import java.sql.SQLException; import ai.chat2db.plugin.mysql.MysqlDBManage; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; +import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.sql.SQLExecutor; public class MariaDBManage extends MysqlDBManage implements DBManage { + String PROCEDURE_SQL = "SELECT COUNT(*)\n" + + "FROM information_schema.ROUTINES\n" + + "WHERE ROUTINE_TYPE = 'PROCEDURE'\n" + + "AND ROUTINE_NAME = '%s'\n" + + "AND ROUTINE_SCHEMA = '%s'"; + + @Override + public void updateProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) throws SQLException { + try { + connection.setAutoCommit(false); + String procedureBody = procedure.getProcedureBody(); + boolean isCreateOrReplace = procedureBody.trim().toUpperCase().startsWith("CREATE OR REPLACE "); + + if (procedureBody == null || !procedureBody.trim().toUpperCase().startsWith("CREATE")) { + throw new IllegalArgumentException("No CREATE statement found."); + } + + String procedureNewName = getSchemaOrProcedureName(procedureBody, databaseName, procedure); + if (!procedureNewName.equals(procedure.getProcedureName())) { + procedureBody = procedureBody.replace(procedure.getProcedureName(), procedureNewName); + } + String checkProcedureSQL = String.format(PROCEDURE_SQL, procedure.getProcedureName().toUpperCase(),schemaName.toUpperCase()); + String finalProcedureBody = procedureBody; + SQLExecutor.getInstance().execute(connection, checkProcedureSQL, resultSet -> { + if (resultSet.next()) { + int count = resultSet.getInt(1); + if (count >= 1) { + if (isCreateOrReplace) { + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet2 -> { + }); + } else { + throw new SQLException("Procedure with the same name already exists."); + } + } + } + }); + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet -> {}); + } catch (Exception e) { + connection.rollback(); + throw new RuntimeException(e); + } finally { + connection.setAutoCommit(true); + } + } + + @Override + public void deleteProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) { + String procedureNewName = getSchemaOrProcedureName(procedure.getProcedureBody(), databaseName, procedure); + String sql = "DROP PROCEDURE " + procedureNewName; + SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); + } + + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } + } diff --git a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlDBManage.java b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlDBManage.java index 56566a0a..ec3fb5b6 100644 --- a/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-mysql/src/main/java/ai/chat2db/plugin/mysql/MysqlDBManage.java @@ -18,9 +18,13 @@ import static cn.hutool.core.date.DatePattern.NORM_DATETIME_PATTERN; @Slf4j public class MysqlDBManage extends DefaultDBManage implements DBManage { + + private static String PROCEDURE_SQL = "SELECT COUNT(*) FROM INFORMATION_SCHEMA.ROUTINES " + + "WHERE ROUTINE_SCHEMA = '%s' AND ROUTINE_NAME = '%s' AND ROUTINE_TYPE = 'PROCEDURE'"; + @Override public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { - asyncContext.write(String.format(EXPORT_TITLE, DateUtil.format(new Date(),NORM_DATETIME_PATTERN))); + asyncContext.write(String.format(EXPORT_TITLE, DateUtil.format(new Date(), NORM_DATETIME_PATTERN))); exportTables(connection, databaseName, schemaName, asyncContext); asyncContext.setProgress(50); exportViews(connection, databaseName, asyncContext); @@ -82,9 +86,9 @@ public class MysqlDBManage extends DefaultDBManage implements DBManage { exportTableData(connection, databaseName, schemaName, tableName, asyncContext); } } - }catch (Exception e){ + } catch (Exception e) { log.error("export table error", e); - asyncContext.error(String.format("export table %s error:%s", tableName,e.getMessage())); + asyncContext.error(String.format("export table %s error:%s", tableName, e.getMessage())); } } @@ -160,9 +164,29 @@ public class MysqlDBManage extends DefaultDBManage implements DBManage { @Override public void updateProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) throws SQLException { try { - String sql = "DROP PROCEDURE " + procedure.getProcedureName(); - SQLExecutor.getInstance().execute(connection, sql, resultSet -> {}); + connection.setAutoCommit(false); String procedureBody = procedure.getProcedureBody(); + if (procedureBody == null || !procedureBody.trim().toUpperCase().startsWith("CREATE")) { + throw new IllegalArgumentException("No CREATE statement found."); + } + + String procedureNewName = getSchemaOrProcedureName(procedureBody, databaseName, procedure); + if (!procedureNewName.equals(procedure.getProcedureName())) { + procedureBody = procedureBody.replace(procedure.getProcedureName(), procedureNewName); + } + String checkProcedureSQL = String.format(PROCEDURE_SQL, databaseName, procedure.getProcedureName()); + SQLExecutor.getInstance().execute(connection, checkProcedureSQL, resultSet -> { + try { + if (resultSet.next()) { + int count = resultSet.getInt(1); + if (count >= 1) { + throw new Exception("Procedure already exists"); + } + } + } catch (Exception e) { + e.printStackTrace(); + } + }); SQLExecutor.getInstance().execute(connection, procedureBody, resultSet -> {}); } catch (Exception e) { connection.rollback(); @@ -192,6 +216,21 @@ public class MysqlDBManage extends DefaultDBManage implements DBManage { SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } + @Override + public void deleteProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) { + String procedureNewName = getSchemaOrProcedureName(procedure.getProcedureBody(), databaseName, procedure); + String sql = "DROP PROCEDURE " + procedureNewName; + SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); + } + + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } + public static String format(String tableName) { return "`" + tableName + "`"; } diff --git a/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseDBManage.java b/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseDBManage.java index ba3a4c89..bb98b990 100644 --- a/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-oceanbase/src/main/java/ai/chat2db/plugin/oceanbase/OceanBaseDBManage.java @@ -3,7 +3,62 @@ package ai.chat2db.plugin.oceanbase; import ai.chat2db.plugin.mysql.MysqlDBManage; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; +import ai.chat2db.spi.model.Procedure; +import ai.chat2db.spi.sql.SQLExecutor; + +import java.sql.Connection; +import java.sql.SQLException; public class OceanBaseDBManage extends MysqlDBManage implements DBManage { + private static String PROCEDURE_SQL = "SELECT COUNT(*) FROM information_schema.routines " + + "WHERE routine_type='PROCEDURE' AND routine_schema='%s' AND routine_name='%s';"; + + + @Override + public void updateProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) throws SQLException { + try { + connection.setAutoCommit(false); + String procedureBody = procedure.getProcedureBody(); + boolean isCreateOrReplace = procedureBody.trim().toUpperCase().startsWith("CREATE OR REPLACE "); + + if (procedureBody == null || !procedureBody.trim().toUpperCase().startsWith("CREATE")) { + throw new IllegalArgumentException("No CREATE statement found."); + } + + String procedureNewName = getSchemaOrProcedureName(procedureBody, schemaName, procedure); + if (!procedureNewName.equals(procedure.getProcedureName())) { + procedureBody = procedureBody.replace(procedure.getProcedureName(), procedureNewName); + } + String checkProcedureSQL = String.format(PROCEDURE_SQL, schemaName.toUpperCase(), procedure.getProcedureName().toUpperCase()); + String finalProcedureBody = procedureBody; + SQLExecutor.getInstance().execute(connection, checkProcedureSQL, resultSet -> { + if (resultSet.next()) { + int count = resultSet.getInt(1); + if (count >= 1) { + if (isCreateOrReplace) { + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet2 -> {}); + } else { + throw new SQLException("Procedure with the same name already exists."); + } + } + } + }); + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet -> {}); + } catch (Exception e) { + connection.rollback(); + throw new RuntimeException(e); + } finally { + connection.setAutoCommit(true); + } + } + + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } + } diff --git a/chat2db-server/chat2db-plugins/chat2db-oracle/src/main/java/ai/chat2db/plugin/oracle/OracleDBManage.java b/chat2db-server/chat2db-plugins/chat2db-oracle/src/main/java/ai/chat2db/plugin/oracle/OracleDBManage.java index fe91d6d6..e14271ce 100644 --- a/chat2db-server/chat2db-plugins/chat2db-oracle/src/main/java/ai/chat2db/plugin/oracle/OracleDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-oracle/src/main/java/ai/chat2db/plugin/oracle/OracleDBManage.java @@ -20,6 +20,9 @@ import java.util.List; @Slf4j public class OracleDBManage extends DefaultDBManage implements DBManage { + private static String PROCEDURE_SQL = "SELECT COUNT(*) FROM ALL_OBJECTS " + + "WHERE OWNER = '%s' AND OBJECT_NAME = '%s' AND OBJECT_TYPE = 'PROCEDURE'"; + public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { exportTables(connection, databaseName, schemaName, asyncContext); exportViews(connection, asyncContext, schemaName); @@ -110,6 +113,45 @@ public class OracleDBManage extends DefaultDBManage implements DBManage { asyncContext.write(function.getFunctionBody() + "\n"); } + @Override + public void updateProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) throws SQLException { + try { + connection.setAutoCommit(false); + String procedureBody = procedure.getProcedureBody(); + boolean isCreateOrReplace = procedureBody.trim().toUpperCase().startsWith("CREATE OR REPLACE "); + + if (procedureBody == null || !procedureBody.trim().toUpperCase().startsWith("CREATE")) { + throw new IllegalArgumentException("No CREATE statement found."); + } + + String procedureNewName = getSchemaOrProcedureName(procedureBody, schemaName, procedure); + if (!procedureNewName.equals(procedure.getProcedureName())) { + procedureBody = procedureBody.replace(procedure.getProcedureName(), procedureNewName); + } + String checkProcedureSQL = String.format(PROCEDURE_SQL, schemaName.toUpperCase(), procedure.getProcedureName().toUpperCase()); + String finalProcedureBody = procedureBody; + SQLExecutor.getInstance().execute(connection, checkProcedureSQL, resultSet -> { + if (resultSet.next()) { + int count = resultSet.getInt(1); + if (count >= 1) { + if (isCreateOrReplace) { + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet2 -> {}); + } else { + throw new SQLException("Procedure with the same name already exists."); + } + } + } + }); + + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet -> {}); + + } catch (Exception e) { + connection.rollback(); + throw new RuntimeException(e); + } finally { + connection.setAutoCommit(true); + } + } @Override public void connectDatabase(Connection connection, String database) { @@ -142,4 +184,12 @@ public class OracleDBManage extends DefaultDBManage implements DBManage { SQLExecutor.getInstance().execute(connection, sql, (resultSet) -> null); } + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } + } diff --git a/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java b/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java index d8411b91..86282870 100644 --- a/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java @@ -3,6 +3,7 @@ package ai.chat2db.plugin.postgresql; import ai.chat2db.spi.DBManage; import ai.chat2db.spi.jdbc.DefaultDBManage; import ai.chat2db.spi.model.AsyncContext; +import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.ConnectInfo; import ai.chat2db.spi.sql.SQLExecutor; @@ -13,6 +14,8 @@ import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; +import java.util.regex.Matcher; +import java.util.regex.Pattern; import static ai.chat2db.plugin.postgresql.consts.SQLConst.ENUM_TYPE_DDL_SQL; @@ -45,13 +48,13 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { tableNames.add(tableName); } for (String tableName : tableNames) { - exportTable(connection, databaseName,schemaName, tableName, asyncContext); + exportTable(connection, databaseName, schemaName, tableName, asyncContext); } } } - public void exportTable(Connection connection,String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { + public void exportTable(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { String sql = String.format("select pg_get_tabledef('%s','%s',true,'COMMENTS') as ddl;", schemaName, tableName); try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) { if (resultSet.next()) { @@ -119,6 +122,33 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { } } + @Override + public void updateProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) throws SQLException { + try { + connection.setAutoCommit(false); + String procedureBody = procedure.getProcedureBody(); + boolean isCreateOrReplace = procedureBody.trim().toUpperCase().startsWith("CREATE OR REPLACE "); + String parameterSignature = extractParameterSignature(procedureBody); + + if (procedureBody == null || !procedureBody.trim().toUpperCase().startsWith("CREATE")) { + throw new IllegalArgumentException("No CREATE statement found."); + } + + String procedureNewName = getSchemaOrProcedureName(procedureBody, schemaName, procedure); + if (!procedureNewName.equals(procedure.getProcedureName())) { + procedureBody = procedureBody.replace(procedure.getProcedureName(), procedureNewName); + } + String dropSql = "DROP PROCEDURE IF EXISTS " + procedureNewName + parameterSignature; + SQLExecutor.getInstance().execute(connection, dropSql, resultSet -> {}); + SQLExecutor.getInstance().execute(connection, procedureBody, resultSet -> {}); + } catch (Exception e) { + connection.rollback(); + throw new RuntimeException(e); + } finally { + connection.setAutoCommit(true); + } + } + @Override public Connection getConnection(ConnectInfo connectInfo) { String url = connectInfo.getUrl(); @@ -160,14 +190,46 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage { } @Override - public void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName,boolean copyData) throws SQLException { + public void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName, boolean copyData) throws SQLException { String sql = ""; - if(copyData){ + if (copyData) { sql = "CREATE TABLE " + newTableName + " AS TABLE " + tableName + " WITH DATA"; - }else { + } else { sql = "CREATE TABLE " + newTableName + " AS TABLE " + tableName + " WITH NO DATA"; } SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } + @Override + public void deleteProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) { + String procedureBody = procedure.getProcedureBody(); + String parameterSignature = extractParameterSignature(procedureBody); + String procedureNewName = getSchemaOrProcedureName(procedureBody, schemaName, procedure); + String sql = "DROP PROCEDURE " + procedureNewName + parameterSignature; + SQLExecutor.getInstance().execute(connection, sql, resultSet -> {}); + } + + private String extractParameterSignature(String input) { + int depth = 0, start = -1; + for (int i = 0; i < input.length(); i++) { + char c = input.charAt(i); + if (c == '(') { + if (depth++ == 0) start = i; + } else if (c == ')' && --depth == 0 && start != -1) { + return "(" + input.substring(start + 1, i) + ")"; + } + } + if (depth == 0) { + return ""; + } + return null; + } + + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } } diff --git a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerDBManage.java b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerDBManage.java index 93d1ed95..e175951b 100644 --- a/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerDBManage.java +++ b/chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerDBManage.java @@ -6,6 +6,7 @@ import ai.chat2db.spi.ValueProcessor; import ai.chat2db.spi.jdbc.DefaultDBManage; import ai.chat2db.spi.model.AsyncContext; import ai.chat2db.spi.model.JDBCDataValue; +import ai.chat2db.spi.model.Procedure; import ai.chat2db.spi.sql.Chat2DBContext; import ai.chat2db.spi.sql.SQLExecutor; import ai.chat2db.spi.util.ResultSetUtils; @@ -51,6 +52,11 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { + "triggerDefinition, CASE WHEN status & 1 = 1 THEN 'Enabled' ELSE 'Disabled' END AS Status FROM sysobjects " + "WHERE xtype = 'TR' "; + private static String PROCEDURE_SQL = "SELECT COUNT(*) AS ProcedureCount\n" + + "FROM sys.procedures\n" + + "WHERE [name] = N'%s'\n" + + "AND schema_id = SCHEMA_ID(N'%s');"; + @Override public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException { @@ -194,6 +200,44 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { } } + @Override + public void updateProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) throws SQLException { + try { + connection.setAutoCommit(false); + String procedureBody = procedure.getProcedureBody(); + boolean isCreateOrAlter = procedureBody.trim().toUpperCase().startsWith("CREATE OR ALTER "); + + if (procedureBody == null || !procedureBody.trim().toUpperCase().startsWith("CREATE")) { + throw new IllegalArgumentException("No CREATE statement found."); + } + + String procedureNewName = getSchemaOrProcedureName(procedureBody, schemaName, procedure); + if (!procedureNewName.equals(procedure.getProcedureName())) { + procedureBody = procedureBody.replace(procedure.getProcedureName(), procedureNewName); + } + String checkProcedureSQL = String.format(PROCEDURE_SQL, procedure.getProcedureName().toUpperCase(), schemaName.toUpperCase()); + String finalProcedureBody = procedureBody; + SQLExecutor.getInstance().execute(connection, checkProcedureSQL, resultSet -> { + if (resultSet.next()) { + int count = resultSet.getInt(1); + if (count >= 1) { + if (isCreateOrAlter) { + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet2 -> {}); + } else { + throw new SQLException("Procedure with the same name already exists."); + } + } + } + }); + SQLExecutor.getInstance().execute(connection, finalProcedureBody, resultSet -> {}); + } catch (Exception e) { + connection.rollback(); + throw new RuntimeException(e); + } finally { + connection.setAutoCommit(true); + } + } + @Override public void connectDatabase(Connection connection, String database) { try { @@ -213,4 +257,12 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage { } SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } + + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } } diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/ProcedureService.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/ProcedureService.java index d4f8053c..0c8cdadb 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/ProcedureService.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/service/ProcedureService.java @@ -20,6 +20,7 @@ public interface ProcedureService { /** * Querying procedure information. + * * @param databaseName * @param schemaName * @param procedureName @@ -28,10 +29,22 @@ public interface ProcedureService { DataResult detail(String databaseName, String schemaName, String procedureName); /** + * Update procedure. + * * @param databaseName * @param schemaName * @param procedure * @return */ ActionResult update(String databaseName, String schemaName, Procedure procedure) throws SQLException; + + /** + * Delete procedure. + * + * @param databaseName + * @param schemaName + * @param procedure + * @return + */ + ActionResult delete(String databaseName, String schemaName, Procedure procedure); } diff --git a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/ProcedureServiceImpl.java b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/ProcedureServiceImpl.java index b8dec55c..326537e1 100644 --- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/ProcedureServiceImpl.java +++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/ProcedureServiceImpl.java @@ -27,4 +27,10 @@ public class ProcedureServiceImpl implements ProcedureService { Chat2DBContext.getDBManage().updateProcedure(Chat2DBContext.getConnection(), databaseName, schemaName, procedure); return ActionResult.isSuccess(); } + + @Override + public ActionResult delete(String databaseName, String schemaName, Procedure procedure) { + Chat2DBContext.getDBManage().deleteProcedure(Chat2DBContext.getConnection(), databaseName, schemaName, procedure); + return ActionResult.isSuccess(); + } } diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/ProcedureController.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/ProcedureController.java index be7980ec..1a7053ee 100644 --- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/ProcedureController.java +++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/ProcedureController.java @@ -41,8 +41,14 @@ public class ProcedureController { } @PostMapping("/update") - public ActionResult update(@Valid @RequestBody ProcedureUpdateRequest request) throws SQLException { + public ActionResult update(@Valid ProcedureUpdateRequest request) throws SQLException { Procedure procedure = procedureConverter.request2param(request); return procedureService.update(request.getDatabaseName(), request.getSchemaName(), procedure); } + + @PostMapping("/delete") + public ActionResult delete(@Valid ProcedureUpdateRequest request) { + Procedure procedure = procedureConverter.request2param(request); + return procedureService.delete(request.getDatabaseName(), request.getSchemaName(), procedure); + } } diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/DBManage.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/DBManage.java index e2410c46..4a5664c8 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/DBManage.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/DBManage.java @@ -163,4 +163,13 @@ public interface DBManage { * @return */ void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName,boolean copyData) throws SQLException; + + /** + * delete procedure + * + * @param databaseName + * @param schemaName + * @param procedure + */ + void deleteProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure); } \ No newline at end of file diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultDBManage.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultDBManage.java index 9efafff3..e0b1e9a7 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultDBManage.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/jdbc/DefaultDBManage.java @@ -181,24 +181,27 @@ public class DefaultDBManage implements DBManage { @Override public void truncateTable(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException { - String sql = "TRUNCATE TABLE " + tableName ; + String sql = "TRUNCATE TABLE " + tableName; SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } @Override - public void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName,boolean copyData) throws SQLException { + public void copyTable(Connection connection, String databaseName, String schemaName, String tableName, String newTableName, boolean copyData) throws SQLException { String sql = ""; - if(copyData){ - sql = "CREATE TABLE " + newTableName + " AS SELECT * FROM " + tableName; - }else { + if (copyData) { + sql = "CREATE TABLE " + newTableName + " AS SELECT * FROM " + tableName; + } else { sql = "CREATE TABLE " + newTableName + " AS SELECT * FROM " + tableName + " WHERE 1=0"; } SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); } -// public void exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException { -// exportTableData(connection, databaseName, schemaName, tableName, asyncContext); -// } + @Override + public void deleteProcedure(Connection connection, String databaseName, String schemaName, Procedure procedure) { + String procedureNewName = getSchemaOrProcedureName(procedure.getProcedureBody(), schemaName, procedure); + String sql = "DROP PROCEDURE " + procedureNewName; + SQLExecutor.getInstance().execute(connection, sql, resultSet -> null); + } @Override public void dropTable(Connection connection, String databaseName, String schemaName, String tableName) { @@ -223,10 +226,18 @@ public class DefaultDBManage implements DBManage { valueList.add(valueString); } String insertSql = sqlBuilder.buildSingleInsertSql(null, null, tableName, columnList, valueList); - asyncContext.write(insertSql+";"); + asyncContext.write(insertSql + ";"); valueList.clear(); } }); } + + private static String getSchemaOrProcedureName(String procedureBody, String schemaName, Procedure procedure) { + if (procedureBody.toLowerCase().contains(schemaName.toLowerCase())) { + return procedure.getProcedureName(); + } else { + return schemaName + "." + procedure.getProcedureName(); + } + } } \ No newline at end of file diff --git a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java index 5d1b7e36..8e249b36 100644 --- a/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java +++ b/chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/SQLExecutor.java @@ -82,10 +82,6 @@ public class SQLExecutor implements CommandExecutor { } } -// public void execute(Connection connection, String sql, Consumer> headerConsumer, -// Consumer> rowConsumer, ValueHandler valueHandler) { -// execute(connection, sql, headerConsumer, rowConsumer, true, valueHandler); -// } public void execute( Connection connection, String sql,