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.NESTED;
+  }
 }

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

@@ -0,0 +1,8 @@
+package org.hypertrace.core.documentstore;
+
+public enum DocumentType {
+  // FLAT documents contain individual columns for each attribute
+  FLAT,
+  // NESTED documents contains attributes as a nested JSON document
+  NESTED
+}

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 = DocumentType.NESTED;
 
   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 {

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;
@@ -1264,6 +1265,10 @@ public PostgresResultIteratorWithBasicTypes(ResultSet resultSet) {
       super(resultSet);
     }
 
+    public PostgresResultIteratorWithBasicTypes(ResultSet resultSet, DocumentType documentType) {
+      super(resultSet, documentType);
+    }
+
     @Override
     public Document next() {
       try {
@@ -1299,7 +1304,7 @@ protected Document prepareDocument() throws SQLException, IOException {
         jsonNode.remove(DOCUMENT_ID);
       }
 
-      return new JSONDocument(MAPPER.writeValueAsString(jsonNode));
+      return new JSONDocument(MAPPER.writeValueAsString(jsonNode), documentType);
     }
 
     private void addColumnToJsonNode(
@@ -1386,17 +1391,29 @@ static class PostgresResultIterator implements CloseableIterator<Document> {
 
     protected final ObjectMapper MAPPER = new ObjectMapper();
     protected ResultSet resultSet;
-    private final boolean removeDocumentId;
     protected boolean cursorMovedForward = false;
     protected boolean hasNext = false;
 
+    private final boolean removeDocumentId;
+    protected DocumentType documentType;
+
     public PostgresResultIterator(ResultSet resultSet) {
       this(resultSet, true);
     }
 
     PostgresResultIterator(ResultSet resultSet, boolean removeDocumentId) {
+      this(resultSet, removeDocumentId, DocumentType.NESTED);
+    }
+
+    public PostgresResultIterator(ResultSet resultSet, DocumentType documentType) {
+      this(resultSet, true, documentType);
+    }
+
+    PostgresResultIterator(
+        ResultSet resultSet, boolean removeDocumentId, DocumentType documentType) {
       this.resultSet = resultSet;
       this.removeDocumentId = removeDocumentId;
+      this.documentType = documentType;
     }
 
     @Override
@@ -1447,7 +1464,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);
     }
 
     protected void closeResultSet() {
@@ -1508,7 +1525,7 @@ protected Document prepareDocument() throws SQLException, IOException {
           }
         }
       }
-      return new JSONDocument(MAPPER.writeValueAsString(jsonNode));
+      return new JSONDocument(MAPPER.writeValueAsString(jsonNode), documentType);
     }
 
     private String getColumnValue(

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

@@ -10,6 +10,7 @@
 import lombok.extern.slf4j.Slf4j;
 import org.hypertrace.core.documentstore.CloseableIterator;
 import org.hypertrace.core.documentstore.Document;
+import org.hypertrace.core.documentstore.DocumentType;
 import org.hypertrace.core.documentstore.postgres.PostgresCollection.PostgresResultIterator;
 import org.hypertrace.core.documentstore.postgres.PostgresCollection.PostgresResultIteratorWithMetaData;
 import org.hypertrace.core.documentstore.postgres.query.v1.transformer.PostgresQueryTransformer;
@@ -18,6 +19,7 @@
 @Slf4j
 @AllArgsConstructor
 public class PostgresQueryExecutor {
+
   private final PostgresTableIdentifier tableIdentifier;
 
   public CloseableIterator<Document> execute(final Connection connection, final Query query) {
@@ -37,7 +39,8 @@ public CloseableIterator<Document> execute(
       final ResultSet resultSet = preparedStatement.executeQuery();
 
       if ((tableIdentifier.getTableName().equals(flatStructureCollectionName))) {
-        return new PostgresCollection.PostgresResultIteratorWithBasicTypes(resultSet);
+        return new PostgresCollection.PostgresResultIteratorWithBasicTypes(
+            resultSet, DocumentType.FLAT);
       }
       return query.getSelections().size() > 0
           ? new PostgresResultIteratorWithMetaData(resultSet)