Make JsonDocument Underlying Storage Type Aware (Document store / SQL store)

document-store/src/main/java/org/hypertrace/core/documentstore/Document.java

@@ -3,4 +3,8 @@
 public interface Document {
 
   String toJson();
+
+  default DocumentType getDocumentType() {
+    return DocumentType.DOCUMENT_STORE;
+  }
 }

document-store/src/main/java/org/hypertrace/core/documentstore/DocumentType.java

@@ -0,0 +1,6 @@
+package org.hypertrace.core.documentstore;
+
+public enum DocumentType {
+  SQL_STORE,
+  DOCUMENT_STORE
+}

document-store/src/main/java/org/hypertrace/core/documentstore/JSONDocument.java

@@ -11,19 +11,40 @@ public class JSONDocument implements Document {
 
   private static ObjectMapper mapper = new ObjectMapper();
   private JsonNode node;
+  private DocumentType documentType;
 
   public JSONDocument(String json) throws IOException {
     node = mapper.readTree(json);
   }
 
+  public JSONDocument(String json, DocumentType documentType) throws IOException {
+    node = mapper.readTree(json);
+    this.documentType = documentType;
+  }
+
   public JSONDocument(Object object) throws IOException {
     node = mapper.readTree(mapper.writeValueAsString(object));
   }
 
+  public JSONDocument(Object object, DocumentType documentType) throws IOException {
+    node = mapper.readTree(mapper.writeValueAsString(object));
+    this.documentType = documentType;
+  }
+
   public JSONDocument(JsonNode node) {
     this.node = node;
   }
 
+  public JSONDocument(JsonNode node, DocumentType documentType) {
+    this.node = node;
+    this.documentType = documentType;
+  }
+
+  @Override
+  public DocumentType getDocumentType() {
+    return this.documentType;
+  }
+
   @Override
   public String toJson() {
     try {
@@ -39,6 +60,12 @@ public static JSONDocument errorDocument(String message) {
     return new JSONDocument(objectNode);
   }
 
+  public static JSONDocument errorDocument(String message, DocumentType documentType) {
+    ObjectNode objectNode = mapper.createObjectNode();
+    objectNode.put("errorMessage", message);
+    return new JSONDocument(objectNode, documentType);
+  }
+
   @Override
   public String toString() {
     return toJson();

document-store/src/main/java/org/hypertrace/core/documentstore/mongo/MongoUtils.java

@@ -29,10 +29,12 @@
 import org.bson.json.JsonMode;
 import org.bson.json.JsonWriterSettings;
 import org.hypertrace.core.documentstore.Document;
+import org.hypertrace.core.documentstore.DocumentType;
 import org.hypertrace.core.documentstore.JSONDocument;
 import org.hypertrace.core.documentstore.model.options.ReturnDocumentType;
 
 public final class MongoUtils {
+
   public static final String FIELD_SEPARATOR = ".";
   public static final String PREFIX = "$";
   private static final String LITERAL = PREFIX + "literal";
@@ -137,10 +139,10 @@ public static Document dbObjectToDocument(BasicDBObject dbObject) {
       jsonString = dbObject.toJson(relaxed);
       JsonNode jsonNode = MAPPER.readTree(jsonString);
       JsonNode decodedJsonNode = recursiveClone(jsonNode, MongoUtils::decodeKey, identity());
-      return new JSONDocument(decodedJsonNode);
+      return new JSONDocument(decodedJsonNode, DocumentType.DOCUMENT_STORE);
     } catch (IOException e) {
       // throwing exception is not very useful here.
-      return JSONDocument.errorDocument(e.getMessage());
+      return JSONDocument.errorDocument(e.getMessage(), DocumentType.DOCUMENT_STORE);
     }
   }
 

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java

@@ -61,6 +61,7 @@
 import org.hypertrace.core.documentstore.Collection;
 import org.hypertrace.core.documentstore.CreateResult;
 import org.hypertrace.core.documentstore.Document;
+import org.hypertrace.core.documentstore.DocumentType;
 import org.hypertrace.core.documentstore.Filter;
 import org.hypertrace.core.documentstore.JSONDocument;
 import org.hypertrace.core.documentstore.Key;
@@ -1276,7 +1277,7 @@ public Document next() {
       } catch (IOException | SQLException e) {
         System.out.println("prepare document failed!");
         closeResultSet();
-        return JSONDocument.errorDocument(e.getMessage());
+        return JSONDocument.errorDocument(e.getMessage(), DocumentType.SQL_STORE);
       }
     }
 
@@ -1299,7 +1300,7 @@ protected Document prepareDocument() throws SQLException, IOException {
         jsonNode.remove(DOCUMENT_ID);
       }
 
-      return new JSONDocument(MAPPER.writeValueAsString(jsonNode));
+      return new JSONDocument(MAPPER.writeValueAsString(jsonNode), DocumentType.SQL_STORE);
     }
 
     private void addColumnToJsonNode(
@@ -1430,7 +1431,7 @@ public Document next() {
         return prepareDocument();
       } catch (IOException | SQLException e) {
         closeResultSet();
-        return JSONDocument.errorDocument(e.getMessage());
+        return JSONDocument.errorDocument(e.getMessage(), DocumentType.SQL_STORE);
       }
     }
 
@@ -1447,7 +1448,7 @@ protected Document prepareDocument() throws SQLException, IOException {
       jsonNode.put(CREATED_AT, String.valueOf(createdAt));
       jsonNode.put(UPDATED_AT, String.valueOf(updatedAt));
 
-      return new JSONDocument(MAPPER.writeValueAsString(jsonNode));
+      return new JSONDocument(MAPPER.writeValueAsString(jsonNode), DocumentType.SQL_STORE);
     }
 
     protected void closeResultSet() {
@@ -1508,7 +1509,7 @@ protected Document prepareDocument() throws SQLException, IOException {
           }
         }
       }
-      return new JSONDocument(MAPPER.writeValueAsString(jsonNode));
+      return new JSONDocument(MAPPER.writeValueAsString(jsonNode), DocumentType.SQL_STORE);
     }
 
     private String getColumnValue(