Quote column names for ADD, RENAME and DROP COLUMN in JDBC connectors

trinodb · Aug 31, 2022 · 3ebd9c2 · hashhar · Sep 9, 2022 · ebyhr
1 parent f6e621d
commit 3ebd9c2
Show file tree

Hide file tree

Showing 10 changed files with 203 additions and 11 deletions.
diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/BaseJdbcClient.java
@@ -795,18 +795,24 @@ protected String renameColumnSql(JdbcTableHandle handle, JdbcColumnHandle jdbcCo
         return format(
                 "ALTER TABLE %s RENAME COLUMN %s TO %s",
                 quoted(handle.asPlainTable().getRemoteTableName()),
-                jdbcColumn.getColumnName(),
-                newRemoteColumnName);
+                quoted(jdbcColumn.getColumnName()),
+                quoted(newRemoteColumnName));
     }
 
     @Override
     public void dropColumn(ConnectorSession session, JdbcTableHandle handle, JdbcColumnHandle column)
     {
-        String sql = format(
-                "ALTER TABLE %s DROP COLUMN %s",
-                quoted(handle.asPlainTable().getRemoteTableName()),
-                column.getColumnName());
-        execute(session, sql);
+        try (Connection connection = connectionFactory.openConnection(session)) {
+            String remoteColumnName = identifierMapping.toRemoteColumnName(connection, column.getColumnName());
+            String sql = format(
+                    "ALTER TABLE %s DROP COLUMN %s",
+                    quoted(handle.asPlainTable().getRemoteTableName()),
+                    quoted(remoteColumnName));
+            execute(connection, sql);
+        }
+        catch (SQLException e) {
+            throw new TrinoException(JDBC_ERROR, e);
+        }
     }
 
     @Override

diff --git a/...rino-clickhouse/src/test/java/io/trino/plugin/clickhouse/BaseClickHouseConnectorTest.java b/...rino-clickhouse/src/test/java/io/trino/plugin/clickhouse/BaseClickHouseConnectorTest.java
@@ -94,6 +94,39 @@ public void testRenameColumn()
         throw new SkipException("TODO: test not implemented yet");
     }
 
+    @Override
+    public void testAddAndDropColumnName(String columnName)
+    {
+        // TODO: Enable this test
+        assertThatThrownBy(() -> super.testAddAndDropColumnName(columnName))
+                .hasMessageContaining("is not supported by storage Log");
+        throw new SkipException("TODO");
+    }
+
+    @Override
+    public void testRenameColumnName(String columnName)
+    {
+        // TODO: Enable this test
+        if (columnName.equals("a.dot")) {
+            assertThatThrownBy(() -> super.testRenameColumnName(columnName))
+                    .hasMessageContaining("Cannot rename column from nested struct to normal column");
+            throw new SkipException("TODO");
+        }
+        assertThatThrownBy(() -> super.testRenameColumnName(columnName))
+                .hasMessageContaining("is not supported by storage Log");
+        throw new SkipException("TODO");
+    }
+
+    @Override
+    protected Optional<String> filterColumnNameTestData(String columnName)
+    {
+        // TODO: Investigate why a\backslash allows creating a table, but it throws an exception when selecting
+        if (columnName.equals("a\\backslash`")) {
+            return Optional.empty();
+        }
+        return Optional.of(columnName);
+    }
+
     @Override
     public void testDropColumn()
     {

diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java
@@ -186,6 +186,17 @@ protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior)
         }
     }
 
+    @Override
+    public void testAddAndDropColumnName(String columnName)
+    {
+        if (columnName.equals("a.dot")) {
+            assertThatThrownBy(() -> super.testAddAndDropColumnName(columnName))
+                    .hasMessageContaining("Cannot add column with ambiguous name");
+            return;
+        }
+        super.testAddAndDropColumnName(columnName);
+    }
+
     @Override
     protected void verifyVersionedQueryFailurePermissible(Exception e)
     {

diff --git a/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java b/plugin/trino-kudu/src/test/java/io/trino/plugin/kudu/TestKuduConnectorTest.java
@@ -140,6 +140,24 @@ public void testRenameTableToUnqualifiedPreservesSchema()
                 .hasMessage("Creating schema in Kudu connector not allowed if schema emulation is disabled.");
     }
 
+    @Override
+    public void testAddAndDropColumnName(String columnName)
+    {
+        // TODO: Enable this test
+        assertThatThrownBy(() -> super.testAddAndDropColumnName(columnName))
+                .hasMessage("Table partitioning must be specified using setRangePartitionColumns or addHashPartitions");
+        throw new SkipException("TODO");
+    }
+
+    @Override
+    public void testRenameColumnName(String columnName)
+    {
+        // TODO: Enable this test
+        assertThatThrownBy(() -> super.testRenameColumnName(columnName))
+                .hasMessage("Table partitioning must be specified using setRangePartitionColumns or addHashPartitions");
+        throw new SkipException("TODO");
+    }
+
     @Test
     @Override
     public void testDescribeTable()

