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

@ -2,6 +2,7 @@ package ai.chat2db.plugin.mysql;
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.SQLExecutor;
import org.springframework.util.StringUtils;
@ -11,113 +12,121 @@ import java.util.Objects;
public class MysqlDBManage extends DefaultDBManage implements DBManage {
@Override
public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException {
StringBuilder sqlBuilder = new StringBuilder();
exportTables(connection, databaseName, sqlBuilder, containData);
exportViews(connection, databaseName, sqlBuilder);
exportProcedures(connection, sqlBuilder);
exportTriggers(connection, sqlBuilder);
exportFunctions(connection, databaseName, sqlBuilder);
return sqlBuilder.toString();
public void exportDatabase(Connection connection, String databaseName, String schemaName, AsyncContext asyncContext) throws SQLException {
exportTables(connection, databaseName, asyncContext);
exportViews(connection, databaseName, asyncContext);
exportProcedures(connection, asyncContext);
exportTriggers(connection, asyncContext);
exportFunctions(connection, databaseName, asyncContext);
}
private void exportFunctions(Connection connection, String databaseName, StringBuilder sqlBuilder) throws SQLException {
private void exportFunctions(Connection connection, String databaseName, AsyncContext asyncContext) throws SQLException {
try (ResultSet resultSet = connection.getMetaData().getFunctions(databaseName, null, null)) {
while (resultSet.next()) {
exportFunction(connection, resultSet.getString("FUNCTION_NAME"), sqlBuilder);
exportFunction(connection, resultSet.getString("FUNCTION_NAME"), asyncContext);
}
}
}
private void exportFunction(Connection connection, String functionName, StringBuilder sqlBuilder) throws SQLException {
private void exportFunction(Connection connection, String functionName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("SHOW CREATE FUNCTION %s;", functionName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
if (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("DROP FUNCTION IF EXISTS ").append(functionName).append(";").append("\n")
.append(resultSet.getString("Create Function")).append(";").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportTables(Connection connection, String databaseName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
private void exportTables(Connection connection, String databaseName, AsyncContext asyncContext) throws SQLException {
try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, null, null, new String[]{"TABLE", "SYSTEM TABLE"})) {
while (resultSet.next()) {
exportTable(connection, resultSet.getString("TABLE_NAME"), sqlBuilder, containData);
exportTable(connection, resultSet.getString("TABLE_NAME"), asyncContext);
}
}
}
private void exportTable(Connection connection, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
private void exportTable(Connection connection, String tableName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("show create table %s ", tableName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
if (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("DROP TABLE IF EXISTS ").append(format(tableName)).append(";").append("\n")
.append(resultSet.getString("Create Table")).append(";").append("\n");
if (containData) {
exportTableData(connection, null,tableName, sqlBuilder);
asyncContext.write(sqlBuilder.toString());
if (asyncContext.isContainsData()) {
exportTableData(connection, null,tableName, asyncContext);
}
}
}
}
private void exportViews(Connection connection, String databaseName, StringBuilder sqlBuilder) throws SQLException {
private void exportViews(Connection connection, String databaseName, AsyncContext asyncContext) throws SQLException {
try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, null, null, new String[]{"VIEW"})) {
while (resultSet.next()) {
exportView(connection, resultSet.getString("TABLE_NAME"), sqlBuilder);
exportView(connection, resultSet.getString("TABLE_NAME"), asyncContext);
}
}
}
private void exportView(Connection connection, String viewName, StringBuilder sqlBuilder) throws SQLException {
private void exportView(Connection connection, String viewName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("show create view %s ", viewName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
if (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("DROP VIEW IF EXISTS ").append(format(viewName)).append(";").append("\n")
.append(resultSet.getString("Create View")).append(";").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}
private void exportProcedures(Connection connection, StringBuilder sqlBuilder) throws SQLException {
private void exportProcedures(Connection connection, AsyncContext asyncContext) throws SQLException {
String sql = "SHOW PROCEDURE STATUS WHERE Db = DATABASE()";
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
exportProcedure(connection, resultSet.getString("Name"), sqlBuilder);
exportProcedure(connection, resultSet.getString("Name"), asyncContext);
}
}
}
private void exportProcedure(Connection connection, String procedureName, StringBuilder sqlBuilder) throws SQLException {
private void exportProcedure(Connection connection, String procedureName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("show create procedure %s ", procedureName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
if (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("DROP PROCEDURE IF EXISTS ").append(format(procedureName)).append(";").append("\n")
.append("delimiter ;;").append("\n").append(resultSet.getString("Create Procedure")).append(";;")
.append("\n").append("delimiter ;").append("\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 = "SHOW TRIGGERS";
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) {
String triggerName = resultSet.getString("Trigger");
exportTrigger(connection, triggerName, sqlBuilder);
exportTrigger(connection, triggerName, asyncContext);
}
}
}
private void exportTrigger(Connection connection, String triggerName, StringBuilder sqlBuilder) throws SQLException {
private void exportTrigger(Connection connection, String triggerName, AsyncContext asyncContext) throws SQLException {
String sql = String.format("show create trigger %s ", triggerName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
if (resultSet.next()) {
StringBuilder sqlBuilder = new StringBuilder();
sqlBuilder.append("DROP TRIGGER IF EXISTS ").append(format(triggerName)).append(";").append("\n")
.append("delimiter ;;").append("\n").append(resultSet.getString("SQL Original Statement")).append(";;")
.append("\n").append("delimiter ;").append("\n");
asyncContext.write(sqlBuilder.toString());
}
}
}

View File

@ -4,9 +4,8 @@ import ai.chat2db.plugin.mysql.builder.MysqlSqlBuilder;
import ai.chat2db.plugin.mysql.type.*;
import ai.chat2db.plugin.mysql.value.MysqlValueProcessor;
import ai.chat2db.spi.MetaData;
import ai.chat2db.spi.ValueProcessor;
import ai.chat2db.spi.SqlBuilder;
import ai.chat2db.spi.ValueHandler;
import ai.chat2db.spi.ValueProcessor;
import ai.chat2db.spi.jdbc.DefaultMetaService;
import ai.chat2db.spi.model.*;
import ai.chat2db.spi.sql.SQLExecutor;
@ -348,10 +347,10 @@ public class MysqlMetaData extends DefaultMetaService implements MetaData {
return Arrays.stream(names).filter(name -> StringUtils.isNotBlank(name)).map(name -> "`" + name + "`").collect(Collectors.joining("."));
}
@Override
public ValueHandler getValueHandler() {
return new MysqlValueHandler();
}
// @Override
// public ValueHandler getValueHandler() {
// return new MysqlValueHandler();
// }
@Override
public ValueProcessor getValueProcessor() {

View File

@ -1,45 +1,45 @@
package ai.chat2db.plugin.mysql;
import ai.chat2db.plugin.mysql.type.MysqlColumnTypeEnum;
import ai.chat2db.plugin.mysql.value.GeometryValueHandler;
import ai.chat2db.spi.ValueHandler;
import ai.chat2db.spi.jdbc.DefaultValueHandler;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Map;
public class MysqlValueHandler extends DefaultValueHandler {
private static final Map<String, ValueHandler> VALUE_HANDLER_MAP = Map.of(
MysqlColumnTypeEnum.GEOMETRY.name(), new GeometryValueHandler()
);
@Override
public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException {
try {
Object obj = rs.getObject(index);
if (obj == null) {
return null;
}
String columnTypeName = rs.getMetaData().getColumnTypeName(index);
if (MysqlColumnTypeEnum.GEOMETRY.name().equalsIgnoreCase(columnTypeName)
|| MysqlColumnTypeEnum.POINT.name().equalsIgnoreCase(columnTypeName)
|| MysqlColumnTypeEnum.LINESTRING.name().equalsIgnoreCase(columnTypeName)
|| MysqlColumnTypeEnum.POLYGON.name().equalsIgnoreCase(columnTypeName)
|| MysqlColumnTypeEnum.MULTIPOINT.name().equalsIgnoreCase(columnTypeName)
|| MysqlColumnTypeEnum.MULTILINESTRING.name().equalsIgnoreCase(columnTypeName)
|| MysqlColumnTypeEnum.MULTIPOLYGON.name().equalsIgnoreCase(columnTypeName)
|| MysqlColumnTypeEnum.GEOMETRYCOLLECTION.name().equalsIgnoreCase(columnTypeName)
) {
ValueHandler handler = VALUE_HANDLER_MAP.get(MysqlColumnTypeEnum.GEOMETRY.name());
return handler.getString(rs, index, limitSize);
} else {
return super.getString(rs, index, limitSize);
}
}catch (Exception e){
return rs.getString(index);
}
}
}
//package ai.chat2db.plugin.mysql;
//
//import ai.chat2db.plugin.mysql.type.MysqlColumnTypeEnum;
//import ai.chat2db.plugin.mysql.value.GeometryValueHandler;
//import ai.chat2db.spi.ValueHandler;
//import ai.chat2db.spi.jdbc.DefaultValueHandler;
//
//import java.sql.ResultSet;
//import java.sql.SQLException;
//import java.util.Map;
//
//public class MysqlValueHandler extends DefaultValueHandler {
//
// private static final Map<String, ValueHandler> VALUE_HANDLER_MAP = Map.of(
// MysqlColumnTypeEnum.GEOMETRY.name(), new GeometryValueHandler()
// );
//
// @Override
// public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException {
// try {
// Object obj = rs.getObject(index);
// if (obj == null) {
// return null;
// }
// String columnTypeName = rs.getMetaData().getColumnTypeName(index);
// if (MysqlColumnTypeEnum.GEOMETRY.name().equalsIgnoreCase(columnTypeName)
// || MysqlColumnTypeEnum.POINT.name().equalsIgnoreCase(columnTypeName)
// || MysqlColumnTypeEnum.LINESTRING.name().equalsIgnoreCase(columnTypeName)
// || MysqlColumnTypeEnum.POLYGON.name().equalsIgnoreCase(columnTypeName)
// || MysqlColumnTypeEnum.MULTIPOINT.name().equalsIgnoreCase(columnTypeName)
// || MysqlColumnTypeEnum.MULTILINESTRING.name().equalsIgnoreCase(columnTypeName)
// || MysqlColumnTypeEnum.MULTIPOLYGON.name().equalsIgnoreCase(columnTypeName)
// || MysqlColumnTypeEnum.GEOMETRYCOLLECTION.name().equalsIgnoreCase(columnTypeName)
// ) {
// ValueHandler handler = VALUE_HANDLER_MAP.get(MysqlColumnTypeEnum.GEOMETRY.name());
// return handler.getString(rs, index, limitSize);
// } else {
// return super.getString(rs, index, limitSize);
// }
// }catch (Exception e){
// return rs.getString(index);
// }
// }
//
//}

View File

@ -1,69 +1,69 @@
package ai.chat2db.plugin.mysql.value;
import ai.chat2db.spi.ValueHandler;
import org.locationtech.jts.geom.Geometry;
import org.locationtech.jts.io.WKBReader;
import java.io.InputStream;
import java.io.ByteArrayOutputStream;
import java.sql.ResultSet;
import java.sql.SQLException;
public class GeometryValueHandler implements ValueHandler {
@Override
public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException {
try {
InputStream inputStream = rs.getBinaryStream(index);
Geometry dbGeometry = null;
if (inputStream != null) {
//convert the stream to a byte[] array
//so it can be passed to the WKBReader
byte[] buffer = new byte[255];
int bytesRead = 0;
ByteArrayOutputStream baos = new ByteArrayOutputStream();
while ((bytesRead = inputStream.read(buffer)) != -1) {
baos.write(buffer, 0, bytesRead);
}
byte[] geometryAsBytes = baos.toByteArray();
if (geometryAsBytes.length < 5) {
throw new Exception("Invalid geometry inputStream - less than five bytes");
}
//first four bytes of the geometry are the SRID,
//followed by the actual WKB. Determine the SRID
//here
byte[] sridBytes = new byte[4];
System.arraycopy(geometryAsBytes, 0, sridBytes, 0, 4);
boolean bigEndian = (geometryAsBytes[4] == 0x00);
int srid = 0;
if (bigEndian) {
for (int i = 0; i < sridBytes.length; i++) {
srid = (srid << 8) + (sridBytes[i] & 0xff);
}
} else {
for (int i = 0; i < sridBytes.length; i++) {
srid += (sridBytes[i] & 0xff) << (8 * i);
}
}
//use the JTS WKBReader for WKB parsing
WKBReader wkbReader = new WKBReader();
//copy the byte array, removing the first four
//SRID bytes
byte[] wkb = new byte[geometryAsBytes.length - 4];
System.arraycopy(geometryAsBytes, 4, wkb, 0, wkb.length);
dbGeometry = wkbReader.read(wkb);
dbGeometry.setSRID(srid);
}
return dbGeometry.toString();
} catch (Exception e) {
return rs.getString(index);
}
}
}
//package ai.chat2db.plugin.mysql.value;
//
//import ai.chat2db.spi.ValueHandler;
//import org.locationtech.jts.geom.Geometry;
//import org.locationtech.jts.io.WKBReader;
//
//import java.io.InputStream;
//import java.io.ByteArrayOutputStream;
//import java.sql.ResultSet;
//import java.sql.SQLException;
//
//public class GeometryValueHandler implements ValueHandler {
// @Override
// public String getString(ResultSet rs, int index, boolean limitSize) throws SQLException {
// try {
// InputStream inputStream = rs.getBinaryStream(index);
// Geometry dbGeometry = null;
// if (inputStream != null) {
//
// //convert the stream to a byte[] array
// //so it can be passed to the WKBReader
// byte[] buffer = new byte[255];
//
// int bytesRead = 0;
// ByteArrayOutputStream baos = new ByteArrayOutputStream();
// while ((bytesRead = inputStream.read(buffer)) != -1) {
// baos.write(buffer, 0, bytesRead);
// }
//
// byte[] geometryAsBytes = baos.toByteArray();
//
// if (geometryAsBytes.length < 5) {
// throw new Exception("Invalid geometry inputStream - less than five bytes");
// }
//
// //first four bytes of the geometry are the SRID,
// //followed by the actual WKB. Determine the SRID
// //here
// byte[] sridBytes = new byte[4];
// System.arraycopy(geometryAsBytes, 0, sridBytes, 0, 4);
// boolean bigEndian = (geometryAsBytes[4] == 0x00);
//
// int srid = 0;
// if (bigEndian) {
// for (int i = 0; i < sridBytes.length; i++) {
// srid = (srid << 8) + (sridBytes[i] & 0xff);
// }
// } else {
// for (int i = 0; i < sridBytes.length; i++) {
// srid += (sridBytes[i] & 0xff) << (8 * i);
// }
// }
//
// //use the JTS WKBReader for WKB parsing
// WKBReader wkbReader = new WKBReader();
//
// //copy the byte array, removing the first four
// //SRID bytes
// byte[] wkb = new byte[geometryAsBytes.length - 4];
// System.arraycopy(geometryAsBytes, 4, wkb, 0, wkb.length);
// dbGeometry = wkbReader.read(wkb);
// dbGeometry.setSRID(srid);
// }
// return dbGeometry.toString();
// } catch (Exception e) {
// return rs.getString(index);
// }
// }
//}