Use RemoteTableName in BaseJdbcClient where reasonable

trinodb · Jul 31, 2020 · fe46bdf · fe46bdf
1 parent a1ef1f2
commit fe46bdf
Showing 1 changed file with 18 additions and 15 deletions.
diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcClient.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/BaseJdbcClient.java
@@ -238,11 +238,7 @@ public Optional<JdbcTableHandle> getTableHandle(JdbcIdentity identity, SchemaTab
             try (ResultSet resultSet = getTables(connection, Optional.of(remoteSchema), Optional.of(remoteTable))) {
                 List<JdbcTableHandle> tableHandles = new ArrayList<>();
                 while (resultSet.next()) {
-                    tableHandles.add(new JdbcTableHandle(
-                            schemaTableName,
-                            resultSet.getString("TABLE_CAT"),
-                            resultSet.getString("TABLE_SCHEM"),
-                            resultSet.getString("TABLE_NAME")));
+                    tableHandles.add(new JdbcTableHandle(schemaTableName, getRemoteTable(resultSet)));
                 }
                 if (tableHandles.isEmpty()) {
                     return Optional.empty();
@@ -267,9 +263,7 @@ public List<JdbcColumnHandle> getColumns(ConnectorSession session, JdbcTableHand
             List<JdbcColumnHandle> columns = new ArrayList<>();
             while (resultSet.next()) {
                 // skip if table doesn't match expected
-                if (!(Objects.equals(tableHandle.getCatalogName(), resultSet.getString("TABLE_CAT"))
-                        && Objects.equals(tableHandle.getSchemaName(), resultSet.getString("TABLE_SCHEM"))
-                        && Objects.equals(tableHandle.getTableName(), resultSet.getString("TABLE_NAME")))) {
+                if (!(Objects.equals(tableHandle.getRemoteTableName(), getRemoteTable(resultSet)))) {
                     continue;
                 }
                 allColumns++;
@@ -318,9 +312,9 @@ protected ResultSet getColumns(JdbcTableHandle tableHandle, DatabaseMetaData met
             throws SQLException
     {
         return metadata.getColumns(
-                tableHandle.getCatalogName(),
-                escapeNamePattern(Optional.ofNullable(tableHandle.getSchemaName()), metadata.getSearchStringEscape()).orElse(null),
-                escapeNamePattern(Optional.ofNullable(tableHandle.getTableName()), metadata.getSearchStringEscape()).orElse(null),
+                tableHandle.getRemoteTableName().getCatalogName().orElse(null),
+                escapeNamePattern(tableHandle.getRemoteTableName().getSchemaName(), metadata.getSearchStringEscape()).orElse(null),
+                escapeNamePattern(Optional.of(tableHandle.getRemoteTableName().getTableName()), metadata.getSearchStringEscape()).orElse(null),
                 null);
     }
 
@@ -615,7 +609,7 @@ public void addColumn(ConnectorSession session, JdbcTableHandle handle, ColumnMe
             }
             String sql = format(
                     "ALTER TABLE %s ADD %s",
-                    quoted(handle.getCatalogName(), handle.getSchemaName(), handle.getTableName()),
+                    quoted(handle.getRemoteTableName()),
                     getColumnSql(session, column, columnName));
             execute(connection, sql);
         }
@@ -633,7 +627,7 @@ public void renameColumn(JdbcIdentity identity, JdbcTableHandle handle, JdbcColu
             }
             String sql = format(
                     "ALTER TABLE %s RENAME COLUMN %s TO %s",
-                    quoted(handle.getCatalogName(), handle.getSchemaName(), handle.getTableName()),
+                    quoted(handle.getRemoteTableName()),
                     jdbcColumn.getColumnName(),
                     newColumnName);
             execute(connection, sql);
@@ -648,15 +642,15 @@ public void dropColumn(JdbcIdentity identity, JdbcTableHandle handle, JdbcColumn
     {
         String sql = format(
                 "ALTER TABLE %s DROP COLUMN %s",
-                quoted(handle.getCatalogName(), handle.getSchemaName(), handle.getTableName()),
+                quoted(handle.getRemoteTableName()),
                 column.getColumnName());
         execute(identity, sql);
     }
 
     @Override
     public void dropTable(JdbcIdentity identity, JdbcTableHandle handle)
     {
-        String sql = "DROP TABLE " + quoted(handle.getCatalogName(), handle.getSchemaName(), handle.getTableName());
+        String sql = "DROP TABLE " + quoted(handle.getRemoteTableName());
         execute(identity, sql);
     }
 
@@ -959,4 +953,13 @@ private static String escapeNamePattern(String name, String escape)
         name = name.replace("%", escape + "%");
         return name;
     }
+
+    private static RemoteTableName getRemoteTable(ResultSet resultSet)
+            throws SQLException
+    {
+        return new RemoteTableName(
+                Optional.ofNullable(resultSet.getString("TABLE_CAT")),
+                Optional.ofNullable(resultSet.getString("TABLE_SCHEM")),
+                resultSet.getString("TABLE_NAME"));
+    }
 }