opensource/Chat2DB
1
0
Fork 0
mirror of https://github.com/CodePhiliaX/Chat2DB.git synced 2025-07-29 02:32:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

batch export

Browse Source
This commit is contained in:
SwallowGG
2024-06-21 17:21:28 +08:00
parent 96e6340f8c
commit 9421c3d3f4
31 changed files with 871 additions and 764 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -2,6 +2,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.sql.Chat2DBContext;
import ai.chat2db.spi.sql.ConnectInfo;
import ai.chat2db.spi.sql.SQLExecutor;
@ -15,24 +16,24 @@ import static ai.chat2db.plugin.postgresql.consts.SQLConst.*;
public class PostgreSQLDBManage extends DefaultDBManage implements DBManage {
public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException {
StringBuilder sqlBuilder = new StringBuilder();
exportTypes(connection, sqlBuilder);
exportTables(connection, databaseName, schemaName, sqlBuilder, containData);
exportViews(connection, schemaName, sqlBuilder);
exportFunctions(connection, schemaName, sqlBuilder);
exportTriggers(connection, sqlBuilder);
return sqlBuilder.toString();
public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
exportTypes(connection, asyncContext);
exportTables(connection, databaseName, schemaName, asyncContext);
exportViews(connection, schemaName, asyncContext);
exportFunctions(connection, schemaName, asyncContext);
exportTriggers(connection, asyncContext);
}
private void exportTypes(Connection connection, StringBuilder sqlBuilder) throws SQLException {
private void exportTypes(Connection connection, AsyncContext asyncContext) throws SQLException {
try (ResultSet resultSet = connection.createStatement().executeQuery(ENUM_TYPE_DDL_SQL)) {
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append(resultSet.getString("ddl")).append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportTables(Connection connection, String databaseName, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
private void exportTables(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, schemaName, null,
new String[]{"TABLE", "SYSTEM TABLE","PARTITIONED TABLE"})) {
ArrayList<String> tableNames = new ArrayList<>();
@ -41,58 +42,66 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage {
tableNames.add(tableName);
}
for (String tableName : tableNames) {
exportTable(connection, schemaName, tableName, sqlBuilder);
exportTable(connection, schemaName, tableName, asyncContext);
}
if (containData) {
if (asyncContext.isContainsData()) {
for (String tableName : tableNames) {
exportTableData(connection, schemaName, tableName, sqlBuilder);
exportTableData(connection, schemaName, tableName, asyncContext);
}
}
}
}
private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder) throws SQLException {
private void exportTable(Connection connection, 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()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("\n").append("DROP TABLE IF EXISTS ").append(tableName).append(";").append("\n")
.append(resultSet.getString("ddl")).append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportViews(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
private void exportViews(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("SELECT table_name, view_definition FROM information_schema.views WHERE table_schema = '%s'",schemaName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
String viewName = resultSet.getString("table_name");
String viewDefinition = resultSet.getString("view_definition");
sqlBuilder.append("CREATE OR REPLACE VIEW ").append(viewName).append(" AS ").append(viewDefinition).append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
private void exportFunctions(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("SELECT proname, pg_get_functiondef(oid) AS function_definition FROM pg_proc " +
"WHERE pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = '%s')", schemaName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
String functionName = resultSet.getString("proname");
String functionDefinition = resultSet.getString("function_definition");
sqlBuilder.append("DROP FUNCTION IF EXISTS ").append(schemaName).append(".").append(functionName).append(";\n");
sqlBuilder.append(functionDefinition).append(";\n\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportTriggers(Connection connection, StringBuilder sqlBuilder) throws SQLException {
private void exportTriggers(Connection connection, AsyncContext asyncContext) throws SQLException {
String sql = "SELECT pg_get_triggerdef(oid) AS trigger_definition FROM pg_trigger";
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(sql)) {
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append(resultSet.getString("trigger_definition")).append(";").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}