Append "data:" after line breaks for SSE JSON data fields

sdeleuze · sdeleuze · commit 2735cba4b353 · 2016-12-08T11:19:46.000+01:00
Issue: SPR-14899
diff --git a/spring-web/src/main/java/org/springframework/http/codec/ServerSentEventHttpMessageWriter.java b/spring-web/src/main/java/org/springframework/http/codec/ServerSentEventHttpMessageWriter.java
@@ -19,6 +19,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -48,6 +49,16 @@
  */
 public class ServerSentEventHttpMessageWriter implements HttpMessageWriter<Object> {
 
+	/**
+	 * Server-Sent Events hint expecting a {@link Boolean} value which when set to true
+	 * will adapt the content in order to comply with Server-Sent Events recommendation.
+	 * For example, it will append "data:" after each line break with data encoders
+	 * supporting it.
+	 * @see <a href="https://www.w3.org/TR/eventsource/">Server-Sent Events W3C recommendation</a>
+	 */
+	public static final String SSE_CONTENT_HINT = ServerSentEventHttpMessageWriter.class.getName() + ".sseContent";
+
+
 	private final List<Encoder<?>> dataEncoders;
 
 
@@ -87,6 +98,8 @@ public Mono<Void> write(Publisher<?> inputStream, ResolvableType elementType, Me
 	private Flux<Publisher<DataBuffer>> encode(Publisher<?> inputStream, DataBufferFactory bufferFactory,
 			ResolvableType type, Map<String, Object> hints) {
 
+		Map<String, Object> hintsWithSse = new HashMap<>(hints);
+		hintsWithSse.put(SSE_CONTENT_HINT, true);
 		return Flux.from(inputStream)
 				.map(o -> toSseEvent(o, type))
 				.map(sse -> {
@@ -107,7 +120,7 @@ private Flux<Publisher<DataBuffer>> encode(Publisher<?> inputStream, DataBufferF
 									return Flux.empty();
 								}
 								else {
-									return applyEncoder(data, bufferFactory, hints);
+									return applyEncoder(data, bufferFactory, hintsWithSse);
 								}
 							}).orElse(Flux.empty());
 
diff --git a/spring-web/src/main/java/org/springframework/http/codec/json/Jackson2JsonEncoder.java b/spring-web/src/main/java/org/springframework/http/codec/json/Jackson2JsonEncoder.java
@@ -22,9 +22,14 @@
 import java.util.List;
 import java.util.Map;
 
+import com.fasterxml.jackson.core.PrettyPrinter;
+import com.fasterxml.jackson.core.util.DefaultIndenter;
+import com.fasterxml.jackson.core.util.DefaultPrettyPrinter;
 import com.fasterxml.jackson.databind.JavaType;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.ObjectWriter;
+import com.fasterxml.jackson.databind.SerializationConfig;
+import com.fasterxml.jackson.databind.SerializationFeature;
 import com.fasterxml.jackson.databind.type.TypeFactory;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -35,6 +40,7 @@
 import org.springframework.core.codec.Encoder;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.http.codec.ServerSentEventHttpMessageWriter;
 import org.springframework.http.converter.json.Jackson2ObjectMapperBuilder;
 import org.springframework.util.Assert;
 import org.springframework.util.MimeType;
@@ -57,12 +63,18 @@ public class Jackson2JsonEncoder extends AbstractJackson2Codec implements Encode
 	private static final ByteBuffer END_ARRAY_BUFFER = ByteBuffer.wrap(new byte[]{']'});
 
 
+	private final PrettyPrinter ssePrettyPrinter;
+
+
 	public Jackson2JsonEncoder() {
-		super(Jackson2ObjectMapperBuilder.json().build());
+		this(Jackson2ObjectMapperBuilder.json().build());
 	}
 
 	public Jackson2JsonEncoder(ObjectMapper mapper) {
 		super(mapper);
+		DefaultPrettyPrinter prettyPrinter = new DefaultPrettyPrinter();
+		prettyPrinter.indentObjectsWith(new DefaultIndenter("  ", "\ndata:"));
+		this.ssePrettyPrinter = prettyPrinter;
 	}
 
 
@@ -123,6 +135,12 @@ private DataBuffer encodeValue(Object value, DataBufferFactory bufferFactory,
 			writer = writer.forType(javaType);
 		}
 
+		Boolean sse = (Boolean)hints.get(ServerSentEventHttpMessageWriter.SSE_CONTENT_HINT);
+		SerializationConfig config = writer.getConfig();
+		if (Boolean.TRUE.equals(sse) && config.isEnabled(SerializationFeature.INDENT_OUTPUT)) {
+			writer = writer.with(this.ssePrettyPrinter);
+		}
+
 		DataBuffer buffer = bufferFactory.allocateBuffer();
 		OutputStream outputStream = buffer.asOutputStream();
 		try {
diff --git a/spring-web/src/main/java/org/springframework/http/converter/json/AbstractJackson2HttpMessageConverter.java b/spring-web/src/main/java/org/springframework/http/converter/json/AbstractJackson2HttpMessageConverter.java
@@ -27,11 +27,14 @@
 import com.fasterxml.jackson.core.JsonEncoding;
 import com.fasterxml.jackson.core.JsonGenerator;
 import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.core.PrettyPrinter;
+import com.fasterxml.jackson.core.util.DefaultIndenter;
 import com.fasterxml.jackson.core.util.DefaultPrettyPrinter;
 import com.fasterxml.jackson.databind.JavaType;
 import com.fasterxml.jackson.databind.JsonMappingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.ObjectWriter;
+import com.fasterxml.jackson.databind.SerializationConfig;
 import com.fasterxml.jackson.databind.SerializationFeature;
 import com.fasterxml.jackson.databind.ser.FilterProvider;
 import com.fasterxml.jackson.databind.type.TypeFactory;
@@ -69,22 +72,29 @@ public abstract class AbstractJackson2HttpMessageConverter extends AbstractGener
 
 	private Boolean prettyPrint;
 
+	private PrettyPrinter ssePrettyPrinter;
+
 
 	protected AbstractJackson2HttpMessageConverter(ObjectMapper objectMapper) {
-		this.objectMapper = objectMapper;
-		setDefaultCharset(DEFAULT_CHARSET);
+		init(objectMapper);
 	}
 
 	protected AbstractJackson2HttpMessageConverter(ObjectMapper objectMapper, MediaType supportedMediaType) {
 		super(supportedMediaType);
-		this.objectMapper = objectMapper;
-		setDefaultCharset(DEFAULT_CHARSET);
+		init(objectMapper);
 	}
 
 	protected AbstractJackson2HttpMessageConverter(ObjectMapper objectMapper, MediaType... supportedMediaTypes) {
 		super(supportedMediaTypes);
+		init(objectMapper);
+	}
+
+	protected void init(ObjectMapper objectMapper) {
 		this.objectMapper = objectMapper;
 		setDefaultCharset(DEFAULT_CHARSET);
+		DefaultPrettyPrinter prettyPrinter = new DefaultPrettyPrinter();
+		prettyPrinter.indentObjectsWith(new DefaultIndenter("  ", "\ndata:"));
+		this.ssePrettyPrinter = prettyPrinter;
 	}
 
 
@@ -234,7 +244,8 @@ private Object readJavaType(JavaType javaType, HttpInputMessage inputMessage) {
 	protected void writeInternal(Object object, Type type, HttpOutputMessage outputMessage)
 			throws IOException, HttpMessageNotWritableException {
 
-		JsonEncoding encoding = getJsonEncoding(outputMessage.getHeaders().getContentType());
+		MediaType contentType = outputMessage.getHeaders().getContentType();
+		JsonEncoding encoding = getJsonEncoding(contentType);
 		JsonGenerator generator = this.objectMapper.getFactory().createGenerator(outputMessage.getBody(), encoding);
 		try {
 			writePrefix(generator, object);
@@ -265,6 +276,11 @@ else if (filters != null) {
 			if (javaType != null && javaType.isContainerType()) {
 				objectWriter = objectWriter.forType(javaType);
 			}
+			SerializationConfig config = objectWriter.getConfig();
+			if (contentType != null && contentType.isCompatibleWith(MediaType.TEXT_EVENT_STREAM) &&
+					config.isEnabled(SerializationFeature.INDENT_OUTPUT)) {
+				objectWriter = objectWriter.with(this.ssePrettyPrinter);
+			}
 			objectWriter.writeValue(generator, value);
 
 			writeSuffix(generator, object);
diff --git a/spring-web/src/test/java/org/springframework/http/codec/ServerSentEventHttpMessageWriterTests.java b/spring-web/src/test/java/org/springframework/http/codec/ServerSentEventHttpMessageWriterTests.java
@@ -20,6 +20,7 @@
 import java.util.Collections;
 import java.util.function.Consumer;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.Test;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -31,6 +32,7 @@
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.http.MediaType;
 import org.springframework.http.codec.json.Jackson2JsonEncoder;
+import org.springframework.http.converter.json.Jackson2ObjectMapperBuilder;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse;
 
 import static org.junit.Assert.*;
@@ -116,7 +118,7 @@ public void encodePojo() {
 				new Pojo("foofoofoo", "barbarbar"));
 		MockServerHttpResponse outputMessage = new MockServerHttpResponse();
 		messageWriter.write(source, ResolvableType.forClass(Pojo.class),
-				new MediaType("text", "event-stream"), outputMessage, Collections.emptyMap());
+				MediaType.TEXT_EVENT_STREAM, outputMessage, Collections.emptyMap());
 
 		Publisher<? extends Publisher<? extends DataBuffer>> result = outputMessage.getBodyWithFlush();
 		StepVerifier.create(result)
@@ -126,6 +128,29 @@ public void encodePojo() {
 				.verify();
 	}
 
+	@Test  // SPR-14899
+	public void encodePojoWithPrettyPrint() {
+		ObjectMapper mapper = Jackson2ObjectMapperBuilder.json().indentOutput(true).build();
+		this.messageWriter = new ServerSentEventHttpMessageWriter(Collections.singletonList(new Jackson2JsonEncoder(mapper)));
+
+		Flux<Pojo> source = Flux.just(new Pojo("foofoo", "barbar"),
+				new Pojo("foofoofoo", "barbarbar"));
+		MockServerHttpResponse outputMessage = new MockServerHttpResponse();
+		messageWriter.write(source, ResolvableType.forClass(Pojo.class),
+				MediaType.TEXT_EVENT_STREAM, outputMessage, Collections.emptyMap());
+
+		Publisher<? extends Publisher<? extends DataBuffer>> result = outputMessage.getBodyWithFlush();
+		StepVerifier.create(result)
+				.consumeNextWith(sseConsumer("data:", "{\n" +
+						"data:  \"foo\" : \"foofoo\",\n" +
+						"data:  \"bar\" : \"barbar\"\n" + "data:}", "\n"))
+				.consumeNextWith(sseConsumer("data:", "{\n" +
+						"data:  \"foo\" : \"foofoofoo\",\n" +
+						"data:  \"bar\" : \"barbarbar\"\n" + "data:}", "\n"))
+				.expectComplete()
+				.verify();
+	}
+
 
 	private Consumer<Publisher<? extends DataBuffer>> sseConsumer(String... expected) {
 		return publisher -> {
diff --git a/spring-web/src/test/java/org/springframework/http/converter/json/MappingJackson2HttpMessageConverterTests.java b/spring-web/src/test/java/org/springframework/http/converter/json/MappingJackson2HttpMessageConverterTests.java
@@ -223,6 +223,20 @@ public void prettyPrint() throws Exception {
 		assertEquals("{" + NEWLINE_SYSTEM_PROPERTY + "  \"name\" : \"Jason\"" + NEWLINE_SYSTEM_PROPERTY + "}", result);
 	}
 
+	@Test
+	public void prettyPrintWithSse() throws Exception {
+		MockHttpOutputMessage outputMessage = new MockHttpOutputMessage();
+		outputMessage.getHeaders().setContentType(MediaType.TEXT_EVENT_STREAM);
+		PrettyPrintBean bean = new PrettyPrintBean();
+		bean.setName("Jason");
+
+		this.converter.setPrettyPrint(true);
+		this.converter.writeInternal(bean, null, outputMessage);
+		String result = outputMessage.getBodyAsString(StandardCharsets.UTF_8);
+
+		assertEquals("{\ndata:  \"name\" : \"Jason\"\ndata:}", result);
+	}
+
 	@Test
 	public void prefixJson() throws Exception {
 		MockHttpOutputMessage outputMessage = new MockHttpOutputMessage();
