Support altering column comments in Hive Glue catalogs

trinodb · Sep 22, 2023 · fe30d7a · fe30d7a
1 parent 726a6ed
commit fe30d7a
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 1 deletion.
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/glue/GlueHiveMetastore.java
@@ -144,6 +144,7 @@
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static io.trino.plugin.hive.HiveErrorCode.HIVE_FILESYSTEM_ERROR;
+import static io.trino.plugin.hive.HiveErrorCode.HIVE_INVALID_METADATA;
 import static io.trino.plugin.hive.HiveErrorCode.HIVE_METASTORE_ERROR;
 import static io.trino.plugin.hive.TableType.MANAGED_TABLE;
 import static io.trino.plugin.hive.TableType.VIRTUAL_VIEW;
@@ -727,7 +728,56 @@ public void setTableOwner(String databaseName, String tableName, HivePrincipal p
     @Override
     public void commentColumn(String databaseName, String tableName, String columnName, Optional<String> comment)
     {
-        throw new TrinoException(NOT_SUPPORTED, "Column comment is not yet supported by Glue service");
+        Table table = getExistingTable(databaseName, tableName);
+        List<Column> dataColumns = table.getDataColumns();
+        List<Column> partitionColumns = table.getPartitionColumns();
+
+        Optional<Integer> matchingDataColumn = indexOfColumnWithName(dataColumns, columnName);
+        Optional<Integer> matchingPartitionColumn = indexOfColumnWithName(partitionColumns, columnName);
+
+        if (matchingDataColumn.isPresent() && matchingPartitionColumn.isPresent()) {
+            throw new TrinoException(HIVE_INVALID_METADATA, "Found two columns with names matching " + columnName);
+        }
+        if (matchingDataColumn.isEmpty() && matchingPartitionColumn.isEmpty()) {
+            throw new ColumnNotFoundException(table.getSchemaTableName(), columnName);
+        }
+
+        Table updatedTable = Table.builder(table)
+                .setDataColumns(matchingDataColumn.map(index -> setColumnCommentForIndex(dataColumns, index, comment)).orElse(dataColumns))
+                .setPartitionColumns(matchingPartitionColumn.map(index -> setColumnCommentForIndex(partitionColumns, index, comment)).orElse(partitionColumns))
+                .build();
+
+        replaceTable(databaseName, tableName, updatedTable, null);
+    }
+
+    private static Optional<Integer> indexOfColumnWithName(List<Column> columns, String columnName)
+    {
+        Optional<Integer> index = Optional.empty();
+        for (int i = 0; i < columns.size(); i++) {
+            // Glue columns are always lowercase
+            if (columns.get(i).getName().equals(columnName)) {
+                index.ifPresent(ignored -> {
+                    throw new TrinoException(HIVE_INVALID_METADATA, "Found two columns with names matching " + columnName);
+                });
+                index = Optional.of(i);
+            }
+        }
+        return index;
+    }
+
+    private static List<Column> setColumnCommentForIndex(List<Column> columns, int indexToUpdate, Optional<String> comment)
+    {
+        ImmutableList.Builder<Column> newColumns = ImmutableList.builder();
+        for (int i = 0; i < columns.size(); i++) {
+            Column originalColumn = columns.get(i);
+            if (i == indexToUpdate) {
+                newColumns.add(new Column(originalColumn.getName(), originalColumn.getType(), comment, originalColumn.getProperties()));
+            }
+            else {
+                newColumns.add(originalColumn);
+            }
+        }
+        return newColumns.build();
     }
 
     @Override

diff --git a/...n/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java b/...n/trino-hive/src/test/java/io/trino/plugin/hive/metastore/glue/TestHiveGlueMetastore.java
@@ -1425,6 +1425,38 @@ public void testGlueObjectsWithoutStorageDescriptor()
         }
     }
 
+    @Test
+    public void testAlterColumnComment()
+            throws Exception
+    {
+        SchemaTableName tableName = temporaryTable("test_alter_column_comment");
+        List<ColumnMetadata> columns = ImmutableList.of(
+                new ColumnMetadata("first_column", BIGINT),
+                new ColumnMetadata("second_column", VARCHAR),
+                new ColumnMetadata("partition_column", BIGINT));
+        createDummyPartitionedTable(tableName, columns, ImmutableList.of("partition_column"), ImmutableList.of());
+        try {
+            metastore.commentColumn(tableName.getSchemaName(), tableName.getTableName(), "second_column", Optional.of("second column comment"));
+            metastore.commentColumn(tableName.getSchemaName(), tableName.getTableName(), "partition_column", Optional.of("partition column comment"));
+
+            Table withComment = metastore.getTable(tableName.getSchemaName(), tableName.getTableName()).orElseThrow();
+            assertThat(withComment.getColumn("first_column").orElseThrow().getComment()).isEmpty();
+            assertThat(withComment.getColumn("second_column").orElseThrow().getComment()).isEqualTo(Optional.of("second column comment"));
+            assertThat(withComment.getColumn("partition_column").orElseThrow().getComment()).isEqualTo(Optional.of("partition column comment"));
+
+            metastore.commentColumn(tableName.getSchemaName(), tableName.getTableName(), "second_column", Optional.empty());
+            withComment = metastore.getTable(tableName.getSchemaName(), tableName.getTableName()).orElseThrow();
+            assertThat(withComment.getColumn("first_column").orElseThrow().getComment()).isEmpty();
+            assertThat(withComment.getColumn("second_column").orElseThrow().getComment()).isEmpty();
+            assertThat(withComment.getColumn("partition_column").orElseThrow().getComment()).isEqualTo(Optional.of("partition column comment"));
+        }
+        finally {
+            glueClient.deleteTable(new DeleteTableRequest()
+                    .withDatabaseName(tableName.getSchemaName())
+                    .withName(tableName.getTableName()));
+        }
+    }
+
     private Block singleValueBlock(long value)
     {
         BlockBuilder blockBuilder = BIGINT.createBlockBuilder(null, 1);