From 5cc6152d8b49da390020690a7004d54af2c7468a Mon Sep 17 00:00:00 2001 From: Marius Grama Date: Sat, 27 Aug 2022 23:33:04 +0200 Subject: [PATCH] Add test for 'COMMENT ON COLUMN' statement --- .../execution/BaseDataDefinitionTaskTest.java | 29 +++++++++++++++++++ .../io/trino/execution/TestCommentTask.java | 19 ++++++++++-- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java b/core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java index e11cc23e13bf..02e9f1d3c0bf 100644 --- a/core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java +++ b/core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java @@ -62,6 +62,7 @@ import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.Objects; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CopyOnWriteArrayList; @@ -69,6 +70,7 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Verify.verifyNotNull; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static com.google.common.util.concurrent.MoreExecutors.directExecutor; import static io.trino.metadata.MetadataManager.createTestMetadataManager; @@ -139,6 +141,11 @@ protected static QualifiedName qualifiedName(String name) return QualifiedName.of(TEST_CATALOG_NAME, SCHEMA, name); } + protected static QualifiedName qualifiedColumnName(String tableName, String columnName) + { + return QualifiedName.of(TEST_CATALOG_NAME, SCHEMA, tableName, columnName); + } + protected static QualifiedObjectName asQualifiedObjectName(QualifiedName viewName) { return QualifiedObjectName.valueOf(viewName.toString()); @@ -420,6 +427,21 @@ public void setViewComment(Session session, QualifiedObjectName viewName, Option view.getRunAsIdentity())); } + @Override + public void setColumnComment(Session session, TableHandle tableHandle, ColumnHandle column, Optional comment) + { + ConnectorTableMetadata tableMetadata = getTableMetadata(tableHandle); + TestingColumnHandle columnHandle = (TestingColumnHandle) column; + ConnectorTableMetadata newTableMetadata = new ConnectorTableMetadata( + tableMetadata.getTable(), + tableMetadata.getColumns().stream() + .map(tableColumn -> Objects.equals(tableColumn.getName(), columnHandle.getName()) ? withComment(tableColumn, comment) : tableColumn) + .collect(toImmutableList()), + tableMetadata.getProperties(), + tableMetadata.getComment()); + tables.put(tableMetadata.getTable(), newTableMetadata); + } + @Override public void renameMaterializedView(Session session, QualifiedObjectName source, QualifiedObjectName target) { @@ -433,5 +455,12 @@ public ResolvedFunction getCoercion(Session session, OperatorType operatorType, { return delegate.getCoercion(session, operatorType, fromType, toType); } + + private static ColumnMetadata withComment(ColumnMetadata tableColumn, Optional comment) + { + return ColumnMetadata.builderFrom(tableColumn) + .setComment(comment) + .build(); + } } } diff --git a/core/trino-main/src/test/java/io/trino/execution/TestCommentTask.java b/core/trino-main/src/test/java/io/trino/execution/TestCommentTask.java index b9ef6ce3b550..b1eaad0f4748 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestCommentTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestCommentTask.java @@ -17,7 +17,9 @@ import com.google.common.util.concurrent.ListenableFuture; import io.trino.execution.warnings.WarningCollector; import io.trino.metadata.QualifiedObjectName; +import io.trino.metadata.TableHandle; import io.trino.security.AllowAllAccessControl; +import io.trino.spi.connector.ConnectorTableMetadata; import io.trino.sql.tree.Comment; import io.trino.sql.tree.QualifiedName; import org.testng.annotations.Test; @@ -26,6 +28,7 @@ import static io.airlift.concurrent.MoreFutures.getFutureValue; import static io.trino.spi.StandardErrorCode.TABLE_NOT_FOUND; +import static io.trino.sql.tree.Comment.Type.COLUMN; import static io.trino.sql.tree.Comment.Type.TABLE; import static io.trino.sql.tree.Comment.Type.VIEW; import static io.trino.testing.TestingHandles.TEST_CATALOG_NAME; @@ -36,8 +39,6 @@ public class TestCommentTask extends BaseDataDefinitionTaskTest { - // TODO: Add test for 'COMMENT ON COLUMN' statements - @Test public void testCommentTable() { @@ -106,6 +107,20 @@ public void testCommentViewOnMaterializedView() .hasMessage("View '%s' does not exist, but a materialized view with that name exists. Setting comments on materialized views is unsupported.", materializedViewName); } + @Test + public void testCommentTableColumn() + { + QualifiedObjectName tableName = qualifiedObjectName("existing_table"); + QualifiedName columnName = qualifiedColumnName("existing_table", "test"); + metadata.createTable(testSession, TEST_CATALOG_NAME, someTable(tableName), false); + + getFutureValue(setComment(COLUMN, columnName, Optional.of("new test column comment"))); + TableHandle tableHandle = metadata.getTableHandle(testSession, tableName).get(); + ConnectorTableMetadata connectorTableMetadata = metadata.getTableMetadata(testSession, tableHandle).getMetadata(); + assertThat(Optional.ofNullable(connectorTableMetadata.getColumns().stream().filter(column -> "test".equals(column.getName())).findFirst().get().getComment())) + .isEqualTo(Optional.of("new test column comment")); + } + private ListenableFuture setComment(Comment.Type type, QualifiedName viewName, Optional comment) { return new CommentTask(metadata, new AllowAllAccessControl()).execute(new Comment(type, viewName, comment), queryStateMachine, ImmutableList.of(), WarningCollector.NOOP);