WIP

Author: suddendust

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

@@ -533,6 +533,52 @@ void testSetNestedJsonbArray(String storeName) throws Exception {
         // Other sales fields preserved
         assertEquals(200, resultJson.get("sales").get("total").asInt());
       }
+
+      @ParameterizedTest(name = "{0}: SET entire JSONB column")
+      @ArgumentsSource(AllStoresProvider.class)
+      void testSetEntireJsonbColumn(String storeName) throws Exception {
+        String docId = generateDocId("set-jsonb-column");
+        insertTestDocument(docId);
+
+        Collection collection = getCollection(storeName);
+        Query query = buildQueryById(docId);
+
+        // Create a completely new object to replace the entire JSONB column
+        ObjectNode newProps = OBJECT_MAPPER.createObjectNode();
+        newProps.put("manufacturer", "NewManufacturer");
+        newProps.put("model", "Model-X");
+        newProps.put("year", 2024);
+        newProps.putArray("features").add("feature1").add("feature2");
+
+        List<SubDocumentUpdate> updates =
+            List.of(
+                SubDocumentUpdate.of("props", SubDocumentValue.of(new JSONDocument(newProps))));
+
+        UpdateOptions options =
+            UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build();
+
+        Optional<Document> result = collection.update(query, updates, options);
+
+        assertTrue(result.isPresent());
+        JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson());
+        JsonNode propsNode = resultJson.get("props");
+        assertNotNull(propsNode);
+
+        // Verify new fields are present
+        assertEquals("NewManufacturer", propsNode.get("manufacturer").asText());
+        assertEquals("Model-X", propsNode.get("model").asText());
+        assertEquals(2024, propsNode.get("year").asInt());
+        assertTrue(propsNode.get("features").isArray());
+        assertEquals(2, propsNode.get("features").size());
+        assertEquals("feature1", propsNode.get("features").get(0).asText());
+        assertEquals("feature2", propsNode.get("features").get(1).asText());
+
+        // Verify old fields are NOT present (entire column replaced)
+        assertFalse(propsNode.has("brand"));
+        assertFalse(propsNode.has("size"));
+        assertFalse(propsNode.has("count"));
+        assertFalse(propsNode.has("colors"));
+      }
     }
 
     @Nested

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

@@ -1466,10 +1466,8 @@ private void addColumnToJsonNode(
           if (jsonString != null) {
             try {
               JsonNode jsonValue = MAPPER.readTree(jsonString);
-              // For FLAT documents, column names with dots are actual column names, not encoded
-              // nested paths. Only apply nesting logic for non-FLAT document types.
-              if (documentType != DocumentType.FLAT
-                  && PostgresUtils.isEncodedNestedField(columnName)) {
+              // Handle like MetaData iterator - check for encoded nested fields
+              if (PostgresUtils.isEncodedNestedField(columnName)) {
                 handleNestedField(
                     PostgresUtils.decodeAliasForNestedField(columnName), jsonNode, jsonValue);
               } else {

document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresSetValueParser.java

@@ -8,6 +8,7 @@
 import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate;
 import org.hypertrace.core.documentstore.postgres.Params;
 import org.hypertrace.core.documentstore.postgres.Params.Builder;
+import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType;
 import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentArrayGetter;
 import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentValueParser;
 
@@ -33,14 +34,21 @@ public String parseNonJsonbField(final UpdateParserInput input) {
               .accept(new PostgresSubDocumentArrayGetter())
               .values();
       input.getParamsBuilder().addObjectParam(values);
+      return String.format("\"%s\" = ?", input.getBaseField());
     } else {
-      // For scalar columns, use standard value parser (ignore returned JSONB expression)
-      input
-          .getUpdate()
-          .getSubDocumentValue()
-          .accept(new PostgresSubDocumentValueParser(input.getParamsBuilder()));
+      // For scalar columns, use value parser which returns proper expression with type cast
+      String valueExpr =
+          input
+              .getUpdate()
+              .getSubDocumentValue()
+              .accept(new PostgresSubDocumentValueParser(input.getParamsBuilder()));
+      // For JSONB columns, use the returned expression (e.g., "?::jsonb" for nested documents)
+      // For other columns, use plain "?"
+      if (input.getColumnType() == PostgresDataType.JSONB) {
+        return String.format("\"%s\" = %s", input.getBaseField(), valueExpr);
+      }
+      return String.format("\"%s\" = ?", input.getBaseField());
     }
-    return String.format("\"%s\" = ?", input.getBaseField());
   }
 
   @Override