batch export

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

View File

@ -3,6 +3,7 @@ package ai.chat2db.plugin.db2;
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.sql.Chat2DBContext;
import ai.chat2db.spi.sql.ConnectInfo;
import ai.chat2db.spi.sql.SQLExecutor;
@ -18,25 +19,23 @@ import java.sql.SQLException;
public class DB2DBManage extends DefaultDBManage implements DBManage {
@Override
public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException {
StringBuilder sqlBuilder = new StringBuilder();
exportTables(connection, schemaName, sqlBuilder, containData);
exportViews(connection, schemaName, sqlBuilder);
exportProceduresAndFunctions(connection, schemaName, sqlBuilder);
exportTriggers(connection, schemaName, sqlBuilder);
return sqlBuilder.toString();
public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
exportTables(connection, schemaName, asyncContext);
exportViews(connection, schemaName, asyncContext);
exportProceduresAndFunctions(connection, schemaName, asyncContext);
exportTriggers(connection, schemaName, asyncContext);
}
private void exportTables(Connection connection, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
private void exportTables(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException {
try (ResultSet resultSet = connection.getMetaData().getTables(null, schemaName, null, new String[]{"TABLE", "SYSTEM TABLE"})) {
while (resultSet.next()) {
exportTable(connection, schemaName, resultSet.getString("TABLE_NAME"), sqlBuilder, containData);
exportTable(connection, schemaName, resultSet.getString("TABLE_NAME"), asyncContext);
}
}
}
private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
private void exportTable(Connection connection, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException {
try {
SQLExecutor.getInstance().execute(connection, SQLConstant.TABLE_DDL_FUNCTION_SQL, resultSet -> null);
} catch (Exception e) {
@ -45,42 +44,50 @@ public class DB2DBManage extends DefaultDBManage implements DBManage {
String sql = String.format("select %s.GENERATE_TABLE_DDL('%s', '%s') as sql from %s;", schemaName, schemaName, tableName, tableName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
if (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append(resultSet.getString("sql")).append("\n");
if (containData) {
exportTableData(connection, schemaName, tableName, sqlBuilder);
asyncContext.write(sqlBuilder.toString());
if (asyncContext.isContainsData()) {
exportTableData(connection, schemaName, tableName, asyncContext);
}
}
}
}
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 TEXT from syscat.views where VIEWSCHEMA='%s';", schemaName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
String ddl = resultSet.getString("TEXT");
sqlBuilder.append(ddl).append(";").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportProceduresAndFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
private void exportProceduresAndFunctions(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("select TEXT from syscat.routines where ROUTINESCHEMA='%s';", schemaName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
String ddl = resultSet.getString("TEXT");
sqlBuilder.append(ddl).append(";").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportTriggers(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
private void exportTriggers(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("select * from SYSCAT.TRIGGERS where TRIGSCHEMA = '%s';", schemaName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
String ddl = resultSet.getString("TEXT");
sqlBuilder.append(ddl).append(";").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}