From 469ead11b32457b067a624f3ae966d6cbc2de956 Mon Sep 17 00:00:00 2001 From: Prashant Pandey Date: Fri, 6 Feb 2026 13:56:23 +0530 Subject: [PATCH] Subdoc update - ADD --- .../FlatCollectionWriteTest.java | 278 ++++++++++++++++++ .../postgres/FlatPostgresCollection.java | 7 +- ...FlatCollectionSubDocAddOperatorParser.java | 165 +++++++++++ 3 files changed, 449 insertions(+), 1 deletion(-) create mode 100644 document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocAddOperatorParser.java diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java index 7bf4e417..27a3ff65 100644 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/FlatCollectionWriteTest.java @@ -1615,6 +1615,284 @@ void testUpdateReturnsBeforeDocument() throws Exception { } } + @Nested + @DisplayName("ADD Operator Tests") + class AddSubdocOperatorTests { + + @Test + @DisplayName("Should increment top-level numeric column with ADD operator") + void testAddTopLevelColumn() throws Exception { + // Row 1 has price = 10 + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + // ADD 5 to price (10 + 5 = 15) + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(5)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + assertEquals(15, resultJson.get("price").asInt()); + + // Verify in database + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"price\" FROM \"%s\" WHERE \"id\" = '1'", FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals(15, rs.getInt("price")); + } + } + + @Test + @DisplayName("Should handle ADD on NULL column (treat as 0)") + void testAddOnNullColumn() throws Exception { + // Create a document with NULL price + String docId = "add-null-test"; + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "NullPriceItem"); + // price is not set, will be NULL + flatCollection.create(key, new JSONDocument(node)); + + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of(key.toString()))) + .build(); + + // ADD 100 to NULL price (COALESCE(NULL, 0) + 100 = 100) + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(100)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + assertEquals(100, resultJson.get("price").asInt()); + } + + @Test + @DisplayName("Should ADD with negative value (decrement)") + void testAddNegativeValue() throws Exception { + // Row 2 has price = 20 + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("2"))) + .build(); + + // ADD -5 to price (20 - 5 = 15) + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(-5)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + assertEquals(15, resultJson.get("price").asInt()); + } + + @Test + @DisplayName("Should ADD with floating point value") + void testAddFloatingPointValue() throws Exception { + // Row 3 has price = 30 + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("3"))) + .build(); + + // ADD 0.5 to price (30 + 0.5 = 30.5, but price is INTEGER so it might truncate) + // Testing with a column that supports decimals - weight is DOUBLE PRECISION + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("weight") + .operator(UpdateOperator.ADD) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(2.5)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + // Initial weight is NULL, so COALESCE(NULL, 0) + 2.5 = 2.5 + assertEquals(2.5, resultJson.get("weight").asDouble(), 0.01); + } + + @Test + @DisplayName("Should ADD to nested JSONB numeric field") + void testAddNestedJsonbField() throws Exception { + // First, set up a document with a JSONB field containing a numeric value + String docId = "add-jsonb-test"; + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "JsonbItem"); + ObjectNode sales = OBJECT_MAPPER.createObjectNode(); + sales.put("total", 100); + sales.put("count", 5); + node.set("sales", sales); + flatCollection.create(key, new JSONDocument(node)); + + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of(key.toString()))) + .build(); + + // ADD 50 to sales.total (100 + 50 = 150) + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("sales.total") + .operator(UpdateOperator.ADD) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(50)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + assertEquals(150, resultJson.get("sales").get("total").asInt()); + // Verify count wasn't affected + assertEquals(5, resultJson.get("sales").get("count").asInt()); + } + + @Test + @DisplayName("Should ADD to nested JSONB field that doesn't exist (creates with value)") + void testAddNestedJsonbFieldNotExists() throws Exception { + // Document with empty JSONB or no such nested key + String docId = "add-jsonb-new-key"; + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "NewKeyItem"); + ObjectNode sales = OBJECT_MAPPER.createObjectNode(); + sales.put("region", "US"); + // No 'total' key + node.set("sales", sales); + flatCollection.create(key, new JSONDocument(node)); + + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of(key.toString()))) + .build(); + + // ADD 75 to sales.total (non-existent, should become 0 + 75 = 75) + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("sales.total") + .operator(UpdateOperator.ADD) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(75)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + assertEquals(75.0, resultJson.get("sales").get("total").asDouble(), 0.01); + // Verify existing key wasn't affected + assertEquals("US", resultJson.get("sales").get("region").asText()); + } + + @Test + @DisplayName("Should throw IllegalArgumentException for non-numeric value") + void testAddNonNumericValue() { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + // ADD with a string value should fail + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + "not-a-number")) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + assertThrows( + IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); + } + } + @Test @DisplayName("Should return empty when no document matches query") void testUpdateNoMatch() throws Exception { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java index 26cc0e7e..238c8bbf 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollection.java @@ -2,6 +2,7 @@ import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.AFTER_UPDATE; import static org.hypertrace.core.documentstore.model.options.ReturnDocumentType.BEFORE_UPDATE; +import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.ADD; import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.SET; import com.fasterxml.jackson.databind.JsonNode; @@ -46,6 +47,7 @@ import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; import org.hypertrace.core.documentstore.postgres.query.v1.transformer.FlatPostgresFieldTransformer; import org.hypertrace.core.documentstore.postgres.update.FlatUpdateContext; +import org.hypertrace.core.documentstore.postgres.update.parser.FlatCollectionSubDocAddOperatorParser; import org.hypertrace.core.documentstore.postgres.update.parser.FlatCollectionSubDocSetOperatorParser; import org.hypertrace.core.documentstore.postgres.update.parser.FlatCollectionSubDocUpdateOperatorParser; import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils; @@ -72,7 +74,10 @@ public class FlatPostgresCollection extends PostgresCollection { private static final String DEFAULT_PRIMARY_KEY_COLUMN = "key"; private static final Map - SUB_DOC_UPDATE_PARSERS = Map.of(SET, new FlatCollectionSubDocSetOperatorParser()); + SUB_DOC_UPDATE_PARSERS = + Map.of( + SET, new FlatCollectionSubDocSetOperatorParser(), + ADD, new FlatCollectionSubDocAddOperatorParser()); private final PostgresLazyilyLoadedSchemaRegistry schemaRegistry; diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocAddOperatorParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocAddOperatorParser.java new file mode 100644 index 00000000..3b10cf6e --- /dev/null +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocAddOperatorParser.java @@ -0,0 +1,165 @@ +package org.hypertrace.core.documentstore.postgres.update.parser; + +import org.hypertrace.core.documentstore.model.subdoc.PrimitiveSubDocumentValue; +import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue; +import org.hypertrace.core.documentstore.model.subdoc.visitor.SubDocumentValueVisitor; +import org.hypertrace.core.documentstore.postgres.update.FlatUpdateContext; + +/** + * Parser for the ADD operator in flat collections. + * + *

ADD increments a numeric field by the given value. Handles two cases: + * + *

    + *
  • Top-level numeric columns: {@code "column" = COALESCE("column", 0) + ?} + *
  • Nested JSONB paths: {@code "column" = jsonb_set(COALESCE("column", '{}'), '{path}', + * (COALESCE("column"->>'path', '0')::float + ?::float)::text::jsonb, true)} + *
+ */ +public class FlatCollectionSubDocAddOperatorParser + implements FlatCollectionSubDocUpdateOperatorParser { + + /** Visitor to extract numeric values from SubDocumentValue. */ + private static final SubDocumentValueVisitor NUMERIC_VALUE_EXTRACTOR = + new SubDocumentValueVisitor<>() { + @Override + public Number visit(PrimitiveSubDocumentValue value) { + Object val = value.getValue(); + if (val instanceof Number) { + return (Number) val; + } + throw new IllegalArgumentException( + "ADD operator requires a numeric value, got: " + val.getClass().getName()); + } + + @Override + public Number visit( + org.hypertrace.core.documentstore.model.subdoc.MultiValuedPrimitiveSubDocumentValue + value) { + throw new IllegalArgumentException("ADD operator does not support multi-valued updates"); + } + + @Override + public Number visit( + org.hypertrace.core.documentstore.model.subdoc.NestedSubDocumentValue value) { + throw new IllegalArgumentException( + "ADD operator does not support nested document values"); + } + + @Override + public Number visit( + org.hypertrace.core.documentstore.model.subdoc.MultiValuedNestedSubDocumentValue + value) { + throw new IllegalArgumentException( + "ADD operator does not support multi-valued nested documents"); + } + + @Override + public Number visit( + org.hypertrace.core.documentstore.model.subdoc.NullSubDocumentValue value) { + throw new IllegalArgumentException("ADD operator does not support null values"); + } + }; + + @Override + public String parse(FlatUpdateContext context) { + validateNumericValue(context.getValue()); + + if (context.isTopLevel()) { + return parseTopLevel(context); + } else { + return parseNestedJsonb(context); + } + } + + private void validateNumericValue(SubDocumentValue value) { + // This will throw if the value is not numeric + value.accept(NUMERIC_VALUE_EXTRACTOR); + } + + /** + * Generates SQL for adding to a top-level numeric column. + * + *

Output: {@code "column" = COALESCE("column", 0) + ?::type} + */ + private String parseTopLevel(FlatUpdateContext context) { + Number value = context.getValue().accept(NUMERIC_VALUE_EXTRACTOR); + context.getParams().add(value); + + String typeCast = getPostgresTypeCast(context); + return String.format( + "\"%s\" = COALESCE(\"%s\", 0) + ?%s", + context.getColumnName(), context.getColumnName(), typeCast); + } + + /** Returns the PostgreSQL type cast for the column type. */ + private String getPostgresTypeCast(FlatUpdateContext context) { + if (context.getColumnType() == null) { + return ""; + } + switch (context.getColumnType()) { + case INTEGER: + return "::integer"; + case BIGINT: + return "::bigint"; + case REAL: + return "::real"; + case DOUBLE_PRECISION: + return "::double precision"; + default: + return ""; + } + } + + /** + * Generates SQL for adding to a numeric field within a JSONB column. Infers the numeric type from + * the value to preserve integer precision when possible. + * + *

Output for integers: {@code "column" = jsonb_set(COALESCE("column", '{}'), ?::text[], + * (COALESCE("column"#>>'{path}', '0')::bigint + ?::bigint)::text::jsonb, true)} + * + *

Output for floats: {@code "column" = jsonb_set(COALESCE("column", '{}'), ?::text[], + * (COALESCE("column"#>>'{path}', '0')::double precision + ?::double precision)::text::jsonb, + * true)} + */ + private String parseNestedJsonb(FlatUpdateContext context) { + String jsonPath = buildJsonPath(context.getNestedPath()); + Number value = context.getValue().accept(NUMERIC_VALUE_EXTRACTOR); + + // Infer type from value to preserve precision + String sqlType = inferSqlTypeFromValue(value); + + // Add params: jsonPath, value + context.getParams().add(jsonPath); + context.getParams().add(value); + + // Extracts nested JSONB value as text, e.g., "metrics"#>>'{sales,total}' traverses + // metrics→sales→total + String fieldAccessor = String.format("\"%s\"#>>'%s'", context.getColumnName(), jsonPath); + + // jsonb_set with arithmetic using inferred type + return String.format( + "\"%s\" = jsonb_set(COALESCE(\"%s\", '{}'), ?::text[], (COALESCE(%s, '0')::%s + ?::%s)::text::jsonb, true)", + context.getColumnName(), context.getColumnName(), fieldAccessor, sqlType, sqlType); + } + + /** Infers PostgreSQL type from the Java Number type. */ + private String inferSqlTypeFromValue(Number value) { + if (value instanceof Integer || value instanceof Short || value instanceof Byte) { + return "integer"; + } else if (value instanceof Long) { + return "bigint"; + } else { + // Float, Double, BigDecimal - use double precision for safety + return "double precision"; + } + } + + /** + * Builds a PostgreSQL text array path from nested path components. For example, ["seller", + * "count"] becomes "{seller,count}" + */ + private String buildJsonPath(String[] nestedPath) { + return "{" + String.join(",", nestedPath) + "}"; + } +}