[PoC] Add TypedDocument for supporting flat collection writes

document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java

@@ -1,7 +1,10 @@
 package org.hypertrace.core.documentstore;
 
 import static org.hypertrace.core.documentstore.utils.Utils.readFileFromResource;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -11,10 +14,16 @@
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import org.hypertrace.core.documentstore.expression.impl.DataType;
+import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
+import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
+import org.hypertrace.core.documentstore.expression.operators.RelationalOperator;
 import org.hypertrace.core.documentstore.postgres.PostgresDatastore;
+import org.hypertrace.core.documentstore.query.Query;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
@@ -169,16 +178,115 @@ public static void shutdown() {
   class UpsertTests {
 
     @Test
-    @DisplayName("Should throw UnsupportedOperationException for upsert")
-    void testUpsertNewDocument() {
+    @DisplayName("Should upsert a new document with TypedDocument")
+    void testUpsertNewDocument() throws IOException {
+      // Create TypedDocument with value and type coupled together
+      TypedDocument typedDoc =
+          TypedDocument.builder()
+              .field("item", "NewItem", DataType.STRING)
+              .field("price", 99, DataType.INTEGER)
+              .field("quantity", 50, DataType.INTEGER)
+              .field("in_stock", true, DataType.BOOLEAN)
+              .arrayField("tags", List.of("new", "test"), DataType.STRING)
+              .jsonbField("props", Map.of("brand", "TestBrand", "size", "medium"))
+              .build();
+
+      Key key = new SingleValueKey("default", "100");
+
+      // Perform upsert
+      boolean result = flatCollection.upsert(key, typedDoc);
+      assertTrue(result, "Upsert should return true for new document");
+
+      // Verify the document was inserted by querying it
+      Query query =
+          Query.builder()
+              .setFilter(
+                  RelationalExpression.of(
+                      IdentifierExpression.of("item"),
+                      RelationalOperator.EQ,
+                      org.hypertrace.core.documentstore.expression.impl.ConstantExpression.of(
+                          "NewItem")))
+              .build();
+
+      Iterator<Document> results = flatCollection.find(query);
+      assertTrue(results.hasNext(), "Should find the inserted document");
+
+      Document foundDoc = results.next();
+      assertNotNull(foundDoc);
+
+      JsonNode foundJson = OBJECT_MAPPER.readTree(foundDoc.toJson());
+      assertEquals("NewItem", foundJson.get("item").asText());
+      assertEquals(99, foundJson.get("price").asInt());
+      assertEquals(50, foundJson.get("quantity").asInt());
+      assertTrue(foundJson.get("in_stock").asBoolean());
+
+      // Verify array field
+      JsonNode tagsNode = foundJson.get("tags");
+      assertNotNull(tagsNode);
+      assertTrue(tagsNode.isArray());
+      assertEquals(2, tagsNode.size());
+
+      // Verify JSONB field
+      JsonNode propsResult = foundJson.get("props");
+      assertNotNull(propsResult);
+      assertEquals("TestBrand", propsResult.get("brand").asText());
+    }
+
+    @Test
+    @DisplayName("Should update existing document on upsert with TypedDocument")
+    void testUpsertExistingDocument() throws IOException {
+      // First, get an existing document ID from the initial data
+      String existingId = "1"; // ID 1 exists in initial data
+
+      // Create TypedDocument with updated values
+      TypedDocument typedDoc =
+          TypedDocument.builder()
+              .field("item", "UpdatedSoap", DataType.STRING)
+              .field("price", 999, DataType.INTEGER)
+              .field("quantity", 100, DataType.INTEGER)
+              .field("in_stock", false, DataType.BOOLEAN)
+              .build();
+
+      Key key = new SingleValueKey("default", existingId);
+
+      // Perform upsert (should update existing)
+      boolean result = flatCollection.upsert(key, typedDoc);
+      assertTrue(result, "Upsert should return true for existing document update");
+
+      // Verify the document was updated
+      Query query =
+          Query.builder()
+              .setFilter(
+                  RelationalExpression.of(
+                      IdentifierExpression.of("item"),
+                      RelationalOperator.EQ,
+                      org.hypertrace.core.documentstore.expression.impl.ConstantExpression.of(
+                          "UpdatedSoap")))
+              .build();
+
+      Iterator<Document> results = flatCollection.find(query);
+      assertTrue(results.hasNext(), "Should find the updated document");
+
+      Document foundDoc = results.next();
+      JsonNode foundJson = OBJECT_MAPPER.readTree(foundDoc.toJson());
+      assertEquals("UpdatedSoap", foundJson.get("item").asText());
+      assertEquals(999, foundJson.get("price").asInt());
+      assertEquals(100, foundJson.get("quantity").asInt());
+    }
+
+    @Test
+    @DisplayName("Should throw IllegalArgumentException for non-TypedDocument")
+    void testUpsertWithoutTypedDocument() {
       ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
-      objectNode.put("_id", 100);
       objectNode.put("item", "NewItem");
       objectNode.put("price", 99);
       Document document = new JSONDocument(objectNode);
       Key key = new SingleValueKey("default", "100");
 
-      assertThrows(UnsupportedOperationException.class, () -> flatCollection.upsert(key, document));
+      assertThrows(
+          IllegalArgumentException.class,
+          () -> flatCollection.upsert(key, document),
+          "Should throw IllegalArgumentException when not using TypedDocument");
     }
 
     @Test

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

@@ -0,0 +1,177 @@
+package org.hypertrace.core.documentstore;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import org.hypertrace.core.documentstore.expression.impl.DataType;
+
+/**
+ * A Document implementation where each field carries its value and type information.
+ *
+ * <p>For flat PostgreSQL tables, type information is needed to correctly set values in prepared
+ * statements, especially for arrays and JSONB columns.
+ *
+ * <p>Usage:
+ *
+ * <pre>{@code
+ * TypedDocument typedDoc = TypedDocument.builder()
+ *     .field("item", "Soap", DataType.STRING)
+ *     .field("price", 99, DataType.INTEGER)
+ *     .field("in_stock", true, DataType.BOOLEAN)
+ *     .arrayField("tags", List.of("hygiene", "bath"), DataType.STRING)
+ *     .jsonbField("props", Map.of("brand", "Dove", "size", "large"))
+ *     .build();
+ * }</pre>
+ */
+public class TypedDocument implements Document {
+
+  private static final ObjectMapper MAPPER = new ObjectMapper();
+
+  private final Map<String, TypedField> fields;
+
+  private TypedDocument(Map<String, TypedField> fields) {
+    this.fields = Collections.unmodifiableMap(fields);
+  }
+
+  public static Builder builder() {
+    return new Builder();
+  }
+
+  @Override
+  public String toJson() {
+    Map<String, Object> jsonMap = new LinkedHashMap<>();
+    for (Map.Entry<String, TypedField> entry : fields.entrySet()) {
+      jsonMap.put(entry.getKey(), entry.getValue().getValue());
+    }
+    try {
+      return MAPPER.writeValueAsString(jsonMap);
+    } catch (JsonProcessingException e) {
+      throw new RuntimeException("Failed to serialize TypedDocument to JSON", e);
+    }
+  }
+
+  @Override
+  public DocumentType getDocumentType() {
+    return DocumentType.FLAT;
+  }
+
+  /** Returns all fields in this document. */
+  public Map<String, TypedField> getFields() {
+    return fields;
+  }
+
+  /** Returns a specific field by name, or null if not present. */
+  public TypedField getField(String fieldName) {
+    return fields.get(fieldName);
+  }
+
+  /** Represents a field's value and type information. */
+  public static class TypedField {
+    private final Object value;
+    private final FieldKind kind;
+    private final DataType dataType; // For SCALAR or ARRAY element type
+
+    private TypedField(Object value, FieldKind kind, DataType dataType) {
+      this.value = value;
+      this.kind = kind;
+      this.dataType = dataType;
+    }
+
+    public static TypedField scalar(Object value, DataType dataType) {
+      return new TypedField(value, FieldKind.SCALAR, dataType);
+    }
+
+    public static TypedField array(Collection<?> value, DataType elementType) {
+      return new TypedField(value, FieldKind.ARRAY, elementType);
+    }
+
+    public static TypedField jsonb(Object value) {
+      return new TypedField(value, FieldKind.JSONB, null);
+    }
+
+    public Object getValue() {
+      return value;
+    }
+
+    public FieldKind getKind() {
+      return kind;
+    }
+
+    /** Returns the DataType for scalar fields, or the element type for array fields. */
+    public DataType getDataType() {
+      return dataType;
+    }
+
+    public boolean isScalar() {
+      return kind == FieldKind.SCALAR;
+    }
+
+    public boolean isArray() {
+      return kind == FieldKind.ARRAY;
+    }
+
+    public boolean isJsonb() {
+      return kind == FieldKind.JSONB;
+    }
+  }
+
+  /** The kind of field (scalar, array, or JSONB). */
+  public enum FieldKind {
+    SCALAR,
+    ARRAY,
+    JSONB
+  }
+
+  /** Builder for TypedDocument */
+  public static class Builder {
+    private final Map<String, TypedField> fields = new LinkedHashMap<>();
+
+    private Builder() {}
+
+    /**
+     * Adds a scalar field with its value and type.
+     *
+     * @param name the column name
+     * @param value the field value
+     * @param type the data type
+     * @return this builder
+     */
+    public Builder field(String name, Object value, DataType type) {
+      fields.put(name, TypedField.scalar(value, type));
+      return this;
+    }
+
+    /**
+     * Adds an array field with its values and element type.
+     *
+     * @param name the column name
+     * @param values the array values
+     * @param elementType the data type of array elements
+     * @return this builder
+     */
+    public Builder arrayField(String name, List<?> values, DataType elementType) {
+      fields.put(name, TypedField.array(values, elementType));
+      return this;
+    }
+
+    /**
+     * Adds a JSONB field with its value (will be serialized as JSON).
+     *
+     * @param name the column name
+     * @param value the value (typically a Map or object that can be serialized to JSON)
+     * @return this builder
+     */
+    public Builder jsonbField(String name, Object value) {
+      fields.put(name, TypedField.jsonb(value));
+      return this;
+    }
+
+    public TypedDocument build() {
+      return new TypedDocument(fields);
+    }
+  }
+}