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

@ -1,6 +1,5 @@
package ai.chat2db.plugin.sqlserver;
import ai.chat2db.spi.ValueHandler;
import ai.chat2db.spi.model.Command;
import ai.chat2db.spi.model.ExecuteResult;
import ai.chat2db.spi.sql.SQLExecutor;
@ -45,8 +44,8 @@ public class SqlServerCommandExecutor extends SQLExecutor {
*
*/
public ExecuteResult execute(final String sql, Connection connection, boolean limitRowSize, Integer offset,
Integer count, ValueHandler valueHandler)
Integer count)
throws SQLException {
return super.execute(removeSpecialGO(sql), connection, limitRowSize, offset, count, valueHandler);
return super.execute(removeSpecialGO(sql), connection, limitRowSize, offset, count);
}
}

View File

@ -2,6 +2,7 @@ package ai.chat2db.plugin.sqlserver;
import ai.chat2db.spi.DBManage;
import ai.chat2db.spi.jdbc.DefaultDBManage;
import ai.chat2db.spi.model.AsyncContext;
import ai.chat2db.spi.sql.SQLExecutor;
import java.sql.*;
@ -42,34 +43,30 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage {
+ "WHERE xtype = 'TR' ";
@Override
public String exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName) throws SQLException {
StringBuilder sqlBuilder = new StringBuilder();
exportTableData(connection, tableName, sqlBuilder);
return sqlBuilder.toString();
public void exportDatabaseData(Connection connection, String databaseName, String schemaName, String tableName, AsyncContext asyncContext) throws SQLException {
exportTableData(connection, tableName, asyncContext);
}
@Override
public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException {
StringBuilder sqlBuilder = new StringBuilder();
exportTables(connection, sqlBuilder, schemaName, containData);
exportViews(connection, databaseName, schemaName, sqlBuilder);
exportFunctions(connection, schemaName, sqlBuilder);
exportProcedures(connection, schemaName, sqlBuilder);
exportTriggers(connection, sqlBuilder);
return sqlBuilder.toString();
public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
exportTables(connection, schemaName, asyncContext);
exportViews(connection, databaseName, schemaName, asyncContext);
exportFunctions(connection, schemaName, asyncContext);
exportProcedures(connection, schemaName, asyncContext);
exportTriggers(connection, asyncContext);
}
private void exportTables(Connection connection, StringBuilder sqlBuilder, String schemaName, boolean containData) throws SQLException {
private void exportTables(Connection connection, String schemaName,AsyncContext asyncContext) throws SQLException {
String sql ="SELECT name FROM SysObjects Where XType='U'";
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
String tableName = resultSet.getString("name");
exportTable(connection, tableName, schemaName, sqlBuilder, containData);
exportTable(connection, tableName, schemaName, asyncContext);
}
}
}
private void exportTable(Connection connection, String tableName, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
private void exportTable(Connection connection, String tableName, String schemaName, AsyncContext asyncContext) throws SQLException {
try {
SQLExecutor.getInstance().execute(connection, tableDDLFunction.replace("tableSchema", schemaName),
resultSet -> null);
@ -79,23 +76,26 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage {
String sql = String.format("SELECT %s.ufn_GetCreateTableScript('%s', '%s') as ddl",schemaName,schemaName,tableName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
if (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("DROP TABLE IF EXISTS ").append(tableName).append(";").append("\n")
.append(resultSet.getString("ddl")).append("\n");
if (containData) {
exportTableData(connection, tableName, sqlBuilder);
asyncContext.write(sqlBuilder.toString());
if (asyncContext.isContainsData()) {
exportTableData(connection, tableName, asyncContext);
} else {
sqlBuilder.append("go").append("\n");
asyncContext.write("go \n");
}
}
}
}
private void exportTableData(Connection connection, String tableName, StringBuilder sqlBuilder) throws SQLException {
private void exportTableData(Connection connection, String tableName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("select * from %s", tableName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
ResultSetMetaData metaData = resultSet.getMetaData();
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("INSERT INTO ").append(tableName).append(" VALUES (");
for (int i = 1; i <= metaData.getColumnCount(); i++) {
String value = resultSet.getString(i);
@ -109,75 +109,84 @@ public class SqlServerDBManage extends DefaultDBManage implements DBManage {
}
}
sqlBuilder.append(");\n");
asyncContext.write(sqlBuilder.toString());
}
sqlBuilder.append("\n");
}
sqlBuilder.append("go").append("\n");
asyncContext.write("go \n");
}
private void exportViews(Connection connection, String databaseName, String schemaName, StringBuilder sqlBuilder) throws SQLException {
private void exportViews(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("SELECT TABLE_NAME, VIEW_DEFINITION FROM INFORMATION_SCHEMA.VIEWS " +
"WHERE TABLE_SCHEMA = '%s' AND TABLE_CATALOG = '%s'; ", schemaName, databaseName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("DROP VIEW IF EXISTS ").append(resultSet.getString("TABLE_NAME")).append(";\n").append("go").append("\n")
.append(resultSet.getString("VIEW_DEFINITION")).append(";").append("\n")
.append("go").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 name FROM sys.objects WHERE type = 'FN' and SCHEMA_ID = SCHEMA_ID('%s')", schemaName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
String functionName = resultSet.getString("name");
exportFunction(connection, functionName, schemaName, sqlBuilder);
exportFunction(connection, functionName, schemaName, asyncContext);
}
}
}
private void exportFunction(Connection connection, String functionName, String schemaName, StringBuilder sqlBuilder) throws SQLException {
private void exportFunction(Connection connection, String functionName, String schemaName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("SELECT OBJECT_DEFINITION(OBJECT_ID('%s.%s')) as ddl", schemaName, functionName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
if (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append(resultSet.getString("ddl")
.replace("CREATE FUNCTION", "CREATE OR ALTER FUNCTION"))
.append("\n").append("go").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportProcedures(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
private void exportProcedures(Connection connection, String schemaName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("SELECT name FROM sys.procedures WHERE SCHEMA_ID = SCHEMA_ID('%s')", schemaName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
String procedureName = resultSet.getString("name");
exportProcedure(connection, procedureName, schemaName, sqlBuilder);
exportProcedure(connection, procedureName, schemaName, asyncContext);
}
}
}
private void exportProcedure(Connection connection, String procedureName, String schemaName, StringBuilder sqlBuilder) throws SQLException {
private void exportProcedure(Connection connection, String procedureName, String schemaName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("SELECT definition FROM sys.sql_modules WHERE object_id = (OBJECT_ID('%s.%s'));", schemaName, procedureName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
if (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append(resultSet.getString("definition")
.replace("CREATE PROCEDURE", "CREATE OR ALTER PROCEDURE"))
.append("\n").append("go").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportTriggers(Connection connection, StringBuilder sqlBuilder) throws SQLException {
private void exportTriggers(Connection connection, AsyncContext asyncContext) throws SQLException {
try (ResultSet resultSet = connection.createStatement().executeQuery(TRIGGER_SQL_LIST)) {
while (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append(resultSet.getString("triggerDefinition")
.replace("CREATE TRIGGER", "CREATE OR ALTER TRIGGER"))
.append("\n").append("go").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}