opensource/Chat2DB
1
0
Fork 0
mirror of https://github.com/CodePhiliaX/Chat2DB.git synced 2025-09-20 11:15:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update and delete operations of procedures

Browse Source
This commit is contained in:
suyue
2024-07-30 18:36:27 +08:00
parent cc00466c8e
commit 052a36b73e
15 changed files with 557 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
chat2db-server/chat2db-plugins/chat2db-sqlserver/src/main/java/ai/chat2db/plugin/sqlserver/SqlServerDBManage.java
View File

@ -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();
}
}
}