Add support for materialized view column comments

core/trino-main/src/main/java/io/trino/execution/CommentTask.java

@@ -17,6 +17,7 @@
 import com.google.inject.Inject;
 import io.trino.Session;
 import io.trino.execution.warnings.WarningCollector;
+import io.trino.metadata.MaterializedViewDefinition;
 import io.trino.metadata.Metadata;
 import io.trino.metadata.QualifiedObjectName;
 import io.trino.metadata.RedirectionAwareTableHandle;
@@ -140,18 +141,14 @@ private void commentOnColumn(Comment statement, Session session)
 
         QualifiedObjectName originalObjectName = createQualifiedObjectName(session, statement, prefix);
         if (metadata.isView(session, originalObjectName)) {
-            String columnName = statement.getName().getSuffix();
             ViewDefinition viewDefinition = metadata.getView(session, originalObjectName).get();
-            ViewColumn viewColumn = viewDefinition.getColumns().stream()
-                    .filter(column -> column.getName().equals(columnName))
-                    .findAny()
-                    .orElseThrow(() -> semanticException(COLUMN_NOT_FOUND, statement, "Column does not exist: %s", columnName));
-
-            accessControl.checkCanSetColumnComment(session.toSecurityContext(), originalObjectName);
+            ViewColumn viewColumn = findAndCheckViewColumn(statement, session, viewDefinition, originalObjectName);
             metadata.setViewColumnComment(session, originalObjectName, viewColumn.getName(), statement.getComment());
         }
         else if (metadata.isMaterializedView(session, originalObjectName)) {
-            throw semanticException(TABLE_NOT_FOUND, statement, "Setting comments on the columns of materialized views is unsupported");
+            MaterializedViewDefinition materializedViewDefinition = metadata.getMaterializedView(session, originalObjectName).get();
+            ViewColumn viewColumn = findAndCheckViewColumn(statement, session, materializedViewDefinition, originalObjectName);
+            metadata.setMaterializedViewColumnComment(session, originalObjectName, viewColumn.getName(), statement.getComment());
         }
         else {
             RedirectionAwareTableHandle redirectionAwareTableHandle = metadata.getRedirectionAwareTableHandle(session, originalObjectName);
@@ -171,4 +168,15 @@ else if (metadata.isMaterializedView(session, originalObjectName)) {
             metadata.setColumnComment(session, tableHandle, columnHandles.get(columnName), statement.getComment());
         }
     }
+
+    private ViewColumn findAndCheckViewColumn(Comment statement, Session session, ViewDefinition viewDefinition, QualifiedObjectName originalObjectName)
+    {
+        String columnName = statement.getName().getSuffix();
+        ViewColumn viewColumn = viewDefinition.getColumns().stream()
+                .filter(column -> column.getName().equals(columnName))
+                .findAny()
+                .orElseThrow(() -> semanticException(COLUMN_NOT_FOUND, statement, "Column does not exist: %s", columnName));
+        accessControl.checkCanSetColumnComment(session.toSecurityContext(), originalObjectName);
+        return viewColumn;
+    }
 }

core/trino-main/src/main/java/io/trino/metadata/MaterializedViewDefinition.java

@@ -92,7 +92,7 @@ public ConnectorMaterializedViewDefinition toConnectorMaterializedViewDefinition
                 getCatalog(),
                 getSchema(),
                 getColumns().stream()
-                        .map(column -> new ConnectorMaterializedViewDefinition.Column(column.getName(), column.getType()))
+                        .map(column -> new ConnectorMaterializedViewDefinition.Column(column.getName(), column.getType(), column.getComment()))
                         .collect(toImmutableList()),
                 getGracePeriod(),
                 getComment(),

core/trino-main/src/main/java/io/trino/metadata/Metadata.java

