From 55d26d4bcfa5d132d0713de8a21f5fe96e0c2f41 Mon Sep 17 00:00:00 2001 From: Tristan Tarrant Date: Mon, 1 Jul 2024 16:18:28 +0200 Subject: [PATCH] ISPN-346 Fix JSON for maps --- .../protostream/ProtobufParser.java | 10 +- .../descriptors/FieldDescriptor.java | 4 + .../descriptors/MapDescriptor.java | 22 +++ .../protostream/impl/JsonUtils.java | 126 ++++++++++++++---- integrationtests/pom.xml | 6 + .../marshaller/GeneratedMarshallerTest.java | 101 +++++++++++++- .../marshaller/model/MapOfMapOfUUID.java | 19 +++ .../marshaller/model/MapOfString.java | 14 ++ .../processor/marshaller/model/MapOfUUID.java | 14 ++ .../processor/marshaller/model/MapSchema.java | 24 ++-- 10 files changed, 304 insertions(+), 36 deletions(-) create mode 100644 integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfMapOfUUID.java create mode 100644 integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfString.java create mode 100644 integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfUUID.java diff --git a/core/src/main/java/org/infinispan/protostream/ProtobufParser.java b/core/src/main/java/org/infinispan/protostream/ProtobufParser.java index 41d57dc9c..828356b79 100644 --- a/core/src/main/java/org/infinispan/protostream/ProtobufParser.java +++ b/core/src/main/java/org/infinispan/protostream/ProtobufParser.java @@ -5,6 +5,7 @@ import org.infinispan.protostream.descriptors.Descriptor; import org.infinispan.protostream.descriptors.FieldDescriptor; +import org.infinispan.protostream.descriptors.MapDescriptor; import org.infinispan.protostream.descriptors.Type; import org.infinispan.protostream.descriptors.WireType; import org.infinispan.protostream.impl.TagReaderImpl; @@ -64,12 +65,19 @@ private void parseMessage(TagHandler tagHandler, Descriptor messageDescriptor, T final int fieldNumber = WireType.getTagFieldNumber(tag); final WireType wireType = WireType.fromTag(tag); final FieldDescriptor fd = messageDescriptor != null ? messageDescriptor.findFieldByNumber(fieldNumber) : null; - switch (wireType) { case LENGTH_DELIMITED: { if (fd == null) { byte[] value = in.readByteArray(); tagHandler.onTag(fieldNumber, null, value); + } else if (fd instanceof MapDescriptor md) { + int length = in.readUInt32(); + int oldLimit = in.pushLimit(length); + tagHandler.onStartNested(fieldNumber, fd); + parseMessage(tagHandler, md.asDescriptor(), in); + tagHandler.onEndNested(fieldNumber, fd); + in.checkLastTagWas(0); + in.popLimit(oldLimit); } else if (fd.getType() == Type.STRING) { String value = in.readString(); tagHandler.onTag(fieldNumber, fd, value); diff --git a/core/src/main/java/org/infinispan/protostream/descriptors/FieldDescriptor.java b/core/src/main/java/org/infinispan/protostream/descriptors/FieldDescriptor.java index 69aa4a3c9..71ff10c45 100644 --- a/core/src/main/java/org/infinispan/protostream/descriptors/FieldDescriptor.java +++ b/core/src/main/java/org/infinispan/protostream/descriptors/FieldDescriptor.java @@ -77,6 +77,10 @@ public boolean isRepeated() { return label == Label.REPEATED; } + public boolean isMap() { + return false; + } + public JavaType getJavaType() { return getType().getJavaType(); } diff --git a/core/src/main/java/org/infinispan/protostream/descriptors/MapDescriptor.java b/core/src/main/java/org/infinispan/protostream/descriptors/MapDescriptor.java index 62c3acb99..ce7d1eca1 100644 --- a/core/src/main/java/org/infinispan/protostream/descriptors/MapDescriptor.java +++ b/core/src/main/java/org/infinispan/protostream/descriptors/MapDescriptor.java @@ -3,11 +3,22 @@ public class MapDescriptor extends FieldDescriptor { private final String keyTypeName; private final Type keyType; + private final Descriptor descriptor; private MapDescriptor(Builder builder) { super(builder); keyTypeName = builder.keyTypeName; keyType = Type.primitiveFromString(keyTypeName); + Descriptor.Builder b = new Descriptor.Builder().withName(name).withFullName(fullName); + FieldDescriptor.Builder kb = new FieldDescriptor.Builder().withNumber(1).withName("key").withTypeName(keyTypeName); + b.addField(kb); + FieldDescriptor.Builder vb = new FieldDescriptor.Builder().withNumber(2).withName("value").withTypeName(typeName); + b.addField(vb); + descriptor = b.build(); + } + + public Descriptor asDescriptor() { + return descriptor; } @Override @@ -15,6 +26,11 @@ public boolean isRepeated() { return true; } + @Override + public boolean isMap() { + return true; + } + public Type getKeyType() { return keyType; } @@ -44,6 +60,12 @@ public Label getLabel() { return Label.OPTIONAL; } + @Override + void setMessageType(Descriptor descriptor) { + super.setMessageType(descriptor); + this.descriptor.getFields().get(1).setMessageType(descriptor); + } + @Override public String toString() { return "MapDescriptor{" + diff --git a/core/src/main/java/org/infinispan/protostream/impl/JsonUtils.java b/core/src/main/java/org/infinispan/protostream/impl/JsonUtils.java index f19583a2e..b7155dfbf 100644 --- a/core/src/main/java/org/infinispan/protostream/impl/JsonUtils.java +++ b/core/src/main/java/org/infinispan/protostream/impl/JsonUtils.java @@ -44,6 +44,7 @@ import org.infinispan.protostream.descriptors.FieldDescriptor; import org.infinispan.protostream.descriptors.GenericDescriptor; import org.infinispan.protostream.descriptors.Label; +import org.infinispan.protostream.descriptors.MapDescriptor; import org.infinispan.protostream.descriptors.Type; import com.fasterxml.jackson.core.JsonFactory; @@ -230,11 +231,15 @@ private static void processObject(ImmutableSerializationContext ctx, JsonParser break; case START_OBJECT: { FieldDescriptor fd = messageDescriptor.findFieldByName(currentField); - Descriptor messageType = fd.getMessageType(); - if (messageType == null) { - throw new IllegalStateException("Field '" + currentField + "' is not an object"); + if (fd.isMap()) { + processMap(ctx, (MapDescriptor) fd, parser, nestedWriter); + } else { + Descriptor messageType = fd.getMessageType(); + if (messageType == null) { + throw new IllegalStateException("Field '" + currentField + "' is not an object"); + } + processObject(ctx, parser, nestedWriter, messageType, fd.getNumber(), false); } - processObject(ctx, parser, nestedWriter, messageType, fd.getNumber(), false); break; } case FIELD_NAME: @@ -280,6 +285,76 @@ private static void processObject(ImmutableSerializationContext ctx, JsonParser writer.flush(); } + private static void processMap(ImmutableSerializationContext ctx, MapDescriptor md, JsonParser parser, TagWriter writer) throws IOException { + while (true) { + JsonToken token = parser.nextToken(); + if (token == JsonToken.END_OBJECT) { + break; + } + if (token != JsonToken.FIELD_NAME) { + throw new IllegalStateException("Unexpected token"); + } + ByteArrayOutputStream baos = new ByteArrayOutputStream(ProtobufUtil.DEFAULT_ARRAY_BUFFER_SIZE); + TagWriter nestedWriter = TagWriterImpl.newInstanceNoBuffer(ctx, baos); + String key = parser.getCurrentName(); + switch (md.getKeyType()) { + case STRING -> nestedWriter.writeString(1, key); + case INT32 -> nestedWriter.writeInt32(1, Integer.parseInt(key)); + case INT64 -> nestedWriter.writeInt64(1, Long.parseLong(key)); + case FIXED32 -> nestedWriter.writeFixed32(1, Integer.parseInt(key)); + case FIXED64 -> nestedWriter.writeFixed64(1, Long.parseLong(key)); + case SINT32 -> nestedWriter.writeSInt32(1, Integer.parseInt(key)); + case SINT64 -> nestedWriter.writeSInt64(1, Long.parseLong(key)); + case SFIXED32 -> nestedWriter.writeSFixed32(1, Integer.parseInt(key)); + case SFIXED64 -> nestedWriter.writeSFixed64(1, Long.parseLong(key)); + case UINT32 -> nestedWriter.writeUInt32(1, Integer.parseInt(key)); + case UINT64 -> nestedWriter.writeUInt64(1, Long.parseLong(key)); + } + processMapValue(ctx, parser, nestedWriter, md); + writer.writeBytes(md.getNumber(), baos.toByteArray()); + writer.flush(); + } + + } + + private static void processMapValue(ImmutableSerializationContext ctx, JsonParser parser, TagWriter writer, MapDescriptor md) throws IOException { + JsonToken token = parser.nextToken(); + if (token == null) { + return; + } + switch (token) { + case START_OBJECT: { + processObject(ctx, parser, writer, md.getMessageType(),2, false); + break; + } + case VALUE_STRING: + writer.writeString(2, parser.getValueAsString()); + break; + case VALUE_NUMBER_INT: + switch (md.getType()) { + case INT32 -> writer.writeInt32(2, parser.getIntValue()); + case INT64 -> writer.writeInt64(2, parser.getLongValue()); + case FIXED32 -> writer.writeFixed32(2, parser.getIntValue()); + case FIXED64 -> writer.writeFixed64(2, parser.getLongValue()); + case SINT32 -> writer.writeSInt32(2, parser.getIntValue()); + case SINT64 -> writer.writeSInt64(2, parser.getLongValue()); + case SFIXED32 -> writer.writeSFixed32(2, parser.getIntValue()); + case SFIXED64 -> writer.writeSFixed64(2, parser.getLongValue()); + case UINT32 -> writer.writeUInt32(2, parser.getIntValue()); + case UINT64 -> writer.writeUInt64(2, parser.getLongValue()); + } + break; + case VALUE_NUMBER_FLOAT: + switch (md.getType()) { + case FLOAT -> writer.writeFloat(2, parser.getFloatValue()); + case DOUBLE -> writer.writeDouble(2, parser.getDoubleValue()); + } + break; + } + writer.flush(); + } + + private static void processPrimitive(JsonParser parser, TagWriter writer, Type fieldType) throws IOException { while (true) { JsonToken token = parser.nextToken(); @@ -546,15 +621,14 @@ public void onStartNested(int fieldNumber, FieldDescriptor fieldDescriptor) { return; } startSlot(fieldDescriptor); - nestingLevel = new JsonNestingLevel(nestingLevel); - if (prettyPrint) { indent(); nestingLevel.indent++; } - - jsonOut.append('{'); + if (!fieldDescriptor.isMap()) { + jsonOut.append('{'); + } } @Override @@ -566,7 +640,9 @@ public void onEndNested(int fieldNumber, FieldDescriptor fieldDescriptor) { nestingLevel.indent--; indent(); } - jsonOut.append('}'); + if (!fieldDescriptor.isMap()) { + jsonOut.append('}'); + } nestingLevel = nestingLevel.previous; } @@ -590,37 +666,40 @@ private void startSlot(FieldDescriptor fieldDescriptor) { if (nestingLevel.repeatedFieldDescriptor != null && nestingLevel.repeatedFieldDescriptor != fieldDescriptor) { endArraySlot(); } - + boolean map = nestingLevel.previous != null && nestingLevel.previous.repeatedFieldDescriptor != null && nestingLevel.previous.repeatedFieldDescriptor.isMap(); if (nestingLevel.isFirstField) { nestingLevel.isFirstField = false; } else { - jsonOut.append(','); + jsonOut.append(map ? ':' : ','); } - if (!fieldDescriptor.isRepeated() || nestingLevel.repeatedFieldDescriptor == null) { - if (prettyPrint) { - indent(); + if (!map) { + if (!fieldDescriptor.isRepeated() || nestingLevel.repeatedFieldDescriptor == null) { + if (prettyPrint) { + indent(); + } + if (fieldDescriptor.getLabel() == Label.ONE_OF) { + jsonOut.append('"').append(JSON_VALUE_FIELD).append("\":"); + } else { + jsonOut.append('"').append(fieldDescriptor.getName()).append("\":"); + } } - if (fieldDescriptor.getLabel() == Label.ONE_OF) { - jsonOut.append('"').append(JSON_VALUE_FIELD).append("\":"); - } else { - jsonOut.append('"').append(fieldDescriptor.getName()).append("\":"); + if (prettyPrint) { + jsonOut.append(' '); } } - if (prettyPrint) { - jsonOut.append(' '); - } if (fieldDescriptor.isRepeated() && nestingLevel.repeatedFieldDescriptor == null) { nestingLevel.repeatedFieldDescriptor = fieldDescriptor; - jsonOut.append('['); + jsonOut.append(fieldDescriptor.isMap() ? '{' : '['); } } private void endArraySlot() { + boolean map = nestingLevel.repeatedFieldDescriptor.isMap(); if (prettyPrint && nestingLevel.repeatedFieldDescriptor.getType() == Type.MESSAGE) { indent(); } nestingLevel.repeatedFieldDescriptor = null; - jsonOut.append(']'); + jsonOut.append(map ? '}' : ']'); } }; @@ -773,6 +852,7 @@ private static void expectField(String expectedFieldName, String actualFieldName } // todo [anistor] do we really need html escaping? so far I'm keeping it so we behave like previous implementation + /** * Escapes a string literal in order to have a valid JSON representation. Optionally it can also escape some html chars. */ diff --git a/integrationtests/pom.xml b/integrationtests/pom.xml index 488261a44..e5c362d99 100644 --- a/integrationtests/pom.xml +++ b/integrationtests/pom.xml @@ -63,6 +63,12 @@ test + + com.fasterxml.jackson.core + jackson-databind + test + + com.google.testing.compile compile-testing diff --git a/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/GeneratedMarshallerTest.java b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/GeneratedMarshallerTest.java index 002479024..bc6f7a7cc 100644 --- a/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/GeneratedMarshallerTest.java +++ b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/GeneratedMarshallerTest.java @@ -1,11 +1,14 @@ package org.infinispan.protostream.integrationtests.processor.marshaller; import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import java.io.IOException; +import java.io.StringReader; import java.util.Collections; import java.util.Map; import java.util.UUID; @@ -16,9 +19,13 @@ import org.infinispan.protostream.annotations.ProtoSchema; import org.infinispan.protostream.annotations.ProtoSyntax; import org.infinispan.protostream.exception.ProtoStreamException; +import org.infinispan.protostream.impl.JsonUtils; import org.infinispan.protostream.integrationtests.processor.marshaller.model.FootballSchema; import org.infinispan.protostream.integrationtests.processor.marshaller.model.FootballSchemaImpl; import org.infinispan.protostream.integrationtests.processor.marshaller.model.FootballTeam; +import org.infinispan.protostream.integrationtests.processor.marshaller.model.MapOfMapOfUUID; +import org.infinispan.protostream.integrationtests.processor.marshaller.model.MapOfString; +import org.infinispan.protostream.integrationtests.processor.marshaller.model.MapOfUUID; import org.infinispan.protostream.integrationtests.processor.marshaller.model.MapSchema; import org.infinispan.protostream.integrationtests.processor.marshaller.model.ModelWithMap; import org.infinispan.protostream.integrationtests.processor.marshaller.model.NullTestModel; @@ -26,7 +33,10 @@ import org.infinispan.protostream.integrationtests.processor.marshaller.model.SimpleEnum; import org.junit.Test; +import com.fasterxml.jackson.databind.ObjectMapper; + public class GeneratedMarshallerTest { + static final UUID uuid = UUID.fromString("00000000-0000-0000-0000-000000000000"); @Test public void testGenericMessage() { @@ -142,6 +152,95 @@ interface NoNullsSchema extends GeneratedSchema { schemaFileName = "allownulls.proto", syntax = ProtoSyntax.PROTO3 ) - interface NullsAllowedSchema extends GeneratedSchema { + interface NullsAllowedSchema extends GeneratedSchema {} + + + @Test + public void testMapOfStringToJson() throws IOException { + var ctx = ProtobufUtil.newSerializationContext(); + MapSchema.INSTANCE.registerSchema(ctx); + MapSchema.INSTANCE.registerMarshallers(ctx); + + var m = new MapOfString(); + m.data = Map.of("1", "2", "3", "4"); + + var bytes = ProtobufUtil.toWrappedByteArray(ctx, m); + String json = JsonUtils.toCanonicalJSON(ctx, bytes, false); + assertJson( + """ + { + "_type":"generic.MapOfString", + "data":{"1":"2","3":"4"} + } + """, + json + ); + byte[] bytes2 = JsonUtils.fromCanonicalJSON(ctx, new StringReader(json)); + assertArrayEquals(bytes, bytes2); + } + + @Test + public void testMapOfUUIDToJson() throws IOException { + var ctx = ProtobufUtil.newSerializationContext(); + MapSchema.INSTANCE.registerSchema(ctx); + MapSchema.INSTANCE.registerMarshallers(ctx); + + var m = new MapOfUUID(); + m.data = Map.of("1", uuid, "3", uuid); + + var bytes = ProtobufUtil.toWrappedByteArray(ctx, m); + assertJson(""" + { + "_type":"generic.MapOfUUID", + "data":{ + "1":{"mostSigBitsFixed":0,"leastSigBitsFixed":0}, + "3":{"mostSigBitsFixed":0,"leastSigBitsFixed":0} + } + } + """, + JsonUtils.toCanonicalJSON(ctx, bytes, false)); + } + + @Test + public void testMapOfMapOfUUIDToJson() throws IOException { + var ctx = ProtobufUtil.newSerializationContext(); + MapSchema.INSTANCE.registerSchema(ctx); + MapSchema.INSTANCE.registerMarshallers(ctx); + + var m = new MapOfUUID(); + m.data = Map.of("1", uuid, "3", uuid); + var m2 = new MapOfMapOfUUID(); + m2.data1 = Map.of("1", "2", "3", "4"); + m2.data2 = m; + m2.data3 = SimpleEnum.SECOND; + + var bytes = ProtobufUtil.toWrappedByteArray(ctx, m2); + MapOfMapOfUUID copy = ProtobufUtil.fromWrappedByteArray(ctx, bytes); + assertNotNull(copy); + String json = JsonUtils.toCanonicalJSON(ctx, bytes, false); + assertJson(""" + { + "_type":"generic.MapOfMapOfUUID", + "data1":{ + "3":"4", + "1":"2" + }, + "data2":{ + "data":{ + "3":{"mostSigBitsFixed":0,"leastSigBitsFixed":0}, + "1":{"mostSigBitsFixed":0,"leastSigBitsFixed":0} + } + }, + "data3":"SECOND" + } + """, json); + + byte[] bytes2 = JsonUtils.fromCanonicalJSON(ctx, new StringReader(json)); + assertArrayEquals(bytes, bytes2); + } + + private static void assertJson(String j1, String j2) throws IOException { + ObjectMapper mapper = new ObjectMapper(); + assertEquals(mapper.readTree(j1), mapper.readTree(j2)); } } diff --git a/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfMapOfUUID.java b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfMapOfUUID.java new file mode 100644 index 000000000..e0245875d --- /dev/null +++ b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfMapOfUUID.java @@ -0,0 +1,19 @@ +package org.infinispan.protostream.integrationtests.processor.marshaller.model; + +import java.util.HashMap; +import java.util.Map; + +import org.infinispan.protostream.annotations.ProtoField; + +public class MapOfMapOfUUID { + + @ProtoField(value = 1, mapImplementation = HashMap.class) + public Map data1; + + @ProtoField(2) + public MapOfUUID data2; + + @ProtoField(3) + public SimpleEnum data3; + +} \ No newline at end of file diff --git a/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfString.java b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfString.java new file mode 100644 index 000000000..5fb800a20 --- /dev/null +++ b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfString.java @@ -0,0 +1,14 @@ + +package org.infinispan.protostream.integrationtests.processor.marshaller.model; + +import java.util.HashMap; +import java.util.Map; + +import org.infinispan.protostream.annotations.ProtoField; + +public class MapOfString { + + @ProtoField(value = 1, mapImplementation = HashMap.class) + public Map data; + +} \ No newline at end of file diff --git a/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfUUID.java b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfUUID.java new file mode 100644 index 000000000..e42456052 --- /dev/null +++ b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapOfUUID.java @@ -0,0 +1,14 @@ +package org.infinispan.protostream.integrationtests.processor.marshaller.model; + +import java.util.HashMap; +import java.util.Map; +import java.util.UUID; + +import org.infinispan.protostream.annotations.ProtoField; + +public class MapOfUUID { + + @ProtoField(value = 1, mapImplementation = HashMap.class) + public Map data; + +} diff --git a/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapSchema.java b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapSchema.java index b8ef6f6ef..e97d7f58f 100644 --- a/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapSchema.java +++ b/integrationtests/src/test/java/org/infinispan/protostream/integrationtests/processor/marshaller/model/MapSchema.java @@ -6,17 +6,19 @@ import org.infinispan.protostream.types.java.CommonTypes; @ProtoSchema( - syntax = ProtoSyntax.PROTO3, - schemaPackageName = "generic", - schemaFilePath = "proto", - schemaFileName = "generic.proto", - dependsOn = CommonTypes.class, - includeClasses = { - ModelWithMap.class, - SimpleEnum.class - } + syntax = ProtoSyntax.PROTO3, + schemaPackageName = "generic", + schemaFilePath = "proto", + schemaFileName = "generic.proto", + dependsOn = CommonTypes.class, + includeClasses = { + ModelWithMap.class, + SimpleEnum.class, + MapOfString.class, + MapOfUUID.class, + MapOfMapOfUUID.class + } ) public interface MapSchema extends GeneratedSchema { - - MapSchema INSTANCE = new MapSchemaImpl(); + MapSchema INSTANCE = new MapSchemaImpl(); }