From 7b7d3f4e93d861905f1a20acdb00d0f88b746808 Mon Sep 17 00:00:00 2001 From: Aaron3S Date: Mon, 13 Jan 2025 17:58:16 +0800 Subject: [PATCH] feat: format pg dbname --- .../postgresql/PostgresqlActuator.java | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 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 index 663041d..aff4855 100644 --- 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,25 +16,28 @@ 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 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 @@ -41,4 +45,16 @@ public SQLExecutePlan createPlan(String schema, String table, SQLQueryParams sql 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); + } }