diff --git a/plugin/trino-mariadb/src/test/java/io/trino/plugin/mariadb/TestMariaDbConnectorTest.java b/plugin/trino-mariadb/src/test/java/io/trino/plugin/mariadb/TestMariaDbConnectorTest.java
@@ -44,6 +44,13 @@ public void testRenameColumn()
                 .hasMessageContaining("Rename column not supported for the MariaDB server version");
     }
 
+    @Override
+    public void testRenameColumnName(String columnName)
+    {
+        assertThatThrownBy(() -> super.testRenameColumnName(columnName))
+                .hasMessageContaining("Rename column not supported for the MariaDB server version");
+    }
+
     @Override
     public void testAlterTableRenameColumnToLongName()
     {

diff --git a/plugin/trino-mysql/src/test/java/io/trino/plugin/mysql/TestMySqlLegacyConnectorTest.java b/plugin/trino-mysql/src/test/java/io/trino/plugin/mysql/TestMySqlLegacyConnectorTest.java
@@ -135,6 +135,14 @@ public void testRenameColumn()
                 .hasStackTraceContaining("RENAME COLUMN x TO before_y");
     }
 
+    @Override
+    public void testRenameColumnName(String columnName)
+    {
+        assertThatThrownBy(() -> super.testRenameColumnName(columnName))
+                .hasMessageContaining("You have an error in your SQL syntax")
+                .hasStackTraceContaining("RENAME COLUMN");
+    }
+
     @Override
     public void testAlterTableRenameColumnToLongName()
     {

diff --git a/plugin/trino-phoenix5/src/main/java/io/trino/plugin/phoenix5/PhoenixMetadata.java b/plugin/trino-phoenix5/src/main/java/io/trino/plugin/phoenix5/PhoenixMetadata.java
@@ -234,7 +234,7 @@ public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle
         phoenixClient.execute(session, format(
                 "ALTER TABLE %s ADD %s %s",
                 getEscapedTableName(handle.getSchemaName(), handle.getTableName()),
-                column.getName(),
+                phoenixClient.quoted(column.getName()),
                 phoenixClient.toWriteMapping(session, column.getType()).getDataType()));
     }
 
@@ -246,7 +246,7 @@ public void dropColumn(ConnectorSession session, ConnectorTableHandle tableHandl
         phoenixClient.execute(session, format(
                 "ALTER TABLE %s DROP COLUMN %s",
                 getEscapedTableName(handle.getSchemaName(), handle.getTableName()),
-                columnHandle.getColumnName()));
+                phoenixClient.quoted(columnHandle.getColumnName())));
     }
 
     @Override

diff --git a/plugin/trino-phoenix5/src/test/java/io/trino/plugin/phoenix5/TestPhoenixConnectorTest.java b/plugin/trino-phoenix5/src/test/java/io/trino/plugin/phoenix5/TestPhoenixConnectorTest.java
@@ -147,6 +147,37 @@ public void testAlterTableRenameColumnToLongName()
         throw new SkipException("Rename column is not yet supported by Phoenix connector");
     }
 
