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 {
@@ -54,6 +75,6 @@ public boolean equals(Object obj) {
     }
 
     JSONDocument other = (JSONDocument) obj;
-    return Objects.equals(node, other.node);
+    return Objects.equals(node, other.node) && documentType == other.documentType;
   }
 }

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)

document-store/src/test/java/org/hypertrace/core/documentstore/JSONDocumentTest.java

@@ -1,11 +1,15 @@
 package org.hypertrace.core.documentstore;
 
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import java.util.Map;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
 public class JSONDocumentTest {
 
+  private static final ObjectMapper mapper = new ObjectMapper();
+
   @Test
   public void testJSONDocument() throws Exception {
     Map<String, String> data = Map.of("key1", "value1", "key2", "value2");
@@ -14,4 +18,92 @@ public void testJSONDocument() throws Exception {
     Assertions.assertEquals(document1, document2);
     Assertions.assertEquals(document1.toJson(), document2.toJson());
   }
+
+  @Test
+  public void testJSONDocumentWithDefaultDocumentType() throws Exception {
+    Map<String, String> data = Map.of("key1", "value1", "key2", "value2");
+    JSONDocument document = new JSONDocument(data);
+    Assertions.assertEquals(DocumentType.NESTED, document.getDocumentType());
+  }
+
+  @Test
+  public void testJSONDocumentWithExplicitDocumentType() throws Exception {
+    Map<String, String> data = Map.of("key1", "value1", "key2", "value2");
+
+    JSONDocument nestedDocument = new JSONDocument(data, DocumentType.NESTED);
+    Assertions.assertEquals(DocumentType.NESTED, nestedDocument.getDocumentType());
+
+    JSONDocument flatDocument = new JSONDocument(data, DocumentType.FLAT);
+    Assertions.assertEquals(DocumentType.FLAT, flatDocument.getDocumentType());
+  }
+
+  @Test
+  public void testJSONDocumentConstructorsWithDocumentType() throws Exception {
+    String json = "{\"key1\":\"value1\",\"key2\":\"value2\"}";
+    JsonNode node = mapper.readTree(json);
+    Map<String, String> data = Map.of("key1", "value1", "key2", "value2");
+
+    // Test string constructor with DocumentType
+    JSONDocument stringDoc = new JSONDocument(json, DocumentType.FLAT);
+    Assertions.assertEquals(DocumentType.FLAT, stringDoc.getDocumentType());
+
+    // Test object constructor with DocumentType
+    JSONDocument objectDoc = new JSONDocument(data, DocumentType.FLAT);
+    Assertions.assertEquals(DocumentType.FLAT, objectDoc.getDocumentType());
+
+    // Test JsonNode constructor with DocumentType
+    JSONDocument nodeDoc = new JSONDocument(node, DocumentType.FLAT);
+    Assertions.assertEquals(DocumentType.FLAT, nodeDoc.getDocumentType());
+  }
+
+  @Test
+  public void testEqualsWithSameContentDifferentDocumentType() throws Exception {
+    Map<String, String> data = Map.of("key1", "value1", "key2", "value2");
+
+    JSONDocument nestedDoc = new JSONDocument(data, DocumentType.NESTED);
+    JSONDocument flatDoc = new JSONDocument(data, DocumentType.FLAT);
+
+    // Current implementation only compares JsonNode, not DocumentType
+    // This test documents the current behavior - documents with same content but different types
+    // are equal
+    Assertions.assertNotEquals(nestedDoc, flatDoc);
+    Assertions.assertEquals(nestedDoc.toJson(), flatDoc.toJson());
+  }
+
+  @Test
+  public void testEqualsWithSameContentSameDocumentType() throws Exception {
+    Map<String, String> data = Map.of("key1", "value1", "key2", "value2");
+
+    JSONDocument doc1 = new JSONDocument(data, DocumentType.FLAT);
+    JSONDocument doc2 = new JSONDocument(data, DocumentType.FLAT);
+
+    Assertions.assertEquals(doc1, doc2);
+    Assertions.assertEquals(doc1.getDocumentType(), doc2.getDocumentType());
+  }
+
+  @Test
+  public void testErrorDocument() throws Exception {
+    String errorMessage = "Test error message";
+    JSONDocument errorDoc = JSONDocument.errorDocument(errorMessage);
+
+    // Verify default document type
+    Assertions.assertEquals(DocumentType.NESTED, errorDoc.getDocumentType());
+
+    // Verify error message is in the JSON
+    String expectedJson = "{\"errorMessage\":\"" + errorMessage + "\"}";
+    Assertions.assertEquals(expectedJson, errorDoc.toJson());
+
+    // Test error document equality
+    JSONDocument anotherErrorDoc = JSONDocument.errorDocument(errorMessage);
+    Assertions.assertEquals(errorDoc, anotherErrorDoc);
+  }
+
+  @Test
+  public void testToStringMethod() throws Exception {
+    Map<String, String> data = Map.of("key1", "value1");
+    JSONDocument document = new JSONDocument(data, DocumentType.FLAT);
+
+    // toString should return the same as toJson
+    Assertions.assertEquals(document.toJson(), document.toString());
+  }
 }