opensource/Chat2DB
1
0
Fork 0
mirror of https://github.com/CodePhiliaX/Chat2DB.git synced 2025-07-29 10:43:06 +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

72
chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java
View File

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