GsonMessageWriter.java
/*
* Copyright 2020 Providence Authors
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package net.morimekta.proto.gson.sio;
import com.google.gson.stream.JsonWriter;
import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import com.google.protobuf.Descriptors;
import com.google.protobuf.Duration;
import com.google.protobuf.Message;
import com.google.protobuf.ProtocolMessageEnum;
import com.google.protobuf.Timestamp;
import net.morimekta.proto.MorimektaOptions;
import net.morimekta.proto.ProtoMessage;
import net.morimekta.proto.gson.ProtoTypeOptions;
import net.morimekta.proto.sio.MessageWriter;
import net.morimekta.proto.utils.JsonNameUtil;
import net.morimekta.proto.utils.ValueUtil;
import java.io.IOException;
import java.io.Writer;
import java.time.format.DateTimeFormatter;
import java.util.Collection;
import java.util.Map;
import java.util.Optional;
import static java.util.Objects.requireNonNull;
import static net.morimekta.proto.ProtoMessage.getMessageClass;
import static net.morimekta.proto.gson.ProtoTypeOptions.Option.WRITE_COMPACT_MESSAGE;
import static net.morimekta.proto.gson.ProtoTypeOptions.Option.WRITE_DURATION_AS_STRING;
import static net.morimekta.proto.gson.ProtoTypeOptions.Option.WRITE_ENUM_AS_NUMBER;
import static net.morimekta.proto.gson.ProtoTypeOptions.Option.WRITE_FIELD_AS_NUMBER;
import static net.morimekta.proto.gson.ProtoTypeOptions.Option.WRITE_TIMESTAMP_AS_ISO;
import static net.morimekta.proto.gson.ProtoTypeOptions.Option.WRITE_UNPACKED_ANY;
import static net.morimekta.proto.gson.ProtoTypeOptions.Value.ANY_TYPE_FIELD_NAME;
import static net.morimekta.proto.gson.ProtoTypeOptions.Value.ANY_TYPE_PREFIX;
import static net.morimekta.proto.utils.ByteStringUtil.toBase64;
import static net.morimekta.proto.utils.FieldUtil.getMapValueDescriptor;
import static net.morimekta.proto.utils.GoogleTypesUtil.simpleDurationString;
import static net.morimekta.proto.utils.GoogleTypesUtil.toInstant;
import static net.morimekta.proto.utils.JsonNameUtil.getJsonEnumName;
import static net.morimekta.strings.Stringable.asString;
/**
* Message writer that uses a Gson for writing the messages to
* stream.
*/
public class GsonMessageWriter implements MessageWriter, AutoCloseable {
private final Writer rawWriter;
private final JsonWriter writer;
private final ProtoTypeOptions options;
/**
* @param writer Writer to write JSON to.
*/
public GsonMessageWriter(Writer writer) {
this(writer, false);
}
/**
* @param writer Writer to write JSON to.
* @param pretty If JSON should be pretty printed.
*/
public GsonMessageWriter(Writer writer, boolean pretty) {
this(writer, makeJsonWriter(writer, pretty), new ProtoTypeOptions());
}
/**
* @param writer Writer to write JSON to.
* @param options Options for the writer.
* @param pretty If JSON should be pretty printed.
*/
public GsonMessageWriter(Writer writer, ProtoTypeOptions options, boolean pretty) {
this(writer, makeJsonWriter(writer, pretty), options);
}
/**
* @param writer JSON Writer to write to.
* @param options Options for the writer.
*/
public GsonMessageWriter(JsonWriter writer, ProtoTypeOptions options) {
this(null, writer, options);
}
/**
* Create a message writer that writes JSON to a writer as
* a stream of messages.
*
* @param rawWriter The raw writer. If this is not null, it will write a newline after
* each message, visually separating them in a stream file of messages.
* @param writer The Gson Json writer.
* @param options Type options for writing GSON.
*/
public GsonMessageWriter(Writer rawWriter,
JsonWriter writer,
ProtoTypeOptions options) {
this.rawWriter = rawWriter;
this.writer = requireNonNull(writer, "writer == null");
this.options = requireNonNull(options, "options == null");
}
@Override
public void write(Message message) throws IOException {
requireNonNull(message, "message == null");
writeMessage(message);
writer.flush();
if (rawWriter != null) {
rawWriter.write("\n");
rawWriter.flush();
}
}
@Override
public void close() throws IOException {
writer.close();
}
private void writeMessage(Message message) throws IOException {
if (message.getDescriptorForType().equals(Any.getDescriptor())) {
if (options.isEnabled(WRITE_UNPACKED_ANY)) {
Any any = (Any) message;
var descriptor = options.getRegistry().messageTypeByTypeUrl(any.getTypeUrl());
if (descriptor != null) {
var type = getMessageClass(descriptor);
var unpacked = any.unpack(type);
writer.beginObject();
writer.name(options.getValue(ANY_TYPE_FIELD_NAME));
writer.value(options.getValue(ANY_TYPE_PREFIX) + descriptor.getFullName());
writeMessageFields(unpacked);
writer.endObject();
return;
}
}
} else if (options.isEnabled(WRITE_TIMESTAMP_AS_ISO) &&
message.getDescriptorForType().equals(Timestamp.getDescriptor())) {
writer.value(DateTimeFormatter.ISO_INSTANT.format(toInstant((Timestamp) message)));
return;
} else if (options.isEnabled(WRITE_DURATION_AS_STRING) &&
message.getDescriptorForType().equals(Duration.getDescriptor())) {
var durationS = simpleDurationString((Duration) message);
if (durationS != null) {
writer.value(durationS);
return;
}
} else if (options.isEnabled(WRITE_COMPACT_MESSAGE) &&
message.getDescriptorForType()
.getOptions()
.getExtension(MorimektaOptions.compact)) {
writer.beginArray();
int skipped = 0;
var msg = new ProtoMessage(message);
for (Descriptors.FieldDescriptor field : message.getDescriptorForType().getFields()) {
Optional<?> val = msg.optional(field);
if (val.isPresent()) {
while (skipped > 0) {
writer.nullValue();
--skipped;
}
writeValue(val.get(), field);
} else {
++skipped;
}
}
writer.endArray();
return;
}
writer.beginObject();
writeMessageFields(message);
writer.endObject();
}
private void writeMessageFields(Message message) throws IOException {
for (var field : message.getAllFields().entrySet()) {
var val = ValueUtil.toJavaValue(field.getKey(), field.getValue());
if (options.isEnabled(WRITE_FIELD_AS_NUMBER)) {
writer.name("" + field.getKey().getNumber());
} else {
writer.name(JsonNameUtil.getJsonFieldName(field.getKey()));
}
writeValue(val, field.getKey());
}
}
@SuppressWarnings("unchecked")
private void writeValue(Object value, Descriptors.FieldDescriptor field) throws IOException {
if (value instanceof Collection) {
writer.beginArray();
for (Object o : (Collection<?>) value) {
writeValue(o, field);
}
writer.endArray();
return;
} else if (value instanceof Map) {
var valueDescriptor = getMapValueDescriptor(field);
writer.beginObject();
for (Map.Entry<Object, Object> entry : ((Map<Object, Object>) value).entrySet()) {
writeMapKey(entry.getKey());
writeValue(entry.getValue(), valueDescriptor);
}
writer.endObject();
return;
}
switch (field.getType().getJavaType()) {
case BOOLEAN: {
writer.value((Boolean) value);
break;
}
case INT:
case LONG:
case FLOAT:
case DOUBLE: {
writer.value((Number) value);
break;
}
case STRING: {
writer.value((String) value);
break;
}
case BYTE_STRING: {
writer.value(toBase64((ByteString) value));
break;
}
case ENUM: {
if (options.isEnabled(WRITE_ENUM_AS_NUMBER)) {
writer.value(((ProtocolMessageEnum) value).getNumber());
} else {
writer.value(getJsonEnumName((ProtocolMessageEnum) value));
}
break;
}
case MESSAGE: {
writeMessage((Message) value);
break;
}
default: {
throw new IOException("Unhandled value type: " + field.getType());
}
}
}
private void writeMapKey(Object value) throws IOException {
if (value instanceof String) {
writer.name((String) value);
} else if (value instanceof Boolean) {
writer.name(value.toString());
} else if (value instanceof Number) {
writer.name(asString(value));
} else {
throw new IOException("Map value with incompatible with map key: " + value.getClass().getSimpleName());
}
}
private static JsonWriter makeJsonWriter(Writer writer, boolean pretty) {
var json = new JsonWriter(writer);
json.setSerializeNulls(true);
if (pretty) {
json.setIndent(" ");
}
return json;
}
}