Override ObjectMapper in the DefaultDataConverter's payload coverter

temporal-sdk/src/main/java/io/temporal/common/converter/DefaultDataConverter.java

@@ -26,6 +26,7 @@
 import io.temporal.api.common.v1.Payloads;
 import java.lang.reflect.Type;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -39,16 +40,16 @@
  */
 public class DefaultDataConverter implements DataConverter {
 
-  private static final AtomicReference<DataConverter> defaultDataConverterInstance =
-      new AtomicReference<>(
-          // Order is important as the first converter that can convert the payload is used
-          new DefaultDataConverter(
-              new NullPayloadConverter(),
-              new ByteArrayPayloadConverter(),
-              new ProtobufJsonPayloadConverter(),
-              new JacksonJsonPayloadConverter()));
+  // Order is important as the first converter that can convert the payload is used
+  private static final PayloadConverter[] DEFAULT_PAYLOAD_CONVERTERS = {
+    new NullPayloadConverter(),
+    new ByteArrayPayloadConverter(),
+    new ProtobufJsonPayloadConverter(),
+    new JacksonJsonPayloadConverter()
+  };
 
-  private static final Object[] EMPTY_OBJECT_ARRAY = new Object[0];
+  private static final AtomicReference<DataConverter> defaultDataConverterInstance =
+      new AtomicReference<>(newDefaultInstance());
 
   private final Map<String, PayloadConverter> converterMap = new ConcurrentHashMap<>();
 
@@ -70,16 +71,44 @@ public static void setDefaultDataConverter(DataConverter converter) {
     defaultDataConverterInstance.set(converter);
   }
 
+  /**
+   * Creates a new instance of {@code DefaultDataConverter} populated with the default list of
+   * payload converters.
+   */
+  public static DefaultDataConverter newDefaultInstance() {
+    return new DefaultDataConverter(DEFAULT_PAYLOAD_CONVERTERS);
+  }
+
   /**
    * Creates instance from ordered array of converters. When converting an object to payload the
-   * array of converters is iterated from the beginning until one of the converters succesfully
+   * array of converters is iterated from the beginning until one of the converters successfully
    * converts the value.
    */
   public DefaultDataConverter(PayloadConverter... converters) {
-    for (PayloadConverter converter : converters) {
-      this.converters.add(converter);
-      this.converterMap.put(converter.getEncodingType(), converter);
+    Collections.addAll(this.converters, converters);
+    updateConverterMap();
+  }
+
+  /**
+   * Modifies this {@code DefaultDataConverter} by overriding some of its {@link PayloadConverter}s.
+   * Every payload converter from {@code overrideConverters} either replaces existing payload
+   * converter with the same encoding type, or is added to the end of payload converters list.
+   */
+  public DefaultDataConverter withPayloadConverterOverrides(
+      PayloadConverter... overrideConverters) {
+    for (PayloadConverter overrideConverter : overrideConverters) {
+      PayloadConverter existingConverter = converterMap.get(overrideConverter.getEncodingType());
+      if (existingConverter != null) {
+        int existingConverterIndex = converters.indexOf(existingConverter);
+        converters.set(existingConverterIndex, overrideConverter);
+      } else {
+        converters.add(overrideConverter);
+      }
     }
+
+    updateConverterMap();
+
+    return this;
   }
 
   @Override
@@ -154,4 +183,11 @@ public <T> T fromPayloads(
     }
     return fromPayload(content.get().getPayloads(index), parameterType, genericParameterType);
   }
+
+  private void updateConverterMap() {
+    converterMap.clear();
+    for (PayloadConverter converter : converters) {
+      converterMap.put(converter.getEncodingType(), converter);
+    }
+  }
 }

temporal-sdk/src/test/java/io/temporal/common/converter/ProtoPayloadConverterTest.java

@@ -21,8 +21,14 @@
 
 import static org.junit.Assert.assertEquals;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.SerializationFeature;
+import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
+import com.google.protobuf.util.JsonFormat;
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.common.v1.WorkflowExecution;
+import java.time.Instant;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
 import org.junit.Test;
@@ -43,6 +49,21 @@ public void testProtoJson() {
     assertEquals(execution, converted);
   }
 
+  @Test
+  public void testCustomJson() {
+    ObjectMapper objectMapper =
+        new ObjectMapper()
+            .configure(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS, true)
+            .registerModule(new JavaTimeModule());
+    DataConverter converter =
+        DefaultDataConverter.newDefaultInstance()
+            .withPayloadConverterOverrides(new JacksonJsonPayloadConverter(objectMapper));
+    TestPayload payload = new TestPayload(1L, Instant.now(), "myPayload");
+    Optional<Payloads> data = converter.toPayloads(payload);
+    TestPayload converted = converter.fromPayloads(0, data, TestPayload.class, TestPayload.class);
+    assertEquals(payload, converted);
+  }
+
   @Test
   public void testProto() {
     DataConverter converter = new DefaultDataConverter(new ProtobufPayloadConverter());
@@ -56,4 +77,92 @@ public void testProto() {
         converter.fromPayloads(0, data, WorkflowExecution.class, WorkflowExecution.class);
     assertEquals(execution, converted);
   }
+
+  @Test
+  public void testCustomProto() {
+    DataConverter converter =
+        DefaultDataConverter.newDefaultInstance()
+            .withPayloadConverterOverrides(
+                new ProtobufJsonPayloadConverter(
+                    JsonFormat.printer().printingEnumsAsInts(), JsonFormat.parser()));
+    WorkflowExecution execution =
+        WorkflowExecution.newBuilder()
+            .setWorkflowId(UUID.randomUUID().toString())
+            .setRunId(UUID.randomUUID().toString())
+            .build();
+    Optional<Payloads> data = converter.toPayloads(execution);
+    WorkflowExecution converted =
+        converter.fromPayloads(0, data, WorkflowExecution.class, WorkflowExecution.class);
+    assertEquals(execution, converted);
+  }
+
+  static class TestPayload {
+    private long id;
+    private Instant timestamp;
+    private String name;
+
+    public TestPayload() {}
+
+    TestPayload(long id, Instant timestamp, String name) {
+      this.id = id;
+      this.timestamp = timestamp;
+      this.name = name;
+    }
+
+    public long getId() {
+      return id;
+    }
+
+    public void setId(long id) {
+      this.id = id;
+    }
+
+    public Instant getTimestamp() {
+      return timestamp;
+    }
+
+    public void setTimestamp(Instant timestamp) {
+      this.timestamp = timestamp;
+    }
+
+    public String getName() {
+      return name;
+    }
+
+    public void setName(String name) {
+      this.name = name;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      TestPayload that = (TestPayload) o;
+      return id == that.id
+          && Objects.equals(timestamp, that.timestamp)
+          && Objects.equals(name, that.name);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(id, timestamp, name);
+    }
+
+    @Override
+    public String toString() {
+      return "TestPayload{"
+          + "id="
+          + id
+          + ", timestamp="
+          + timestamp
+          + ", name='"
+          + name
+          + '\''
+          + '}';
+    }
+  }
 }