Extend EventData with AmqpAnnotatedMessage

sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventData.java

@@ -3,8 +3,12 @@
 
 package com.azure.messaging.eventhubs;
 
+import com.azure.core.amqp.models.AmqpAnnotatedMessage;
+import com.azure.core.amqp.models.AmqpMessageBody;
+import com.azure.core.amqp.models.AmqpMessageId;
 import com.azure.core.util.BinaryData;
 import com.azure.core.util.Context;
+import com.azure.core.util.logging.ClientLogger;
 
 import java.nio.ByteBuffer;
 import java.time.Instant;
@@ -54,8 +58,8 @@ public class EventData {
     static final Set<String> RESERVED_SYSTEM_PROPERTIES;
 
     private final Map<String, Object> properties;
-    private final BinaryData body;
     private final SystemProperties systemProperties;
+    private final AmqpAnnotatedMessage annotatedMessage;
     private Context context;
 
     static {
@@ -73,16 +77,24 @@ public class EventData {
      * Creates an event containing the {@code body}.
      *
      * @param body The data to set for this event.
+     *
      * @throws NullPointerException if {@code body} is {@code null}.
      */
     public EventData(byte[] body) {
-        this(BinaryData.fromBytes(Objects.requireNonNull(body, "'body' cannot be null.")));
+        this.systemProperties = new SystemProperties();
+        this.context = Context.NONE;
+
+        final AmqpMessageBody messageBody = AmqpMessageBody.fromData(
+            Objects.requireNonNull(body, "'body' cannot be null."));
+        this.annotatedMessage = new AmqpAnnotatedMessage(messageBody);
+        this.properties = annotatedMessage.getApplicationProperties();
     }
 
     /**
      * Creates an event containing the {@code body}.
      *
      * @param body The data to set for this event.
+     *
      * @throws NullPointerException if {@code body} is {@code null}.
      */
     public EventData(ByteBuffer body) {
@@ -93,6 +105,7 @@ public EventData(ByteBuffer body) {
      * Creates an event by encoding the {@code body} using UTF-8 charset.
      *
      * @param body The string that will be UTF-8 encoded to create an event.
+     *
      * @throws NullPointerException if {@code body} is {@code null}.
      */
     public EventData(String body) {
@@ -105,22 +118,38 @@ public EventData(String body) {
      * @param body The {@link BinaryData} payload for this event.
      */
     public EventData(BinaryData body) {
-        this(body, new SystemProperties(), Context.NONE);
+        this(Objects.requireNonNull(body, "'body' cannot be null.").toBytes());
     }
 
     /**
      * Creates an event with the given {@code body}, system properties and context.
      *
-     * @param body The data to set for this event.
      * @param systemProperties System properties set by message broker for this event.
      * @param context A specified key-value pair of type {@link Context}.
-     * @throws NullPointerException if {@code body}, {@code systemProperties}, or {@code context} is {@code null}.
+     * @param amqpAnnotatedMessage Backing annotated message.
+     *
+     * @throws NullPointerException if {@code amqpAnnotatedMessage}, {@code systemProperties}, or {@code context} is
+     *     {@code null}.
+     * @throws IllegalArgumentException if {@code amqpAnnotatedMessage}'s body type is unknown.
      */
-    EventData(BinaryData body, SystemProperties systemProperties, Context context) {
-        this.body = Objects.requireNonNull(body, "'body' cannot be null.");
+    EventData(AmqpAnnotatedMessage amqpAnnotatedMessage, SystemProperties systemProperties, Context context) {
         this.context = Objects.requireNonNull(context, "'context' cannot be null.");
-        this.systemProperties =  Objects.requireNonNull(systemProperties, "'systemProperties' cannot be null.");
-        this.properties = new HashMap<>();
+        this.systemProperties = Objects.requireNonNull(systemProperties, "'systemProperties' cannot be null.");
+        this.properties = Collections.unmodifiableMap(amqpAnnotatedMessage.getApplicationProperties());
+        this.annotatedMessage = Objects.requireNonNull(amqpAnnotatedMessage, "'amqpAnnotatedMessage' cannot be null.");
+
+        switch (annotatedMessage.getBody().getBodyType()) {
+            case DATA:
+                break;
+            case SEQUENCE:
+            case VALUE:
+                new ClientLogger(EventData.class).warning("Message body type '{}' is not supported in EH. "
+                    + " Getting contents of body may throw.", annotatedMessage.getBody().getBodyType());
+                break;
+            default:
+                throw new ClientLogger(EventData.class).logExceptionAsError(new IllegalArgumentException(
+                    "Body type not valid " + annotatedMessage.getBody().getBodyType()));
+        }
     }
 
     /**
@@ -133,7 +162,8 @@ public EventData(BinaryData body) {
      *
      * {@codesnippet com.azure.messaging.eventhubs.eventdata.getProperties}
      *
-     * @return Application properties associated with this {@link EventData}.
+     * @return Application properties associated with this {@link EventData}. For received {@link EventData}, the map is
+     *     a read-only view.
      */
     public Map<String, Object> getProperties() {
         return properties;
@@ -143,8 +173,8 @@ public Map<String, Object> getProperties() {
      * Properties that are populated by Event Hubs service. As these are populated by the Event Hubs service, they are
      * only present on a <b>received</b> {@link EventData}.
      *
-     * @return An encapsulation of all system properties appended by EventHubs service into {@link EventData}.
-     *     {@code null} if the {@link EventData} is not received from the Event Hubs service.
+     * @return An encapsulation of all system properties appended by EventHubs service into {@link EventData}. {@code
+     *     null} if the {@link EventData} is not received from the Event Hubs service.
      */
     public Map<String, Object> getSystemProperties() {
         return systemProperties;
@@ -162,7 +192,7 @@ public Map<String, Object> getSystemProperties() {
      * @return A byte array representing the data.
      */
     public byte[] getBody() {
-        return body.toBytes();
+        return annotatedMessage.getBody().getFirstData();
     }
 
     /**
@@ -171,7 +201,7 @@ public byte[] getBody() {
      * @return UTF-8 decoded string representation of the event data.
      */
     public String getBodyAsString() {
-        return new String(body.toBytes(), UTF_8);
+        return new String(annotatedMessage.getBody().getFirstData(), UTF_8);
     }
 
     /**
@@ -180,7 +210,7 @@ public String getBodyAsString() {
      * @return the {@link BinaryData} payload associated with this event.
      */
     public BinaryData getBodyAsBinaryData() {
-        return body;
+        return BinaryData.fromBytes(annotatedMessage.getBody().getFirstData());
     }
 
     /**
@@ -219,8 +249,8 @@ public Instant getEnqueuedTime() {
 
     /**
      * Gets the sequence number assigned to the event when it was enqueued in the associated Event Hub partition. This
-     * is unique for every message received in the Event Hub partition. This is only present on a <b>received</b>
-     * {@link EventData}.
+     * is unique for every message received in the Event Hub partition. This is only present on a <b>received</b> {@link
+     * EventData}.
      *
      * @return The sequence number for this event. {@code null} if the {@link EventData} was not received from Event
      *     Hubs service.
@@ -229,6 +259,84 @@ public Long getSequenceNumber() {
         return systemProperties.getSequenceNumber();
     }
 
+    /**
+     * Gets the underlying AMQP message.
+     *
+     * @return The underlying AMQP message.
+     */
+    public AmqpAnnotatedMessage getRawAmqpMessage() {
+        return annotatedMessage;
+    }
+
+    /**
+     * Gets the content type.
+     *
+     * @return The content type.
+     */
+    public String getContentType() {
+        return annotatedMessage.getProperties().getContentType();
+    }
+
+    /**
+     * Sets the content type.
+     *
+     * @param contentType The content type.
+     *
+     * @return The updated {@link EventData}.
+     */
+    public EventData setContentType(String contentType) {
+        annotatedMessage.getProperties().setContentType(contentType);
+        return this;
+    }
+
+    /**
+     * Gets the correlation id.
+     *
+     * @return The correlation id. {@code null} if there is none set.
+     */
+    public String getCorrelationId() {
+        final AmqpMessageId messageId = annotatedMessage.getProperties().getCorrelationId();
+        return messageId != null ? messageId.toString() : null;
+    }
+
+    /**
+     * Sets the correlation id.
+     *
+     * @param correlationId The correlation id.
+     *
+     * @return The updated {@link EventData}.
+     */
+    public EventData setCorrelationId(String correlationId) {
+        final AmqpMessageId id = correlationId != null ? new AmqpMessageId(correlationId) : null;
+
+        annotatedMessage.getProperties().setCorrelationId(id);
+        return this;
+    }
+
+    /**
+     * Gets the message id.
+     *
+     * @return The message id. {@code null} if there is none set.
+     */
+    public String getMessageId() {
+        final AmqpMessageId messageId = annotatedMessage.getProperties().getMessageId();
+        return messageId != null ? messageId.toString() : null;
+    }
+
+    /**
+     * Sets the message id.
+     *
+     * @param messageId The message id.
+     *
+     * @return The updated {@link EventData}.
+     */
+    public EventData setMessageId(String messageId) {
+        final AmqpMessageId id = messageId != null ? new AmqpMessageId(messageId) : null;
+
+        annotatedMessage.getProperties().setMessageId(id);
+        return this;
+    }
+
     /**
      * {@inheritDoc}
      */
@@ -242,15 +350,16 @@ public boolean equals(Object o) {
         }
 
         EventData eventData = (EventData) o;
-        return Arrays.equals(body.toBytes(), eventData.body.toBytes());
+        return Arrays.equals(annotatedMessage.getBody().getFirstData(),
+            eventData.annotatedMessage.getBody().getFirstData());
     }
 
     /**
      * {@inheritDoc}
      */
     @Override
     public int hashCode() {
-        return Arrays.hashCode(body.toBytes());
+        return Arrays.hashCode(annotatedMessage.getBody().getFirstData());
     }
 
     /**
@@ -267,8 +376,10 @@ Context getContext() {
      *
      * @param key The key for this context object
      * @param value The value for this context object.
-     * @throws NullPointerException if {@code key} or {@code value} is null.
+     *
      * @return The updated {@link EventData}.
+     *
+     * @throws NullPointerException if {@code key} or {@code value} is null.
      */
     public EventData addContext(String key, Object value) {
         Objects.requireNonNull(key, "The 'key' parameter cannot be null.");
@@ -300,27 +411,29 @@ static class SystemProperties extends HashMap<String, Object> {
             super(map);
             this.partitionKey = removeSystemProperty(PARTITION_KEY_ANNOTATION_NAME.getValue());
 
-            final String offset = removeSystemProperty(OFFSET_ANNOTATION_NAME.getValue());
+            final Long offset = removeSystemPropertyAsLong(OFFSET_ANNOTATION_NAME.getValue());
             if (offset == null) {
                 throw new IllegalStateException(String.format(Locale.US,
                     "offset: %s should always be in map.", OFFSET_ANNOTATION_NAME.getValue()));
             }
-            this.offset = Long.valueOf(offset);
+
+            this.offset = offset;
             put(OFFSET_ANNOTATION_NAME.getValue(), this.offset);
 
-            final Date enqueuedTimeValue = removeSystemProperty(ENQUEUED_TIME_UTC_ANNOTATION_NAME.getValue());
+            final Instant enqueuedTimeValue = removeSystemPropertyAsInstant(ENQUEUED_TIME_UTC_ANNOTATION_NAME.getValue());
             if (enqueuedTimeValue == null) {
                 throw new IllegalStateException(String.format(Locale.US,
                     "enqueuedTime: %s should always be in map.", ENQUEUED_TIME_UTC_ANNOTATION_NAME.getValue()));
             }
-            this.enqueuedTime = enqueuedTimeValue.toInstant();
+            this.enqueuedTime = enqueuedTimeValue;
             put(ENQUEUED_TIME_UTC_ANNOTATION_NAME.getValue(), this.enqueuedTime);
 
-            final Long sequenceNumber = removeSystemProperty(SEQUENCE_NUMBER_ANNOTATION_NAME.getValue());
+            final Long sequenceNumber = removeSystemPropertyAsLong(SEQUENCE_NUMBER_ANNOTATION_NAME.getValue());
             if (sequenceNumber == null) {
                 throw new IllegalStateException(String.format(Locale.US,
                     "sequenceNumber: %s should always be in map.", SEQUENCE_NUMBER_ANNOTATION_NAME.getValue()));
             }
+
             this.sequenceNumber = sequenceNumber;
             put(SEQUENCE_NUMBER_ANNOTATION_NAME.getValue(), this.sequenceNumber);
         }
@@ -358,8 +471,9 @@ private Instant getEnqueuedTime() {
          * Event Hub.
          *
          * @return Sequence number for this event.
-         * @throws IllegalStateException if {@link SystemProperties} does not contain the sequence number in a retrieved
-         * event.
+         *
+         * @throws IllegalStateException if {@link SystemProperties} does not contain the sequence number in a
+         *     retrieved event.
          */
         private Long getSequenceNumber() {
             return sequenceNumber;
@@ -373,5 +487,37 @@ private <T> T removeSystemProperty(final String key) {
 
             return null;
         }
+
+        private Long removeSystemPropertyAsLong(final String key) {
+            if (!this.containsKey(key)) {
+                return null;
+            }
+
+            final Object value = this.remove(key);
+            if (value instanceof String) {
+                return Long.valueOf((String) value);
+            } else if (value instanceof Long) {
+                return (Long) value;
+            } else {
+                throw new ClientLogger(EventData.class).logExceptionAsError(new IllegalStateException(
+                    String.format(Locale.US, "Key: %s Value %s is not of type String or Long.", key, value)));
+            }
+        }
+
+        private Instant removeSystemPropertyAsInstant(final String key) {
+            if (!this.containsKey(key)) {
+                return null;
+            }
+
+            final Object value = this.remove(key);
+            if (value instanceof Date) {
+                return ((Date) value).toInstant();
+            } else if (value instanceof Instant) {
+                return (Instant) value;
+            } else {
+                throw new ClientLogger(EventData.class).logExceptionAsError(new IllegalStateException(
+                    String.format(Locale.US, "Key: %s Value %s is not of type Date or Instant.", key, value)));
+            }
+        }
+    }
     }
-}