+    @Override
+    public void testRenameColumnName(String columnName)
+    {
+        // The column name is rejected when creating a table
+        if (columnName.equals("a\"quote")) {
+            super.testRenameColumnName(columnName);
+            return;
+        }
+        assertThatThrownBy(() -> super.testRenameColumnName(columnName))
+                // TODO (https://github.com/trinodb/trino/issues/7205) support column rename in Phoenix
+                .hasMessageContaining("Syntax error. Encountered \"RENAME\"");
+        throw new SkipException("Rename column is not yet supported by Phoenix connector");
+    }
+
+    @Override
+    public void testAddAndDropColumnName(String columnName)
+    {
+        // TODO: Investigate why these two case fail
+        if (columnName.equals("an'apostrophe")) {
+            assertThatThrownBy(() -> super.testAddAndDropColumnName(columnName))
+                    .hasMessageContaining("Syntax error. Mismatched input");
+            throw new SkipException("TODO");
+        }
+        if (columnName.equals("a\\backslash`")) {
+            assertThatThrownBy(() -> super.testAddAndDropColumnName(columnName))
+                    .hasMessageContaining("Undefined column");
+            throw new SkipException("TODO");
+        }
+        super.testAddAndDropColumnName(columnName);
+    }
+
     @Override
     public void testInsert()
     {

diff --git a/plugin/trino-sqlserver/src/main/java/io/trino/plugin/sqlserver/SqlServerClient.java b/plugin/trino-sqlserver/src/main/java/io/trino/plugin/sqlserver/SqlServerClient.java
@@ -323,8 +323,13 @@ protected String renameColumnSql(JdbcTableHandle handle, JdbcColumnHandle jdbcCo
     {
         return format(
                 "sp_rename %s, %s, 'COLUMN'",
-                singleQuote(handle.getCatalogName(), handle.getSchemaName(), handle.getTableName(), jdbcColumn.getColumnName()),
-                singleQuote(newRemoteColumnName));
+                singleQuote(handle.getCatalogName(), handle.getSchemaName(), handle.getTableName(), "[" + escape(jdbcColumn.getColumnName()) + "]"),
+                "[" + newRemoteColumnName + "]");
+    }
+
+    private static String escape(String name)
+    {
+        return name.replace("'", "''");
     }
 
     @Override

diff --git a/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java b/testing/trino-testing/src/main/java/io/trino/testing/BaseConnectorTest.java
@@ -3831,6 +3831,79 @@ protected void testColumnName(String columnName, boolean delimited)
         }
     }
 
+    @Test(dataProvider = "testColumnNameDataProvider")
+    public void testAddAndDropColumnName(String columnName)
+    {
+        skipTestUnless(hasBehavior(SUPPORTS_ADD_COLUMN) && hasBehavior(SUPPORTS_DROP_COLUMN));
+
+        if (!requiresDelimiting(columnName)) {
+            testAddAndDropColumnName(columnName, false);
+        }
+        testAddAndDropColumnName(columnName, true);
+    }
+
+    protected void testAddAndDropColumnName(String columnName, boolean delimited)
+    {
+        String nameInSql = toColumnNameInSql(columnName, delimited);
+        String tableName = "tcn_" + nameInSql.toLowerCase(ENGLISH).replaceAll("[^a-z0-9]", "") + randomTableSuffix();
+
+        try {
+            assertUpdate("CREATE TABLE " + tableName + "(" + nameInSql + " varchar(50), value varchar(50))");
+        }
+        catch (RuntimeException e) {
+            if (isColumnNameRejected(e, columnName, delimited)) {
+                // It is OK if give column name is not allowed and is clearly rejected by the connector.
+                return;
+            }
+            throw e;
+        }
+        assertTableColumnNames(tableName, columnName.toLowerCase(ENGLISH), "value");
+
+        assertUpdate("ALTER TABLE " + tableName + " DROP COLUMN " + nameInSql);
+        assertTableColumnNames(tableName, "value");
+
+        assertUpdate("ALTER TABLE " + tableName + " ADD COLUMN " + nameInSql + " varchar(50)");
+        assertTableColumnNames(tableName, "value", columnName.toLowerCase(ENGLISH));
+
+        assertUpdate("DROP TABLE " + tableName);
+    }
+
+    @Test(dataProvider = "testColumnNameDataProvider")
+    public void testRenameColumnName(String columnName)
+    {
+        skipTestUnless(hasBehavior(SUPPORTS_RENAME_COLUMN));
+
+        if (!requiresDelimiting(columnName)) {
+            testRenameColumnName(columnName, false);
+        }
+        testRenameColumnName(columnName, true);
+    }
+
+    protected void testRenameColumnName(String columnName, boolean delimited)
+    {
+        String nameInSql = toColumnNameInSql(columnName, delimited);
+        String tableName = "tcn_" + nameInSql.replaceAll("[^a-z0-9]", "") + randomTableSuffix();
+        // Use complex identifier to test a source column name when renaming columns
+        String sourceColumnName = "a;b$c";
+
+        try {
+            assertUpdate("CREATE TABLE " + tableName + "(\"" + sourceColumnName + "\" varchar(50))");
+            assertTableColumnNames(tableName, sourceColumnName);
+
+            assertUpdate("ALTER TABLE " + tableName + " RENAME COLUMN \"" + sourceColumnName + "\" TO " + nameInSql);
+            assertTableColumnNames(tableName, columnName.toLowerCase(ENGLISH));
+        }
+        catch (RuntimeException e) {
+            if (isColumnNameRejected(e, columnName, delimited)) {
+                // It is OK if give column name is not allowed and is clearly rejected by the connector.
+                return;
+            }
+            throw e;
+        }
+
+        assertUpdate("DROP TABLE " + tableName);
+    }
+
     private static String toColumnNameInSql(String columnName, boolean delimited)
     {
         String nameInSql = columnName;