SNOW-1842220 Add comments to columns created by schema evolution

snowflakedb · Dec 3, 2024 · 2dce516 · 2dce516
1 parent 85db567
commit 2dce516
Show file tree

Hide file tree

Showing 7 changed files with 52 additions and 3 deletions.
diff --git a/src/main/java/com/snowflake/kafka/connector/internal/DescribeTableRow.java b/src/main/java/com/snowflake/kafka/connector/internal/DescribeTableRow.java
@@ -7,9 +7,18 @@ public class DescribeTableRow {
   private final String column;
   private final String type;
 
+  private final String comment;
+
+  public DescribeTableRow(String column, String type, String comment) {
+    this.column = column;
+    this.type = type;
+    this.comment = comment;
+  }
+
   public DescribeTableRow(String column, String type) {
     this.column = column;
     this.type = type;
+    this.comment = null;
   }
 
   public String getColumn() {
@@ -20,6 +29,10 @@ public String getType() {
     return type;
   }
 
+  public String getComment() {
+    return comment;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) return true;

diff --git a/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeConnectionServiceV1.java b/src/main/java/com/snowflake/kafka/connector/internal/SnowflakeConnectionServiceV1.java
@@ -1179,7 +1179,8 @@ public Optional<List<DescribeTableRow>> describeTable(String tableName) {
       while (result.next()) {
         String columnName = result.getString("name");
         String type = result.getString("type");
-        rows.add(new DescribeTableRow(columnName, type));
+        String comment = result.getString("comment");
+        rows.add(new DescribeTableRow(columnName, type, comment));
       }
       return Optional.of(rows);
     } catch (Exception e) {

diff --git a/...owflake/kafka/connector/internal/streaming/schemaevolution/iceberg/IcebergColumnTree.java b/...owflake/kafka/connector/internal/streaming/schemaevolution/iceberg/IcebergColumnTree.java
@@ -4,6 +4,7 @@
 class IcebergColumnTree {
 
   private final IcebergFieldNode rootNode;
+  private final String comment;
 
   String getColumnName() {
     return rootNode.name;
@@ -13,7 +14,17 @@ IcebergFieldNode getRootNode() {
     return rootNode;
   }
 
+  String getComment() {
+    return comment;
+  }
+
+  public IcebergColumnTree(IcebergFieldNode rootNode, String comment) {
+    this.rootNode = rootNode;
+    this.comment = comment;
+  }
+
   IcebergColumnTree(IcebergFieldNode rootNode) {
     this.rootNode = rootNode;
+    this.comment = null;
   }
 }
diff --git a/.../kafka/connector/internal/streaming/schemaevolution/iceberg/IcebergColumnTreeFactory.java b/.../kafka/connector/internal/streaming/schemaevolution/iceberg/IcebergColumnTreeFactory.java
@@ -44,7 +44,7 @@ IcebergColumnTree fromConnectSchema(Field kafkaConnectField) {
             + kafkaConnectField.name());
     IcebergFieldNode rootNode =
         createNode(kafkaConnectField.name().toUpperCase(), kafkaConnectField.schema());
-    return new IcebergColumnTree(rootNode);
+    return new IcebergColumnTree(rootNode, kafkaConnectField.schema().doc());
   }
 
   // -- parse tree from Iceberg schema logic --

diff --git a/...a/connector/internal/streaming/schemaevolution/iceberg/IcebergSchemaEvolutionService.java b/...a/connector/internal/streaming/schemaevolution/iceberg/IcebergSchemaEvolutionService.java
@@ -170,7 +170,8 @@ private Map<String, ColumnInfos> toColumnInfos(List<IcebergColumnTree> columnTre
         .map(
             columnTree ->
                 Maps.immutableEntry(
-                    columnTree.getColumnName(), new ColumnInfos(typeBuilder.buildType(columnTree))))
+                    columnTree.getColumnName(),
+                    new ColumnInfos(typeBuilder.buildType(columnTree), columnTree.getComment())))
         .collect(
             Collectors.toMap(
                 Map.Entry::getKey, Map.Entry::getValue, (oldValue, newValue) -> newValue));

diff --git a/...va/com/snowflake/kafka/connector/streaming/iceberg/IcebergIngestionSchemaEvolutionIT.java b/...va/com/snowflake/kafka/connector/streaming/iceberg/IcebergIngestionSchemaEvolutionIT.java
@@ -483,4 +483,15 @@ private static Stream<Arguments> testEvolutionOfComplexTypes_dataSource() {
             twoObjectsExtendedWithMapAndArrayPayload(),
             false));
   }
+
+  @Test
+  // @Disabled
+  void shouldAppendedCommentTest() throws Exception {
+    // when
+    insertWithRetry(schemaAndPayloadWithComment(), 0, true);
+    waitForOffset(1);
+    // then
+    List<DescribeTableRow> columns = describeTable(tableName);
+    assertEquals("Test comment", columns.get(1).getComment());
+  }
 }
diff --git a/src/test/java/com/snowflake/kafka/connector/streaming/iceberg/TestJsons.java b/src/test/java/com/snowflake/kafka/connector/streaming/iceberg/TestJsons.java
@@ -476,6 +476,15 @@ static String twoObjectsExtendedWithMapAndArrayPayload() {
         + "    }";
   }
 
+  static String schemaAndPayloadWithComment() {
+    return SCHEMA_BEGINNING
+        + COMMENTED_SCHEMA
+        + SCHEMA_END
+        + "\"payload\": {"
+        + STRING_PAYLOAD
+        + "}}";
+  }
+
   static String BOOL_SCHEMA = " {  \"field\" : \"test_boolean\", \"type\" : \"boolean\"} ";
 
   static String INT64_SCHEMA = "{  \"field\" : \"test_int64\", \"type\" : \"int64\" }";
@@ -489,6 +498,9 @@ static String twoObjectsExtendedWithMapAndArrayPayload() {
 
   static String STRING_SCHEMA = "{  \"field\" : \"test_string\",  \"type\" : \"string\"  }";
 
+  static String COMMENTED_SCHEMA =
+      "{  \"field\" : \"test_string\",  \"type\" : \"string\", \"doc\": \"Test comment\"  }";
+
   static final String BOOL_PAYLOAD = "\"test_boolean\" : true ";
   static final String INT64_PAYLOAD = "\"test_int64\" : 2137324241343241 ";
   static final String INT32_PAYLOAD = "\"test_int32\" : 2137 ";