[apache#1171] bugfix(postgresql): Fix bug where the same table name c…

…an be loaded under multiple schemas.
Clearvive · Dec 18, 2023 · 3cf66f0 · 3cf66f0
1 parent f3e1bae
commit 3cf66f0
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 13 deletions.
diff --git a/...java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlTableOperations.java b/...java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlTableOperations.java
@@ -49,17 +49,17 @@ public class PostgreSqlTableOperations extends JdbcTableOperations {
           + "    pg_namespace AS n ON n.oid = c.relnamespace\n"
           + "WHERE \n"
           + "    a.attnum > 0 \n"
-          + "    AND c.relname = ?";
+          + "    AND c.relname = ? AND n.nspname = ?";
 
   private static final String SHOW_COLUMN_INFO_SQL =
-      "select * FROM information_schema.columns WHERE table_name = ? order by ordinal_position";
+      "select * FROM information_schema.columns WHERE table_name = ? AND table_schema = ? order by ordinal_position";
 
   private static final String SHOW_TABLE_COMMENT_SQL =
       "SELECT tb.table_name, d.description\n"
           + "FROM information_schema.tables tb\n"
           + "         JOIN pg_class c ON c.relname = tb.table_name\n"
           + "         LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = '0'\n"
-          + "WHERE tb.table_name = ?;";
+          + "WHERE tb.table_name = ? AND table_schema = ?;";
 
   private String database;
 
@@ -80,15 +80,18 @@ public void initialize(
   public JdbcTable load(String schema, String tableName) throws NoSuchTableException {
     try (Connection connection = getConnection(schema)) {
       // The first step is to obtain the comment information of the column.
-      Map<String, String> columnCommentMap = selectColumnComment(tableName, connection);
+      Map<String, String> columnCommentMap = selectColumnComment(schema, tableName, connection);
 
       // The second step is to obtain the column information of the table.
       List<JdbcColumn> jdbcColumns =
           selectColumnInfoAndExecute(
-              tableName, connection, (builder, s) -> builder.withComment(columnCommentMap.get(s)));
+              schema,
+              tableName,
+              connection,
+              (builder, s) -> builder.withComment(columnCommentMap.get(s)));
 
       // The third step is to obtain the comment information of the table.
-      String comment = selectTableComment(tableName, connection);
+      String comment = selectTableComment(schema, tableName, connection);
       return new JdbcTable.Builder()
           .withName(tableName)
           .withColumns(jdbcColumns.toArray(new JdbcColumn[0]))
@@ -102,13 +105,15 @@ public JdbcTable load(String schema, String tableName) throws NoSuchTableExcepti
   }
 
   private List<JdbcColumn> selectColumnInfoAndExecute(
+      String schemaName,
       String tableName,
       Connection connection,
       BiConsumer<JdbcColumn.Builder, String> builderConsumer)
       throws SQLException {
     List<JdbcColumn> jdbcColumns = new ArrayList<>();
     try (PreparedStatement preparedStatement = connection.prepareStatement(SHOW_COLUMN_INFO_SQL)) {
       preparedStatement.setString(1, tableName);
+      preparedStatement.setString(2, schemaName);
       try (ResultSet resultSet = preparedStatement.executeQuery()) {
         while (resultSet.next()) {
           ColDataType colDataType = new ColDataType();
@@ -153,10 +158,12 @@ private static List<String> getArgList(ResultSet resultSet) throws SQLException
     return result;
   }
 
-  private String selectTableComment(String tableName, Connection connection) throws SQLException {
+  private String selectTableComment(String schema, String tableName, Connection connection)
+      throws SQLException {
     try (PreparedStatement preparedStatement =
         connection.prepareStatement(SHOW_TABLE_COMMENT_SQL)) {
       preparedStatement.setString(1, tableName);
+      preparedStatement.setString(2, schema);
       try (ResultSet resultSet = preparedStatement.executeQuery()) {
         if (resultSet.next()) {
           return resultSet.getString("description");
@@ -170,13 +177,14 @@ private String selectTableComment(String tableName, Connection connection) throw
    * @return Returns the column names and comments of the table
    * @throws SQLException
    */
-  private Map<String, String> selectColumnComment(String tableName, Connection connection)
-      throws SQLException {
+  private Map<String, String> selectColumnComment(
+      String schema, String tableName, Connection connection) throws SQLException {
     Map<String, String> columnCommentMap = new HashMap<>();
 
     try (PreparedStatement preparedStatement =
         connection.prepareStatement(SHOW_COLUMN_COMMENT_SQL)) {
       preparedStatement.setString(1, tableName);
+      preparedStatement.setString(2, schema);
       try (ResultSet resultSet = preparedStatement.executeQuery()) {
         while (resultSet.next()) {
           String comment = resultSet.getString("comment");

diff --git a/...ato/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlTableOperations.java b/...ato/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlTableOperations.java
@@ -350,6 +350,14 @@ public void testCreateMultipleTable() throws SQLException {
     List<String> tableNames = TABLE_OPERATIONS.listTables(TEST_DB_NAME);
     Assertions.assertFalse(tableNames.contains(table_1));
 
+    Assertions.assertThrows(
+        NoSuchTableException.class, () -> TABLE_OPERATIONS.load(TEST_DB_NAME, table_1));
+
+    Assertions.assertThrows(
+        NoSuchTableException.class, () -> TABLE_OPERATIONS.load("other_schema", table_1));
+    Assertions.assertThrows(
+        NoSuchTableException.class, () -> postgreSqlTableOperations.load("other_schema", table_1));
+
     String table_2 = "table_multiple_2";
     TABLE_OPERATIONS.create(
         TEST_DB_NAME,
@@ -369,10 +377,11 @@ public void testCreateMultipleTable() throws SQLException {
     Assertions.assertFalse(tableNames.contains(table_2));
 
     Assertions.assertThrows(
-        NoSuchTableException.class,
-        () -> {
-          postgreSqlTableOperations.load(TEST_DB_NAME, table_2);
-        });
+        NoSuchTableException.class, () -> postgreSqlTableOperations.load(TEST_DB_NAME, table_2));
+    Assertions.assertThrows(
+        NoSuchTableException.class, () -> postgreSqlTableOperations.load("other_schema", table_2));
+    Assertions.assertThrows(
+        NoSuchTableException.class, () -> TABLE_OPERATIONS.load("other_schema", table_2));
     postgreSqlTableOperations.purge(TEST_DB_NAME, table_1);
 
     Assertions.assertThrows(