use SdkJsonGenerator

Author: RanVaknin

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/document/StrategyJsonSerializer.java

@@ -1,54 +1,50 @@
 package software.amazon.awssdk.enhanced.dynamodb.internal.document;
 
-import java.io.IOException;
-import java.io.StringWriter;
+import java.nio.charset.StandardCharsets;
 import java.util.Map;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.SdkBytes;
+import software.amazon.awssdk.protocols.json.SdkJsonGenerator;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.thirdparty.jackson.core.JsonFactory;
-import software.amazon.awssdk.thirdparty.jackson.core.JsonGenerator;
 
 @SdkInternalApi
 public final class StrategyJsonSerializer {
 
-    private StrategyJsonSerializer() {
-    }
-
     private enum JsonSerializationStrategy {
         NULL {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
                 generator.writeNull();
             }
         },
         STRING {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
-                generator.writeString(av.s());
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
+                generator.writeValue(av.s());
             }
         },
         NUMBER {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
                 generator.writeNumber(av.n());
             }
         },
         BOOLEAN {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
-                generator.writeBoolean(av.bool());
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
+                generator.writeValue(av.bool());
             }
         },
         BYTES {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
-                generator.writeBinary(av.b().asByteArray());
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
+                generator.writeValue(av.b().asByteBuffer());
             }
         },
         LIST {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
                 generator.writeStartArray();
                 for (AttributeValue item : av.l()) {
                     serializeAttributeValue(generator, item);
@@ -58,7 +54,7 @@ public void serialize(JsonGenerator generator, AttributeValue av) throws IOExcep
         },
         MAP {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
                 generator.writeStartObject();
                 for (Map.Entry<String, AttributeValue> entry : av.m().entrySet()) {
                     generator.writeFieldName(entry.getKey());
@@ -69,17 +65,17 @@ public void serialize(JsonGenerator generator, AttributeValue av) throws IOExcep
         },
         STRING_SET {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
                 generator.writeStartArray();
                 for (String s : av.ss()) {
-                    generator.writeString(s);
+                    generator.writeValue(s);
                 }
                 generator.writeEndArray();
             }
         },
         NUMBER_SET {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
                 generator.writeStartArray();
                 for (String n : av.ns()) {
                     generator.writeNumber(n);
@@ -89,34 +85,32 @@ public void serialize(JsonGenerator generator, AttributeValue av) throws IOExcep
         },
         BYTES_SET {
             @Override
-            public void serialize(JsonGenerator generator, AttributeValue av) throws IOException {
+            public void serialize(SdkJsonGenerator generator, AttributeValue av) {
                 generator.writeStartArray();
                 for (SdkBytes b : av.bs()) {
-                    generator.writeBinary(b.asByteArray());
+                    generator.writeValue(b.asByteBuffer());
                 }
                 generator.writeEndArray();
             }
         };
 
-        public abstract void serialize(JsonGenerator generator, AttributeValue av) throws IOException;
+        public abstract void serialize(SdkJsonGenerator generator, AttributeValue av);
     }
 
     public static String serializeAttributeValueMap(Map<String, AttributeValue> map) {
-        StringWriter writer = new StringWriter();
-        try (JsonGenerator generator = new JsonFactory().createGenerator(writer)) {
-            generator.writeStartObject();
-            for (Map.Entry<String, AttributeValue> entry : map.entrySet()) {
-                generator.writeFieldName(entry.getKey());
-                serializeAttributeValue(generator, entry.getValue());
-            }
-            generator.writeEndObject();
-        } catch (IOException e) {
-            throw new IllegalStateException("Failed to serialize AttributeValue map to JSON", e);
-        }
-        return writer.toString();
+        SdkJsonGenerator jsonGen = new SdkJsonGenerator(new JsonFactory(), "application/json");
+
+        jsonGen.writeStartObject();
+        map.forEach((key, value) -> {
+            jsonGen.writeFieldName(key);
+            serializeAttributeValue(jsonGen, value);
+        });
+        jsonGen.writeEndObject();
+
+        return new String(jsonGen.getBytes(), StandardCharsets.UTF_8);
     }
 
-    private static void serializeAttributeValue(JsonGenerator generator, AttributeValue av) throws IOException {
+    public static void serializeAttributeValue(SdkJsonGenerator generator, AttributeValue av) {
         JsonSerializationStrategy strategy = getStrategy(av);
         strategy.serialize(generator, av);
     }