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 138229ee3..474dbbe0b 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 @@ -36,6 +36,7 @@ import org.hypertrace.core.documentstore.model.options.ReturnDocumentType; import org.hypertrace.core.documentstore.model.options.UpdateOptions; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; +import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue; import org.hypertrace.core.documentstore.model.subdoc.UpdateOperator; import org.hypertrace.core.documentstore.postgres.PostgresDatastore; import org.hypertrace.core.documentstore.query.Query; @@ -1905,9 +1906,543 @@ class SubDocUpdateTests { class SetOperatorTests { @Test - @DisplayName("Should update top-level column with SET operator") - void testUpdateTopLevelColumn() throws Exception { - // Update the price of item with id = 1 + @DisplayName("Should update multiple top-level columns in single update") + void testSetMultipleColumns() throws Exception { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("2"))) + .build(); + + List updates = + List.of(SubDocumentUpdate.of("price", 555), SubDocumentUpdate.of("quantity", 100)); + + 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(555, resultJson.get("price").asInt()); + assertEquals(100, resultJson.get("quantity").asInt()); + + // Verify in database + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"price\", \"quantity\" FROM \"%s\" WHERE \"id\" = '2'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals(555, rs.getInt("price")); + assertEquals(100, rs.getInt("quantity")); + } + } + + @Test + @DisplayName("Should update nested path in JSONB column") + void testUpdateNestedJsonbPath() throws Exception { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("3"))) + .build(); + + // Update props.brand nested path + List updates = + List.of(SubDocumentUpdate.of("props.brand", "UpdatedBrand")); + + 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()); + assertNotNull(resultJson.get("props")); + assertEquals("UpdatedBrand", resultJson.get("props").get("brand").asText()); + + // Verify in database + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"props\"->>'brand' as brand FROM \"%s\" WHERE \"id\" = '3'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("UpdatedBrand", rs.getString("brand")); + } + } + + @Test + @DisplayName("Should return BEFORE_UPDATE document") + void testUpdateReturnsBeforeDocument() throws Exception { + // First get the current price + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("4"))) + .build(); + + List updates = List.of(SubDocumentUpdate.of("price", 777)); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.BEFORE_UPDATE).build(); + + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + // Should return the old price (5 from initial data), not the new one (777) + assertEquals(5, resultJson.get("price").asInt()); + + // But database should have the new value + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"price\" FROM \"%s\" WHERE \"id\" = '4'", FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals(777, rs.getInt("price")); + } + } + + @Test + @DisplayName("Case 1: SET on field not in schema should skip (default SKIP strategy)") + void testSetFieldNotInSchema() throws Exception { + // Update a field that doesn't exist in the schema + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + SubDocumentUpdate update = + SubDocumentUpdate.builder() + .subDocument("nonexistent_column.some_key") + .operator(UpdateOperator.SET) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of("new_value")) + .build(); + + // With default SKIP strategy, this should not throw but skip the update + Optional result = + flatCollection.update( + query, + List.of(update), + UpdateOptions.builder() + .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) + .build()); + + // Document should still be returned (unchanged since update was skipped) + assertTrue(result.isPresent()); + + // Verify the document wasn't modified (item should still be "Soap") + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"item\" FROM \"%s\" WHERE \"id\" = '1'", FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("Soap", rs.getString("item")); + } + } + + @Test + @DisplayName("Case 2: SET on JSONB column that is NULL should create the structure") + void testSetJsonbColumnIsNull() throws Exception { + // Row 2 has props = NULL + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("2"))) + .build(); + + SubDocumentUpdate update = + SubDocumentUpdate.builder() + .subDocument("props.newKey") + .operator(UpdateOperator.SET) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of("newValue")) + .build(); + + Optional result = + flatCollection.update( + query, + List.of(update), + UpdateOptions.builder() + .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) + .build()); + + assertTrue(result.isPresent()); + + // Verify props now has the new key + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"props\"->>'newKey' as newKey FROM \"%s\" WHERE \"id\" = '2'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("newValue", rs.getString("newKey")); + } + } + + @Test + @DisplayName("Case 3: SET on JSONB path that exists should update the value") + void testSetJsonbPathExists() throws Exception { + // Row 1 has props.brand = "Dettol" + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + SubDocumentUpdate update = + SubDocumentUpdate.builder() + .subDocument("props.brand") + .operator(UpdateOperator.SET) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + "UpdatedBrand")) + .build(); + + Optional result = + flatCollection.update( + query, + List.of(update), + UpdateOptions.builder() + .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) + .build()); + + assertTrue(result.isPresent()); + + // Verify props.brand was updated + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"props\"->>'brand' as brand FROM \"%s\" WHERE \"id\" = '1'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("UpdatedBrand", rs.getString("brand")); + } + } + + @Test + @DisplayName("Case 4: SET on JSONB path that doesn't exist should create the key") + void testSetJsonbPathDoesNotExist() throws Exception { + // Row 1 has props but no "newAttribute" key + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + SubDocumentUpdate update = + SubDocumentUpdate.builder() + .subDocument("props.newAttribute") + .operator(UpdateOperator.SET) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + "brandNewValue")) + .build(); + + Optional result = + flatCollection.update( + query, + List.of(update), + UpdateOptions.builder() + .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) + .build()); + + assertTrue(result.isPresent()); + + // Verify props.newAttribute was created + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"props\"->>'newAttribute' as newAttr, \"props\"->>'brand' as brand FROM \"%s\" WHERE \"id\" = '1'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("brandNewValue", rs.getString("newAttr")); + // Verify existing data wasn't lost + assertEquals("Dettol", rs.getString("brand")); + } + } + + @Test + @DisplayName("SET on top-level column should update the value directly") + void testSetTopLevelColumn() throws Exception { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + SubDocumentUpdate update = + SubDocumentUpdate.builder() + .subDocument("item") + .operator(UpdateOperator.SET) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + "UpdatedSoap")) + .build(); + + Optional result = + flatCollection.update( + query, + List.of(update), + UpdateOptions.builder() + .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) + .build()); + + assertTrue(result.isPresent()); + + // Verify item was updated + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"item\" FROM \"%s\" WHERE \"id\" = '1'", FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("UpdatedSoap", rs.getString("item")); + } + } + + @Test + @DisplayName("SET with empty object value") + void testSetWithEmptyObjectValue() throws Exception { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + // SET a JSON object containing an empty object + SubDocumentUpdate update = + SubDocumentUpdate.builder() + .subDocument("props.newProperty") + .operator(UpdateOperator.SET) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + new JSONDocument( + Map.of("hello", "world", "emptyObject", Collections.emptyMap())))) + .build(); + + Optional result = + flatCollection.update( + query, + List.of(update), + UpdateOptions.builder() + .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) + .build()); + + assertTrue(result.isPresent()); + + // Verify the JSON object was set correctly + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"props\"->'newProperty' as newProp FROM \"%s\" WHERE \"id\" = '1'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + String jsonStr = rs.getString("newProp"); + assertNotNull(jsonStr); + assertTrue(jsonStr.contains("hello")); + assertTrue(jsonStr.contains("emptyObject")); + } + } + + @Test + @DisplayName("SET with JSON document as value") + void testSetWithJsonDocumentValue() throws Exception { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + SubDocumentUpdate update = + SubDocumentUpdate.builder() + .subDocument("props.nested") + .operator(UpdateOperator.SET) + .subDocumentValue( + SubDocumentValue.of(new JSONDocument(Map.of("key1", "value1", "key2", 123)))) + .build(); + + Optional result = + flatCollection.update( + query, + List.of(update), + UpdateOptions.builder() + .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) + .build()); + + assertTrue(result.isPresent()); + + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"props\"->'nested'->>'key1' as key1, \"props\"->'nested'->>'key2' as key2 FROM \"%s\" WHERE \"id\" = '1'", + FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals("value1", rs.getString("key1")); + assertEquals("123", rs.getString("key2")); + } + } + } + + @Nested + @DisplayName("UNSET Operator Tests") + class UnsetOperatorTests { + + @Test + @DisplayName("Should UNSET top-level column (set to NULL)") + void testUnsetTopLevelColumn() throws Exception { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("item") + .operator(UpdateOperator.UNSET) + .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()); + + JsonNode itemNode = resultJson.get("item"); + assertTrue(itemNode == null || itemNode.isNull()); + + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"item\" FROM \"%s\" WHERE \"id\" = '1'", FLAT_COLLECTION_NAME)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertNull(rs.getString("item")); + } + } + + @Test + @DisplayName("Should UNSET nested JSONB field (remove key)") + void testUnsetNestedJsonbField() throws Exception { + String docId = getRandomDocId(4); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "JsonbItem"); + ObjectNode props = OBJECT_MAPPER.createObjectNode(); + props.put("brand", "TestBrand"); + props.put("color", "Red"); + node.set("props", props); + flatCollection.create(key, new JSONDocument(node)); + + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of(key.toString()))) + .build(); + + // UNSET props.brand + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.brand") + .operator(UpdateOperator.UNSET) + .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()); + assertFalse(resultJson.get("props").has("brand")); + assertEquals("Red", resultJson.get("props").get("color").asText()); + } + } + + @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( @@ -1917,7 +2452,15 @@ void testUpdateTopLevelColumn() throws Exception { ConstantExpression.of("1"))) .build(); - List updates = List.of(SubDocumentUpdate.of("price", 999)); + // 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(); @@ -1926,9 +2469,8 @@ void testUpdateTopLevelColumn() throws Exception { assertTrue(result.isPresent()); JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - assertEquals(999, resultJson.get("price").asInt()); + assertEquals(15, resultJson.get("price").asInt()); - // Verify in database PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; try (Connection conn = pgDatastore.getPostgresClient(); PreparedStatement ps = @@ -1937,13 +2479,54 @@ void testUpdateTopLevelColumn() throws Exception { "SELECT \"price\" FROM \"%s\" WHERE \"id\" = '1'", FLAT_COLLECTION_NAME)); ResultSet rs = ps.executeQuery()) { assertTrue(rs.next()); - assertEquals(999, rs.getInt("price")); + assertEquals(15, rs.getInt("price")); } } @Test - @DisplayName("Should update multiple top-level columns in single update") - void testUpdateMultipleColumns() throws Exception { + @DisplayName("Should handle ADD on NULL column (treat as 0)") + void testAddOnNullColumn() throws Exception { + // Create a document with NULL price + String docId = getRandomDocId(4); + 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( @@ -1953,8 +2536,15 @@ void testUpdateMultipleColumns() throws Exception { ConstantExpression.of("2"))) .build(); + // ADD -5 to price (20 - 5 = 15) List updates = - List.of(SubDocumentUpdate.of("price", 555), SubDocumentUpdate.of("quantity", 100)); + 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(); @@ -1963,13 +2553,13 @@ void testUpdateMultipleColumns() throws Exception { assertTrue(result.isPresent()); JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - assertEquals(555, resultJson.get("price").asInt()); - assertEquals(100, resultJson.get("quantity").asInt()); + assertEquals(15, resultJson.get("price").asInt()); } @Test - @DisplayName("Should update nested path in JSONB column") - void testUpdateNestedJsonbPath() throws Exception { + @DisplayName("Should ADD with floating point value") + void testAddFloatingPointValue() throws Exception { + // Row 3 has price = 30 Query query = Query.builder() .setFilter( @@ -1979,9 +2569,16 @@ void testUpdateNestedJsonbPath() throws Exception { ConstantExpression.of("3"))) .build(); - // Update props.brand nested path + // 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.of("props.brand", "UpdatedBrand")); + 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(); @@ -1990,57 +2587,193 @@ void testUpdateNestedJsonbPath() throws Exception { assertTrue(result.isPresent()); JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - assertNotNull(resultJson.get("props")); - assertEquals("UpdatedBrand", resultJson.get("props").get("brand").asText()); + // Initial weight is NULL, so COALESCE(NULL, 0) + 2.5 = 2.5 + assertEquals(2.5, resultJson.get("weight").asDouble(), 0.01); } @Test - @DisplayName("Should return BEFORE_UPDATE document") - void testUpdateReturnsBeforeDocument() throws Exception { - // First get the current price + @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 = getRandomDocId(4); + 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("4"))) + ConstantExpression.of(key.toString()))) .build(); - List updates = List.of(SubDocumentUpdate.of("price", 777)); + // 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 = getRandomDocId(4); + 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 throw IllegalArgumentException for multi-valued primitive value") + void testAddMultiValuedPrimitiveValue() { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + // ADD with an array of numbers should fail + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + new Integer[] {1, 2, 3})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + assertThrows( + IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); + } + + @Test + @DisplayName("Should throw IllegalArgumentException for nested document value") + void testAddNestedDocumentValue() throws Exception { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("1"))) + .build(); + + // ADD with a nested document should fail + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue( + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + new JSONDocument("{\"nested\": 123}"))) + .build()); UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.BEFORE_UPDATE).build(); - - Optional result = flatCollection.update(query, updates, options); - - assertTrue(result.isPresent()); - JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - // Should return the old price (5 from initial data), not the new one (777) - assertEquals(5, resultJson.get("price").asInt()); + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - // But database should have the new value - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"price\" FROM \"%s\" WHERE \"id\" = '4'", FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals(777, rs.getInt("price")); - } + assertThrows( + IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); } - } - - @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 + @DisplayName("Should throw IllegalArgumentException for multi-valued nested document value") + void testAddMultiValuedNestedDocumentValue() throws Exception { Query query = Query.builder() .setFilter( @@ -2050,47 +2783,35 @@ void testAddTopLevelColumn() throws Exception { ConstantExpression.of("1"))) .build(); - // ADD 5 to price (10 + 5 = 15) + // ADD with an array of documents should fail List updates = List.of( SubDocumentUpdate.builder() .subDocument("price") .operator(UpdateOperator.ADD) .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(5)) + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + new Document[] { + new JSONDocument("{\"a\": 1}"), new JSONDocument("{\"b\": 2}") + })) .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")); - } + assertThrows( + IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); } @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"; + @DisplayName("Should ADD to BIGINT column with correct type cast") + void testAddBigintColumn() throws Exception { + // Create a document with big_number set + String docId = getRandomDocId(4); Key key = new SingleValueKey(DEFAULT_TENANT, docId); ObjectNode node = OBJECT_MAPPER.createObjectNode(); - node.put("item", "NullPriceItem"); - // price is not set, will be NULL + node.put("item", "BigintItem"); + node.put("big_number", 1000000000000L); flatCollection.create(key, new JSONDocument(node)); Query query = @@ -2102,14 +2823,14 @@ void testAddOnNullColumn() throws Exception { ConstantExpression.of(key.toString()))) .build(); - // ADD 100 to NULL price (COALESCE(NULL, 0) + 100 = 100) + // ADD 500 to big_number List updates = List.of( SubDocumentUpdate.builder() - .subDocument("price") + .subDocument("big_number") .operator(UpdateOperator.ADD) .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(100)) + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(500L)) .build()); UpdateOptions options = @@ -2119,30 +2840,37 @@ void testAddOnNullColumn() throws Exception { assertTrue(result.isPresent()); JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - assertEquals(100, resultJson.get("price").asInt()); + assertEquals(1000000000500L, resultJson.get("big_number").asLong()); } @Test - @DisplayName("Should ADD with negative value (decrement)") - void testAddNegativeValue() throws Exception { - // Row 2 has price = 20 + @DisplayName("Should ADD to REAL column with correct type cast") + void testAddRealColumn() throws Exception { + // Create a document with rating set + String docId = getRandomDocId(4); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "RealItem"); + node.put("rating", 3.5); + flatCollection.create(key, new JSONDocument(node)); + Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of("2"))) + ConstantExpression.of(key.toString()))) .build(); - // ADD -5 to price (20 - 5 = 15) + // ADD 1.0 to rating (3.5 + 1.0 = 4.5) List updates = List.of( SubDocumentUpdate.builder() - .subDocument("price") + .subDocument("rating") .operator(UpdateOperator.ADD) .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(-5)) + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(1.0)) .build()); UpdateOptions options = @@ -2151,32 +2879,55 @@ void testAddNegativeValue() throws Exception { 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 directly + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + try (Connection conn = pgDatastore.getPostgresClient(); + PreparedStatement ps = + conn.prepareStatement( + String.format( + "SELECT \"rating\" FROM \"%s\" WHERE \"id\" = '%s'", + FLAT_COLLECTION_NAME, key)); + ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + assertEquals(4.5f, rs.getFloat("rating"), 0.01f); + } } + } + + @Nested + @DisplayName("APPEND_TO_LIST Operator Tests") + class AppendToListOperatorTests { @Test - @DisplayName("Should ADD with floating point value") - void testAddFloatingPointValue() throws Exception { - // Row 3 has price = 30 + @DisplayName("Should append values to top-level array column") + void testAppendToTopLevelArray() throws Exception { + // Create a document with known tags for predictable testing + String docId = getRandomDocId(4); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "TestItem"); + node.putArray("tags").add("tag1").add("tag2"); + flatCollection.create(key, new JSONDocument(node)); + Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of("3"))) + ConstantExpression.of(key.toString()))) .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 + // Append new tags List updates = List.of( SubDocumentUpdate.builder() - .subDocument("weight") - .operator(UpdateOperator.ADD) + .subDocument("tags") + .operator(UpdateOperator.APPEND_TO_LIST) .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(2.5)) + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + new String[] {"newTag1", "newTag2"})) .build()); UpdateOptions options = @@ -2186,22 +2937,24 @@ void testAddFloatingPointValue() throws Exception { 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); + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(4, tagsNode.size()); + assertEquals("newTag1", tagsNode.get(2).asText()); + assertEquals("newTag2", tagsNode.get(3).asText()); } @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"; + @DisplayName("Should append values to nested JSONB array") + void testAppendToNestedJsonbArray() throws Exception { + // Set up a document with JSONB containing an array + String docId = getRandomDocId(4); 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); + node.put("item", "JsonbArrayItem"); + ObjectNode props = OBJECT_MAPPER.createObjectNode(); + props.putArray("colors").add("red").add("blue"); + node.set("props", props); flatCollection.create(key, new JSONDocument(node)); Query query = @@ -2213,14 +2966,15 @@ void testAddNestedJsonbField() throws Exception { ConstantExpression.of(key.toString()))) .build(); - // ADD 50 to sales.total (100 + 50 = 150) + // Append to props.colors List updates = List.of( SubDocumentUpdate.builder() - .subDocument("sales.total") - .operator(UpdateOperator.ADD) + .subDocument("props.colors") + .operator(UpdateOperator.APPEND_TO_LIST) .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(50)) + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + new String[] {"green", "yellow"})) .build()); UpdateOptions options = @@ -2230,23 +2984,23 @@ void testAddNestedJsonbField() throws Exception { 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()); + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(4, colorsNode.size()); } @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"; + @DisplayName("Should create list when appending to non-existent JSONB array") + void testAppendToNonExistentJsonbArray() throws Exception { + // Create a document with props but NO colors array + String docId = getRandomDocId(4); 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); + node.put("item", "ItemWithoutColors"); + ObjectNode props = OBJECT_MAPPER.createObjectNode(); + props.put("brand", "TestBrand"); + // Note: no colors array in props + node.set("props", props); flatCollection.create(key, new JSONDocument(node)); Query query = @@ -2258,14 +3012,13 @@ void testAddNestedJsonbFieldNotExists() throws Exception { ConstantExpression.of(key.toString()))) .build(); - // ADD 75 to sales.total (non-existent, should become 0 + 75 = 75) + // Append to props.colors which doesn't exist List updates = List.of( SubDocumentUpdate.builder() - .subDocument("sales.total") - .operator(UpdateOperator.ADD) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(75)) + .subDocument("props.colors") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"green", "yellow"})) .build()); UpdateOptions options = @@ -2275,142 +3028,225 @@ void testAddNestedJsonbFieldNotExists() throws Exception { 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()); + + // Should create the array with the appended values + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertNotNull(colorsNode, "colors array should be created"); + assertTrue(colorsNode.isArray()); + assertEquals(2, colorsNode.size()); + assertEquals("green", colorsNode.get(0).asText()); + assertEquals("yellow", colorsNode.get(1).asText()); + + assertEquals("TestBrand", resultJson.get("props").get("brand").asText()); } + } + + @Nested + @DisplayName("ADD_TO_LIST_IF_ABSENT Operator Tests") + class AddToListIfAbsentOperatorTests { @Test - @DisplayName("Should throw IllegalArgumentException for non-numeric value") - void testAddNonNumericValue() { + @DisplayName("Should add unique values to top-level array column") + void testAddToListIfAbsentTopLevel() throws Exception { + String docId = getRandomDocId(4); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "TestItem"); + node.putArray("tags").add("existing1").add("existing2"); + flatCollection.create(key, new JSONDocument(node)); + Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of("1"))) + ConstantExpression.of(key.toString()))) .build(); - // ADD with a string value should fail + // Add tags - 'existing1' already exists, 'newTag' is new List updates = List.of( SubDocumentUpdate.builder() - .subDocument("price") - .operator(UpdateOperator.ADD) + .subDocument("tags") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) .subDocumentValue( org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - "not-a-number")) + new String[] {"existing1", "newTag"})) .build()); UpdateOptions options = UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - assertThrows( - IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(3, tagsNode.size()); // original 2 + 1 new unique + + // Verify 'newTag' was added + boolean hasNewTag = false; + for (JsonNode tag : tagsNode) { + if ("newTag".equals(tag.asText())) { + hasNewTag = true; + break; + } + } + assertTrue(hasNewTag); } @Test - @DisplayName("Should throw IllegalArgumentException for multi-valued primitive value") - void testAddMultiValuedPrimitiveValue() { + @DisplayName("Should add unique values to nested JSONB array") + void testAddToListIfAbsentNestedJsonb() throws Exception { + // Set up a document with JSONB containing an array + String docId = getRandomDocId(4); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "JsonbArrayItem"); + ObjectNode props = OBJECT_MAPPER.createObjectNode(); + props.putArray("colors").add("red").add("blue"); + node.set("props", props); + flatCollection.create(key, new JSONDocument(node)); + Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of("1"))) + ConstantExpression.of(key.toString()))) .build(); - // ADD with an array of numbers should fail + // Add colors - 'red' already exists, 'green' is new List updates = List.of( SubDocumentUpdate.builder() - .subDocument("price") - .operator(UpdateOperator.ADD) + .subDocument("props.colors") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) .subDocumentValue( org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new Integer[] {1, 2, 3})) + new String[] {"red", "green"})) .build()); UpdateOptions options = UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - assertThrows( - IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(3, colorsNode.size()); + assertEquals("red", colorsNode.get(0).asText()); + assertEquals("blue", colorsNode.get(1).asText()); + assertEquals("green", colorsNode.get(2).asText()); } @Test - @DisplayName("Should throw IllegalArgumentException for nested document value") - void testAddNestedDocumentValue() throws Exception { + @DisplayName("Should not add duplicates when all values already exist") + void testAddToListIfAbsentNoDuplicates() throws Exception { + String docId = getRandomDocId(4); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "TestItem"); + node.putArray("tags").add("tag1").add("tag2"); + flatCollection.create(key, new JSONDocument(node)); + Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of("1"))) + ConstantExpression.of(key.toString()))) .build(); - // ADD with a nested document should fail + // Add tags that already exist List updates = List.of( SubDocumentUpdate.builder() - .subDocument("price") - .operator(UpdateOperator.ADD) + .subDocument("tags") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) .subDocumentValue( org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new JSONDocument("{\"nested\": 123}"))) + new String[] {"tag1", "tag2"})) .build()); UpdateOptions options = UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - assertThrows( - IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(2, tagsNode.size()); + assertEquals("tag1", tagsNode.get(0).asText()); + assertEquals("tag2", tagsNode.get(1).asText()); } + } + + @Nested + @DisplayName("REMOVE_ALL_FROM_LIST Operator Tests") + class RemoveAllFromListOperatorTests { @Test - @DisplayName("Should throw IllegalArgumentException for multi-valued nested document value") - void testAddMultiValuedNestedDocumentValue() throws Exception { + @DisplayName("Should remove values from top-level array column") + void testRemoveAllFromTopLevelArray() throws Exception { + String docId = getRandomDocId(4); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "TestItem"); + node.putArray("tags").add("tag1").add("tag2").add("tag3"); + flatCollection.create(key, new JSONDocument(node)); + Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of("1"))) + ConstantExpression.of(key.toString()))) .build(); - // ADD with an array of documents should fail + // Remove 'tag1' from tags List updates = List.of( SubDocumentUpdate.builder() - .subDocument("price") - .operator(UpdateOperator.ADD) + .subDocument("tags") + .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) .subDocumentValue( org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new Document[] { - new JSONDocument("{\"a\": 1}"), new JSONDocument("{\"b\": 2}") - })) + new String[] {"tag1"})) .build()); UpdateOptions options = UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - assertThrows( - IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(2, tagsNode.size()); // 'tag2' and 'tag3' remain } @Test - @DisplayName("Should ADD to BIGINT column with correct type cast") - void testAddBigintColumn() throws Exception { - // Create a document with big_number set - String docId = "add-bigint-test"; + @DisplayName("Should remove values from nested JSONB array") + void testRemoveAllFromNestedJsonbArray() throws Exception { + // Set up a document with JSONB containing an array + String docId = getRandomDocId(4); Key key = new SingleValueKey(DEFAULT_TENANT, docId); ObjectNode node = OBJECT_MAPPER.createObjectNode(); - node.put("item", "BigintItem"); - node.put("big_number", 1000000000000L); + node.put("item", "JsonbArrayItem"); + ObjectNode props = OBJECT_MAPPER.createObjectNode(); + props.putArray("colors").add("red").add("blue").add("green"); + node.set("props", props); flatCollection.create(key, new JSONDocument(node)); Query query = @@ -2422,14 +3258,15 @@ void testAddBigintColumn() throws Exception { ConstantExpression.of(key.toString()))) .build(); - // ADD 500 to big_number + // Remove 'red' and 'blue' from props.colors List updates = List.of( SubDocumentUpdate.builder() - .subDocument("big_number") - .operator(UpdateOperator.ADD) + .subDocument("props.colors") + .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(500L)) + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + new String[] {"red", "blue"})) .build()); UpdateOptions options = @@ -2439,18 +3276,19 @@ void testAddBigintColumn() throws Exception { assertTrue(result.isPresent()); JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - assertEquals(1000000000500L, resultJson.get("big_number").asLong()); + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(1, colorsNode.size()); // Only 'green' remains } @Test - @DisplayName("Should ADD to REAL column with correct type cast") - void testAddRealColumn() throws Exception { - // Create a document with rating set - String docId = "add-real-test"; + @DisplayName("Should handle removing non-existent values (no-op)") + void testRemoveNonExistentValues() throws Exception { + String docId = getRandomDocId(4); Key key = new SingleValueKey(DEFAULT_TENANT, docId); ObjectNode node = OBJECT_MAPPER.createObjectNode(); - node.put("item", "RealItem"); - node.put("rating", 3.5); + node.put("item", "TestItem"); + node.putArray("tags").add("tag1").add("tag2"); flatCollection.create(key, new JSONDocument(node)); Query query = @@ -2462,14 +3300,15 @@ void testAddRealColumn() throws Exception { ConstantExpression.of(key.toString()))) .build(); - // ADD 1.0 to rating (3.5 + 1.0 = 4.5) + // Try to remove values that don't exist List updates = List.of( SubDocumentUpdate.builder() - .subDocument("rating") - .operator(UpdateOperator.ADD) + .subDocument("tags") + .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(1.0)) + org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( + new String[] {"nonexistent1", "nonexistent2"})) .build()); UpdateOptions options = @@ -2478,19 +3317,10 @@ void testAddRealColumn() throws Exception { Optional result = flatCollection.update(query, updates, options); assertTrue(result.isPresent()); - - // Verify in database directly - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"rating\" FROM \"%s\" WHERE \"id\" = '%s'", - FLAT_COLLECTION_NAME, key)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals(4.5f, rs.getFloat("rating"), 0.01f); - } + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(2, tagsNode.size()); // No change } } @@ -2557,33 +3387,6 @@ void testUpdateNestedPathOnNonJsonbColumn() { assertThrows(IOException.class, () -> flatCollection.update(query, updates, options)); } - - @Test - @DisplayName("Should throw IOException for unsupported operator") - void testUpdateUnsupportedOperator() { - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("_id"), - RelationalOperator.EQ, - ConstantExpression.of(1))) - .build(); - - // UNSET is not supported yet - List updates = - List.of( - SubDocumentUpdate.builder() - .subDocument("price") - .operator(UpdateOperator.UNSET) - .build()); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - assertThrows( - IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); - } } @Nested @@ -2930,60 +3733,6 @@ void testBulkUpdateNonExistentColumnWithThrowStrategy() { } } - @Nested - @DisplayName("Drop Operations") - class DropTests { - - @Test - @DisplayName("Should throw UnsupportedOperationException for drop") - void testDrop() { - assertThrows(UnsupportedOperationException.class, () -> flatCollection.drop()); - } - } - - @Nested - @DisplayName("Sub-Document Operations") - class SubDocumentTests { - - @Test - @DisplayName("Should throw UnsupportedOperationException for updateSubDoc") - void testSubDocumentUpdate() { - Key docKey = new SingleValueKey("default", "1"); - ObjectNode subDoc = OBJECT_MAPPER.createObjectNode(); - subDoc.put("newField", "newValue"); - Document subDocument = new JSONDocument(subDoc); - - assertThrows( - UnsupportedOperationException.class, - () -> flatCollection.updateSubDoc(docKey, "props.nested", subDocument)); - } - - @Test - @DisplayName("Should throw UnsupportedOperationException for deleteSubDoc") - void testSubDocumentDelete() { - Key docKey = new SingleValueKey("default", "1"); - - assertThrows( - UnsupportedOperationException.class, - () -> flatCollection.deleteSubDoc(docKey, "props.brand")); - } - - @Test - @DisplayName("Should throw UnsupportedOperationException for bulkUpdateSubDocs") - void testBulkUpdateSubDocs() { - Map> documents = new HashMap<>(); - Key key1 = new SingleValueKey("default", "1"); - Map subDocs1 = new HashMap<>(); - ObjectNode subDoc1 = OBJECT_MAPPER.createObjectNode(); - subDoc1.put("updated", true); - subDocs1.put("props.status", new JSONDocument(subDoc1)); - documents.put(key1, subDocs1); - - assertThrows( - UnsupportedOperationException.class, () -> flatCollection.bulkUpdateSubDocs(documents)); - } - } - @Nested @DisplayName("Bulk Array Value Operations") class BulkArrayValueOperationTests { @@ -3072,318 +3821,13 @@ void testCreateOrReplaceRefreshesSchemaOnDroppedColumn() throws Exception { } @Nested - @DisplayName("Update SET Operator Tests") - class UpdateSetOperatorTests { - - @Test - @DisplayName("Case 1: SET on field not in schema should skip (default SKIP strategy)") - void testSetFieldNotInSchema() throws Exception { - // Update a field that doesn't exist in the schema - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("1"))) - .build(); - - SubDocumentUpdate update = - SubDocumentUpdate.builder() - .subDocument("nonexistent_column.some_key") - .operator(UpdateOperator.SET) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of("new_value")) - .build(); - - // With default SKIP strategy, this should not throw but skip the update - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build()); - - // Document should still be returned (unchanged since update was skipped) - assertTrue(result.isPresent()); - - // Verify the document wasn't modified (item should still be "Soap") - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"item\" FROM \"%s\" WHERE \"id\" = '1'", FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals("Soap", rs.getString("item")); - } - } - - @Test - @DisplayName("Case 2: SET on JSONB column that is NULL should create the structure") - void testSetJsonbColumnIsNull() throws Exception { - // Row 2 has props = NULL - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("2"))) - .build(); - - SubDocumentUpdate update = - SubDocumentUpdate.builder() - .subDocument("props.newKey") - .operator(UpdateOperator.SET) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of("newValue")) - .build(); - - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build()); - - assertTrue(result.isPresent()); - - // Verify props now has the new key - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"props\"->>'newKey' as newKey FROM \"%s\" WHERE \"id\" = '2'", - FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals("newValue", rs.getString("newKey")); - } - } - - @Test - @DisplayName("Case 3: SET on JSONB path that exists should update the value") - void testSetJsonbPathExists() throws Exception { - // Row 1 has props.brand = "Dettol" - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("1"))) - .build(); - - SubDocumentUpdate update = - SubDocumentUpdate.builder() - .subDocument("props.brand") - .operator(UpdateOperator.SET) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - "UpdatedBrand")) - .build(); - - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build()); - - assertTrue(result.isPresent()); - - // Verify props.brand was updated - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"props\"->>'brand' as brand FROM \"%s\" WHERE \"id\" = '1'", - FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals("UpdatedBrand", rs.getString("brand")); - } - } - - @Test - @DisplayName("Case 4: SET on JSONB path that doesn't exist should create the key") - void testSetJsonbPathDoesNotExist() throws Exception { - // Row 1 has props but no "newAttribute" key - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("1"))) - .build(); - - SubDocumentUpdate update = - SubDocumentUpdate.builder() - .subDocument("props.newAttribute") - .operator(UpdateOperator.SET) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - "brandNewValue")) - .build(); - - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build()); - - assertTrue(result.isPresent()); - - // Verify props.newAttribute was created - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"props\"->>'newAttribute' as newAttr, \"props\"->>'brand' as brand FROM \"%s\" WHERE \"id\" = '1'", - FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals("brandNewValue", rs.getString("newAttr")); - // Verify existing data wasn't lost - assertEquals("Dettol", rs.getString("brand")); - } - } - - @Test - @DisplayName("SET on top-level column should update the value directly") - void testSetTopLevelColumn() throws Exception { - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("1"))) - .build(); - - SubDocumentUpdate update = - SubDocumentUpdate.builder() - .subDocument("item") - .operator(UpdateOperator.SET) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of("UpdatedSoap")) - .build(); - - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build()); - - assertTrue(result.isPresent()); - - // Verify item was updated - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"item\" FROM \"%s\" WHERE \"id\" = '1'", FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals("UpdatedSoap", rs.getString("item")); - } - } - - @Test - @DisplayName("SET with empty object value") - void testSetWithEmptyObjectValue() throws Exception { - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("1"))) - .build(); - - // SET a JSON object containing an empty object - SubDocumentUpdate update = - SubDocumentUpdate.builder() - .subDocument("props.newProperty") - .operator(UpdateOperator.SET) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new JSONDocument( - Map.of("hello", "world", "emptyObject", Collections.emptyMap())))) - .build(); - - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build()); - - assertTrue(result.isPresent()); - - // Verify the JSON object was set correctly - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"props\"->'newProperty' as newProp FROM \"%s\" WHERE \"id\" = '1'", - FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - String jsonStr = rs.getString("newProp"); - assertNotNull(jsonStr); - assertTrue(jsonStr.contains("hello")); - assertTrue(jsonStr.contains("emptyObject")); - } - } + @DisplayName("Drop Operations") + class DropTests { @Test - @DisplayName("SET with JSON document as value") - void testSetWithJsonDocumentValue() throws Exception { - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("1"))) - .build(); - - // SET a JSON document as value - SubDocumentUpdate update = - SubDocumentUpdate.builder() - .subDocument("props.nested") - .operator(UpdateOperator.SET) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new JSONDocument(Map.of("key1", "value1", "key2", 123)))) - .build(); - - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build()); - - assertTrue(result.isPresent()); - - // Verify the JSON document was set correctly - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"props\"->'nested'->>'key1' as key1, \"props\"->'nested'->>'key2' as key2 FROM \"%s\" WHERE \"id\" = '1'", - FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals("value1", rs.getString("key1")); - assertEquals("123", rs.getString("key2")); - } + @DisplayName("Should throw UnsupportedOperationException for drop") + void testDrop() { + assertThrows(UnsupportedOperationException.class, () -> flatCollection.drop()); } } } diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoFlatPgConsistencyTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoFlatPgConsistencyTest.java new file mode 100644 index 000000000..9461cf69d --- /dev/null +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoFlatPgConsistencyTest.java @@ -0,0 +1,749 @@ +package org.hypertrace.core.documentstore; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import java.io.IOException; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.stream.Stream; +import org.hypertrace.core.documentstore.expression.impl.ConstantExpression; +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.model.options.ReturnDocumentType; +import org.hypertrace.core.documentstore.model.options.UpdateOptions; +import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; +import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue; +import org.hypertrace.core.documentstore.model.subdoc.UpdateOperator; +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.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.extension.ExtensionContext; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.ArgumentsProvider; +import org.junit.jupiter.params.provider.ArgumentsSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.wait.strategy.Wait; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.utility.DockerImageName; + +@Testcontainers +public class MongoFlatPgConsistencyTest { + + private static final Logger LOGGER = LoggerFactory.getLogger(MongoFlatPgConsistencyTest.class); + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + private static final String COLLECTION_NAME = "consistency_test"; + private static final String DEFAULT_TENANT = "default"; + private static final String MONGO_STORE = "Mongo"; + private static final String POSTGRES_FLAT_STORE = "PostgresFlat"; + + private static Map datastoreMap; + private static Map collectionMap; + + private static GenericContainer mongo; + private static GenericContainer postgres; + + @BeforeAll + public static void init() throws IOException { + datastoreMap = new HashMap<>(); + collectionMap = new HashMap<>(); + + // Start MongoDB + mongo = + new GenericContainer<>(DockerImageName.parse("mongo:8.0.1")) + .withExposedPorts(27017) + .waitingFor(Wait.forListeningPort()); + mongo.start(); + + Map mongoConfig = new HashMap<>(); + mongoConfig.put("host", "localhost"); + mongoConfig.put("port", mongo.getMappedPort(27017).toString()); + Config mongoCfg = ConfigFactory.parseMap(mongoConfig); + + Datastore mongoDatastore = DatastoreProvider.getDatastore("Mongo", mongoCfg); + datastoreMap.put(MONGO_STORE, mongoDatastore); + + // Start PostgreSQL + postgres = + new GenericContainer<>(DockerImageName.parse("postgres:13.1")) + .withEnv("POSTGRES_PASSWORD", "postgres") + .withEnv("POSTGRES_USER", "postgres") + .withExposedPorts(5432) + .waitingFor(Wait.forListeningPort()); + postgres.start(); + + String postgresConnectionUrl = + String.format("jdbc:postgresql://localhost:%s/", postgres.getMappedPort(5432)); + + Map postgresConfig = new HashMap<>(); + postgresConfig.put("url", postgresConnectionUrl); + postgresConfig.put("user", "postgres"); + postgresConfig.put("password", "postgres"); + + Datastore postgresDatastore = + DatastoreProvider.getDatastore("Postgres", ConfigFactory.parseMap(postgresConfig)); + datastoreMap.put(POSTGRES_FLAT_STORE, postgresDatastore); + + // Create Postgres flat collection schema + createFlatCollectionSchema((PostgresDatastore) postgresDatastore); + + // Create collections + mongoDatastore.deleteCollection(COLLECTION_NAME); + mongoDatastore.createCollection(COLLECTION_NAME, null); + collectionMap.put(MONGO_STORE, mongoDatastore.getCollection(COLLECTION_NAME)); + collectionMap.put( + POSTGRES_FLAT_STORE, + postgresDatastore.getCollectionForType(COLLECTION_NAME, DocumentType.FLAT)); + + LOGGER.info("Test setup complete. Collections ready for both Mongo and PostgresFlat."); + } + + private static void createFlatCollectionSchema(PostgresDatastore pgDatastore) { + String createTableSQL = + String.format( + "CREATE TABLE \"%s\" (" + + "\"id\" TEXT PRIMARY KEY," + + "\"item\" TEXT," + + "\"price\" INTEGER," + + "\"quantity\" INTEGER," + + "\"in_stock\" BOOLEAN," + + "\"tags\" TEXT[]," + + "\"props\" JSONB" + + ");", + COLLECTION_NAME); + + try (Connection connection = pgDatastore.getPostgresClient(); + PreparedStatement statement = connection.prepareStatement(createTableSQL)) { + statement.execute(); + LOGGER.info("Created flat collection table: {}", COLLECTION_NAME); + } catch (Exception e) { + LOGGER.error("Failed to create flat collection schema: {}", e.getMessage(), e); + throw new RuntimeException("Failed to create flat collection schema", e); + } + } + + @BeforeEach + public void clearCollections() { + Collection mongoCollection = collectionMap.get(MONGO_STORE); + mongoCollection.deleteAll(); + + PostgresDatastore pgDatastore = (PostgresDatastore) datastoreMap.get(POSTGRES_FLAT_STORE); + String deleteSQL = String.format("DELETE FROM \"%s\"", COLLECTION_NAME); + try (Connection connection = pgDatastore.getPostgresClient(); + PreparedStatement statement = connection.prepareStatement(deleteSQL)) { + statement.executeUpdate(); + } catch (Exception e) { + LOGGER.error("Failed to clear Postgres table: {}", e.getMessage(), e); + } + } + + @AfterAll + public static void shutdown() { + if (mongo != null) { + mongo.stop(); + } + if (postgres != null) { + postgres.stop(); + } + } + + private static class AllStoresProvider implements ArgumentsProvider { + + @Override + public Stream provideArguments(final ExtensionContext context) { + return Stream.of(Arguments.of(MONGO_STORE), Arguments.of(POSTGRES_FLAT_STORE)); + } + } + + private Collection getCollection(String storeName) { + return collectionMap.get(storeName); + } + + private static String generateDocId(String prefix) { + return prefix + "-" + System.currentTimeMillis() + "-" + (int) (Math.random() * 10000); + } + + private static String getKeyString(String docId) { + return new SingleValueKey(DEFAULT_TENANT, docId).toString(); + } + + private Query buildQueryById(String docId) { + return Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of(getKeyString(docId)))) + .build(); + } + + private void insertMinimalTestDocument(String docId) throws IOException { + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + String keyStr = key.toString(); + + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", keyStr); + objectNode.put("item", "Minimal Item"); + + Document document = new JSONDocument(objectNode); + + for (Collection collection : collectionMap.values()) { + collection.upsert(key, document); + } + } + + @Nested + @DisplayName("SubDocument Compatibility Tests") + class SubDocCompatibilityTest { + + @Nested + @DisplayName( + "Non-Existent Fields in JSONB Column. Subdoc updates on non-existent JSONB fields should create those fields in both Mongo and PG") + class JsonbNonExistentFieldTests { + + @ParameterizedTest(name = "{0}: SET on non-existent nested field should create field") + @ArgumentsSource(AllStoresProvider.class) + void testSet(String storeName) throws Exception { + String docId = generateDocId("set-nonexistent"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + + Query query = buildQueryById(docId); + + // SET props.brand which doesn't exist + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.brand") + .operator(UpdateOperator.SET) + .subDocumentValue(SubDocumentValue.of("NewBrand")) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent(), storeName + ": Should return updated document"); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + JsonNode propsNode = resultJson.get("props"); + assertNotNull(propsNode, storeName + ": props should be created"); + assertEquals( + "NewBrand", propsNode.get("brand").asText(), storeName + ": brand should be set"); + } + + @ParameterizedTest(name = "{0}: ADD on non-existent nested field behavior") + @ArgumentsSource(AllStoresProvider.class) + void testAdd(String storeName) throws Exception { + String docId = generateDocId("add-nonexistent"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + + Query query = buildQueryById(docId); + + // ADD to props.count which doesn't exist + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.count") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(10)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent(), storeName + ": Should return updated document"); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + // ADD on non-existent field should treat it as 0 and add, resulting in the value + JsonNode propsNode = resultJson.get("props"); + assertNotNull(propsNode, storeName + ": props should be created"); + assertEquals( + 10, propsNode.get("count").asInt(), storeName + ": count should be 10 (0 + 10)"); + } + + @ParameterizedTest(name = "{0}: UNSET on non-existent nested field behavior") + @ArgumentsSource(AllStoresProvider.class) + void testUnset(String storeName) throws Exception { + String docId = generateDocId("unset-nonexistent"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + + Query query = buildQueryById(docId); + + // UNSET props.brand which doesn't exist + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.brand") + .operator(UpdateOperator.UNSET) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + // Should succeed without error - UNSET on non-existent is a no-op + assertTrue(result.isPresent(), storeName + ": Should return updated document"); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + // Document should still exist with original fields + assertEquals("Minimal Item", resultJson.get("item").asText()); + } + + @ParameterizedTest(name = "{0}: APPEND_TO_LIST on non-existent nested array behavior") + @ArgumentsSource(AllStoresProvider.class) + void testAppendToList(String storeName) throws Exception { + String docId = generateDocId("append-nonexistent"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + + Query query = buildQueryById(docId); + + // APPEND_TO_LIST on props.colors which doesn't exist + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.colors") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"red", "blue"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent(), storeName + ": Should return updated document"); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + // Should create the array with the appended values + JsonNode propsNode = resultJson.get("props"); + assertNotNull(propsNode, storeName + ": props should be created"); + JsonNode colorsNode = propsNode.get("colors"); + assertNotNull(colorsNode, storeName + ": colors should be created"); + assertTrue(colorsNode.isArray(), storeName + ": colors should be an array"); + assertEquals(2, colorsNode.size(), storeName + ": colors should have 2 elements"); + } + + @ParameterizedTest(name = "{0}: ADD_TO_LIST_IF_ABSENT on non-existent nested array behavior") + @ArgumentsSource(AllStoresProvider.class) + void testAddToListIfAbsent(String storeName) throws Exception { + String docId = generateDocId("addifabsent-nonexistent"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + + Query query = buildQueryById(docId); + + // ADD_TO_LIST_IF_ABSENT on props.tags which doesn't exist + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.tags") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) + .subDocumentValue(SubDocumentValue.of(new String[] {"tag1", "tag2"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent(), storeName + ": Should return updated document"); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + // Should create the array with the values + JsonNode propsNode = resultJson.get("props"); + assertNotNull(propsNode, storeName + ": props should be created"); + JsonNode tagsNode = propsNode.get("tags"); + assertNotNull(tagsNode, storeName + ": tags should be created"); + assertTrue(tagsNode.isArray(), storeName + ": tags should be an array"); + assertEquals(2, tagsNode.size(), storeName + ": tags should have 2 elements"); + } + + @ParameterizedTest(name = "{0}: REMOVE_ALL_FROM_LIST on non-existent nested array behavior") + @ArgumentsSource(AllStoresProvider.class) + void testRemoveAllFromList(String storeName) throws Exception { + String docId = generateDocId("removeall-nonexistent"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + + Query query = buildQueryById(docId); + + // REMOVE_ALL_FROM_LIST on props.colors which doesn't exist + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.colors") + .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"red"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + // Should succeed - removing from non-existent list is a no-op or results in empty array + assertTrue(result.isPresent(), storeName + ": Should return updated document"); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + // Document should still exist + assertEquals("Minimal Item", resultJson.get("item").asText()); + } + + @ParameterizedTest(name = "{0}: SET on deep nested path should create intermediate objects") + @ArgumentsSource(AllStoresProvider.class) + void testSetDeepNested(String storeName) throws Exception { + String docId = generateDocId("set-deep"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // SET props.brand.category.name - all intermediate objects don't exist + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.brand.category.name") + .operator(UpdateOperator.SET) + .subDocumentValue(SubDocumentValue.of("Electronics")) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + // Verify deep nested structure was created + JsonNode propsNode = resultJson.get("props"); + assertNotNull(propsNode, storeName + ": props should be created"); + JsonNode brandNode = propsNode.get("brand"); + assertNotNull(brandNode, storeName + ": props.brand should be created"); + JsonNode categoryNode = brandNode.get("category"); + assertNotNull(categoryNode, storeName + ": props.brand.category should be created"); + assertEquals("Electronics", categoryNode.get("name").asText()); + } + + @ParameterizedTest(name = "{0}: ADD on deep nested path should create intermediate objects") + @ArgumentsSource(AllStoresProvider.class) + void testAddDeepNested(String storeName) throws Exception { + String docId = generateDocId("add-deep"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // ADD to props.stats.sales.count - all intermediate objects don't exist + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.stats.sales.count") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(5)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + JsonNode propsNode = resultJson.get("props"); + assertNotNull(propsNode, storeName + ": props should be created"); + JsonNode statsNode = propsNode.get("stats"); + assertNotNull(statsNode, storeName + ": props.stats should be created"); + JsonNode salesNode = statsNode.get("sales"); + assertNotNull(salesNode, storeName + ": props.stats.sales should be created"); + assertEquals(5, salesNode.get("count").asInt()); + } + + @ParameterizedTest( + name = "{0}: APPEND_TO_LIST on deep nested path should create intermediate objects") + @ArgumentsSource(AllStoresProvider.class) + void testAppendToListDeepNested(String storeName) throws Exception { + String docId = generateDocId("append-deep"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // APPEND_TO_LIST to props.metadata.tags.items - all intermediate objects don't exist + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.metadata.tags.items") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"tag1", "tag2"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + JsonNode propsNode = resultJson.get("props"); + assertNotNull(propsNode); + JsonNode metadataNode = propsNode.get("metadata"); + assertNotNull(metadataNode); + JsonNode tagsNode = metadataNode.get("tags"); + assertNotNull(tagsNode); + JsonNode itemsNode = tagsNode.get("items"); + assertNotNull(itemsNode); + assertTrue(itemsNode.isArray()); + assertEquals(2, itemsNode.size()); + } + } + + @Nested + @DisplayName("Top-Level Fields Not In PG Schema (Mongo creates, PG skips)") + class TopLevelSchemaMissingFieldTests { + + @ParameterizedTest(name = "{0}: SET on field not in PG schema") + @ArgumentsSource(AllStoresProvider.class) + void testSet(String storeName) throws Exception { + String docId = generateDocId("set-schema-missing"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // SET unknownField which doesn't exist in PG schema + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("unknownField") + .operator(UpdateOperator.SET) + .subDocumentValue(SubDocumentValue.of("newValue")) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent(), storeName + ": Should return updated document"); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + if (MONGO_STORE.equals(storeName)) { + // Mongo creates the field + assertNotNull( + resultJson.get("unknownField"), storeName + ": unknownField should be created"); + assertEquals("newValue", resultJson.get("unknownField").asText()); + } else { + // Postgres SKIP strategy: field not created, no-op + assertTrue( + resultJson.get("unknownField") == null || resultJson.get("unknownField").isNull()); + } + } + + @ParameterizedTest(name = "{0}: ADD on field not in PG schema") + @ArgumentsSource(AllStoresProvider.class) + void testAdd(String storeName) throws Exception { + String docId = generateDocId("add-schema-missing"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // ADD to unknownCount which doesn't exist in PG schema + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("unknownCount") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(10)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent(), storeName + ": Should return updated document"); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + if (MONGO_STORE.equals(storeName)) { + // Mongo creates the field with value + assertNotNull( + resultJson.get("unknownCount"), storeName + ": unknownCount should be created"); + assertEquals(10, resultJson.get("unknownCount").asInt()); + } else { + // Postgres SKIP strategy: field not created, no-op + assertTrue( + resultJson.get("unknownCount") == null || resultJson.get("unknownCount").isNull()); + } + } + + @ParameterizedTest(name = "{0}: UNSET on field not in PG schema") + @ArgumentsSource(AllStoresProvider.class) + void testUnset(String storeName) throws Exception { + String docId = generateDocId("unset-schema-missing"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // UNSET unknownField which doesn't exist in schema or document + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("unknownField") + .operator(UpdateOperator.UNSET) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + // Both Mongo and Postgres: UNSET on non-existent field is a no-op + assertTrue(result.isPresent(), storeName + ": Should return updated document"); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + assertEquals("Minimal Item", resultJson.get("item").asText()); + } + + @ParameterizedTest(name = "{0}: APPEND_TO_LIST on field not in PG schema") + @ArgumentsSource(AllStoresProvider.class) + void testAppendToList(String storeName) throws Exception { + String docId = generateDocId("append-schema-missing"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // APPEND_TO_LIST on unknownList which doesn't exist in PG schema + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("unknownList") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"item1", "item2"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + JsonNode unknownList = resultJson.get("unknownList"); + if (MONGO_STORE.equals(storeName)) { + // Mongo creates the array + assertNotNull(unknownList); + assertTrue(unknownList.isArray()); + assertEquals(2, unknownList.size()); + } else { + // Postgres SKIP strategy: field not created, no-op + assertTrue(unknownList == null || unknownList.isNull()); + } + } + + @ParameterizedTest(name = "{0}: ADD_TO_LIST_IF_ABSENT on field not in PG schema") + @ArgumentsSource(AllStoresProvider.class) + void testAddToList(String storeName) throws Exception { + String docId = generateDocId("addifabsent-schema-missing"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // ADD_TO_LIST_IF_ABSENT on unknownSet which doesn't exist in PG schema + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("unknownSet") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) + .subDocumentValue(SubDocumentValue.of(new String[] {"val1", "val2"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + JsonNode unknownSet = resultJson.get("unknownSet"); + if (MONGO_STORE.equals(storeName)) { + // Mongo creates the array + assertNotNull(unknownSet); + assertTrue(unknownSet.isArray()); + assertEquals(2, unknownSet.size()); + } else { + // Postgres SKIP strategy: field not created, no-op + assertTrue(unknownSet == null || unknownSet.isNull()); + } + } + + @ParameterizedTest(name = "{0}: REMOVE_ALL_FROM_LIST on field not in PG schema") + @ArgumentsSource(AllStoresProvider.class) + void testRemoveAllFromList(String storeName) throws Exception { + String docId = generateDocId("removeall-schema-missing"); + insertMinimalTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // REMOVE_ALL_FROM_LIST on unknownList which doesn't exist in schema or document + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("unknownList") + .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"item1"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = collection.update(query, updates, options); + + // Both Mongo and Postgres: REMOVE_ALL from non-existent is a no-op + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + assertEquals("Minimal Item", resultJson.get("item").asText()); + } + } + } +} 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 cf4835a57..d9d0cc4f5 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 @@ -1,9 +1,14 @@ package org.hypertrace.core.documentstore.postgres; +import static java.util.Map.entry; 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.ADD_TO_LIST_IF_ABSENT; +import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.APPEND_TO_LIST; +import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.REMOVE_ALL_FROM_LIST; import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.SET; +import static org.hypertrace.core.documentstore.model.subdoc.UpdateOperator.UNSET; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; @@ -47,10 +52,14 @@ 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.query.v1.transformer.LegacyFilterToQueryFilterTransformer; -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.update.parser.PostgresAddToListIfAbsentParser; +import org.hypertrace.core.documentstore.postgres.update.parser.PostgresAddValueParser; +import org.hypertrace.core.documentstore.postgres.update.parser.PostgresAppendToListParser; +import org.hypertrace.core.documentstore.postgres.update.parser.PostgresRemoveAllFromListParser; +import org.hypertrace.core.documentstore.postgres.update.parser.PostgresSetValueParser; +import org.hypertrace.core.documentstore.postgres.update.parser.PostgresUnsetPathParser; +import org.hypertrace.core.documentstore.postgres.update.parser.PostgresUpdateOperationParser; +import org.hypertrace.core.documentstore.postgres.update.parser.PostgresUpdateOperationParser.UpdateParserInput; import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils; import org.hypertrace.core.documentstore.query.Query; import org.postgresql.util.PSQLException; @@ -74,11 +83,14 @@ public class FlatPostgresCollection extends PostgresCollection { private static final String MISSING_COLUMN_STRATEGY_CONFIG = "missingColumnStrategy"; private static final String DEFAULT_PRIMARY_KEY_COLUMN = "key"; - private static final Map - SUB_DOC_UPDATE_PARSERS = - Map.of( - SET, new FlatCollectionSubDocSetOperatorParser(), - ADD, new FlatCollectionSubDocAddOperatorParser()); + private static final Map UPDATE_PARSER_MAP = + Map.ofEntries( + entry(SET, new PostgresSetValueParser()), + entry(UNSET, new PostgresUnsetPathParser()), + entry(ADD, new PostgresAddValueParser()), + entry(REMOVE_ALL_FROM_LIST, new PostgresRemoveAllFromListParser()), + entry(ADD_TO_LIST_IF_ABSENT, new PostgresAddToListIfAbsentParser()), + entry(APPEND_TO_LIST, new PostgresAppendToListParser())); private final PostgresLazyilyLoadedSchemaRegistry schemaRegistry; @@ -624,7 +636,7 @@ private Map resolvePathsToColumns( UpdateOperator operator = update.getOperator(); Preconditions.checkArgument( - SUB_DOC_UPDATE_PARSERS.containsKey(operator), "Unsupported UPDATE operator: " + operator); + UPDATE_PARSER_MAP.containsKey(operator), "Unsupported UPDATE operator: " + operator); String path = update.getSubDocument().getPath(); Optional columnName = resolveColumnName(path, tableName); @@ -744,20 +756,40 @@ private void executeUpdate( PostgresColumnMetadata colMeta = schemaRegistry.getColumnOrRefresh(tableName, columnName).orElseThrow(); - FlatUpdateContext context = - FlatUpdateContext.builder() - .columnName(columnName) - // get the nested path. So for example, if colName is `customAttr` and full path is - // `customAttr.props`, then the nested path is `props`. - .nestedPath(getNestedPath(path, columnName)) - .columnType(colMeta.getPostgresType()) - .value(update.getSubDocumentValue()) - .params(params) - .build(); - - FlatCollectionSubDocUpdateOperatorParser operatorParser = - SUB_DOC_UPDATE_PARSERS.get(update.getOperator()); - String fragment = operatorParser.parse(context); + String[] nestedPath = getNestedPath(path, columnName); + boolean isTopLevel = nestedPath.length == 0; + UpdateOperator operator = update.getOperator(); + + Params.Builder paramsBuilder = Params.newBuilder(); + PostgresUpdateOperationParser unifiedParser = UPDATE_PARSER_MAP.get(operator); + + String fragment; + + if (isTopLevel) { + UpdateParserInput input = + UpdateParserInput.builder() + .baseField(columnName) + .path(new String[0]) + .update(update) + .paramsBuilder(paramsBuilder) + .columnType(colMeta.getPostgresType()) + .build(); + fragment = unifiedParser.parseNonJsonbField(input); + } else { + // parseInternal() returns just the value expression + UpdateParserInput jsonbInput = + UpdateParserInput.builder() + .baseField(String.format("\"%s\"", columnName)) + .path(nestedPath) + .update(update) + .paramsBuilder(paramsBuilder) + .columnType(colMeta.getPostgresType()) + .build(); + String valueExpr = unifiedParser.parseInternal(jsonbInput); + fragment = String.format("\"%s\" = %s", columnName, valueExpr); + } + // Transfer params from builder to our list + params.addAll(paramsBuilder.build().getObjectParams().values()); setFragments.add(fragment); } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresInRelationalFilterParserJsonArray.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresInRelationalFilterParserJsonArray.java index 798a26f67..e47d5b60f 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresInRelationalFilterParserJsonArray.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/PostgresInRelationalFilterParserJsonArray.java @@ -6,6 +6,7 @@ import org.hypertrace.core.documentstore.expression.impl.JsonIdentifierExpression; import org.hypertrace.core.documentstore.expression.impl.RelationalExpression; import org.hypertrace.core.documentstore.postgres.Params; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; /** * Optimized parser for IN operations on JSON array fields with type-specific casting. @@ -119,7 +120,7 @@ private String prepareFilterStringForArrayInOperator( final Params.Builder paramsBuilder) { // Determine the appropriate type cast for jsonb_build_array elements - String typeCast = getTypeCastForArray(fieldType); + String typeCast = PostgresDataType.getJsonArrayElementTypeCast(fieldType); // For JSON arrays, we use the @> containment operator // Check if ANY of the RHS values is contained in the LHS array @@ -137,26 +138,4 @@ private String prepareFilterStringForArrayInOperator( ? String.format("(%s)", orConditions) : orConditions; } - - /** - * Returns the PostgreSQL type cast string for jsonb_build_array elements based on array type. - * - * @param fieldType The JSON field type (must not be null) - * @return Type cast string (e.g., "::text", "::numeric") - */ - private String getTypeCastForArray(JsonFieldType fieldType) { - switch (fieldType) { - case STRING_ARRAY: - return "::text"; - case NUMBER_ARRAY: - return "::numeric"; - case BOOLEAN_ARRAY: - return "::boolean"; - case OBJECT_ARRAY: - return "::jsonb"; - default: - throw new IllegalArgumentException( - "Unsupported array type: " + fieldType + ". Expected *_ARRAY types."); - } - } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataType.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataType.java index c920473fb..18063b370 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataType.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataType.java @@ -1,6 +1,7 @@ package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field; import org.hypertrace.core.documentstore.expression.impl.DataType; +import org.hypertrace.core.documentstore.expression.impl.JsonFieldType; /** * PostgreSQL-specific data types with their SQL type strings. @@ -40,6 +41,31 @@ public String getArraySqlType() { return sqlType + "[]"; } + public String getTypeCast() { + return sqlType == null ? "" : "::" + sqlType; + } + + public String getArrayTypeCast() { + return sqlType == null ? "" : "::" + sqlType + "[]"; + } + + public static PostgresDataType fromJavaValue(Object value) { + if (value instanceof String) { + return TEXT; + } else if (value instanceof Integer) { + return INTEGER; + } else if (value instanceof Long) { + return BIGINT; + } else if (value instanceof Float) { + return REAL; + } else if (value instanceof Double) { + return DOUBLE_PRECISION; + } else if (value instanceof Boolean) { + return BOOLEAN; + } + return UNKNOWN; + } + /** * Maps a generic DataType to its PostgreSQL equivalent. * @@ -70,4 +96,27 @@ public static PostgresDataType fromDataType(DataType dataType) { throw new IllegalArgumentException("Unknown DataType: " + dataType); } } + + /** + * Returns the PostgreSQL type cast string for JSONB array element types. + * + * @param fieldType the JSON field type (must be an array type) + * @return Type cast string (e.g., "::text", "::numeric", "::boolean", "::jsonb") + * @throws IllegalArgumentException if fieldType is not a supported array type + */ + public static String getJsonArrayElementTypeCast(JsonFieldType fieldType) { + switch (fieldType) { + case STRING_ARRAY: + return "::text"; + case NUMBER_ARRAY: + return "::numeric"; + case BOOLEAN_ARRAY: + return "::boolean"; + case OBJECT_ARRAY: + return "::jsonb"; + default: + throw new IllegalArgumentException( + "Unsupported array type: " + fieldType + ". Expected *_ARRAY types."); + } + } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java index 41ffebf15..cab7713b7 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresFilterTypeExpressionVisitor.java @@ -30,6 +30,7 @@ import org.hypertrace.core.documentstore.postgres.query.v1.parser.builder.PostgresSelectExpressionParserBuilderImpl; import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.PostgresRelationalFilterParser.PostgresRelationalFilterContext; import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.PostgresRelationalFilterParserFactoryImpl; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; public class PostgresFilterTypeExpressionVisitor implements FilterTypeExpressionVisitor { @@ -260,26 +261,20 @@ private String inferArrayTypeCastFromFilter(FilterTypeExpression filter) { if (filter instanceof RelationalExpression) { RelationalExpression relExpr = (RelationalExpression) filter; - // The visitor returns a string representation, but we need the actual value // Try to get the constant value directly if it's a ConstantExpression if (relExpr.getRhs() instanceof ConstantExpression) { ConstantExpression constExpr = (ConstantExpression) relExpr.getRhs(); Object value = constExpr.getValue(); - if (value instanceof String) { - return "::text[]"; - } else if (value instanceof Integer || value instanceof Long) { - return "::bigint[]"; - } else if (value instanceof Double || value instanceof Float) { - return "::double precision[]"; - } else if (value instanceof Boolean) { - return "::boolean[]"; + PostgresDataType pgType = PostgresDataType.fromJavaValue(value); + if (pgType != PostgresDataType.UNKNOWN) { + return pgType.getArrayTypeCast(); } } } // Default to text[] if we can't infer the type - return "::text[]"; + return PostgresDataType.TEXT.getArrayTypeCast(); } private String getFilterStringForAnyOperator(final DocumentArrayFilterExpression expression) { diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/FlatUpdateContext.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/FlatUpdateContext.java deleted file mode 100644 index 5537c9740..000000000 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/FlatUpdateContext.java +++ /dev/null @@ -1,43 +0,0 @@ -package org.hypertrace.core.documentstore.postgres.update; - -import java.util.List; -import lombok.Builder; -import lombok.Value; -import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue; -import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; - -/** - * Context object containing all information needed to generate SQL for a single field update in - * flat collections. - */ -@Value -@Builder -public class FlatUpdateContext { - /** The column name in the database (e.g., "price", "props") */ - String columnName; - - /** - * The nested path within a JSONB column, empty array for top-level columns. For example, for - * "props.seller.name", columnName would be "props" and nestedPath would be ["seller", "name"]. - */ - String[] nestedPath; - - /** The PostgreSQL data type of the column */ - PostgresDataType columnType; - - /** The value to set/update */ - SubDocumentValue value; - - /** Accumulator for prepared statement parameters (mutable) */ - List params; - - /** Returns true if this is a top-level column update (no nested path) */ - public boolean isTopLevel() { - return nestedPath == null || nestedPath.length == 0; - } - - /** Returns true if the column is a JSONB type */ - public boolean isJsonbColumn() { - return columnType == PostgresDataType.JSONB; - } -} 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 deleted file mode 100644 index 3b10cf6e4..000000000 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocAddOperatorParser.java +++ /dev/null @@ -1,165 +0,0 @@ -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) + "}"; - } -} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocSetOperatorParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocSetOperatorParser.java deleted file mode 100644 index 40cc11f6f..000000000 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocSetOperatorParser.java +++ /dev/null @@ -1,123 +0,0 @@ -package org.hypertrace.core.documentstore.postgres.update.parser; - -import org.hypertrace.core.documentstore.model.subdoc.MultiValuedNestedSubDocumentValue; -import org.hypertrace.core.documentstore.model.subdoc.MultiValuedPrimitiveSubDocumentValue; -import org.hypertrace.core.documentstore.model.subdoc.NestedSubDocumentValue; -import org.hypertrace.core.documentstore.model.subdoc.NullSubDocumentValue; -import org.hypertrace.core.documentstore.model.subdoc.PrimitiveSubDocumentValue; -import org.hypertrace.core.documentstore.model.subdoc.visitor.SubDocumentValueVisitor; -import org.hypertrace.core.documentstore.postgres.update.FlatUpdateContext; - -/** - * Parser for the SET operator in flat collections. - * - *

Handles two cases: - * - *

    - *
  • Top-level columns: {@code SET "column" = ?} - *
  • Nested JSONB paths: {@code SET "column" = jsonb_set(COALESCE("column", '{}'), '{path}', - * to_jsonb(?))} - *
- */ -public class FlatCollectionSubDocSetOperatorParser - implements FlatCollectionSubDocUpdateOperatorParser { - - /** Visitor to extract raw values from SubDocumentValue for use in prepared statements. */ - private static final SubDocumentValueVisitor VALUE_EXTRACTOR = - new SubDocumentValueVisitor<>() { - @Override - public Object visit(PrimitiveSubDocumentValue value) { - return value.getValue(); - } - - @Override - public Object visit(MultiValuedPrimitiveSubDocumentValue value) { - return value.getValues(); - } - - @Override - public Object visit(NestedSubDocumentValue value) { - return value.getJsonValue(); - } - - @Override - public Object visit(MultiValuedNestedSubDocumentValue value) { - return value.getJsonValues(); - } - - @Override - public Object visit(NullSubDocumentValue value) { - return null; - } - }; - - /** - * Visitor that returns the appropriate SQL value expression for jsonb_set. JSON document values - * use ?::jsonb to parse the JSON string directly. Primitive values use to_jsonb(?) to convert to - * proper JSONB format. - */ - private static final SubDocumentValueVisitor VALUE_EXPR_VISITOR = - new SubDocumentValueVisitor<>() { - @Override - public String visit(PrimitiveSubDocumentValue value) { - return "to_jsonb(?)"; - } - - @Override - public String visit(MultiValuedPrimitiveSubDocumentValue value) { - return "to_jsonb(?)"; - } - - @Override - public String visit(NestedSubDocumentValue value) { - return "?::jsonb"; - } - - @Override - public String visit(MultiValuedNestedSubDocumentValue value) { - return "?::jsonb"; - } - - @Override - public String visit(NullSubDocumentValue value) { - return "to_jsonb(?)"; - } - }; - - @Override - public String parse(FlatUpdateContext context) { - if (context.isTopLevel()) { - return parseTopLevel(context); - } else { - return parseNestedJsonb(context); - } - } - - private String parseTopLevel(FlatUpdateContext context) { - context.getParams().add(context.getValue().accept(VALUE_EXTRACTOR)); - return String.format("\"%s\" = ?", context.getColumnName()); - } - - private String parseNestedJsonb(FlatUpdateContext context) { - String jsonPath = buildJsonPath(context.getNestedPath()); - Object value = context.getValue().accept(VALUE_EXTRACTOR); - - context.getParams().add(jsonPath); - context.getParams().add(value); - - // Use jsonb_set with COALESCE to handle null columns - // 4th param (true) creates the key if it doesn't exist - String valueExpr = context.getValue().accept(VALUE_EXPR_VISITOR); - return String.format( - "\"%s\" = jsonb_set(COALESCE(\"%s\", '{}'), ?::text[], %s, true)", - context.getColumnName(), context.getColumnName(), valueExpr); - } - - /** - * Builds a PostgreSQL text array path from nested path components. For example, ["seller", - * "name"] becomes "{seller,name}" - */ - private String buildJsonPath(String[] nestedPath) { - return "{" + String.join(",", nestedPath) + "}"; - } -} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocUpdateOperatorParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocUpdateOperatorParser.java deleted file mode 100644 index 38ef44efa..000000000 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/FlatCollectionSubDocUpdateOperatorParser.java +++ /dev/null @@ -1,26 +0,0 @@ -package org.hypertrace.core.documentstore.postgres.update.parser; - -import org.hypertrace.core.documentstore.postgres.update.FlatUpdateContext; - -/** - * Parser interface for converting SubDocumentUpdate operations to SQL fragments for flat - * collections. - * - *

Each implementation handles a specific {@link - * org.hypertrace.core.documentstore.model.subdoc.UpdateOperator} and generates the appropriate SQL - * SET clause fragment. - */ -public interface FlatCollectionSubDocUpdateOperatorParser { - - /** - * Generates SQL SET clause fragment for this operator. - * - *

For top-level columns, this typically produces: {@code "column" = ?} - * - *

For nested JSONB paths, this produces: {@code "column" = jsonb_set(...)} - * - * @param context The update context containing column info, value, and parameter accumulator - * @return SQL fragment to be used in SET clause - */ - String parse(FlatUpdateContext context); -} diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAddToListIfAbsentParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAddToListIfAbsentParser.java index 2f621c058..57fcbc430 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAddToListIfAbsentParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAddToListIfAbsentParser.java @@ -2,10 +2,28 @@ import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue; import org.hypertrace.core.documentstore.postgres.Params; +import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentArrayGetter; import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentValueAddToListIfAbsentParser; public class PostgresAddToListIfAbsentParser implements PostgresUpdateOperationParser { + @Override + public String parseNonJsonbField(final UpdateParserInput input) { + final SubDocumentValue value = input.getUpdate().getSubDocumentValue(); + + // Extract array values directly for top-level array columns + final PostgresSubDocumentArrayGetter arrayGetter = new PostgresSubDocumentArrayGetter(); + Object[] arrayValues = value.accept(arrayGetter).values(); + input.getParamsBuilder().addObjectParam(arrayValues); + + // For top-level array columns: add unique values using ARRAY(SELECT DISTINCT unnest(...)) + String arrayType = + input.getColumnType() != null ? input.getColumnType().getArraySqlType() : "text[]"; + return String.format( + "\"%s\" = ARRAY(SELECT DISTINCT unnest(COALESCE(\"%s\", '{}') || ?::%s))", + input.getBaseField(), input.getBaseField(), arrayType); + } + @Override public String parseInternal(final UpdateParserInput input) { return new PostgresSetValueParser(this, 0).parseInternal(input); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAddValueParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAddValueParser.java index e03579bb7..965a13ae0 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAddValueParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAddValueParser.java @@ -3,10 +3,75 @@ import static org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.formatSubDocPath; import static org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.prepareFieldDataAccessorExpr; +import org.hypertrace.core.documentstore.model.subdoc.MultiValuedNestedSubDocumentValue; +import org.hypertrace.core.documentstore.model.subdoc.MultiValuedPrimitiveSubDocumentValue; +import org.hypertrace.core.documentstore.model.subdoc.NestedSubDocumentValue; +import org.hypertrace.core.documentstore.model.subdoc.NullSubDocumentValue; +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.Params; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentValueParser; public class PostgresAddValueParser implements PostgresUpdateOperationParser { + + /** Visitor to validate and extract numeric values from SubDocumentValue. */ + private static final SubDocumentValueVisitor NUMERIC_VALUE_VALIDATOR = + 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(MultiValuedPrimitiveSubDocumentValue value) { + throw new IllegalArgumentException("ADD operator does not support multi-valued updates"); + } + + @Override + public Number visit(NestedSubDocumentValue value) { + throw new IllegalArgumentException( + "ADD operator does not support nested document values"); + } + + @Override + public Number visit(MultiValuedNestedSubDocumentValue value) { + throw new IllegalArgumentException( + "ADD operator does not support multi-valued nested documents"); + } + + @Override + public Number visit(NullSubDocumentValue value) { + throw new IllegalArgumentException("ADD operator does not support null values"); + } + }; + + @Override + public String parseNonJsonbField(UpdateParserInput input) { + // Validate that the value is numeric + SubDocumentValue value = input.getUpdate().getSubDocumentValue(); + value.accept(NUMERIC_VALUE_VALIDATOR); + + final Params.Builder paramsBuilder = input.getParamsBuilder(); + final PostgresSubDocumentValueParser valueParser = + new PostgresSubDocumentValueParser(paramsBuilder); + + // Add the numeric value to params + value.accept(valueParser); + + // Generate: "column" = COALESCE("column", 0) + ?::type + PostgresDataType columnType = input.getColumnType(); + String typeCast = (columnType != null) ? columnType.getTypeCast() : ""; + return String.format( + "\"%s\" = COALESCE(\"%s\", 0) + ?%s", input.getBaseField(), input.getBaseField(), typeCast); + } + @Override public String parseInternal(UpdateParserInput input) { return new PostgresSetValueParser(this, 1).parseInternal(input); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAppendToListParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAppendToListParser.java index 01bbf78af..5c07f00fa 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAppendToListParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresAppendToListParser.java @@ -1,10 +1,28 @@ package org.hypertrace.core.documentstore.postgres.update.parser; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue; +import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentArrayGetter; import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentValueParser; public class PostgresAppendToListParser implements PostgresUpdateOperationParser { + @Override + public String parseNonJsonbField(final UpdateParserInput input) { + final SubDocumentValue value = input.getUpdate().getSubDocumentValue(); + + // Extract array values directly for top-level array columns + final PostgresSubDocumentArrayGetter arrayGetter = new PostgresSubDocumentArrayGetter(); + Object[] arrayValues = value.accept(arrayGetter).values(); + input.getParamsBuilder().addObjectParam(arrayValues); + + // For top-level array columns: "column" = COALESCE("column", '{}') || ?::arrayType + String arrayType = + input.getColumnType() != null ? input.getColumnType().getArraySqlType() : "text[]"; + return String.format( + "\"%s\" = COALESCE(\"%s\", '{}') || ?::%s", + input.getBaseField(), input.getBaseField(), arrayType); + } + @Override public String parseInternal(final UpdateParserInput input) { return new PostgresSetValueParser(this, 0).parseInternal(input); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresRemoveAllFromListParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresRemoveAllFromListParser.java index a618c29df..eded52341 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresRemoveAllFromListParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresRemoveAllFromListParser.java @@ -14,6 +14,24 @@ public class PostgresRemoveAllFromListParser implements PostgresUpdateOperationParser { + @Override + public String parseNonJsonbField(final UpdateParserInput input) { + final PostgresSubDocumentArrayGetter subDocArrayGetter = new PostgresSubDocumentArrayGetter(); + final SubDocumentArray array = + input.getUpdate().getSubDocumentValue().accept(subDocArrayGetter); + final Object[] values = array.values(); + + // Add array as single param (not individual values) + input.getParamsBuilder().addObjectParam(values); + + // For top-level array columns: remove values using array_agg with filter + String arrayType = + input.getColumnType() != null ? input.getColumnType().getArraySqlType() : "text[]"; + return String.format( + "\"%s\" = (SELECT array_agg(elem) FROM unnest(\"%s\") AS elem WHERE NOT (elem = ANY(?::%s)))", + input.getBaseField(), input.getBaseField(), arrayType); + } + @Override public String parseInternal(final UpdateParserInput input) { final String baseField = input.getBaseField(); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresSetValueParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresSetValueParser.java index d4e7f4bc5..d3763a60e 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresSetValueParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresSetValueParser.java @@ -20,6 +20,17 @@ public PostgresSetValueParser() { leafNodePathSize = 1; } + @Override + public String parseNonJsonbField(final UpdateParserInput input) { + final Params.Builder paramsBuilder = input.getParamsBuilder(); + final PostgresSubDocumentValueParser valueParser = + new PostgresSubDocumentValueParser(paramsBuilder); + + // For top-level columns, just set the value directly: "column" = ? + input.getUpdate().getSubDocumentValue().accept(valueParser); + return String.format("\"%s\" = ?", input.getBaseField()); + } + @Override public String parseInternal(final UpdateParserInput input) { final String baseField = input.getBaseField(); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresUnsetPathParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresUnsetPathParser.java index a9ee2e400..a82c3d911 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresUnsetPathParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresUnsetPathParser.java @@ -5,6 +5,11 @@ public class PostgresUnsetPathParser implements PostgresUpdateOperationParser { + @Override + public String parseNonJsonbField(final UpdateParserInput input) { + return String.format("\"%s\" = NULL", input.getBaseField()); + } + @Override public String parseInternal(final UpdateParserInput input) { return parse(input); diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresUpdateOperationParser.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresUpdateOperationParser.java index fe4daa1fa..249491004 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresUpdateOperationParser.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/update/parser/PostgresUpdateOperationParser.java @@ -4,8 +4,22 @@ import lombok.Value; import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; import org.hypertrace.core.documentstore.postgres.Params; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; public interface PostgresUpdateOperationParser { + + /** + * Parses an update operation for a top-level column in flat collections. + * + *

For example, for SET on a top-level "price" column: {@code "price" = ?} + * + * @param input the update parser input containing column info and value + * @return SQL fragment for the SET clause + */ + default String parseNonJsonbField(final UpdateParserInput input) { + throw new UnsupportedOperationException("parseNonJsonbField not implemented for this operator"); + } + String parseInternal(final UpdateParserInput input); String parseLeaf(final UpdateParserInput input); @@ -17,5 +31,7 @@ class UpdateParserInput { String[] path; SubDocumentUpdate update; Params.Builder paramsBuilder; + // only for flat collections + PostgresDataType columnType; } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataTypeTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataTypeTest.java new file mode 100644 index 000000000..de803f7aa --- /dev/null +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/query/v1/parser/filter/nonjson/field/PostgresDataTypeTest.java @@ -0,0 +1,181 @@ +package org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.hypertrace.core.documentstore.expression.impl.JsonFieldType; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; + +class PostgresDataTypeTest { + + @Nested + @DisplayName("getTypeCast Tests") + class GetTypeCastTests { + + @Test + @DisplayName("TEXT should return ::text") + void testTextTypeCast() { + assertEquals("::text", PostgresDataType.TEXT.getTypeCast()); + } + + @Test + @DisplayName("INTEGER should return ::int4") + void testIntegerTypeCast() { + assertEquals("::int4", PostgresDataType.INTEGER.getTypeCast()); + } + + @Test + @DisplayName("BIGINT should return ::int8") + void testBigintTypeCast() { + assertEquals("::int8", PostgresDataType.BIGINT.getTypeCast()); + } + + @Test + @DisplayName("REAL should return ::float4") + void testRealTypeCast() { + assertEquals("::float4", PostgresDataType.REAL.getTypeCast()); + } + + @Test + @DisplayName("DOUBLE_PRECISION should return ::float8") + void testDoublePrecisionTypeCast() { + assertEquals("::float8", PostgresDataType.DOUBLE_PRECISION.getTypeCast()); + } + + @Test + @DisplayName("BOOLEAN should return ::bool") + void testBooleanTypeCast() { + assertEquals("::bool", PostgresDataType.BOOLEAN.getTypeCast()); + } + + @Test + @DisplayName("JSONB should return ::jsonb") + void testJsonbTypeCast() { + assertEquals("::jsonb", PostgresDataType.JSONB.getTypeCast()); + } + + @Test + @DisplayName("TIMESTAMPTZ should return ::timestamptz") + void testTimestamptzTypeCast() { + assertEquals("::timestamptz", PostgresDataType.TIMESTAMPTZ.getTypeCast()); + } + + @Test + @DisplayName("DATE should return ::date") + void testDateTypeCast() { + assertEquals("::date", PostgresDataType.DATE.getTypeCast()); + } + + @Test + @DisplayName("UNKNOWN (null sqlType) should return empty string") + void testUnknownTypeCastReturnsEmpty() { + assertEquals("", PostgresDataType.UNKNOWN.getTypeCast()); + } + } + + @Nested + @DisplayName("getArrayTypeCast Tests") + class GetArrayTypeCastTests { + + @Test + @DisplayName("TEXT should return ::text[]") + void testTextArrayTypeCast() { + assertEquals("::text[]", PostgresDataType.TEXT.getArrayTypeCast()); + } + + @Test + @DisplayName("INTEGER should return ::int4[]") + void testIntegerArrayTypeCast() { + assertEquals("::int4[]", PostgresDataType.INTEGER.getArrayTypeCast()); + } + + @Test + @DisplayName("BIGINT should return ::int8[]") + void testBigintArrayTypeCast() { + assertEquals("::int8[]", PostgresDataType.BIGINT.getArrayTypeCast()); + } + + @Test + @DisplayName("REAL should return ::float4[]") + void testRealArrayTypeCast() { + assertEquals("::float4[]", PostgresDataType.REAL.getArrayTypeCast()); + } + + @Test + @DisplayName("DOUBLE_PRECISION should return ::float8[]") + void testDoublePrecisionArrayTypeCast() { + assertEquals("::float8[]", PostgresDataType.DOUBLE_PRECISION.getArrayTypeCast()); + } + + @Test + @DisplayName("BOOLEAN should return ::bool[]") + void testBooleanArrayTypeCast() { + assertEquals("::bool[]", PostgresDataType.BOOLEAN.getArrayTypeCast()); + } + + @Test + @DisplayName("UNKNOWN (null sqlType) should return empty string") + void testUnknownArrayTypeCastReturnsEmpty() { + assertEquals("", PostgresDataType.UNKNOWN.getArrayTypeCast()); + } + } + + @Nested + @DisplayName("getJsonArrayElementTypeCast Tests") + class GetJsonArrayElementTypeCastTests { + + @Test + @DisplayName("STRING_ARRAY should return ::text") + void testStringArrayTypeCast() { + assertEquals( + "::text", PostgresDataType.getJsonArrayElementTypeCast(JsonFieldType.STRING_ARRAY)); + } + + @Test + @DisplayName("NUMBER_ARRAY should return ::numeric") + void testNumberArrayTypeCast() { + assertEquals( + "::numeric", PostgresDataType.getJsonArrayElementTypeCast(JsonFieldType.NUMBER_ARRAY)); + } + + @Test + @DisplayName("BOOLEAN_ARRAY should return ::boolean") + void testBooleanArrayTypeCast() { + assertEquals( + "::boolean", PostgresDataType.getJsonArrayElementTypeCast(JsonFieldType.BOOLEAN_ARRAY)); + } + + @Test + @DisplayName("OBJECT_ARRAY should return ::jsonb") + void testObjectArrayTypeCast() { + assertEquals( + "::jsonb", PostgresDataType.getJsonArrayElementTypeCast(JsonFieldType.OBJECT_ARRAY)); + } + + @Test + @DisplayName("Non-array type should throw IllegalArgumentException") + void testNonArrayTypeThrowsException() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> PostgresDataType.getJsonArrayElementTypeCast(JsonFieldType.STRING)); + + assertEquals( + "Unsupported array type: STRING. Expected *_ARRAY types.", exception.getMessage()); + } + + @Test + @DisplayName("NUMBER type should throw IllegalArgumentException") + void testNumberTypeThrowsException() { + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> PostgresDataType.getJsonArrayElementTypeCast(JsonFieldType.NUMBER)); + + assertEquals( + "Unsupported array type: NUMBER. Expected *_ARRAY types.", exception.getMessage()); + } + } +}