PubSubMessageSerDe for query pointcuts (#99)

Author: akshaisarma

pom.xml

@@ -4,7 +4,7 @@
 
     <groupId>com.yahoo.bullet</groupId>
     <artifactId>bullet-core</artifactId>
-    <version>1.3.3-SNAPSHOT</version>
+    <version>1.4.0-SNAPSHOT</version>
     <packaging>jar</packaging>
     <name>bullet-core</name>
     <description>

src/main/java/com/yahoo/bullet/common/BulletConfig.java

@@ -63,6 +63,7 @@ public class BulletConfig extends Config {
 
     public static final String PUBSUB_CONTEXT_NAME = "bullet.pubsub.context.name";
     public static final String PUBSUB_CLASS_NAME = "bullet.pubsub.class.name";
+    public static final String PUBSUB_MESSAGE_SERDE_CLASS_NAME = "bullet.pubsub.message.serde.class.name";
 
     public static final String STORAGE_CLASS_NAME = "bullet.storage.class.name";
 
@@ -142,6 +143,7 @@ public class BulletConfig extends Config {
 
     public static final String DEFAULT_PUBSUB_CONTEXT_NAME = Context.QUERY_PROCESSING.name();
     public static final String DEFAULT_PUBSUB_CLASS_NAME = "com.yahoo.bullet.pubsub.MockPubSub";
+    public static final String DEFAULT_PUBSUB_MESSAGE_SERDE_CLASS_NAME = "com.yahoo.bullet.pubsub.IdentityPubSubMessageSerDe";
 
     public static final String DEFAULT_RECORD_PROVIDER_CLASS_NAME = "com.yahoo.bullet.record.avro.TypedAvroBulletRecordProvider";
 
@@ -278,6 +280,9 @@ public class BulletConfig extends Config {
         VALIDATOR.define(PUBSUB_CLASS_NAME)
                  .defaultTo(DEFAULT_PUBSUB_CLASS_NAME)
                  .checkIf(Validator::isClassName);
+        VALIDATOR.define(PUBSUB_MESSAGE_SERDE_CLASS_NAME)
+                 .defaultTo(DEFAULT_PUBSUB_MESSAGE_SERDE_CLASS_NAME)
+                 .checkIf(Validator::isClassName);
 
         VALIDATOR.define(RECORD_PROVIDER_CLASS_NAME)
                  .defaultTo(DEFAULT_RECORD_PROVIDER_CLASS_NAME)

src/main/java/com/yahoo/bullet/pubsub/IdentityPubSubMessageSerDe.java

@@ -0,0 +1,31 @@
+/*
+ *  Copyright 2021 Yahoo Inc.
+ *  Licensed under the terms of the Apache License, Version 2.0.
+ *  See the LICENSE file associated with the project for terms.
+ */
+package com.yahoo.bullet.pubsub;
+
+import com.yahoo.bullet.common.BulletConfig;
+
+public class IdentityPubSubMessageSerDe extends PubSubMessageSerDe {
+    private static final long serialVersionUID = -1709000962888195381L;
+
+    /**
+     * Constructor.
+     *
+     * @param config The {@link BulletConfig} to configure this class.
+     */
+    public IdentityPubSubMessageSerDe(BulletConfig config) {
+        super(config);
+    }
+
+    @Override
+    public PubSubMessage toMessage(PubSubMessage message) {
+        return message;
+    }
+
+    @Override
+    public PubSubMessage fromMessage(PubSubMessage message) {
+        return message;
+    }
+}

src/main/java/com/yahoo/bullet/pubsub/Metadata.java

@@ -58,7 +58,7 @@ public Metadata copy() {
     }
 
     /**
-     * Set a serializable content for this metadata.
+     * Set a {@link Serializable} content for this metadata.
      *
      * @param content The content for this metadata.
      */

src/main/java/com/yahoo/bullet/pubsub/PubSubMessage.java

@@ -5,14 +5,18 @@
  */
 package com.yahoo.bullet.pubsub;
 
+import com.google.gson.Gson;
+import com.yahoo.bullet.common.SerializerDeserializer;
 import com.yahoo.bullet.pubsub.Metadata.Signal;
+import com.yahoo.bullet.query.Query;
 import com.yahoo.bullet.result.JSONFormatter;
 import lombok.Getter;
 import lombok.Setter;
 
 import java.io.Serializable;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
+import java.util.Base64;
 import java.util.Objects;
 
 /**
@@ -21,11 +25,12 @@
  */
 @Getter
 public class PubSubMessage implements Serializable, JSONFormatter {
-    private static final long serialVersionUID = -5068189058170874687L;
     public static final Charset CHARSET = StandardCharsets.UTF_8;
+    private static final long serialVersionUID = 5096747716667851530L;
 
     private String id;
-    private byte[] content;
+    // Serializable enforced through the constructors, getter and setter. Is Object so GSON can reify an instance.
+    private Object content;
     @Setter
     private Metadata metadata;
 
@@ -52,50 +57,29 @@ public PubSubMessage(String id, Signal signal) {
      * @param id The ID associated with the message.
      * @param content The content of the message.
      */
-    public PubSubMessage(String id, byte[] content) {
+    public PubSubMessage(String id, Serializable content) {
         this(id, content, (Metadata) null);
     }
 
-    /**
-     * Constructor for a message having only content as a String.
-     *
-     * @param id The ID associated with the message.
-     * @param content The content of the message as a String.
-     */
-    public PubSubMessage(String id, String content) {
-        this(id, content, null);
-    }
-
     /**
      * Constructor for a message having content and a {@link Metadata.Signal}.
      *
      * @param id The ID associated with the message.
      * @param content The content of the message.
      * @param signal The Signal to be sent with the message.
      */
-    public PubSubMessage(String id, byte[] content, Signal signal) {
+    public PubSubMessage(String id, Serializable content, Signal signal) {
         this(id, content, new Metadata(signal, null));
     }
 
-    /**
-     * Constructor for a message having content as a String and {@link Metadata}.
-     *
-     * @param id The ID associated with the message.
-     * @param content The content of the message as a String.
-     * @param metadata The Metadata associated with the message.
-     */
-    public PubSubMessage(String id, String content, Metadata metadata) {
-        this(id, content == null ? null : content.getBytes(CHARSET), metadata);
-    }
-
     /**
      * Constructor for a message having content and {@link Metadata}.
      *
      * @param id The ID associated with the message.
      * @param content The content of the message.
      * @param metadata The Metadata associated with the message.
      */
-    public PubSubMessage(String id, byte[] content, Metadata metadata) {
+    public PubSubMessage(String id, Serializable content, Metadata metadata) {
         this.id = Objects.requireNonNull(id, "ID cannot be null");
         this.content = content;
         this.metadata = metadata;
@@ -138,14 +122,52 @@ public boolean hasSignal() {
         return hasMetadata() && metadata.hasSignal();
     }
 
+    /**
+     * Returns the {@link Serializable} content stored in the message.
+     *
+     * @return The content stored.
+     */
+    public Serializable getContent() {
+        return (Serializable) content;
+    }
+
+    /**
+     * Returns the content stored in the message as a byte[]. You should use this to read the byte[] back from the
+     * message if you provided it originally to the message as a byte[].
+     *
+     * @return The content stored as a byte[].
+     */
+    public byte[] getContentAsByteArray() {
+        return (byte[]) content;
+    }
+
     /**
      * Returns the content stored in the message as a String. You should use this to read the String back from the
      * message if you provided it originally to the message as a String.
      *
-     * @return The content stored as a String using the {@link PubSubMessage#CHARSET}.
+     * @return The content stored as a String.
      */
     public String getContentAsString() {
-        return content == null ? null : new String(content, CHARSET);
+        return (String) content;
+    }
+
+    /**
+     * Returns the content stored in the message as a {@link Query}. You should use this to read the {@link Query} back
+     * if you originally provided to the message as a {@link Serializable}.
+     *
+     * @return The content stored as a {@link Query}.
+     */
+    public Query getContentAsQuery() {
+        return (Query) content;
+    }
+
+    /**
+     * Set a {@link Serializable} content for this message.
+     *
+     * @param content The content for this message.
+     */
+    public void setContent(Serializable content) {
+        this.content = content;
     }
 
     @Override
@@ -169,16 +191,38 @@ public String toString() {
 
     @Override
     public String asJSON() {
-        return JSONFormatter.asJSON(this);
+        String data = Base64.getEncoder().encodeToString(SerializerDeserializer.toBytes((Serializable) content));
+        PubSubMessage message = new PubSubMessage(id, data, metadata);
+        return JSONFormatter.asJSON(message);
     }
 
     /**
-     * Converts a json representation back to an instance.
+     * Converts a json representation back to an instance. Is the inverse of {@link #asJSON()}.
      *
      * @param json The string representation of the JSON.
      * @return An instance of this class.
      */
     public static PubSubMessage fromJSON(String json) {
-        return JSONFormatter.fromJSON(json, PubSubMessage.class);
+        return fromJSON(json, GSON);
+    }
+
+    /**
+     * Converts a JSON representation back to an instance using a specific {@link Gson} converter.
+     * Is the inverse of {@link #asJSON()}.
+     *
+     * @param json The string representation of the JSON.
+     * @param gson The {@link Gson} converter to use.
+     * @return An instance of this class.
+     */
+    public static PubSubMessage fromJSON(String json, Gson gson) {
+        return fromJSON(gson.fromJson(json, PubSubMessage.class));
+    }
+
+    private static PubSubMessage fromJSON(PubSubMessage message) {
+        if (message == null || message.getContent() == null) {
+            return message;
+        }
+        message.content = SerializerDeserializer.fromBytes(Base64.getDecoder().decode(message.getContentAsString()));
+        return message;
     }
 }

src/main/java/com/yahoo/bullet/pubsub/PubSubMessageSerDe.java

@@ -0,0 +1,81 @@
+/*
+ *  Copyright 2021 Yahoo Inc.
+ *  Licensed under the terms of the Apache License, Version 2.0.
+ *  See the LICENSE file associated with the project for terms.
+ */
+package com.yahoo.bullet.pubsub;
+
+import com.yahoo.bullet.common.BulletConfig;
+import com.yahoo.bullet.query.Query;
+
+import java.io.Serializable;
+
+/**
+ * This allows you to hook into the {@link PubSubMessage} serialization and deserialization - for instance, if you want
+ * to customize what is done to the message before it is sent to or read in the backend. You can implement a custom
+ * version of this to change what the {@link PubSubMessage} contains (see {@link #toMessage(PubSubMessage)}) and invert
+ * it when it is read back (see {@link #fromMessage(PubSubMessage)}).
+ *
+ * Please note that this is not intended to be applied to messages read from the backend or for signals sent to the
+ * backend. You can, for example, use this if you want to customize what is stored in storage or to lazily create a
+ * {@link Query} in the backend.
+ */
+public abstract class PubSubMessageSerDe implements Serializable {
+    private static final long serialVersionUID = 5352288558800960763L;
+
+    protected BulletConfig config;
+
+    /**
+     * Constructor.
+     *
+     * @param config The {@link BulletConfig} to configure this class.
+     */
+    public PubSubMessageSerDe(BulletConfig config) {
+        this.config = config;
+    }
+
+    /**
+     * A helper to submit a query by converting it to a standard {@link PubSubMessage} and running it through
+     * {@link #toMessage(PubSubMessage)}.
+     *
+     * @param id The ID for the query.
+     * @param query The {@link Query} object to create a {@link PubSubMessage}.
+     * @param queryString The BQL for the query.
+     * @return A converted {@link PubSubMessage}.
+     */
+    public PubSubMessage toMessage(String id, Query query, String queryString) {
+        return toMessage(new PubSubMessage(id, query, new Metadata(null, queryString)));
+    }
+
+    /**
+     * Takes a {@link PubSubMessage} and returns the new format of the message. Note that it is allowed to modify the
+     * original message. See {@link #fromMessage(PubSubMessage)} for the inverse.
+     *
+     * @param message The {@link PubSubMessage} to convert.
+     * @return A converted {@link PubSubMessage}.
+     */
+    public abstract PubSubMessage toMessage(PubSubMessage message);
+
+    /**
+     * Takes a converted {@link PubSubMessage} and returns the original format of the message. Note that it is allowed
+     * to modify the original message. See {@link #toMessage(PubSubMessage)} for the inverse.
+     *
+     * @param message The {@link PubSubMessage} to revert.
+     * @return A reverted {@link PubSubMessage}.
+     */
+    public abstract PubSubMessage fromMessage(PubSubMessage message);
+
+    /**
+     * Create a {@link PubSubMessageSerDe} instance using the class specified in the config file.
+     *
+     * @param config The non-null {@link BulletConfig} containing the class name and its settings.
+     * @return An instance of specified class initialized with settings from the input file and defaults.
+     */
+    public static PubSubMessageSerDe from(BulletConfig config) {
+        try {
+            return config.loadConfiguredClass(BulletConfig.PUBSUB_MESSAGE_SERDE_CLASS_NAME);
+        } catch (RuntimeException e) {
+            throw new RuntimeException("Cannot create PubSubMessageSerDe instance.", e.getCause());
+        }
+    }
+}

src/main/java/com/yahoo/bullet/pubsub/Publisher.java

@@ -5,6 +5,8 @@
  */
 package com.yahoo.bullet.pubsub;
 
+import java.io.Serializable;
+
 public interface Publisher extends AutoCloseable {
     /**
      * Send a message with an ID and content.
@@ -14,7 +16,7 @@ public interface Publisher extends AutoCloseable {
      * @return The sent {@link PubSubMessage}.
      * @throws PubSubException if the messaging system throws an error.
      */
-    default PubSubMessage send(String id, byte[] content) throws PubSubException {
+    default PubSubMessage send(String id, Serializable content) throws PubSubException {
         return send(new PubSubMessage(id, content));
     }
 

src/main/java/com/yahoo/bullet/pubsub/rest/RESTPubSub.java

@@ -12,6 +12,9 @@
 import com.yahoo.bullet.pubsub.Subscriber;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.http.impl.client.HttpClients;
+
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
@@ -21,7 +24,7 @@
 public class RESTPubSub extends PubSub {
     public static final int OK_200 = 200;
     public static final int NO_CONTENT_204 = 204;
-    public static final String UTF_8 = "UTF-8";
+    public static final Charset UTF_8 = StandardCharsets.UTF_8;
 
     /**
      * Create a RESTPubSub from a {@link BulletConfig}.

src/main/java/com/yahoo/bullet/pubsub/rest/RESTSubscriber.java

@@ -72,7 +72,7 @@ public List<PubSubMessage> getMessages() {
                     HttpEntity httpEntity = response.getEntity();
                     String message = EntityUtils.toString(httpEntity, RESTPubSub.UTF_8);
                     log.debug("Received message from url: {}. Message was {}", url, message);
-                    messages.add(GSON.fromJson(message, PubSubMessage.class));
+                    messages.add(PubSubMessage.fromJSON(message, GSON));
                     EntityUtils.consume(httpEntity);
                 } else if (statusCode != RESTPubSub.NO_CONTENT_204) {
                     // NO_CONTENT_204 indicates there are no new messages - anything else indicates a problem

src/main/resources/bullet_defaults.yaml

@@ -249,3 +249,5 @@ bullet.query.partitioner.equality.delimiter: "|"
 bullet.pubsub.class.name: "com.yahoo.bullet.pubsub.rest.RESTPubSub"
 # The current context. This can be QUERY_PROCESSING or QUERY_SUBMISSION. The PubSub implementation should use this to generate appropriate Publishers and Subscribers.
 bullet.pubsub.context.name: "QUERY_PROCESSING"
+# The class to use for converting and reading PubSubMessage queries sent to the backend.
+bullet.pubsub.message.serde.class.name: "com.yahoo.bullet.pubsub.IdentityPubSubMessageSerDe"