@@ -227,6 +227,11 @@ Optional<TableExecuteHandle> getTableHandleForExecute(
      */
     void setViewColumnComment(Session session, QualifiedObjectName viewName, String columnName, Optional<String> comment);
 
+    /**
+     * Comments to the specified materialized view column.
+     */
+    void setMaterializedViewColumnComment(Session session, QualifiedObjectName viewName, String columnName, Optional<String> comment);
+
     /**
      * Comments to the specified column.
      */

core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java

@@ -609,7 +609,11 @@ public List<TableColumnsMetadata> listTableColumns(Session session, QualifiedTab
                     ImmutableList.Builder<ColumnMetadata> columns = ImmutableList.builder();
                     for (ViewColumn column : entry.getValue().getColumns()) {
                         try {
-                            columns.add(new ColumnMetadata(column.getName(), typeManager.getType(column.getType())));
+                            columns.add(ColumnMetadata.builder()
+                                    .setName(column.getName())
+                                    .setType(typeManager.getType(column.getType()))
+                                    .setComment(column.getComment())
+                                    .build());
                         }
                         catch (TypeNotFoundException e) {
                             throw new TrinoException(INVALID_VIEW, format("Unknown type '%s' for column '%s' in materialized view: %s", column.getType(), column.getName(), entry.getKey()));
@@ -736,6 +740,15 @@ public void setViewColumnComment(Session session, QualifiedObjectName viewName,
         metadata.setViewColumnComment(session.toConnectorSession(catalogHandle), viewName.asSchemaTableName(), columnName, comment);
     }
 
+    @Override
+    public void setMaterializedViewColumnComment(Session session, QualifiedObjectName viewName, String columnName, Optional<String> comment)
+    {
+        CatalogMetadata catalogMetadata = getCatalogMetadataForWrite(session, viewName.getCatalogName());
+        CatalogHandle catalogHandle = catalogMetadata.getCatalogHandle();
+        ConnectorMetadata metadata = catalogMetadata.getMetadata(session);
+        metadata.setMaterializedViewColumnComment(session.toConnectorSession(catalogHandle), viewName.asSchemaTableName(), columnName, comment);
+    }
+
     @Override
     public void setColumnComment(Session session, TableHandle tableHandle, ColumnHandle column, Optional<String> comment)
     {

core/trino-main/src/main/java/io/trino/metadata/ViewInfo.java

@@ -44,7 +44,7 @@ public ViewInfo(ConnectorMaterializedViewDefinition viewDefinition)
     {
         this.originalSql = viewDefinition.getOriginalSql();
         this.columns = viewDefinition.getColumns().stream()
-                .map(column -> new ViewColumn(column.getName(), column.getType(), Optional.empty()))
+                .map(column -> new ViewColumn(column.getName(), column.getType(), column.getComment()))
                 .collect(toImmutableList());
         this.comment = viewDefinition.getComment();
         this.storageTable = viewDefinition.getStorageTable();

core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java

@@ -423,6 +423,15 @@ public void setViewColumnComment(ConnectorSession session, SchemaTableName viewN
         }
     }
 
+    @Override
+    public void setMaterializedViewColumnComment(ConnectorSession session, SchemaTableName viewName, String columnName, Optional<String> comment)
+    {
+        Span span = startSpan("setMaterializedViewColumnComment", viewName);
+        try (var ignored = scopedSpan(span)) {
+            delegate.setMaterializedViewColumnComment(session, viewName, columnName, comment);
+        }
+    }
+
     @Override
     public void setColumnComment(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column, Optional<String> comment)
     {

core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java

@@ -422,6 +422,15 @@ public void setViewColumnComment(Session session, QualifiedObjectName viewName,
         }
     }
 
+    @Override
+    public void setMaterializedViewColumnComment(Session session, QualifiedObjectName viewName, String columnName, Optional<String> comment)
+    {
+        Span span = startSpan("setMaterializedViewColumnComment", viewName);
+        try (var ignored = scopedSpan(span)) {
+            delegate.setMaterializedViewColumnComment(session, viewName, columnName, comment);
+        }
+    }
+
     @Override
     public void setColumnComment(Session session, TableHandle tableHandle, ColumnHandle column, Optional<String> comment)
     {

core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java

@@ -576,6 +576,26 @@ public void setViewColumnComment(Session session, QualifiedObjectName viewName,
                             view.getRunAsIdentity()));
         }
 
+        @Override
+        public void setMaterializedViewColumnComment(Session session, QualifiedObjectName viewName, String columnName, Optional<String> comment)
+        {
+            MaterializedViewDefinition view = materializedViews.get(viewName.asSchemaTableName());
+            materializedViews.put(
+                    viewName.asSchemaTableName(),
+                    new MaterializedViewDefinition(
+                            view.getOriginalSql(),
+                            view.getCatalog(),
+                            view.getSchema(),
+                            view.getColumns().stream()
+                                    .map(currentViewColumn -> columnName.equals(currentViewColumn.getName()) ? new ViewColumn(currentViewColumn.getName(), currentViewColumn.getType(), comment) : currentViewColumn)
+                                    .collect(toImmutableList()),
+                            view.getGracePeriod(),
+                            view.getComment(),
+                            view.getRunAsIdentity().get(),
+                            view.getStorageTable(),
+                            view.getProperties()));
+        }
+
         @Override
         public void renameMaterializedView(Session session, QualifiedObjectName source, QualifiedObjectName target)
         {

core/trino-main/src/test/java/io/trino/execution/TestCommentTask.java

@@ -141,6 +141,25 @@ public void testCommentViewColumn()
                 .hasMessage("Column does not exist: %s", missingColumnName.getSuffix());
     }
 
+    @Test
+    public void testCommentMaterializedViewColumn()
+    {
+        QualifiedObjectName materializedViewName = qualifiedObjectName("existing_materialized_view");
+        metadata.createMaterializedView(testSession, QualifiedObjectName.valueOf(materializedViewName.toString()), someMaterializedView(), false, false);
+        assertThat(metadata.isMaterializedView(testSession, materializedViewName)).isTrue();
+
+        QualifiedName columnName = qualifiedColumnName("existing_materialized_view", "test");
+        QualifiedName missingColumnName = qualifiedColumnName("existing_materialized_view", "missing");
+
+        getFutureValue(setComment(COLUMN, columnName, Optional.of("new test column comment")));
+        assertThat(metadata.getMaterializedView(testSession, materializedViewName).get().getColumns().stream().filter(column -> "test".equals(column.getName())).collect(onlyElement()).getComment())
+                .isEqualTo(Optional.of("new test column comment"));
+
+        assertTrinoExceptionThrownBy(() -> getFutureValue(setComment(COLUMN, missingColumnName, Optional.of("comment for missing column"))))
+                .hasErrorCode(COLUMN_NOT_FOUND)
+                .hasMessage("Column does not exist: %s", missingColumnName.getSuffix());
+    }
+
     private ListenableFuture<Void> setComment(Comment.Type type, QualifiedName viewName, Optional<String> comment)
     {
         return new CommentTask(metadata, new AllowAllAccessControl()).execute(new Comment(type, viewName, comment), queryStateMachine, ImmutableList.of(), WarningCollector.NOOP);

core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java

@@ -294,6 +294,12 @@ public void setViewColumnComment(Session session, QualifiedObjectName viewName,
         throw new UnsupportedOperationException();
     }
 
+    @Override
+    public void setMaterializedViewColumnComment(Session session, QualifiedObjectName viewName, String columnName, Optional<String> comment)
+    {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public void setColumnComment(Session session, TableHandle tableHandle, ColumnHandle column, Optional<String> comment)
     {

core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMaterializedViewDefinition.java

@@ -182,11 +182,19 @@ public static final class Column
     {
         private final String name;
         private final TypeId type;
+        private final Optional<String> comment;
 
+        @Deprecated
         public Column(String name, TypeId type)
+        {
+            this(name, type, Optional.empty());
+        }
+
+        public Column(String name, TypeId type, Optional<String> comment)
         {
             this.name = requireNonNull(name, "name is null");
             this.type = requireNonNull(type, "type is null");
+            this.comment = requireNonNull(comment, "comment is null");
         }
 
         public String getName()
@@ -199,6 +207,11 @@ public TypeId getType()
             return type;
         }
 
+        public Optional<String> getComment()
+        {
+            return comment;
+        }
+
         @Override
         public String toString()
         {
@@ -216,13 +229,14 @@ public boolean equals(Object o)
             }
             Column column = (Column) o;
             return Objects.equals(name, column.name) &&
-                    Objects.equals(type, column.type);
+                    Objects.equals(type, column.type) &&
+                    Objects.equals(comment, column.comment);
         }
 
         @Override
         public int hashCode()
         {
-            return Objects.hash(name, type);
+            return Objects.hash(name, type, comment);
         }
     }
 }

core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java

@@ -434,6 +434,14 @@ default void setViewColumnComment(ConnectorSession session, SchemaTableName view
         throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting view column comments");
     }
 
+    /**
+     * Comments to the specified materialized view column.
+     */
+    default void setMaterializedViewColumnComment(ConnectorSession session, SchemaTableName viewName, String columnName, Optional<String> comment)
+    {
+        throw new TrinoException(NOT_SUPPORTED, "This connector does not support setting materialized view column comments");
+    }
+
     /**
      * Comments to the specified column
      */

docs/src/main/sphinx/connector/iceberg.rst

@@ -1277,6 +1277,7 @@ The Iceberg connector supports setting comments on the following objects:
 - tables
 - views
 - table columns
+- materialized view columns
 
 The ``COMMENT`` option is supported on both the table and the table columns for
 the :doc:`/sql/create-table` operation.

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java

@@ -487,6 +487,14 @@ public void setViewColumnComment(ConnectorSession session, SchemaTableName viewN
         }
     }
 
+    @Override
+    public void setMaterializedViewColumnComment(ConnectorSession session, SchemaTableName viewName, String columnName, Optional<String> comment)
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
+            delegate.setMaterializedViewColumnComment(session, viewName, columnName, comment);
+        }
+    }
+
     @Override
     public void setColumnComment(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column, Optional<String> comment)
     {

plugin/trino-blackhole/pom.xml

@@ -104,6 +104,12 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>io.trino</groupId>
+            <artifactId>trino-testing-services</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>io.trino</groupId>
             <artifactId>trino-tpch</artifactId>