Merge pull request #79 from jumpserver/pr@dev@feat_format_pg_dbname

feat: format pg dbname
jumpserver · Jan 13, 2025 · 347675e · 347675e
2 parents 8c8244a + 7b7d3f4
commit 347675e
Showing 1 changed file with 20 additions and 4 deletions.
diff --git a/backend/modules/src/main/java/org.jumpserver.chen.modules/postgresql/PostgresqlActuator.java b/backend/modules/src/main/java/org.jumpserver.chen.modules/postgresql/PostgresqlActuator.java
@@ -1,5 +1,6 @@
 package org.jumpserver.chen.modules.postgresql;
 
+import org.apache.commons.lang3.StringUtils;
 import org.jumpserver.chen.framework.datasource.ConnectionManager;
 import org.jumpserver.chen.framework.datasource.base.BaseSQLActuator;
 import org.jumpserver.chen.framework.datasource.sql.SQL;
@@ -15,30 +16,45 @@ public PostgresqlActuator(ConnectionManager connectionManager) {
         super(connectionManager);
     }
 
+    private String dbName;
+
     public PostgresqlActuator(PostgresqlActuator sqlActuator, Connection connection) {
         super(sqlActuator, connection);
     }
 
     @Override
     public String getCurrentSchema() throws SQLException {
-        var result = this.execute(SQL.of("SELECT CURRENT_SCHEMA()"));
-        return (String) result.getData().get(0).get(0);
+        var result = this.execute(SQL.of("SELECT current_schema()"));
+        return this.formatSchemaName((String) result.getData().get(0).get(0));
     }
 
     @Override
     public List<String> getSchemas() throws SQLException {
         var result = this.execute(SQL.of("SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA"));
-        return result.getData().stream().map(row -> (String) row.get(0)).toList();
+        return result.getData().stream().map(row -> (String) row.get(0)).toList().stream().map(this::formatSchemaName).toList();
     }
 
     @Override
     public void changeSchema(String schema) throws SQLException {
-        this.execute(SQL.of("SET SEARCH_PATH TO '?';", schema));
+        var ss = schema.split("\\.");
+        this.execute(SQL.of("SET SEARCH_PATH TO '?';", ss[1]));
     }
 
     @Override
     public SQLExecutePlan createPlan(String schema, String table, SQLQueryParams sqlQueryParams) throws SQLException {
         var sql = SQL.of("select * from \"?\".\"?\"", schema, table);
         return this.createPlan(sql, sqlQueryParams);
     }
+
+    private String formatSchemaName(String schema) {
+        try {
+            if (StringUtils.isEmpty(this.dbName)) {
+                var result = this.execute(SQL.of("SELECT current_database();"));
+                this.dbName = (String) result.getData().get(0).get(0);
+            }
+        } catch (SQLException e) {
+            throw new RuntimeException(e);
+        }
+        return String.format("%s.%s", this.dbName, schema);
+    }
 }