From 843622f916abd224797edb8673c62fe778347508 Mon Sep 17 00:00:00 2001
From: DimaVilda <dimavilda@gmail.com>
Date: Fri, 29 Nov 2024 16:47:39 +0100
Subject: [PATCH] fixes: avro schema ser for nullable enums

---
 .../jsonschema/AvroJsonSchemaConverter.java   | 13 +++---
 .../AvroJsonSchemaConverterTest.java          | 42 +++++++++++++++++++
 .../jsonschema/JsonAvroConversionTest.java    | 37 ++++++++++++++++
 3 files changed, 84 insertions(+), 8 deletions(-)

diff --git a/api/src/main/java/io/kafbat/ui/util/jsonschema/AvroJsonSchemaConverter.java b/api/src/main/java/io/kafbat/ui/util/jsonschema/AvroJsonSchemaConverter.java
index 55c75c950..ae2d378c4 100644
--- a/api/src/main/java/io/kafbat/ui/util/jsonschema/AvroJsonSchemaConverter.java
+++ b/api/src/main/java/io/kafbat/ui/util/jsonschema/AvroJsonSchemaConverter.java
@@ -80,14 +80,11 @@ private FieldSchema createUnionSchema(Schema schema, Map<String, FieldSchema> de
     final Map<String, FieldSchema> fields = schema.getTypes().stream()
         .filter(t -> !t.getType().equals(Schema.Type.NULL))
         .map(f -> {
-          String oneOfFieldName;
-          if (f.getType().equals(Schema.Type.RECORD)) {
-            // for records using full record name
-            oneOfFieldName = f.getFullName();
-          } else {
-            // for primitive types - using type name
-            oneOfFieldName = f.getType().getName().toLowerCase();
-          }
+          String oneOfFieldName = switch (f.getType()) {
+            case RECORD -> f.getFullName();
+            case ENUM -> f.getName();
+            default -> f.getType().getName().toLowerCase();
+          };
           return Tuples.of(oneOfFieldName, convertSchema(f, definitions, false));
         }).collect(Collectors.toMap(
             Tuple2::getT1,
diff --git a/api/src/test/java/io/kafbat/ui/util/jsonschema/AvroJsonSchemaConverterTest.java b/api/src/test/java/io/kafbat/ui/util/jsonschema/AvroJsonSchemaConverterTest.java
index 299283aed..35c230f56 100644
--- a/api/src/test/java/io/kafbat/ui/util/jsonschema/AvroJsonSchemaConverterTest.java
+++ b/api/src/test/java/io/kafbat/ui/util/jsonschema/AvroJsonSchemaConverterTest.java
@@ -244,6 +244,48 @@ void testRecordReferences() {
     convertAndCompare(expectedJsonSchema, avroSchema);
   }
 
+  @Test
+  void testNullableUnionEnum()  {
+    String avroSchema =
+        " {"
+            + "     \"type\": \"record\","
+            + "     \"name\": \"Message\","
+            + "     \"namespace\": \"com.provectus.kafka\","
+            + "     \"fields\": ["
+            + "                     {"
+            + "                         \"name\": \"enum_nullable_union\","
+            + "                         \"type\": [\"null\", {"
+            + "                             \"type\": \"enum\","
+            + "                             \"name\": \"Suit\","
+            + "                             \"symbols\": [\"SPADES\",\"HEARTS\",\"DIAMONDS\",\"CLUBS\"]"
+            + "                         }]"
+            + "                     }"
+            + "     ]"
+            + " }";
+
+    String expectedJsonSchema =
+        "{\"$id\":\"http://example.com/Message\","
+            + "\"$schema\":\"https://json-schema.org/draft/2020-12/schema\","
+            + "\"type\":\"object\","
+            + "\"properties\":{"
+            + "\"enum_nullable_union\":{"
+            + "\"oneOf\":["
+            + "{\"type\":\"null\"},"
+            + "{\"type\":\"object\","
+            + "\"properties\":{"
+            + "\"Suit\":{"
+            + "\"type\":\"string\","
+            + "\"enum\":[\"SPADES\",\"HEARTS\",\"DIAMONDS\",\"CLUBS\"]"
+            + "}}}"
+            + "]"
+            + "}},"
+            + "\"definitions\":{"
+            + "\"com.provectus.kafka.Message\":{\"$ref\":\"#\"}"
+            + "}}";
+
+    convertAndCompare(expectedJsonSchema, avroSchema);
+  }
+
   @SneakyThrows
   private void convertAndCompare(String expectedJsonSchema, String sourceAvroSchema) {
     var parseAvroSchema = new Schema.Parser().parse(sourceAvroSchema);
diff --git a/api/src/test/java/io/kafbat/ui/util/jsonschema/JsonAvroConversionTest.java b/api/src/test/java/io/kafbat/ui/util/jsonschema/JsonAvroConversionTest.java
index 5d3daf08e..5ac977fb9 100644
--- a/api/src/test/java/io/kafbat/ui/util/jsonschema/JsonAvroConversionTest.java
+++ b/api/src/test/java/io/kafbat/ui/util/jsonschema/JsonAvroConversionTest.java
@@ -699,6 +699,43 @@ void unionFieldWithInnerTypesNamesClash() {
 
   }
 
+  @Test
+  void unionNullableEnumField() {
+    var schema = createSchema(
+        """
+             {
+               "type": "record",
+               "namespace": "com.test",
+               "name": "TestAvroRecord",
+               "fields": [
+                 {
+                   "name": "enum_nullable_union",
+                   "type" : [ "null", {
+                               "type" : "enum",
+                               "name" : "Suit",
+                               "symbols" : ["SPADES", "HEARTS", "DIAMONDS", "CLUBS"]
+                             } ]
+                 }
+               ]
+            }"""
+    );
+
+    GenericData.Record inputRecord = new GenericData.Record(schema);
+    inputRecord.put("enum_nullable_union",
+        new GenericData.EnumSymbol(
+            schema.getField("enum_nullable_union").schema().getTypes().get(1), "SPADES"));
+    String expectedJsonWithEnum = """
+          {
+            "enum_nullable_union": { "Suit": "SPADES"}\s
+          }
+         \s""";
+    assertJsonsEqual(expectedJsonWithEnum, convertAvroToJson(inputRecord, schema));
+
+    GenericData.Record inputNullRecord  = new GenericData.Record(schema);
+    inputNullRecord.put("enum_nullable_union", null);
+    assertJsonsEqual("{}", convertAvroToJson(inputNullRecord, schema));
+  }
+
   private Schema createSchema(String schema) {
     return new AvroSchema(schema).rawSchema();
   }