From 9120a4581af90b665e256816c403c486258aa33b Mon Sep 17 00:00:00 2001
From: Marius Grama <findinpath@gmail.com>
Date: Tue, 22 Feb 2022 10:02:18 +0100
Subject: [PATCH] Add support for editing column comments on Iceberg tables

---
 .../java/io/trino/plugin/iceberg/IcebergMetadata.java    | 6 ++++++
 .../main/java/io/trino/plugin/iceberg/TrinoCatalog.java  | 2 ++
 .../java/io/trino/plugin/iceberg/TrinoHiveCatalog.java   | 9 +++++++++
 .../trino/plugin/iceberg/BaseIcebergConnectorTest.java   | 1 -
 .../plugin/iceberg/TestIcebergConnectorSmokeTest.java    | 1 -
 5 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java
index cbc1d7252e46..6ce2661cf3e1 100644
--- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java
+++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java
@@ -1341,6 +1341,12 @@ else if (strings.size() != 2) {
         return new MaterializedViewFreshness(true);
     }
 
+    @Override
+    public void setColumnComment(ConnectorSession session, ConnectorTableHandle tableHandle, ColumnHandle column, Optional<String> comment)
+    {
+        catalog.updateColumnComment(session, ((IcebergTableHandle) tableHandle).getSchemaTableName(), ((IcebergColumnHandle) column).getColumnIdentity(), comment);
+    }
+
     private Map<String, Optional<TableToken>> getMaterializedViewToken(ConnectorSession session, SchemaTableName name)
     {
         Map<String, Optional<TableToken>> viewToken = new HashMap<>();
diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TrinoCatalog.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TrinoCatalog.java
index ca29c826e880..189ea9d4a042 100644
--- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TrinoCatalog.java
+++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TrinoCatalog.java
@@ -115,4 +115,6 @@ void createMaterializedView(
     Optional<ConnectorMaterializedViewDefinition> getMaterializedView(ConnectorSession session, SchemaTableName schemaViewName);
 
     void renameMaterializedView(ConnectorSession session, SchemaTableName source, SchemaTableName target);
+
+    void updateColumnComment(ConnectorSession session, SchemaTableName schemaTableName, ColumnIdentity columnIdentity, Optional<String> comment);
 }
diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TrinoHiveCatalog.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TrinoHiveCatalog.java
index 4b81bca2c479..f95daf5f76f5 100644
--- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TrinoHiveCatalog.java
+++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/TrinoHiveCatalog.java
@@ -650,4 +650,13 @@ private List<String> listNamespaces(ConnectorSession session, Optional<String> n
         }
         return listNamespaces(session);
     }
+
+    @Override
+    public void updateColumnComment(ConnectorSession session, SchemaTableName schemaTableName, ColumnIdentity columnIdentity, Optional<String> comment)
+    {
+        metastore.commentColumn(schemaTableName.getSchemaName(), schemaTableName.getTableName(), columnIdentity.getName(), comment);
+
+        Table icebergTable = loadTable(session, schemaTableName);
+        icebergTable.updateSchema().updateColumnDoc(columnIdentity.getName(), comment.orElse(null)).commit();
+    }
 }
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
index 3c1f34baee6b..aeaac2a75b68 100644
--- 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
@@ -146,7 +146,6 @@ protected QueryRunner createQueryRunner()
     protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior)
     {
         switch (connectorBehavior) {
-            case SUPPORTS_COMMENT_ON_COLUMN:
             case SUPPORTS_TOPN_PUSHDOWN:
                 return false;
 
diff --git a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergConnectorSmokeTest.java b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergConnectorSmokeTest.java
index a6f7dd68d61e..2ee4dee126a2 100644
--- a/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergConnectorSmokeTest.java
+++ b/plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/TestIcebergConnectorSmokeTest.java
@@ -42,7 +42,6 @@ protected QueryRunner createQueryRunner()
     protected boolean hasBehavior(TestingConnectorBehavior connectorBehavior)
     {
         switch (connectorBehavior) {
-            case SUPPORTS_COMMENT_ON_COLUMN:
             case SUPPORTS_TOPN_PUSHDOWN:
                 return false;