Skip to content

Commit 7db68d2

Browse files
Chat2DB-ProChat2DB-Pro
authored
Merge pull request OtterMind#1195 from openai0229/pgsql-schema-export
pgsql schema export
2 parents 0852517 + 1469908 commit 7db68d2

4 files changed

Lines changed: 797 additions & 155 deletions

File tree

‎chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLDBManage.java‎

Lines changed: 109 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,122 @@
11
package ai.chat2db.plugin.postgresql;
22

3-
import java.sql.Connection;
4-
53
import ai.chat2db.spi.DBManage;
64
import ai.chat2db.spi.jdbc.DefaultDBManage;
75
import ai.chat2db.spi.sql.Chat2DBContext;
86
import ai.chat2db.spi.sql.ConnectInfo;
97
import ai.chat2db.spi.sql.SQLExecutor;
108
import org.apache.commons.lang3.StringUtils;
119

10+
import java.sql.*;
11+
12+
import static ai.chat2db.plugin.postgresql.consts.SQLConst.*;
13+
1214
public class PostgreSQLDBManage extends DefaultDBManage implements DBManage {
15+
16+
17+
public String exportDatabase(Connection connection, String databaseName, String schemaName, boolean containData) throws SQLException {
18+
StringBuilder sqlBuilder = new StringBuilder();
19+
exportTypes(connection, schemaName, sqlBuilder);
20+
exportTables(connection, schemaName, sqlBuilder, containData);
21+
exportViews(connection, schemaName, sqlBuilder);
22+
exportFunctions(connection, schemaName, sqlBuilder);
23+
exportTriggers(connection, schemaName, sqlBuilder);
24+
return sqlBuilder.toString();
25+
}
26+
27+
private void exportTypes(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
28+
try (Statement statement = connection.createStatement(); ResultSet ddl = statement.executeQuery(ENUM_TYPE_DDL_SQL)) {
29+
while (ddl.next()) {
30+
sqlBuilder.append(ddl.getString(1)).append("\n");
31+
}
32+
}
33+
}
34+
private void exportTables(Connection connection, String schemaName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
35+
String tablesQuery = "SELECT table_name FROM information_schema.tables WHERE table_schema = '" + schemaName + "' AND table_type = 'BASE TABLE'";
36+
try (Statement statement = connection.createStatement(); ResultSet tables = statement.executeQuery(tablesQuery)) {
37+
while (tables.next()) {
38+
String tableName = tables.getString(1);
39+
exportTable(connection, schemaName, tableName, sqlBuilder, containData);
40+
}
41+
}
42+
}
43+
44+
private void exportTable(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder, boolean containData) throws SQLException {
45+
String tableQuery = "select pg_get_tabledef" + "(" + "'" + schemaName + "'" + "," + "'" + tableName + "'" + "," + "true" + "," + "'" + "COMMENTS" + "'" + ")" + ";";
46+
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(tableQuery)) {
47+
sqlBuilder.append("\n").append("DROP TABLE IF EXISTS ").append(schemaName).append(".").append(tableName).append(";\n");
48+
if (resultSet.next()) {
49+
sqlBuilder.append(resultSet.getString(1)).append("\n");
50+
}
51+
if (containData) {
52+
exportTableData(connection, schemaName, tableName, sqlBuilder);
53+
}
54+
}
55+
}
56+
57+
private void exportTableData(Connection connection, String schemaName, String tableName, StringBuilder sqlBuilder) throws SQLException {
58+
StringBuilder insertSql = new StringBuilder();
59+
String dataQuery = "SELECT * FROM " + schemaName + "." + tableName;
60+
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(dataQuery)) {
61+
ResultSetMetaData metaData = resultSet.getMetaData();
62+
int columnCount = metaData.getColumnCount();
63+
while (resultSet.next()) {
64+
insertSql.append("INSERT INTO ").append(tableName).append(" VALUES (");
65+
for (int i = 1; i <= columnCount; i++) {
66+
String value = resultSet.getString(i);
67+
if (value != null) {
68+
insertSql.append("'").append(value).append("'");
69+
} else {
70+
insertSql.append("NULL");
71+
}
72+
if (i < columnCount) {
73+
insertSql.append(", ");
74+
}
75+
}
76+
insertSql.append(");\n");
77+
}
78+
insertSql.append("\n");
79+
sqlBuilder.append(insertSql);
80+
}
81+
}
82+
83+
84+
private void exportViews(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
85+
String viewsQuery = "SELECT table_name, view_definition FROM information_schema.views WHERE table_schema = '" + schemaName + "'";
86+
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(viewsQuery)) {
87+
while (resultSet.next()) {
88+
String viewName = resultSet.getString("table_name");
89+
String viewDefinition = resultSet.getString("view_definition");
90+
sqlBuilder.append("DROP VIEW IF EXISTS ").append(schemaName).append(".").append(viewName).append(";\n");
91+
sqlBuilder.append("CREATE VIEW ").append(schemaName).append(".").append(viewName).append(" AS ").append(viewDefinition).append(";\n\n");
92+
}
93+
}
94+
}
95+
96+
private void exportFunctions(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
97+
String functionsQuery = "SELECT proname, pg_get_functiondef(oid) AS function_definition FROM pg_proc WHERE pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = '" + schemaName + "')";
98+
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(functionsQuery)) {
99+
while (resultSet.next()) {
100+
String functionName = resultSet.getString("proname");
101+
String functionDefinition = resultSet.getString("function_definition");
102+
sqlBuilder.append("DROP FUNCTION IF EXISTS ").append(schemaName).append(".").append(functionName).append(";\n");
103+
sqlBuilder.append(functionDefinition).append(";\n\n");
104+
}
105+
}
106+
}
107+
108+
private void exportTriggers(Connection connection, String schemaName, StringBuilder sqlBuilder) throws SQLException {
109+
String triggersQuery = "SELECT tgname, pg_get_triggerdef(oid) AS trigger_definition FROM pg_trigger";
110+
try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(triggersQuery)) {
111+
while (resultSet.next()) {
112+
String triggerName = resultSet.getString("tgname");
113+
String triggerDefinition = resultSet.getString("trigger_definition");
114+
sqlBuilder.append("DROP TRIGGER IF EXISTS ").append(schemaName).append(".").append(triggerName).append(";\n");
115+
sqlBuilder.append(triggerDefinition).append(";\n\n");
116+
}
117+
}
118+
}
119+
13120
@Override
14121
public void connectDatabase(Connection connection, String database) {
15122
try {

‎chat2db-server/chat2db-plugins/chat2db-postgresql/src/main/java/ai/chat2db/plugin/postgresql/PostgreSQLMetaData.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -102,13 +102,12 @@ public List<Trigger> triggers(Connection connection, String databaseName, String
102102

103103
@Override
104104
public String tableDDL(Connection connection, String databaseName, String schemaName, String tableName) {
105-
SQLExecutor.getInstance().execute(connection, FUNCTION_SQL.replaceFirst("tableSchema", schemaName),
106-
resultSet -> null);
107-
String ddlSql = "select showcreatetable('" + schemaName + "','" + tableName + "') as sql";
105+
SQLExecutor.getInstance().execute(connection, FUNCTION_SQL, resultSet -> null);
106+
String ddlSql = "select pg_get_tabledef" + "(" + "'" + schemaName + "'" + "," + "'" + tableName + "'" + "," + "false" + "," + "'" + "COMMENTS" + "'" + ")" + ";";
108107
return SQLExecutor.getInstance().execute(connection, ddlSql, resultSet -> {
109108
try {
110109
if (resultSet.next()) {
111-
return resultSet.getString("sql");
110+
return resultSet.getString(1);
112111
}
113112
} catch (SQLException e) {
114113
throw new RuntimeException(e);

0 commit comments

Comments
 (0)