Merge pull request #1239 from openai0229/pgsql-schema-export

Pgsql schema export
This commit is contained in:
ji
2024-03-22 11:17:00 +08:00
committed by GitHub
2 changed files with 52 additions and 48 deletions

View File

@ -8,6 +8,8 @@ import ai.chat2db.spi.sql.SQLExecutor;
import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.StringUtils;
import java.sql.*; import java.sql.*;
import java.util.ArrayList;
import java.util.Objects;
import static ai.chat2db.plugin.postgresql.consts.SQLConst.*; import static ai.chat2db.plugin.postgresql.consts.SQLConst.*;
@ -16,86 +18,90 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage {
public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException { public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException {
StringBuilder sqlBuilder = new StringBuilder(); StringBuilder sqlBuilder = new StringBuilder();
exportTypes(connection, schemaName, sqlBuilder); exportTypes(connection, sqlBuilder);
exportTables(connection, schemaName, sqlBuilder, containData); exportTables(connection, databaseName, schemaName, sqlBuilder, containData);
exportViews(connection, schemaName, sqlBuilder); exportViews(connection, schemaName, sqlBuilder);
exportFunctions(connection, schemaName, sqlBuilder); exportFunctions(connection, schemaName, sqlBuilder);
exportTriggers(connection, schemaName, sqlBuilder); exportTriggers(connection, sqlBuilder);
return sqlBuilder.toString(); return sqlBuilder.toString();
} }
private void exportTypes(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { private void exportTypes(Connection connection, StringBuilder sqlBuilder) throws SQLException {
try (Statement statement = connection.createStatement(); ResultSet ddl = statement.executeQuery(ENUM_TYPE_DDL_SQL)) { try (ResultSet resultSet = connection.createStatement().executeQuery(ENUM_TYPE_DDL_SQL)) {
while (ddl.next()) { while (resultSet.next()) {
sqlBuilder.append(ddl.getString(1)).append("\n"); sqlBuilder.append(resultSet.getString("ddl")).append("\n");
} }
} }
} }
private void exportTables(Connection connection, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException { private void exportTables(Connection connection, String databaseName, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
String tablesQuery = "SELECT table_name FROM information_schema.tables WHERE table_schema = '" + schemaName + "' AND table_type = 'BASE TABLE'"; try (ResultSet resultSet = connection.getMetaData().getTables(databaseName, schemaName, null,
try (Statement statement = connection.createStatement(); ResultSet tables = statement.executeQuery(tablesQuery)) { new String[]{"TABLE", "SYSTEM TABLE","PARTITIONED TABLE"})) {
while (tables.next()) { ArrayList<String> tableNames = new ArrayList<>();
String tableName = tables.getString(1); while (resultSet.next()) {
exportTable(connection, schemaName, tableName, sqlBuilder, containData); String tableName = resultSet.getString("TABLE_NAME");
tableNames.add(tableName);
}
for (String tableName : tableNames) {
exportTable(connection, schemaName, tableName, sqlBuilder);
}
if (containData) {
for (String tableName : tableNames) {
exportTableData(connection, schemaName, tableName, sqlBuilder);
}
} }
} }
} }
private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException { private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder) throws SQLException {
String tableQuery = "select pg_get_tabledef" + "(" + "'" + schemaName + "'" + "," + "'" + tableName + "'" + "," + "true" + "," + "'" + "COMMENTS" + "'" + ")" + ";"; String sql =String.format( "select pg_get_tabledef('%s','%s',true,'COMMENTS') as ddl;", schemaName,tableName);
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(tableQuery)) { try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
sqlBuilder.append("\n").append("DROP TABLE IF EXISTS ").append(schemaName).append(".").append(tableName).append(";\n");
if (resultSet.next()) { if (resultSet.next()) {
sqlBuilder.append(resultSet.getString(1)).append("\n"); sqlBuilder.append("\n").append("DROP TABLE IF EXISTS ").append(tableName).append(";").append("\n")
} .append(resultSet.getString("ddl")).append("\n");
if (containData) {
exportTableData(connection, schemaName, tableName, sqlBuilder);
} }
} }
} }
private void exportTableData(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder) throws SQLException { private void exportTableData(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder) throws SQLException {
StringBuilder insertSql = new StringBuilder(); String sql =String.format("select * from %s.%s", schemaName,tableName);
String dataQuery = "SELECT * FROM " + schemaName + "." + tableName; try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(dataQuery)) {
ResultSetMetaData metaData = resultSet.getMetaData(); ResultSetMetaData metaData = resultSet.getMetaData();
int columnCount = metaData.getColumnCount();
while (resultSet.next()) { while (resultSet.next()) {
insertSql.append("INSERT INTO ").append(tableName).append(" VALUES ("); sqlBuilder.append("INSERT INTO ").append(tableName).append(" VALUES (");
for (int i = 1; i <= columnCount; i++) { for (int i = 1; i <= metaData.getColumnCount(); i++) {
String value = resultSet.getString(i); String value = resultSet.getString(i);
if (value != null) { if (Objects.isNull(value)) {
insertSql.append("'").append(value).append("'"); sqlBuilder.append("NULL");
} else { } else {
insertSql.append("NULL"); sqlBuilder.append("'").append(value).append("'");
} }
if (i < columnCount) { if (i < metaData.getColumnCount()) {
insertSql.append(", "); sqlBuilder.append(", ");
} }
} }
insertSql.append(");\n"); sqlBuilder.append(");\n");
} }
insertSql.append("\n"); sqlBuilder.append("\n");
sqlBuilder.append(insertSql);
} }
} }
private void exportViews(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { private void exportViews(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
String viewsQuery = "SELECT table_name, view_definition FROM information_schema.views WHERE table_schema = '" + schemaName + "'";
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(viewsQuery)) { 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()) { while (resultSet.next()) {
String viewName = resultSet.getString("table_name"); String viewName = resultSet.getString("table_name");
String viewDefinition = resultSet.getString("view_definition"); String viewDefinition = resultSet.getString("view_definition");
sqlBuilder.append("DROP VIEW IF EXISTS ").append(schemaName).append(".").append(viewName).append(";\n"); sqlBuilder.append("CREATE OR REPLACE VIEW ").append(viewName).append(" AS ").append(viewDefinition).append("\n");
sqlBuilder.append("CREATE VIEW ").append(schemaName).append(".").append(viewName).append(" AS ").append(viewDefinition).append(";\n\n");
} }
} }
} }
private void exportFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { private void exportFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
String functionsQuery = "SELECT proname, pg_get_functiondef(oid) AS function_definition FROM pg_proc WHERE pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = '" + schemaName + "')"; String sql = String.format("SELECT proname, pg_get_functiondef(oid) AS function_definition FROM pg_proc " +
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(functionsQuery)) { "WHERE pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = '%s')", schemaName);
try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
while (resultSet.next()) { while (resultSet.next()) {
String functionName = resultSet.getString("proname"); String functionName = resultSet.getString("proname");
String functionDefinition = resultSet.getString("function_definition"); String functionDefinition = resultSet.getString("function_definition");
@ -105,14 +111,11 @@ public class PostgreSQLDBManage extends DefaultDBManage implements DBManage {
} }
} }
private void exportTriggers(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException { private void exportTriggers(Connection connection, StringBuilder sqlBuilder) throws SQLException {
String triggersQuery = "SELECT tgname, pg_get_triggerdef(oid) AS trigger_definition FROM pg_trigger"; String sql = "SELECT pg_get_triggerdef(oid) AS trigger_definition FROM pg_trigger";
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(triggersQuery)) { try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(sql)) {
while (resultSet.next()) { while (resultSet.next()) {
String triggerName = resultSet.getString("tgname"); sqlBuilder.append(resultSet.getString("trigger_definition")).append(";").append("\n");
String triggerDefinition = resultSet.getString("trigger_definition");
sqlBuilder.append("DROP TRIGGER IF EXISTS ").append(schemaName).append(".").append(triggerName).append(";\n");
sqlBuilder.append(triggerDefinition).append(";\n\n");
} }
} }
} }

View File

@ -610,6 +610,7 @@ public class SQLConst {
if bSkip THEN CONTINUE; END IF; if bSkip THEN CONTINUE; END IF;
-- Add IF NOT EXISTS clause so partition index additions will not be created if declarative partition in effect and index already created on parent -- Add IF NOT EXISTS clause so partition index additions will not be created if declarative partition in effect and index already created on parent
v_indexrec.indexdef := REPLACE(v_indexrec.indexdef, 'CREATE UNIQUE INDEX', 'CREATE UNIQUE INDEX IF NOT EXISTS');
v_indexrec.indexdef := REPLACE(v_indexrec.indexdef, 'CREATE INDEX', 'CREATE INDEX IF NOT EXISTS'); v_indexrec.indexdef := REPLACE(v_indexrec.indexdef, 'CREATE INDEX', 'CREATE INDEX IF NOT EXISTS');
-- RAISE INFO 'DEBUG8: adding index, %', v_indexrec.indexname; -- RAISE INFO 'DEBUG8: adding index, %', v_indexrec.indexname;