fixup! Support creating table with column comment in Delta Lake

trinodb · Jun 2, 2022 · 04d0d3f · 04d0d3f
1 parent c9f46be
commit 04d0d3f
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 4 deletions.
diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java
@@ -969,7 +969,7 @@ public void addColumn(ConnectorSession session, ConnectorTableHandle tableHandle
                     handle.getMetadataEntry().getId(),
                     columnsBuilder.build(),
                     partitionColumns,
-                    ImmutableMap.of(),
+                    getColumnComments(handle.getMetadataEntry()),
                     buildDeltaMetadataConfiguration(checkpointInterval),
                     ADD_COLUMN_OPERATION,
                     session,

diff --git a/...ava/io/trino/tests/product/deltalake/TestDeltaLakeDatabricksCreateTableCompatibility.java b/...ava/io/trino/tests/product/deltalake/TestDeltaLakeDatabricksCreateTableCompatibility.java
@@ -206,7 +206,7 @@ private static String getTableCommentOnDelta(String schemaName, String tableName
     }
 
     @Test(groups = {DELTA_LAKE_DATABRICKS, PROFILE_SPECIFIC_TESTS})
-    public void testCreateTableWithColumnComment()
+    public void testCreateTableWithColumnCommentOnTrino()
     {
         String tableName = "test_dl_create_column_comment_" + randomTableSuffix();
         String tableDirectory = "databricks-compatibility-test-" + tableName;
@@ -217,16 +217,44 @@ public void testCreateTableWithColumnComment()
                 tableDirectory));
 
         try {
-            assertThat(onTrino().executeQuery("SELECT comment FROM information_schema.columns WHERE table_schema = 'default' AND table_name = '" + tableName + "' AND column_name = 'col'"))
-                    .containsOnly(row("test comment"));
+            assertEquals(getColumnCommentOnTrino("default", tableName, "col"), "test comment");
+            assertEquals(getColumnCommentOnDelta("default", tableName, "col"), "test comment");
 
+            // Verify that adding a new column doesn't remove existing column comments
+            onTrino().executeQuery("ALTER TABLE delta.default." + tableName + " ADD COLUMN new_col INT");
+            assertEquals(getColumnCommentOnTrino("default", tableName, "col"), "test comment");
             assertEquals(getColumnCommentOnDelta("default", tableName, "col"), "test comment");
         }
         finally {
             onTrino().executeQuery("DROP TABLE delta.default." + tableName);
         }
     }
 
+    @Test(groups = {DELTA_LAKE_DATABRICKS, PROFILE_SPECIFIC_TESTS})
+    public void testCreateTableWithColumnCommentOnDelta()
+    {
+        String tableName = "test_dl_create_column_comment_" + randomTableSuffix();
+        String tableDirectory = "databricks-compatibility-test-" + tableName;
+
+        onDelta().executeQuery(format("CREATE TABLE default.%s (col INT COMMENT 'test comment') USING DELTA LOCATION 's3://%s/%s'",
+                tableName,
+                bucketName,
+                tableDirectory));
+
+        try {
+            assertEquals(getColumnCommentOnTrino("default", tableName, "col"), "test comment");
+        }
+        finally {
+            onDelta().executeQuery("DROP TABLE default." + tableName);
+        }
+    }
+
+    private static String getColumnCommentOnTrino(String schemaName, String tableName, String columnName)
+    {
+        QueryResult result = onTrino().executeQuery("SELECT comment FROM information_schema.columns WHERE table_schema = '" + schemaName + "' AND table_name = '" + tableName + "' AND column_name = '" + columnName + "'");
+        return (String) result.row(0).get(0);
+    }
+
     private static String getColumnCommentOnDelta(String schemaName, String tableName, String columnName)
     {
         QueryResult result = onDelta().executeQuery(format("DESCRIBE %s.%s %s", schemaName, tableName, columnName));