diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/BaseWriteTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/BaseWriteTest.java new file mode 100644 index 000000000..a7ab70e63 --- /dev/null +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/BaseWriteTest.java @@ -0,0 +1,215 @@ +package org.hypertrace.core.documentstore; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.typesafe.config.ConfigFactory; +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.nio.charset.StandardCharsets; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.util.HashMap; +import java.util.Map; +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.MissingColumnStrategy; +import org.hypertrace.core.documentstore.postgres.PostgresDatastore; +import org.hypertrace.core.documentstore.query.Query; +import org.junit.jupiter.api.extension.ExtensionContext; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.ArgumentsProvider; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.containers.wait.strategy.Wait; +import org.testcontainers.utility.DockerImageName; + +/** Base class for write tests */ +public abstract class BaseWriteTest { + + protected static final Logger LOGGER = LoggerFactory.getLogger(BaseWriteTest.class); + protected static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + protected static final String DEFAULT_TENANT = "default"; + + // MongoDB container and datastore - shared by all subclasses + protected static GenericContainer mongoContainer; + protected static Datastore mongoDatastore; + + // PostgreSQL container and datastore - shared by all subclasses + protected static GenericContainer postgresContainer; + protected static Datastore postgresDatastore; + + // Maps for multi-store tests + protected static Map datastoreMap = new HashMap<>(); + protected static Map collectionMap = new HashMap<>(); + + protected Collection getCollection(String storeName) { + return collectionMap.get(storeName); + } + + private static final String FLAT_COLLECTION_SCHEMA_PATH = + "schema/flat_collection_test_schema.sql"; + + protected static String loadFlatCollectionSchema() { + try (InputStream is = + BaseWriteTest.class.getClassLoader().getResourceAsStream(FLAT_COLLECTION_SCHEMA_PATH); + BufferedReader reader = + new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) { + StringBuilder sb = new StringBuilder(); + String line; + while ((line = reader.readLine()) != null) { + sb.append(line).append(" "); + } + return sb.toString().trim(); + } catch (Exception e) { + throw new RuntimeException("Failed to load schema from " + FLAT_COLLECTION_SCHEMA_PATH, e); + } + } + + protected static void initMongo() { + mongoContainer = + new GenericContainer<>(DockerImageName.parse("mongo:8.0.1")) + .withExposedPorts(27017) + .waitingFor(Wait.forListeningPort()); + mongoContainer.start(); + + Map mongoConfig = new HashMap<>(); + mongoConfig.put("host", "localhost"); + mongoConfig.put("port", mongoContainer.getMappedPort(27017).toString()); + + mongoDatastore = DatastoreProvider.getDatastore("Mongo", ConfigFactory.parseMap(mongoConfig)); + LOGGER.info("Mongo datastore initialized"); + } + + protected static void shutdownMongo() { + if (mongoContainer != null) { + mongoContainer.stop(); + } + } + + protected static void initPostgres() { + postgresContainer = + new GenericContainer<>(DockerImageName.parse("postgres:13.1")) + .withEnv("POSTGRES_PASSWORD", "postgres") + .withEnv("POSTGRES_USER", "postgres") + .withExposedPorts(5432) + .waitingFor(Wait.forListeningPort()); + postgresContainer.start(); + + String postgresConnectionUrl = + String.format("jdbc:postgresql://localhost:%s/", postgresContainer.getMappedPort(5432)); + + Map postgresConfig = new HashMap<>(); + postgresConfig.put("url", postgresConnectionUrl); + postgresConfig.put("user", "postgres"); + postgresConfig.put("password", "postgres"); + + postgresDatastore = + DatastoreProvider.getDatastore("Postgres", ConfigFactory.parseMap(postgresConfig)); + LOGGER.info("Postgres datastore initialized"); + } + + protected static void shutdownPostgres() { + if (postgresContainer != null) { + postgresContainer.stop(); + } + } + + protected static void createFlatCollectionSchema( + PostgresDatastore pgDatastore, String tableName) { + String createTableSQL = String.format(loadFlatCollectionSchema(), tableName); + + try (Connection connection = pgDatastore.getPostgresClient(); + PreparedStatement statement = connection.prepareStatement(createTableSQL)) { + statement.execute(); + LOGGER.info("Created flat collection table: {}", tableName); + } catch (Exception e) { + LOGGER.error("Failed to create flat collection schema: {}", e.getMessage(), e); + throw new RuntimeException("Failed to create flat collection schema", e); + } + } + + protected static String generateDocId(String prefix) { + return prefix + "-" + System.currentTimeMillis() + "-" + (int) (Math.random() * 10000); + } + + protected static String getKeyString(String docId) { + return new SingleValueKey(DEFAULT_TENANT, docId).toString(); + } + + protected Query buildQueryById(String docId) { + return Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of(getKeyString(docId)))) + .build(); + } + + protected Document createTestDocument(String docId) { + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + String keyStr = key.toString(); + + ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); + objectNode.put("id", keyStr); + objectNode.put("item", "TestItem"); + objectNode.put("price", 100); + objectNode.put("quantity", 50); + objectNode.put("in_stock", true); + objectNode.put("big_number", 1000000000000L); + objectNode.put("rating", 3.5); + objectNode.put("weight", 50.0); + objectNode.putArray("tags").add("tag1").add("tag2"); + objectNode.putArray("numbers").add(1).add(2).add(3); + ObjectNode props = OBJECT_MAPPER.createObjectNode(); + props.put("brand", "TestBrand"); + props.put("size", "M"); + props.put("count", 10); + props.putArray("colors").add("red").add("blue"); + objectNode.set("props", props); + ObjectNode sales = OBJECT_MAPPER.createObjectNode(); + sales.put("total", 200); + sales.put("count", 10); + objectNode.set("sales", sales); + + return new JSONDocument(objectNode); + } + + protected Key createKey(String docId) { + return new SingleValueKey(DEFAULT_TENANT, docId); + } + + protected static void clearTable(String tableName) { + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + String deleteSQL = String.format("DELETE FROM \"%s\"", tableName); + try (Connection connection = pgDatastore.getPostgresClient(); + PreparedStatement statement = connection.prepareStatement(deleteSQL)) { + statement.executeUpdate(); + } catch (Exception e) { + LOGGER.error("Failed to clear table {}: {}", tableName, e.getMessage(), e); + } + } + + protected void insertTestDocument(String docId, Collection collection) throws IOException { + Key key = createKey(docId); + Document document = createTestDocument(docId); + collection.upsert(key, document); + } + + /** Provides all MissingColumnStrategy values for parameterized tests */ + protected static class MissingColumnStrategyProvider implements ArgumentsProvider { + @Override + public Stream provideArguments(ExtensionContext context) { + return Stream.of( + Arguments.of(MissingColumnStrategy.SKIP), + Arguments.of(MissingColumnStrategy.THROW), + Arguments.of(MissingColumnStrategy.IGNORE_DOCUMENT)); + } + } +} 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 ed79c0f71..8bcce2256 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 @@ -9,7 +9,6 @@ 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.google.common.base.Preconditions; import com.typesafe.config.ConfigFactory; @@ -17,16 +16,15 @@ import java.sql.Connection; import java.sql.PreparedStatement; import java.sql.ResultSet; -import java.sql.Timestamp; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; -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; @@ -48,108 +46,28 @@ import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; -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; - -/** - * Integration tests for write operations on flat PostgreSQL collections. - * - *

Flat collections are PostgreSQL tables with explicit column schemas (not JSONB-based nested - * documents). This test class verifies that Collection interface write operations work correctly on - * such collections. - */ + @Testcontainers -public class FlatCollectionWriteTest { +public class FlatCollectionWriteTest extends BaseWriteTest { - private static final Logger LOGGER = LoggerFactory.getLogger(FlatCollectionWriteTest.class); - private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); private static final String FLAT_COLLECTION_NAME = "myTestFlat"; private static final String INSERT_STATEMENTS_FILE = "query/pg_flat_collection_insert.json"; - private static final String DEFAULT_TENANT = "default"; - private static Datastore postgresDatastore; private static Collection flatCollection; - private static GenericContainer postgres; @BeforeAll public static void init() throws IOException { - 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"); - postgresConfig.put( - "postgres.collectionConfigs." + FLAT_COLLECTION_NAME + ".timestampFields.created", - "createdTime"); - postgresConfig.put( - "postgres.collectionConfigs." + FLAT_COLLECTION_NAME + ".timestampFields.lastUpdated", - "lastUpdateTime"); - - postgresDatastore = - DatastoreProvider.getDatastore("Postgres", ConfigFactory.parseMap(postgresConfig)); + initPostgres(); LOGGER.info("Postgres datastore initialized: {}", postgresDatastore.listCollections()); - createFlatCollectionSchema(); + createFlatCollectionSchema((PostgresDatastore) postgresDatastore, FLAT_COLLECTION_NAME); flatCollection = postgresDatastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT); } - private static void createFlatCollectionSchema() { - String createTableSQL = - String.format( - "CREATE TABLE \"%s\" (" - + "\"id\" TEXT PRIMARY KEY," - + "\"item\" TEXT," - + "\"price\" INTEGER," - + "\"quantity\" INTEGER," - + "\"date\" TIMESTAMPTZ," - + "\"in_stock\" BOOLEAN," - + "\"tags\" TEXT[]," - + "\"categoryTags\" TEXT[]," - + "\"props\" JSONB," - + "\"sales\" JSONB," - + "\"numbers\" INTEGER[]," - + "\"scores\" DOUBLE PRECISION[]," - + "\"flags\" BOOLEAN[]," - + "\"big_number\" BIGINT," - + "\"rating\" REAL," - + "\"created_date\" DATE," - + "\"weight\" DOUBLE PRECISION," - + "\"createdTime\" BIGINT," - + "\"lastUpdateTime\" TIMESTAMP WITH TIME ZONE" - + ");", - FLAT_COLLECTION_NAME); - - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - - try (Connection connection = pgDatastore.getPostgresClient(); - PreparedStatement statement = connection.prepareStatement(createTableSQL)) { - statement.execute(); - LOGGER.info("Created flat collection table: {}", FLAT_COLLECTION_NAME); - } catch (Exception e) { - LOGGER.error("Failed to create flat collection schema: {}", e.getMessage(), e); - } - } - private static void executeInsertStatements() { PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; try { @@ -183,22 +101,10 @@ private static void executeInsertStatements() { @BeforeEach public void setupData() { // Clear and repopulate with initial data before each test - clearTable(); + clearTable(FLAT_COLLECTION_NAME); executeInsertStatements(); } - private static void clearTable() { - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - String deleteSQL = String.format("DELETE FROM \"%s\"", FLAT_COLLECTION_NAME); - try (Connection connection = pgDatastore.getPostgresClient(); - PreparedStatement statement = connection.prepareStatement(deleteSQL)) { - statement.executeUpdate(); - LOGGER.info("Cleared table: {}", FLAT_COLLECTION_NAME); - } catch (Exception e) { - LOGGER.error("Failed to clear table: {}", e.getMessage(), e); - } - } - @AfterEach public void cleanup() { // Data is cleared in @BeforeEach, but cleanup here for safety @@ -206,7 +112,7 @@ public void cleanup() { @AfterAll public static void shutdown() { - postgres.stop(); + shutdownPostgres(); } @Nested @@ -216,7 +122,7 @@ class UpsertTests { @Test @DisplayName("Should create new document when key doesn't exist and return true") void testUpsertNewDocument() throws Exception { - String docId = getRandomDocId(4); + String docId = generateDocId("test"); ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); objectNode.put("id", docId); @@ -227,7 +133,7 @@ void testUpsertNewDocument() throws Exception { boolean isNew = flatCollection.upsert(key, document); - assertTrue(isNew, "Should return true for new document"); + assertTrue(isNew); queryAndAssert( key, @@ -241,7 +147,7 @@ void testUpsertNewDocument() throws Exception { @Test @DisplayName("Should merge with existing document preserving unspecified fields") void testUpsertMergesWithExistingDocument() throws Exception { - String docId = getRandomDocId(4); + String docId = generateDocId("test"); // First, create a document with multiple fields ObjectNode initialNode = OBJECT_MAPPER.createObjectNode(); @@ -254,7 +160,7 @@ void testUpsertMergesWithExistingDocument() throws Exception { Key key = new SingleValueKey(DEFAULT_TENANT, docId); boolean firstResult = flatCollection.upsert(key, initialDoc); - assertTrue(firstResult, "First upsert should create new document"); + assertTrue(firstResult); // Now upsert with only some fields - others should be PRESERVED (merge behavior) ObjectNode mergeNode = OBJECT_MAPPER.createObjectNode(); @@ -264,7 +170,7 @@ void testUpsertMergesWithExistingDocument() throws Exception { Document mergeDoc = new JSONDocument(mergeNode); boolean secondResult = flatCollection.upsert(key, mergeDoc); - assertFalse(secondResult, "Second upsert should update existing document"); + assertTrue(secondResult); // Verify merge behavior: item updated, price/quantity/in_stock preserved queryAndAssert( @@ -282,8 +188,8 @@ void testUpsertMergesWithExistingDocument() throws Exception { @Test @DisplayName("Upsert vs CreateOrReplace: upsert preserves, createOrReplace resets to default") void testUpsertVsCreateOrReplaceBehavior() throws Exception { - String docId1 = getRandomDocId(4); - String docId2 = getRandomDocId(4); + String docId1 = generateDocId("test"); + String docId2 = generateDocId("test"); // Setup: Create two identical documents ObjectNode initialNode = OBJECT_MAPPER.createObjectNode(); @@ -341,7 +247,7 @@ void testUpsertVsCreateOrReplaceBehavior() throws Exception { @Test @DisplayName("Should skip unknown fields in upsert (default SKIP strategy)") void testUpsertSkipsUnknownFields() throws Exception { - String docId = getRandomDocId(4); + String docId = generateDocId("test"); ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); objectNode.put("id", docId); @@ -386,7 +292,7 @@ class CreateTests { @DisplayName("Should create document with all supported data types") void testCreateWithAllDataTypes() throws Exception { ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - String docId = getRandomDocId(4); + String docId = generateDocId("test"); objectNode.put("id", docId); objectNode.put("item", "Comprehensive Test Item"); @@ -489,7 +395,7 @@ void testCreateWithAllDataTypes() throws Exception { @DisplayName("Should throw DuplicateDocumentException when creating with existing key") void testCreateDuplicateDocument() throws Exception { - String docId = getRandomDocId(4); + String docId = generateDocId("test"); ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); objectNode.put("id", "dup-doc-200"); objectNode.put("item", "First Item"); @@ -516,7 +422,7 @@ void testCreateDuplicateDocument() throws Exception { void testUnknownFieldsAsPerMissingColumnStrategy(MissingColumnStrategy missingColumnStrategy) throws Exception { - String docId = getRandomDocId(4); + String docId = generateDocId("test"); ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); objectNode.put("id", docId); @@ -539,7 +445,7 @@ void testUnknownFieldsAsPerMissingColumnStrategy(MissingColumnStrategy missingCo FLAT_COLLECTION_NAME, key)); ResultSet rs = ps.executeQuery()) { assertTrue(rs.next()); - assertEquals(0, rs.getInt(1), "Document should not exist in DB after exception"); + assertEquals(0, rs.getInt(1)); } } else { CreateResult result = flatCollection.create(key, document); @@ -565,7 +471,7 @@ void testEmptyMissingColumnStrategyConfigUsesDefault() throws Exception { Collection collectionWithEmptyStrategy = getFlatCollectionWithStrategy(""); // Test that it uses default SKIP strategy (unknown fields are skipped, not thrown) - String docId = getRandomDocId(4); + String docId = generateDocId("test"); ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); objectNode.put("id", docId); objectNode.put("item", "Test Item"); @@ -586,7 +492,7 @@ void testEmptyMissingColumnStrategyConfigUsesDefault() throws Exception { void testInvalidMissingColumnStrategyConfigUsesDefault() throws Exception { Collection collectionWithInvalidStrategy = getFlatCollectionWithStrategy("INVALID_STRATEGY"); - String docId = getRandomDocId(4); + String docId = generateDocId("test"); ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); objectNode.put("id", docId); objectNode.put("item", "Test Item"); @@ -705,7 +611,7 @@ void testCreateRefreshesSchemaOnUndefinedColumnError() throws Exception { void testUnparsableValuesAsPerMissingColStrategy(MissingColumnStrategy missingColumnStrategy) throws Exception { - String docId = getRandomDocId(4); + String docId = generateDocId("test"); // Try to insert a string value into an integer column with wrong type // The unparseable column should be skipped, not throw an exception @@ -759,14 +665,9 @@ void testUnparsableValuesAsPerMissingColStrategy(MissingColumnStrategy missingCo } } - private String getRandomDocId(int len) { - return org.testcontainers.shaded.org.apache.commons.lang3.RandomStringUtils.random( - len, true, false); - } - private static Collection getFlatCollectionWithStrategy(String strategy) { String postgresConnectionUrl = - String.format("jdbc:postgresql://localhost:%s/", postgres.getMappedPort(5432)); + String.format("jdbc:postgresql://localhost:%s/", postgresContainer.getMappedPort(5432)); Map configWithStrategy = new HashMap<>(); configWithStrategy.put("url", postgresConnectionUrl); @@ -798,19 +699,6 @@ interface ResultSetConsumer { void accept(ResultSet rs) throws Exception; } - static class MissingColumnStrategyProvider implements ArgumentsProvider { - - @Override - public Stream provideArguments(ExtensionContext context) { - return Stream.of(MissingColumnStrategy.values()) - .filter( - strategy -> - (strategy == MissingColumnStrategy.THROW) - || (strategy == MissingColumnStrategy.SKIP)) - .map(Arguments::of); - } - } - @Nested @DisplayName("CreateOrReplace Operations") class CreateOrReplaceTests { @@ -820,7 +708,7 @@ class CreateOrReplaceTests { "Should create new document and return true. Cols not specified should be set of default NULL") void testCreateOrReplaceNewDocument() throws Exception { - String docId = getRandomDocId(4); + String docId = generateDocId("test"); ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); objectNode.put("id", "upsert-new-doc-100"); @@ -851,7 +739,7 @@ void testCreateOrReplaceNewDocument() throws Exception { @Test @DisplayName("Should replace existing document and return false") void testCreateOrReplaceExistingDocument() throws Exception { - String docId = getRandomDocId(4); + String docId = generateDocId("test"); ObjectNode initialNode = OBJECT_MAPPER.createObjectNode(); initialNode.put("id", docId); initialNode.put("item", "Original Item"); @@ -913,7 +801,7 @@ void testCreateOrReplaceSkipsUnknownFields() throws Exception { @Test @DisplayName("Should handle JSONB fields in createOrReplace") void testCreateOrReplaceWithJsonbField() throws Exception { - String docId = getRandomDocId(4); + String docId = generateDocId("test"); ObjectNode initialNode = OBJECT_MAPPER.createObjectNode(); initialNode.put("id", docId); initialNode.put("item", "Item with props"); @@ -1904,6 +1792,70 @@ void testUpdateWithCondition() { UnsupportedOperationException.class, () -> flatCollection.update(key, document, condition)); } + + @Test + @DisplayName("Should return empty when no document matches query") + void testUpdateNoMatch() throws Exception { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("id"), + RelationalOperator.EQ, + ConstantExpression.of("9999"))) + .build(); + + List updates = List.of(SubDocumentUpdate.of("price", 100)); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + Optional result = flatCollection.update(query, updates, options); + + assertTrue(result.isEmpty()); + } + + @Test + @DisplayName("Should throw IOException when column does not exist") + void testUpdateNonExistentColumn() { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("_id"), + RelationalOperator.EQ, + ConstantExpression.of(1))) + .build(); + + List updates = + List.of(SubDocumentUpdate.of("nonexistent_column", "value")); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + assertThrows(IOException.class, () -> flatCollection.update(query, updates, options)); + } + + @Test + @DisplayName("Should throw IOException when nested path on non-JSONB column") + void testUpdateNestedPathOnNonJsonbColumn() { + Query query = + Query.builder() + .setFilter( + RelationalExpression.of( + IdentifierExpression.of("_id"), + RelationalOperator.EQ, + ConstantExpression.of(1))) + .build(); + + // "item" is TEXT, not JSONB - nested path should fail + List updates = List.of(SubDocumentUpdate.of("item.nested", "value")); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + assertThrows(IOException.class, () -> flatCollection.update(query, updates, options)); + } } @Nested @@ -1915,127 +1867,139 @@ class SubDocUpdateTests { class SetOperatorTests { @Test - @DisplayName("Should update multiple top-level columns in single update") - void testSetMultipleColumns() throws Exception { + @DisplayName("Cases 1-4: SET all field types via bulkUpdate") + void testSetAllFieldTypes() throws Exception { Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of("2"))) + ConstantExpression.of("1"))) .build(); List updates = - List.of(SubDocumentUpdate.of("price", 555), SubDocumentUpdate.of("quantity", 100)); + List.of( + // Case 1: Top-level primitives + SubDocumentUpdate.of("item", "UpdatedItem"), + SubDocumentUpdate.of("price", 999), + SubDocumentUpdate.of("quantity", 50), + SubDocumentUpdate.of("in_stock", false), + SubDocumentUpdate.of("big_number", 9999999999L), + SubDocumentUpdate.of("rating", 4.5f), + SubDocumentUpdate.of("weight", 123.456), + // Case 2: Top-level arrays + SubDocumentUpdate.of("tags", new String[] {"tag4", "tag5", "tag6"}), + SubDocumentUpdate.of("numbers", new Integer[] {10, 20, 30}), + SubDocumentUpdate.of("scores", new Double[] {1.1, 2.2, 3.3}), + SubDocumentUpdate.of("flags", new Boolean[] {true, false, true}), + // Case 3 & 4: One nested path in JSONB (props) - tests nested primitive + SubDocumentUpdate.of("props.brand", "NewBrand"), + // Use 'sales' JSONB column for nested array test + SubDocumentUpdate.of( + "sales.regions", SubDocumentValue.of(new String[] {"US", "EU", "APAC"}))); UpdateOptions options = UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - Optional result = flatCollection.update(query, updates, options); + // Read expected values from JSON file + String expectedJsonContent = + readFileFromResource("expected/set_all_field_types_expected.json").orElseThrow(); + ObjectNode expectedJson = (ObjectNode) OBJECT_MAPPER.readTree(expectedJsonContent); - assertTrue(result.isPresent()); - JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - assertEquals(555, resultJson.get("price").asInt()); - assertEquals(100, resultJson.get("quantity").asInt()); + try (CloseableIterator results = + flatCollection.bulkUpdate(query, updates, options)) { + assertTrue(results.hasNext()); + Document resultDoc = results.next(); + ObjectNode resultJson = (ObjectNode) OBJECT_MAPPER.readTree(resultDoc.toJson()); - // 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")); + // Remove 'date' field from comparison - it's timezone-dependent and not updated in this + // test + expectedJson.remove("date"); + resultJson.remove("date"); + + assertEquals(expectedJson, resultJson); } } @Test - @DisplayName("Should update nested path in JSONB column") - void testUpdateNestedJsonbPath() throws Exception { + @DisplayName("Case 6: SET on non-existent top-level column should skip by default") + void testSetNonExistentTopLevelColumnSkips() throws Exception { Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of("3"))) + ConstantExpression.of("1"))) .build(); - // Update props.brand nested path List updates = - List.of(SubDocumentUpdate.of("props.brand", "UpdatedBrand")); - + List.of( + SubDocumentUpdate.of("nonexistent_column1", "some_value"), + SubDocumentUpdate.of("nonexistent_column2.value", "some_value")); UpdateOptions options = UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); Optional result = flatCollection.update(query, updates, options); + // Document returned (unchanged since update was skipped) 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 + // Verify original data is intact 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)); + "SELECT \"item\" FROM \"%s\" WHERE \"id\" = '1'", FLAT_COLLECTION_NAME)); ResultSet rs = ps.executeQuery()) { assertTrue(rs.next()); - assertEquals("UpdatedBrand", rs.getString("brand")); + assertEquals("Soap", rs.getString("item")); } } @Test - @DisplayName("Should return BEFORE_UPDATE document") - void testUpdateReturnsBeforeDocument() throws Exception { - // First get the current price + @DisplayName("Case 7b: SET nested path in NULL JSONB column should create structure") + void testSetNestedPathInNullJsonbColumn() throws Exception { + // Row 2 has props = NULL Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of("4"))) + ConstantExpression.of("2"))) .build(); - List updates = List.of(SubDocumentUpdate.of("price", 777)); - + // In this case, props is NULL + List updates = List.of(SubDocumentUpdate.of("props.newKey", "newValue")); UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.BEFORE_UPDATE).build(); + 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()); - // Should return the old price (5 from initial data), not the new one (777) - assertEquals(5, resultJson.get("price").asInt()); + assertEquals("newValue", resultJson.get("props").get("newKey").asText()); - // But database should have the new value + // Verify in database 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)); + "SELECT \"props\"->>'newKey' as newKey FROM \"%s\" WHERE \"id\" = '2'", + FLAT_COLLECTION_NAME)); ResultSet rs = ps.executeQuery()) { assertTrue(rs.next()); - assertEquals(777, rs.getInt("price")); + assertEquals("newValue", rs.getString("newKey")); } } @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 + @DisplayName("Case 7c: SET non-existent nested path in existing JSONB should create key") + void testSetNonExistentNestedPathInExistingJsonb() throws Exception { Query query = Query.builder() .setFilter( @@ -2045,134 +2009,94 @@ void testSetFieldNotInSchema() throws Exception { 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(); + // In this case, props exists but props.newAttribute doesn't exist. + List updates = + List.of(SubDocumentUpdate.of("props.newAttribute", "brandNewValue")); + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).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()); + Optional result = flatCollection.update(query, updates, options); - // 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")); - } + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + assertEquals("brandNewValue", resultJson.get("props").get("newAttribute").asText()); + // Existing data should be preserved + assertEquals("Dettol", resultJson.get("props").get("brand").asText()); } @Test - @DisplayName("Case 2: SET on JSONB column that is NULL should create the structure") - void testSetJsonbColumnIsNull() throws Exception { - // Row 2 has props = NULL + @DisplayName("SET should return correct document based on ReturnDocumentType") + void testSetReturnDocumentTypes() throws Exception { + PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; 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")) + ConstantExpression.of("4"))) .build(); - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder() - .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) - .build()); + // Test BEFORE_UPDATE - returns old value + List updates1 = List.of(SubDocumentUpdate.of("price", 777)); + UpdateOptions beforeOptions = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.BEFORE_UPDATE).build(); - assertTrue(result.isPresent()); + Optional beforeResult = flatCollection.update(query, updates1, beforeOptions); + assertTrue(beforeResult.isPresent()); + JsonNode beforeJson = OBJECT_MAPPER.readTree(beforeResult.get().toJson()); + assertEquals(5, beforeJson.get("price").asInt()); // Old value - // Verify props now has the new key - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + // Verify database has new value try (Connection conn = pgDatastore.getPostgresClient(); PreparedStatement ps = conn.prepareStatement( String.format( - "SELECT \"props\"->>'newKey' as newKey FROM \"%s\" WHERE \"id\" = '2'", - FLAT_COLLECTION_NAME)); + "SELECT \"price\" FROM \"%s\" WHERE \"id\" = '4'", FLAT_COLLECTION_NAME)); ResultSet rs = ps.executeQuery()) { assertTrue(rs.next()); - assertEquals("newValue", rs.getString("newKey")); + assertEquals(777, rs.getInt("price")); } - } - @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(); + // Test AFTER_UPDATE - returns new value + List updates2 = List.of(SubDocumentUpdate.of("price", 888)); + UpdateOptions afterOptions = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - SubDocumentUpdate update = - SubDocumentUpdate.builder() - .subDocument("props.brand") - .operator(UpdateOperator.SET) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - "UpdatedBrand")) - .build(); + Optional afterResult = flatCollection.update(query, updates2, afterOptions); + assertTrue(afterResult.isPresent()); + JsonNode afterJson = OBJECT_MAPPER.readTree(afterResult.get().toJson()); + assertEquals(888, afterJson.get("price").asInt()); // New value - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder() - .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) - .build()); + // Test NONE - returns empty + List updates3 = List.of(SubDocumentUpdate.of("price", 999)); + UpdateOptions noneOptions = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.NONE).build(); - assertTrue(result.isPresent()); + Optional noneResult = flatCollection.update(query, updates3, noneOptions); + assertFalse(noneResult.isPresent()); - // Verify props.brand was updated - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; + // Verify database has the final value try (Connection conn = pgDatastore.getPostgresClient(); PreparedStatement ps = conn.prepareStatement( String.format( - "SELECT \"props\"->>'brand' as brand FROM \"%s\" WHERE \"id\" = '1'", - FLAT_COLLECTION_NAME)); + "SELECT \"price\" FROM \"%s\" WHERE \"id\" = '4'", FLAT_COLLECTION_NAME)); ResultSet rs = ps.executeQuery()) { assertTrue(rs.next()); - assertEquals("UpdatedBrand", rs.getString("brand")); + assertEquals(999, rs.getInt("price")); } } + } + + @Nested + @DisplayName("UNSET Operator Tests") + class UnsetOperatorTests { @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 + @DisplayName("Should UNSET top-level column and nested JSONB field via bulkUpdate") + void testUnsetTopLevelAndNestedFields() throws Exception { + // Row 1 has item="Soap" and props.brand="Dettol" Query query = Query.builder() .setFilter( @@ -2182,1035 +2106,295 @@ void testSetJsonbPathDoesNotExist() throws Exception { 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) + // UNSET both top-level column and nested JSONB field in one operation + List updates = + List.of( + // Top-level: sets column to NULL + SubDocumentUpdate.builder() + .subDocument("item") + .operator(UpdateOperator.UNSET) + .build(), + // Nested JSONB: removes key from JSON object + SubDocumentUpdate.builder() + .subDocument("props.brand") + .operator(UpdateOperator.UNSET) + .build(), + // non existent columns. Shouldn't fail + SubDocumentUpdate.builder() + .subDocument("nonexistentCol") + .operator(UpdateOperator.UNSET) + .build(), + SubDocumentUpdate.builder() + .subDocument("nonexistentCol.key") + .operator(UpdateOperator.UNSET) .build()); - assertTrue(result.isPresent()); + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + try (CloseableIterator results = + flatCollection.bulkUpdate(query, updates, options)) { + assertTrue(results.hasNext()); + Document resultDoc = results.next(); + JsonNode resultJson = OBJECT_MAPPER.readTree(resultDoc.toJson()); + + // Verify top-level column is NULL + JsonNode itemNode = resultJson.get("item"); + assertTrue(itemNode == null || itemNode.isNull()); - // Verify props.newAttribute was created + // Verify nested JSONB key is removed, but other keys preserved + assertFalse(resultJson.get("props").has("brand")); + assertEquals("M", resultJson.get("props").get("size").asText()); + } + + // Verify in database 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'", + "SELECT \"item\", \"props\" 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")); + assertNull(rs.getString("item")); + JsonNode propsJson = OBJECT_MAPPER.readTree(rs.getString("props")); + assertFalse(propsJson.has("brand")); + assertEquals("M", propsJson.get("size").asText()); } } + } + + @Nested + @DisplayName("ADD Operator Tests") + class AddSubdocOperatorTests { @Test - @DisplayName("SET on top-level column should update the value directly") - void testSetTopLevelColumn() throws Exception { + @DisplayName("Should ADD to all numeric types via bulkUpdate") + void testAddAllNumericTypes() throws Exception { + String docId = generateDocId("test"); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "NumericTestItem"); + node.put("price", 100); // INT (positive ADD) + node.put("quantity", 50); // INT (negative ADD - decrement) + node.put("big_number", 1000000000000L); // BIGINT + node.put("rating", 3.5); // REAL + node.put("weight", 50.0); // DOUBLE PRECISION + ObjectNode sales = OBJECT_MAPPER.createObjectNode(); + sales.put("total", 200); // Nested JSONB numeric + sales.put("count", 10); + node.set("sales", sales); + flatCollection.create(key, new JSONDocument(node)); + 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")) + ConstantExpression.of(key.toString()))) .build(); - Optional result = - flatCollection.update( - query, - List.of(update), - UpdateOptions.builder() - .returnDocumentType(ReturnDocumentType.AFTER_UPDATE) + List updates = + List.of( + // Top-level INT: 100 + 5 = 105 + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(5)) + .build(), + // Top-level INT (negative): 50 + (-15) = 35 + SubDocumentUpdate.builder() + .subDocument("quantity") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(-15)) + .build(), + // Top-level BIGINT: 1000000000000 + 500 = 1000000000500 + SubDocumentUpdate.builder() + .subDocument("big_number") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(500L)) + .build(), + // Top-level REAL: 3.5 + 1.0 = 4.5 + SubDocumentUpdate.builder() + .subDocument("rating") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(1.0f)) + .build(), + // Top-level DOUBLE: 50.0 + 2.5 = 52.5 + SubDocumentUpdate.builder() + .subDocument("weight") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(2.5)) + .build(), + // Nested JSONB: 200 + 50 = 250 + SubDocumentUpdate.builder() + .subDocument("sales.total") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(50)) .build()); - assertTrue(result.isPresent()); + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + String expectedJsonContent = + readFileFromResource("expected/add_all_numeric_types_expected.json").orElseThrow(); + JsonNode expectedJson = OBJECT_MAPPER.readTree(expectedJsonContent); + + try (CloseableIterator results = + flatCollection.bulkUpdate(query, updates, options)) { + assertTrue(results.hasNext()); + JsonNode resultJson = OBJECT_MAPPER.readTree(results.next().toJson()); + + ((ObjectNode) resultJson).remove("id"); + assertEquals(expectedJson, resultJson); + } - // Verify item was updated + // Verify in database 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)); + "SELECT \"price\", \"quantity\", \"big_number\", \"rating\", \"weight\", \"sales\" " + + "FROM \"%s\" WHERE \"id\" = '%s'", + FLAT_COLLECTION_NAME, key)); ResultSet rs = ps.executeQuery()) { assertTrue(rs.next()); - assertEquals("UpdatedSoap", rs.getString("item")); + assertEquals(expectedJson.get("price").asInt(), rs.getInt("price")); + assertEquals(expectedJson.get("quantity").asInt(), rs.getInt("quantity")); + assertEquals(expectedJson.get("big_number").asLong(), rs.getLong("big_number")); + assertEquals(expectedJson.get("rating").floatValue(), rs.getFloat("rating"), 0.01f); + assertEquals(expectedJson.get("weight").asDouble(), rs.getDouble("weight"), 0.01); + JsonNode salesJson = OBJECT_MAPPER.readTree(rs.getString("sales")); + assertEquals( + expectedJson.get("sales").get("total").asInt(), salesJson.get("total").asInt()); + assertEquals( + expectedJson.get("sales").get("count").asInt(), salesJson.get("count").asInt()); } } @Test - @DisplayName("SET with empty object value") - void testSetWithEmptyObjectValue() throws Exception { + @DisplayName("Should handle ADD on NULL column (treat as 0)") + void testAddOnNullColumn() throws Exception { + // Create a document with NULL numeric columns + String docId = generateDocId("test"); + Key key = new SingleValueKey(DEFAULT_TENANT, docId); + ObjectNode node = OBJECT_MAPPER.createObjectNode(); + node.put("item", "NullPriceItem"); + // price, weight are 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("1"))) + ConstantExpression.of(key.toString()))) .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( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("1"))) - .build(); - - // ADD 5 to price (10 + 5 = 15) - List updates = - List.of( - SubDocumentUpdate.builder() - .subDocument("price") - .operator(UpdateOperator.ADD) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(5)) - .build()); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - Optional result = flatCollection.update(query, updates, options); - - assertTrue(result.isPresent()); - JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - assertEquals(15, resultJson.get("price").asInt()); - - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"price\" FROM \"%s\" WHERE \"id\" = '1'", FLAT_COLLECTION_NAME)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - assertEquals(15, rs.getInt("price")); - } - } - - @Test - @DisplayName("Should handle ADD on NULL column (treat as 0)") - void testAddOnNullColumn() throws Exception { - // Create a document with NULL price - String docId = 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( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("2"))) - .build(); - - // ADD -5 to price (20 - 5 = 15) + // ADD to NULL columns - COALESCE(NULL, 0) + value List updates = List.of( SubDocumentUpdate.builder() .subDocument("price") .operator(UpdateOperator.ADD) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(-5)) - .build()); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - Optional result = flatCollection.update(query, updates, options); - - assertTrue(result.isPresent()); - JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - assertEquals(15, resultJson.get("price").asInt()); - } - - @Test - @DisplayName("Should ADD with floating point value") - void testAddFloatingPointValue() throws Exception { - // Row 3 has price = 30 - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("3"))) - .build(); - - // ADD 0.5 to price (30 + 0.5 = 30.5, but price is INTEGER so it might truncate) - // Testing with a column that supports decimals - weight is DOUBLE PRECISION - List updates = - List.of( + .subDocumentValue(SubDocumentValue.of(100)) + .build(), SubDocumentUpdate.builder() .subDocument("weight") .operator(UpdateOperator.ADD) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(2.5)) - .build()); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - Optional result = flatCollection.update(query, updates, options); - - assertTrue(result.isPresent()); - JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - // Initial weight is NULL, so COALESCE(NULL, 0) + 2.5 = 2.5 - assertEquals(2.5, resultJson.get("weight").asDouble(), 0.01); - } - - @Test - @DisplayName("Should ADD to nested JSONB numeric field") - void testAddNestedJsonbField() throws Exception { - // First, set up a document with a JSONB field containing a numeric value - String docId = 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(key.toString()))) - .build(); - - // ADD 50 to sales.total (100 + 50 = 150) - List updates = - List.of( - SubDocumentUpdate.builder() - .subDocument("sales.total") - .operator(UpdateOperator.ADD) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(50)) + .subDocumentValue(SubDocumentValue.of(25.5)) .build()); UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - Optional result = flatCollection.update(query, updates, options); - - assertTrue(result.isPresent()); - JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); - assertEquals(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.AFTER_UPDATE).build(); - - assertThrows( - IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); - } - - @Test - @DisplayName("Should throw IllegalArgumentException for multi-valued nested document value") - void testAddMultiValuedNestedDocumentValue() throws Exception { - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("1"))) - .build(); - - // 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( - new Document[] { - new JSONDocument("{\"a\": 1}"), new JSONDocument("{\"b\": 2}") - })) - .build()); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - assertThrows( - IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); - } - - @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", "BigintItem"); - node.put("big_number", 1000000000000L); - flatCollection.create(key, new JSONDocument(node)); - - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of(key.toString()))) - .build(); - - // ADD 500 to big_number - List updates = - List.of( - SubDocumentUpdate.builder() - .subDocument("big_number") - .operator(UpdateOperator.ADD) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(500L)) - .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(1000000000500L, resultJson.get("big_number").asLong()); - } - - @Test - @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(key.toString()))) - .build(); - - // ADD 1.0 to rating (3.5 + 1.0 = 4.5) - List updates = - List.of( - SubDocumentUpdate.builder() - .subDocument("rating") - .operator(UpdateOperator.ADD) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of(1.0)) - .build()); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - 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); - } - } - } - - @Nested - @DisplayName("APPEND_TO_LIST Operator Tests") - class AppendToListOperatorTests { - - @Test - @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(key.toString()))) - .build(); - - // Append new tags - List updates = - List.of( - SubDocumentUpdate.builder() - .subDocument("tags") - .operator(UpdateOperator.APPEND_TO_LIST) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new String[] {"newTag1", "newTag2"})) - .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 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 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", "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(key.toString()))) - .build(); - - // Append to props.colors - List updates = - List.of( - SubDocumentUpdate.builder() - .subDocument("props.colors") - .operator(UpdateOperator.APPEND_TO_LIST) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new String[] {"green", "yellow"})) - .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 colorsNode = resultJson.get("props").get("colors"); - assertTrue(colorsNode.isArray()); - assertEquals(4, colorsNode.size()); - } - - @Test - @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", "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 = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of(key.toString()))) - .build(); - - // Append to 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[] {"green", "yellow"})) - .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()); - - // 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 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(key.toString()))) - .build(); - - // Add tags - 'existing1' already exists, 'newTag' is new - List updates = - List.of( - SubDocumentUpdate.builder() - .subDocument("tags") - .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new String[] {"existing1", "newTag"})) - .build()); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - Optional result = flatCollection.update(query, updates, options); + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - 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; - } + try (CloseableIterator results = + flatCollection.bulkUpdate(query, updates, options)) { + assertTrue(results.hasNext()); + JsonNode resultJson = OBJECT_MAPPER.readTree(results.next().toJson()); + assertEquals(100, resultJson.get("price").asInt()); + assertEquals(25.5, resultJson.get("weight").asDouble(), 0.01); } - assertTrue(hasNewTag); } @Test - @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)); - + @DisplayName("Should throw IllegalArgumentException for non-numeric value") + void testAddNonNumericValue() { Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of(key.toString()))) + ConstantExpression.of("1"))) .build(); - // Add colors - 'red' already exists, 'green' is new List updates = List.of( SubDocumentUpdate.builder() - .subDocument("props.colors") - .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new String[] {"red", "green"})) + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of("not-a-number")) .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 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()); + assertThrows( + IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); } @Test - @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)); - + @DisplayName("Should throw IllegalArgumentException for array value") + void testAddArrayValue() { Query query = Query.builder() .setFilter( RelationalExpression.of( IdentifierExpression.of("id"), RelationalOperator.EQ, - ConstantExpression.of(key.toString()))) + ConstantExpression.of("1"))) .build(); - // Add tags that already exist List updates = List.of( SubDocumentUpdate.builder() - .subDocument("tags") - .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new String[] {"tag1", "tag2"})) + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(new Integer[] {1, 2, 3})) .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 tagsNode = resultJson.get("tags"); - assertTrue(tagsNode.isArray()); - assertEquals(2, tagsNode.size()); - assertEquals("tag1", tagsNode.get(0).asText()); - assertEquals("tag2", tagsNode.get(1).asText()); + assertThrows( + IllegalArgumentException.class, () -> flatCollection.update(query, updates, options)); } } @Nested - @DisplayName("REMOVE_ALL_FROM_LIST Operator Tests") - class RemoveAllFromListOperatorTests { + @DisplayName("APPEND_TO_LIST Operator Tests") + class AppendToListOperatorTests { @Test - @DisplayName("Should remove values from top-level array column") - void testRemoveAllFromTopLevelArray() throws Exception { - String docId = getRandomDocId(4); + @DisplayName("Should APPEND_TO_LIST for top-level and nested arrays via bulkUpdate") + void testAppendToListAllCases() throws Exception { + String docId = generateDocId("test"); 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"); + node.put("item", "AppendTestItem"); + node.putArray("tags").add("tag1").add("tag2"); // Top-level array (existing) + ObjectNode props = OBJECT_MAPPER.createObjectNode(); + props.putArray("colors").add("red").add("blue"); // Nested JSONB array (existing) + props.put("brand", "TestBrand"); + node.set("props", props); + ObjectNode sales = OBJECT_MAPPER.createObjectNode(); + sales.put("total", 100); // Nested JSONB without array + node.set("sales", sales); flatCollection.create(key, new JSONDocument(node)); Query query = @@ -3222,39 +2406,85 @@ void testRemoveAllFromTopLevelArray() throws Exception { ConstantExpression.of(key.toString()))) .build(); - // Remove 'tag1' from tags List updates = List.of( + // Top-level array: append to existing tags SubDocumentUpdate.builder() .subDocument("tags") - .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new String[] {"tag1"})) + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"newTag1", "newTag2"})) + .build(), + // Nested JSONB array: append to existing props.colors + SubDocumentUpdate.builder() + .subDocument("props.colors") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"green", "yellow"})) + .build(), + // Nested JSONB: append to non-existent array (creates it) + SubDocumentUpdate.builder() + .subDocument("sales.regions") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"US", "EU"})) .build()); UpdateOptions options = UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - Optional result = flatCollection.update(query, updates, options); + try (CloseableIterator results = + flatCollection.bulkUpdate(query, updates, options)) { + assertTrue(results.hasNext()); + JsonNode resultJson = OBJECT_MAPPER.readTree(results.next().toJson()); + + // Verify top-level array append + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(4, tagsNode.size()); + assertEquals("tag1", tagsNode.get(0).asText()); + assertEquals("tag2", tagsNode.get(1).asText()); + assertEquals("newTag1", tagsNode.get(2).asText()); + assertEquals("newTag2", tagsNode.get(3).asText()); + + // Verify nested JSONB array append + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(4, colorsNode.size()); + assertEquals("red", colorsNode.get(0).asText()); + assertEquals("blue", colorsNode.get(1).asText()); + assertEquals("green", colorsNode.get(2).asText()); + assertEquals("yellow", colorsNode.get(3).asText()); + + // Verify non-existent array was created + JsonNode regionsNode = resultJson.get("sales").get("regions"); + assertNotNull(regionsNode); + assertTrue(regionsNode.isArray()); + assertEquals(2, regionsNode.size()); + assertEquals("US", regionsNode.get(0).asText()); + assertEquals("EU", regionsNode.get(1).asText()); + + // Verify other fields preserved + assertEquals("TestBrand", resultJson.get("props").get("brand").asText()); + assertEquals(100, resultJson.get("sales").get("total").asInt()); + } - 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 + // todo: Add negative test cases based on Mongo's behaviour } + } + + @Nested + @DisplayName("ADD_TO_LIST_IF_ABSENT Operator Tests") + class AddToListIfAbsentOperatorTests { @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); + @DisplayName("Should ADD_TO_LIST_IF_ABSENT for top-level and nested arrays via bulkUpdate") + void testAddToListIfAbsentAllCases() throws Exception { + String docId = generateDocId("test"); Key key = new SingleValueKey(DEFAULT_TENANT, docId); ObjectNode node = OBJECT_MAPPER.createObjectNode(); - node.put("item", "JsonbArrayItem"); + node.put("item", "AddIfAbsentTestItem"); + node.putArray("tags").add("existing1").add("existing2"); // Top-level array + node.putArray("numbers").add(1).add(2); // Top-level (all duplicates test) ObjectNode props = OBJECT_MAPPER.createObjectNode(); - props.putArray("colors").add("red").add("blue").add("green"); + props.putArray("colors").add("red").add("blue"); // Nested JSONB array node.set("props", props); flatCollection.create(key, new JSONDocument(node)); @@ -3267,37 +2497,72 @@ void testRemoveAllFromNestedJsonbArray() throws Exception { ConstantExpression.of(key.toString()))) .build(); - // Remove 'red' and 'blue' from props.colors List updates = List.of( + // Top-level: 'existing1' exists, 'newTag' is new → adds only 'newTag' + SubDocumentUpdate.builder() + .subDocument("tags") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) + .subDocumentValue(SubDocumentValue.of(new String[] {"existing1", "newTag"})) + .build(), + // Nested JSONB: 'red' exists, 'green' is new → adds only 'green' SubDocumentUpdate.builder() .subDocument("props.colors") - .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new String[] {"red", "blue"})) + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) + .subDocumentValue(SubDocumentValue.of(new String[] {"red", "green"})) .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 colorsNode = resultJson.get("props").get("colors"); - assertTrue(colorsNode.isArray()); - assertEquals(1, colorsNode.size()); // Only 'green' remains + try (CloseableIterator results = + flatCollection.bulkUpdate(query, updates, options)) { + assertTrue(results.hasNext()); + JsonNode resultJson = OBJECT_MAPPER.readTree(results.next().toJson()); + + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(3, tagsNode.size()); + Set tagValues = new HashSet<>(); + tagsNode.forEach(n -> tagValues.add(n.asText())); + assertTrue(tagValues.contains("existing1")); + assertTrue(tagValues.contains("existing2")); + assertTrue(tagValues.contains("newTag")); + + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(3, colorsNode.size()); + Set colorValues = new HashSet<>(); + colorsNode.forEach(n -> colorValues.add(n.asText())); + assertTrue(colorValues.contains("red")); + assertTrue(colorValues.contains("blue")); + assertTrue(colorValues.contains("green")); + } } + // todo: Add a negative case to check what happens to Mongo when this operator is applied to + // non-array columns + } + + @Nested + @DisplayName("REMOVE_ALL_FROM_LIST Operator Tests") + class RemoveAllFromListOperatorTests { @Test - @DisplayName("Should handle removing non-existent values (no-op)") - void testRemoveNonExistentValues() throws Exception { - String docId = getRandomDocId(4); + @DisplayName("Should REMOVE_ALL_FROM_LIST for top-level and nested arrays via bulkUpdate") + void testRemoveAllFromListAllCases() throws Exception { + String docId = generateDocId("test"); Key key = new SingleValueKey(DEFAULT_TENANT, docId); ObjectNode node = OBJECT_MAPPER.createObjectNode(); - node.put("item", "TestItem"); - node.putArray("tags").add("tag1").add("tag2"); + node.put("item", "RemoveTestItem"); + node.putArray("tags").add("tag1").add("tag2").add("tag3"); // Top-level: remove existing + node.putArray("numbers").add(1).add(2).add(3); // Top-level: remove non-existent (no-op) + ObjectNode props = OBJECT_MAPPER.createObjectNode(); + props + .putArray("colors") + .add("red") + .add("blue") + .add("green"); // Nested JSONB: remove multiple + node.set("props", props); flatCollection.create(key, new JSONDocument(node)); Query query = @@ -3309,93 +2574,49 @@ void testRemoveNonExistentValues() throws Exception { ConstantExpression.of(key.toString()))) .build(); - // Try to remove values that don't exist List updates = List.of( + // Top-level: remove 'tag1' → leaves tag2, tag3 SubDocumentUpdate.builder() .subDocument("tags") .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) - .subDocumentValue( - org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue.of( - new String[] {"nonexistent1", "nonexistent2"})) + .subDocumentValue(SubDocumentValue.of(new String[] {"tag1"})) + .build(), + // Nested JSONB: remove 'red' and 'blue' → leaves green + SubDocumentUpdate.builder() + .subDocument("props.colors") + .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"red", "blue"})) .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 tagsNode = resultJson.get("tags"); - assertTrue(tagsNode.isArray()); - assertEquals(2, tagsNode.size()); // No change + try (CloseableIterator results = + flatCollection.bulkUpdate(query, updates, options)) { + assertTrue(results.hasNext()); + JsonNode resultJson = OBJECT_MAPPER.readTree(results.next().toJson()); + + // Verify top-level: tag1 removed, tag2 and tag3 remain + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(2, tagsNode.size()); + assertEquals("tag2", tagsNode.get(0).asText()); + assertEquals("tag3", tagsNode.get(1).asText()); + + // Verify nested JSONB: red and blue removed, green remains + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(1, colorsNode.size()); + assertEquals("green", colorsNode.get(0).asText()); + + // Verify numbers unchanged (no-op since we didn't update it) + JsonNode numbersNode = resultJson.get("numbers"); + assertTrue(numbersNode.isArray()); + assertEquals(3, numbersNode.size()); + } } } - - @Test - @DisplayName("Should return empty when no document matches query") - void testUpdateNoMatch() throws Exception { - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("id"), - RelationalOperator.EQ, - ConstantExpression.of("9999"))) - .build(); - - List updates = List.of(SubDocumentUpdate.of("price", 100)); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - Optional result = flatCollection.update(query, updates, options); - - assertTrue(result.isEmpty()); - } - - @Test - @DisplayName("Should throw IOException when column does not exist") - void testUpdateNonExistentColumn() { - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("_id"), - RelationalOperator.EQ, - ConstantExpression.of(1))) - .build(); - - List updates = - List.of(SubDocumentUpdate.of("nonexistent_column", "value")); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - assertThrows(IOException.class, () -> flatCollection.update(query, updates, options)); - } - - @Test - @DisplayName("Should throw IOException when nested path on non-JSONB column") - void testUpdateNestedPathOnNonJsonbColumn() { - Query query = - Query.builder() - .setFilter( - RelationalExpression.of( - IdentifierExpression.of("_id"), - RelationalOperator.EQ, - ConstantExpression.of(1))) - .build(); - - // "item" is TEXT, not JSONB - nested path should fail - List updates = List.of(SubDocumentUpdate.of("item.nested", "value")); - - UpdateOptions options = - UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); - - assertThrows(IOException.class, () -> flatCollection.update(query, updates, options)); - } } @Nested @@ -3429,11 +2650,11 @@ void testBulkUpdateWithAfterUpdateReturn() throws Exception { } resultIterator.close(); - assertTrue(results.size() > 1, "Should return multiple updated documents"); + assertTrue(results.size() > 1); for (Document doc : results) { JsonNode json = OBJECT_MAPPER.readTree(doc.toJson()); - assertEquals(999, json.get("quantity").asInt(), "All docs should have updated quantity"); + assertEquals(999, json.get("quantity").asInt()); } PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; @@ -3445,7 +2666,7 @@ void testBulkUpdateWithAfterUpdateReturn() throws Exception { FLAT_COLLECTION_NAME)); ResultSet rs = ps.executeQuery()) { assertTrue(rs.next()); - assertEquals(results.size(), rs.getInt(1), "DB should have same number of updated rows"); + assertEquals(results.size(), rs.getInt(1)); } } @@ -3497,11 +2718,8 @@ void testBulkUpdateWithBeforeUpdateReturn() throws Exception { String id = json.get("id").asText(); int returnedQuantity = json.get("quantity").asInt(); - assertTrue(originalQuantities.containsKey(id), "Returned doc ID should be in original set"); - assertEquals( - originalQuantities.get(id).intValue(), - returnedQuantity, - "Returned quantity should be the ORIGINAL value"); + assertTrue(originalQuantities.containsKey(id)); + assertEquals(originalQuantities.get(id).intValue(), returnedQuantity); } // But database should have the NEW value @@ -3513,7 +2731,7 @@ void testBulkUpdateWithBeforeUpdateReturn() throws Exception { FLAT_COLLECTION_NAME)); ResultSet rs = ps.executeQuery()) { while (rs.next()) { - assertEquals(888, rs.getInt("quantity"), "DB should have the updated value"); + assertEquals(888, rs.getInt("quantity")); } } } @@ -3539,7 +2757,7 @@ void testBulkUpdateWithNoneReturn() throws Exception { flatCollection.bulkUpdate(query, updates, options); // Should return empty iterator - assertFalse(resultIterator.hasNext(), "Should return empty iterator for NONE return type"); + assertFalse(resultIterator.hasNext()); resultIterator.close(); // But database should still be updated @@ -3575,7 +2793,7 @@ void testBulkUpdateNoMatchingDocuments() throws Exception { CloseableIterator resultIterator = flatCollection.bulkUpdate(query, updates, options); - assertFalse(resultIterator.hasNext(), "Should return empty iterator when no docs match"); + assertFalse(resultIterator.hasNext()); resultIterator.close(); } @@ -3608,7 +2826,7 @@ void testBulkUpdateMultipleFields() throws Exception { } resultIterator.close(); - assertEquals(3, results.size(), "Should return 3 Soap items"); + assertEquals(3, results.size()); for (Document doc : results) { JsonNode json = OBJECT_MAPPER.readTree(doc.toJson()); @@ -3712,8 +2930,8 @@ void testBulkUpdateNonExistentColumnWithSkipStrategy() throws Exception { assertEquals(1, results.size()); JsonNode json = OBJECT_MAPPER.readTree(results.get(0).toJson()); - assertEquals(111, json.get("price").asInt(), "Valid column should be updated"); - assertFalse(json.has("nonExistentColumn"), "Non-existent column should not appear"); + assertEquals(111, json.get("price").asInt()); + assertFalse(json.has("nonExistentColumn")); } @Test @@ -3742,27 +2960,6 @@ void testBulkUpdateNonExistentColumnWithThrowStrategy() { } } - @Nested - @DisplayName("Bulk Array Value Operations") - class BulkArrayValueOperationTests { - - @Test - @DisplayName("Should throw UnsupportedOperationException for bulkOperationOnArrayValue") - void testBulkOperationOnArrayValue() throws IOException { - Set keys = - Set.of(new SingleValueKey("default", "1"), new SingleValueKey("default", "2")); - List subDocs = - List.of(new JSONDocument("\"newTag1\""), new JSONDocument("\"newTag2\"")); - BulkArrayValueUpdateRequest request = - new BulkArrayValueUpdateRequest( - keys, "tags", BulkArrayValueUpdateRequest.Operation.SET, subDocs); - - assertThrows( - UnsupportedOperationException.class, - () -> flatCollection.bulkOperationOnArrayValue(request)); - } - } - @Nested @DisplayName("CreateOrReplace Schema Refresh Tests") class CreateOrReplaceSchemaRefreshTests { @@ -3839,209 +3036,4 @@ void testDrop() { assertThrows(UnsupportedOperationException.class, () -> flatCollection.drop()); } } - - @Nested - @DisplayName("Timestamp Auto-Population Tests") - class TimestampTests { - - @Test - @DisplayName( - "Should auto-populate createdTime (BIGINT) and lastUpdateTime (TIMESTAMPTZ) on create") - void testTimestampsOnCreate() throws Exception { - long beforeCreate = System.currentTimeMillis(); - - ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("id", "ts-test-1"); - objectNode.put("item", "TimestampTestItem"); - objectNode.put("price", 100); - Document document = new JSONDocument(objectNode); - Key key = new SingleValueKey(DEFAULT_TENANT, "ts-test-1"); - - CreateResult result = flatCollection.create(key, document); - assertTrue(result.isSucceed()); - - long afterCreate = System.currentTimeMillis(); - - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"createdTime\", \"lastUpdateTime\" FROM \"%s\" WHERE \"id\" = '%s'", - FLAT_COLLECTION_NAME, key)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - - long createdTime = rs.getLong("createdTime"); - assertFalse(rs.wasNull(), "createdTime should not be NULL"); - assertTrue( - createdTime >= beforeCreate && createdTime <= afterCreate, - "createdTime should be within test execution window"); - - Timestamp lastUpdateTime = rs.getTimestamp("lastUpdateTime"); - assertNotNull(lastUpdateTime, "lastUpdateTime should not be NULL"); - assertTrue( - lastUpdateTime.getTime() >= beforeCreate && lastUpdateTime.getTime() <= afterCreate, - "lastUpdateTime should be within test execution window"); - } - } - - @Test - @DisplayName("Should preserve createdTime and update lastUpdateTime on upsert") - void testTimestampsOnUpsert() throws Exception { - // First create - ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("id", "ts-test-2"); - objectNode.put("item", "UpsertTimestampTest"); - objectNode.put("price", 100); - Document document = new JSONDocument(objectNode); - Key key = new SingleValueKey(DEFAULT_TENANT, "ts-test-2"); - - flatCollection.create(key, document); - - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - long originalCreatedTime; - long originalLastUpdateTime; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"createdTime\", \"lastUpdateTime\" FROM \"%s\" WHERE \"id\" = '%s'", - FLAT_COLLECTION_NAME, key.toString())); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - originalCreatedTime = rs.getLong("createdTime"); - originalLastUpdateTime = rs.getTimestamp("lastUpdateTime").getTime(); - } - - // Wait a bit to ensure time difference - Thread.sleep(50); - - // Upsert (update existing) - long beforeUpsert = System.currentTimeMillis(); - objectNode.put("price", 200); - Document updatedDoc = new JSONDocument(objectNode); - flatCollection.createOrReplace(key, updatedDoc); - long afterUpsert = System.currentTimeMillis(); - - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"createdTime\", \"lastUpdateTime\" FROM \"%s\" WHERE \"id\" = '%s'", - FLAT_COLLECTION_NAME, key.toString())); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - - long newCreatedTime = rs.getLong("createdTime"); - assertEquals( - originalCreatedTime, newCreatedTime, "createdTime should be preserved on upsert"); - - long newLastUpdateTime = rs.getTimestamp("lastUpdateTime").getTime(); - assertTrue(newLastUpdateTime > originalLastUpdateTime, "lastUpdateTime should be updated"); - assertTrue( - newLastUpdateTime >= beforeUpsert && newLastUpdateTime <= afterUpsert, - "lastUpdateTime should be within upsert execution window"); - } - } - - @Test - @DisplayName( - "Should not throw exception when timestampFields config is missing - cols remain NULL") - void testNoExceptionWhenTimestampConfigMissing() throws Exception { - // Create a collection WITHOUT timestampFields config - String postgresConnectionUrl = - String.format("jdbc:postgresql://localhost:%s/", postgres.getMappedPort(5432)); - - Map configWithoutTimestamps = new HashMap<>(); - configWithoutTimestamps.put("url", postgresConnectionUrl); - configWithoutTimestamps.put("user", "postgres"); - configWithoutTimestamps.put("password", "postgres"); - // Note: NO customParams.timestampFields config - - Datastore datastoreWithoutTimestamps = - DatastoreProvider.getDatastore( - "Postgres", ConfigFactory.parseMap(configWithoutTimestamps)); - Collection collectionWithoutTimestamps = - datastoreWithoutTimestamps.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT); - - // Create a document - should NOT throw exception - ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("id", "ts-test-no-config"); - objectNode.put("item", "NoTimestampConfigTest"); - objectNode.put("price", 100); - Document document = new JSONDocument(objectNode); - Key key = new SingleValueKey(DEFAULT_TENANT, "ts-test-no-config"); - - CreateResult result = collectionWithoutTimestamps.create(key, document); - assertTrue(result.isSucceed()); - - // Verify timestamp columns are NULL - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"createdTime\", \"lastUpdateTime\" FROM \"%s\" WHERE \"id\" = '%s'", - FLAT_COLLECTION_NAME, key)); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - - assertNull( - rs.getObject("createdTime"), "createdTime should be NULL when config is missing"); - assertNull( - rs.getObject("lastUpdateTime"), "lastUpdateTime should be NULL when config is missing"); - } - } - - @Test - @DisplayName( - "Should not throw exception when timestampFields config is invalid JSON - cols remain NULL") - void testNoExceptionWhenTimestampConfigInvalidJson() throws Exception { - // Create a collection with INVALID JSON in timestampFields config - String postgresConnectionUrl = - String.format("jdbc:postgresql://localhost:%s/", postgres.getMappedPort(5432)); - - Map configWithInvalidJson = new HashMap<>(); - configWithInvalidJson.put("url", postgresConnectionUrl); - configWithInvalidJson.put("user", "postgres"); - configWithInvalidJson.put("password", "postgres"); - // Invalid JSON - missing quotes, malformed - configWithInvalidJson.put("customParams.timestampFields", "not valid json {{{"); - - Datastore datastoreWithInvalidConfig = - DatastoreProvider.getDatastore("Postgres", ConfigFactory.parseMap(configWithInvalidJson)); - Collection collectionWithInvalidConfig = - datastoreWithInvalidConfig.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT); - - // Create a document - should NOT throw exception - ObjectNode objectNode = OBJECT_MAPPER.createObjectNode(); - objectNode.put("id", "ts-test-invalid-json"); - objectNode.put("item", "InvalidJsonConfigTest"); - objectNode.put("price", 100); - Document document = new JSONDocument(objectNode); - Key key = new SingleValueKey(DEFAULT_TENANT, "ts-test-invalid-json"); - - CreateResult result = collectionWithInvalidConfig.create(key, document); - assertTrue(result.isSucceed()); - - // Verify timestamp columns are NULL (config parsing failed gracefully) - PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore; - try (Connection conn = pgDatastore.getPostgresClient(); - PreparedStatement ps = - conn.prepareStatement( - String.format( - "SELECT \"createdTime\", \"lastUpdateTime\" FROM \"%s\" WHERE \"id\" = '%s'", - FLAT_COLLECTION_NAME, key.toString())); - ResultSet rs = ps.executeQuery()) { - assertTrue(rs.next()); - - rs.getLong("createdTime"); - assertTrue(rs.wasNull(), "createdTime should be NULL when config JSON is invalid"); - - rs.getTimestamp("lastUpdateTime"); - assertTrue(rs.wasNull(), "lastUpdateTime should be NULL when config JSON is invalid"); - } - } - } } 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 deleted file mode 100644 index 9461cf69d..000000000 --- a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoFlatPgConsistencyTest.java +++ /dev/null @@ -1,749 +0,0 @@ -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/integrationTest/java/org/hypertrace/core/documentstore/MongoPostgresWriteConsistencyTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoPostgresWriteConsistencyTest.java new file mode 100644 index 000000000..8c46dc712 --- /dev/null +++ b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/MongoPostgresWriteConsistencyTest.java @@ -0,0 +1,1089 @@ +package org.hypertrace.core.documentstore; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.io.IOException; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Stream; +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.Disabled; +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.testcontainers.junit.jupiter.Testcontainers; + +/*Validates write consistency b/w Mongo and Postgres*/ +@Testcontainers +public class MongoPostgresWriteConsistencyTest extends BaseWriteTest { + + private static final String COLLECTION_NAME = "consistency_test"; + private static final String MONGO_STORE = "Mongo"; + private static final String POSTGRES_FLAT_STORE = "PostgresFlat"; + + @BeforeAll + public static void init() throws IOException { + // Start MongoDB and PostgreSQL using BaseWriteTest setup + initMongo(); + initPostgres(); + + datastoreMap.put(MONGO_STORE, mongoDatastore); + datastoreMap.put(POSTGRES_FLAT_STORE, postgresDatastore); + + // Create Postgres flat collection schema + createFlatCollectionSchema((PostgresDatastore) postgresDatastore, COLLECTION_NAME); + + // 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."); + } + + @BeforeEach + public void clearCollections() { + collectionMap.get(MONGO_STORE).deleteAll(); + clearTable(COLLECTION_NAME); + } + + @AfterAll + public static void shutdown() { + shutdownMongo(); + shutdownPostgres(); + } + + 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)); + } + } + + /** Inserts a test document into all collections (Mongo and PG) */ + private void insertTestDocument(String docId) throws IOException { + for (Collection collection : collectionMap.values()) { + insertTestDocument(docId, collection); + } + } + + @Nested + @DisplayName("Upsert Consistency Tests") + class UpsertConsistencyTests { + + @ParameterizedTest(name = "{0}: upsert with all field types") + @ArgumentsSource(AllStoresProvider.class) + void testUpsertNewDoc(String storeName) throws Exception { + String docId = generateDocId("upsert-all"); + Key key = createKey(docId); + + Collection collection = getCollection(storeName); + + // Create document with all field types + Document document = createTestDocument(docId); + boolean isNew = collection.upsert(key, document); + assertTrue(isNew); + + // Verify by upserting again (returns true again if the operation succeeds) + boolean secondUpsert = collection.upsert(key, document); + assertTrue(secondUpsert); + + // Query the collection to get the document back + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + assertTrue(iterator.hasNext()); + Document retrievedDoc = iterator.next(); + JsonNode resultJson = OBJECT_MAPPER.readTree(retrievedDoc.toJson()); + + // Verify primitives + assertEquals("TestItem", resultJson.get("item").asText()); + assertEquals(100, resultJson.get("price").asInt()); + assertEquals(50, resultJson.get("quantity").asInt()); + assertTrue(resultJson.get("in_stock").asBoolean()); + assertEquals(1000000000000L, resultJson.get("big_number").asLong()); + assertEquals(3.5, resultJson.get("rating").asDouble(), 0.01); + assertEquals(50.0, resultJson.get("weight").asDouble(), 0.01); + + // Verify arrays + JsonNode tagsNode = resultJson.get("tags"); + assertNotNull(tagsNode); + assertTrue(tagsNode.isArray()); + assertEquals(2, tagsNode.size()); + assertEquals("tag1", tagsNode.get(0).asText()); + assertEquals("tag2", tagsNode.get(1).asText()); + + JsonNode numbersNode = resultJson.get("numbers"); + assertNotNull(numbersNode); + assertTrue(numbersNode.isArray()); + assertEquals(3, numbersNode.size()); + + // Verify JSONB - props + JsonNode propsNode = resultJson.get("props"); + assertNotNull(propsNode); + assertEquals("TestBrand", propsNode.get("brand").asText()); + assertEquals("M", propsNode.get("size").asText()); + assertEquals(10, propsNode.get("count").asInt()); + JsonNode colorsNode = propsNode.get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(2, colorsNode.size()); + + // Verify JSONB - sales + JsonNode salesNode = resultJson.get("sales"); + assertNotNull(salesNode); + assertEquals(200, salesNode.get("total").asInt()); + assertEquals(10, salesNode.get("count").asInt()); + } + } + + @ParameterizedTest(name = "{0}: upsert preserves existing values (merge behavior)") + @ArgumentsSource(AllStoresProvider.class) + void testUpsertExistingDoc(String storeName) throws Exception { + String docId = generateDocId("upsert-merge"); + Key key = createKey(docId); + + Collection collection = getCollection(storeName); + + Document initialDoc = createTestDocument(docId); + collection.upsert(key, initialDoc); + + ObjectNode partialNode = OBJECT_MAPPER.createObjectNode(); + partialNode.put("id", getKeyString(docId)); + partialNode.put("item", "UpdatedItem"); + partialNode.put("price", 999); + Document partialDoc = new JSONDocument(partialNode); + + collection.upsert(key, partialDoc); + + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + assertTrue(iterator.hasNext()); + Document retrievedDoc = iterator.next(); + JsonNode resultJson = OBJECT_MAPPER.readTree(retrievedDoc.toJson()); + + // Updated fields + assertEquals("UpdatedItem", resultJson.get("item").asText()); + assertEquals(999, resultJson.get("price").asInt()); + + // Non-updated fields + assertEquals(50, resultJson.get("quantity").asInt()); + assertTrue(resultJson.get("in_stock").asBoolean()); + assertEquals(1000000000000L, resultJson.get("big_number").asLong()); + assertEquals(3.5, resultJson.get("rating").asDouble(), 0.01); + + JsonNode tagsNode = resultJson.get("tags"); + assertNotNull(tagsNode); + assertEquals(2, tagsNode.size()); + + JsonNode propsNode = resultJson.get("props"); + assertNotNull(propsNode); + assertEquals("TestBrand", propsNode.get("brand").asText()); + + JsonNode salesNode = resultJson.get("sales"); + assertNotNull(salesNode); + assertEquals(200, salesNode.get("total").asInt()); + } + } + + @ParameterizedTest(name = "{0}: bulkUpsert multiple documents") + @ArgumentsSource(AllStoresProvider.class) + void testBulkUpsert(String storeName) throws Exception { + String docId1 = generateDocId("bulk-1"); + String docId2 = generateDocId("bulk-2"); + + Collection collection = getCollection(storeName); + + Map documents = new HashMap<>(); + documents.put(createKey(docId1), createTestDocument(docId1)); + documents.put(createKey(docId2), createTestDocument(docId2)); + + boolean result = collection.bulkUpsert(documents); + assertTrue(result); + + for (String docId : List.of(docId1, docId2)) { + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + assertTrue(iterator.hasNext()); + Document doc = iterator.next(); + JsonNode json = OBJECT_MAPPER.readTree(doc.toJson()); + + assertEquals("TestItem", json.get("item").asText()); + assertEquals(100, json.get("price").asInt()); + assertEquals(50, json.get("quantity").asInt()); + assertTrue(json.get("in_stock").asBoolean()); + + JsonNode tagsNode = json.get("tags"); + assertNotNull(tagsNode); + assertEquals(2, tagsNode.size()); + } + } + } + + @ParameterizedTest(name = "{0}: upsert with non-existing fields (schema mismatch)") + @ArgumentsSource(AllStoresProvider.class) + void testUpsertNonExistingFields(String storeName) throws Exception { + String docId = generateDocId("upsert-unknown"); + Key key = createKey(docId); + + Collection collection = getCollection(storeName); + + // Create document with fields that don't exist in the PG schema + ObjectNode docNode = OBJECT_MAPPER.createObjectNode(); + docNode.put("id", getKeyString(docId)); + docNode.put("item", "TestItem"); + docNode.put("price", 100); + docNode.put("unknown_field_1", "unknown_value"); + docNode.put("unknown_field_2", 999); + Document document = new JSONDocument(docNode); + + // Upsert should succeed (PG skips unknown fields with default strategy) + boolean result = collection.upsert(key, document); + assertTrue(result); + + // Verify document exists with known fields + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + assertTrue(iterator.hasNext()); + Document retrievedDoc = iterator.next(); + JsonNode json = OBJECT_MAPPER.readTree(retrievedDoc.toJson()); + + // Known fields should exist + assertEquals("TestItem", json.get("item").asText()); + assertEquals(100, json.get("price").asInt()); + + // For Mongo, unknown fields will be stored; for PG with SKIP strategy, they won't + if (storeName.equals("Mongo")) { + assertNotNull(json.get("unknown_field_1")); + assertEquals("unknown_value", json.get("unknown_field_1").asText()); + assertNotNull(json.get("unknown_field_2")); + assertEquals(999, json.get("unknown_field_2").asInt()); + } + } + } + } + + @Nested + @DisplayName("CreateOrReplace Consistency Tests") + class CreateOrReplaceConsistencyTests { + + @ParameterizedTest(name = "{0}: createOrReplace new document") + @ArgumentsSource(AllStoresProvider.class) + void testCreateOrReplaceNewDoc(String storeName) throws Exception { + String docId = generateDocId("cor-new"); + Key key = createKey(docId); + + Collection collection = getCollection(storeName); + + Document document = createTestDocument(docId); + boolean isNew = collection.createOrReplace(key, document); + assertTrue(isNew); + + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + assertTrue(iterator.hasNext()); + Document retrievedDoc = iterator.next(); + JsonNode json = OBJECT_MAPPER.readTree(retrievedDoc.toJson()); + + assertEquals("TestItem", json.get("item").asText()); + assertEquals(100, json.get("price").asInt()); + assertEquals(50, json.get("quantity").asInt()); + } + } + + @ParameterizedTest(name = "{0}: createOrReplace replaces entire document") + @ArgumentsSource(AllStoresProvider.class) + void testCreateOrReplaceExistingDoc(String storeName) throws Exception { + String docId = generateDocId("cor-replace"); + Key key = createKey(docId); + + Collection collection = getCollection(storeName); + + // First create with all fields + Document initialDoc = createTestDocument(docId); + collection.createOrReplace(key, initialDoc); + + // Replace with partial document - unlike upsert, this should REPLACE entirely + ObjectNode replacementNode = OBJECT_MAPPER.createObjectNode(); + replacementNode.put("id", getKeyString(docId)); + replacementNode.put("item", "ReplacedItem"); + replacementNode.put("price", 777); + // Note: quantity, in_stock, tags, props, sales are NOT specified + Document replacementDoc = new JSONDocument(replacementNode); + + boolean isNew = collection.createOrReplace(key, replacementDoc); + assertFalse(isNew); + + Query query = buildQueryById(docId); + try (CloseableIterator iterator = collection.find(query)) { + assertTrue(iterator.hasNext()); + Document retrievedDoc = iterator.next(); + JsonNode json = OBJECT_MAPPER.readTree(retrievedDoc.toJson()); + + // Replaced fields should have new values + assertEquals("ReplacedItem", json.get("item").asText()); + assertEquals(777, json.get("price").asInt()); + + // Note that PG should return null for non-specified fields. However, the iterator + // specifically excludes null fields + // from the result set, so we expect these to be missing. + assertNull(json.get("quantity")); + assertNull(json.get("in_stock")); + assertNull(json.get("tags")); + assertNull(json.get("props")); + assertNull(json.get("sales")); + } + } + + @ParameterizedTest(name = "{0}: createOrReplaceAndReturn") + @ArgumentsSource(AllStoresProvider.class) + @Disabled("Not implemented for PG") + void testCreateOrReplaceAndReturn(String storeName) throws Exception { + String docId = generateDocId("cor-return"); + Key key = createKey(docId); + + Collection collection = getCollection(storeName); + + Document document = createTestDocument(docId); + Document returned = collection.createOrReplaceAndReturn(key, document); + + assertNotNull(returned); + JsonNode json = OBJECT_MAPPER.readTree(returned.toJson()); + + assertEquals("TestItem", json.get("item").asText()); + assertEquals(100, json.get("price").asInt()); + assertEquals(50, json.get("quantity").asInt()); + } + } + + @Nested + class SubdocUpdateConsistencyTests { + + @Nested + @DisplayName("SET Operator Tests") + class SetOperatorTests { + + @ParameterizedTest(name = "{0}: SET top-level primitives") + @ArgumentsSource(AllStoresProvider.class) + void testSetTopLevelPrimitives(String storeName) throws Exception { + String docId = generateDocId("set-primitives"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + List updates = + List.of( + SubDocumentUpdate.of("item", "UpdatedItem"), + SubDocumentUpdate.of("price", 999), + SubDocumentUpdate.of("quantity", 50), + SubDocumentUpdate.of("in_stock", false), + SubDocumentUpdate.of("big_number", 9999999999L), + SubDocumentUpdate.of("rating", 4.5f), + SubDocumentUpdate.of("weight", 123.456)); + + 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()); + assertEquals("UpdatedItem", resultJson.get("item").asText()); + assertEquals(999, resultJson.get("price").asInt()); + assertFalse(resultJson.get("in_stock").asBoolean()); + assertEquals(9999999999L, resultJson.get("big_number").asLong()); + assertEquals(4.5, resultJson.get("rating").asDouble(), 0.01); + assertEquals(123.456, resultJson.get("weight").asDouble(), 0.01); + } + + @ParameterizedTest(name = "{0}: SET top-level array") + @ArgumentsSource(AllStoresProvider.class) + void testSetTopLevelArray(String storeName) throws Exception { + String docId = generateDocId("set-array"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + List updates = + List.of(SubDocumentUpdate.of("tags", new String[] {"tag4", "tag5", "tag6"})); + + 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 tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(3, tagsNode.size()); + assertEquals("tag4", tagsNode.get(0).asText()); + assertEquals("tag5", tagsNode.get(1).asText()); + assertEquals("tag6", tagsNode.get(2).asText()); + } + + @ParameterizedTest(name = "{0}: SET top-level array") + @ArgumentsSource(AllStoresProvider.class) + void testSetTopLevelEmptyArray(String storeName) throws Exception { + String docId = generateDocId("set-array"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + List updates = List.of(SubDocumentUpdate.of("tags", new String[] {})); + + 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 tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(0, tagsNode.size()); + } + + @ParameterizedTest(name = "{0}: SET nested JSONB primitive") + @ArgumentsSource(AllStoresProvider.class) + void testSetNestedJsonbPrimitive(String storeName) throws Exception { + String docId = generateDocId("set-nested"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + 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()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + assertEquals("NewBrand", resultJson.get("props").get("brand").asText()); + // Other props fields preserved + assertEquals("M", resultJson.get("props").get("size").asText()); + assertEquals(10, resultJson.get("props").get("count").asInt()); + } + + @ParameterizedTest(name = "{0}: SET nested JSONB array") + @ArgumentsSource(AllStoresProvider.class) + void testSetNestedJsonbArray(String storeName) throws Exception { + String docId = generateDocId("set-nested-array"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("sales.regions") + .operator(UpdateOperator.SET) + .subDocumentValue(SubDocumentValue.of(new String[] {"US", "EU", "APAC"})) + .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 regionsNode = resultJson.get("sales").get("regions"); + assertTrue(regionsNode.isArray()); + assertEquals(3, regionsNode.size()); + // Other sales fields preserved + assertEquals(200, resultJson.get("sales").get("total").asInt()); + } + + @ParameterizedTest(name = "{0}: SET entire JSONB column") + @ArgumentsSource(AllStoresProvider.class) + void testSetEntireJsonbColumn(String storeName) throws Exception { + String docId = generateDocId("set-jsonb-column"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // Create a completely new object to replace the entire JSONB column + ObjectNode newProps = OBJECT_MAPPER.createObjectNode(); + newProps.put("manufacturer", "NewManufacturer"); + newProps.put("model", "Model-X"); + newProps.put("year", 2024); + newProps.putArray("features").add("feature1").add("feature2"); + + List updates = + List.of(SubDocumentUpdate.of("props", SubDocumentValue.of(new JSONDocument(newProps)))); + + 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); + + // Verify new fields are present + assertEquals("NewManufacturer", propsNode.get("manufacturer").asText()); + assertEquals("Model-X", propsNode.get("model").asText()); + assertEquals(2024, propsNode.get("year").asInt()); + assertTrue(propsNode.get("features").isArray()); + assertEquals(2, propsNode.get("features").size()); + assertEquals("feature1", propsNode.get("features").get(0).asText()); + assertEquals("feature2", propsNode.get("features").get(1).asText()); + + // Verify old fields are NOT present (entire column replaced) + assertFalse(propsNode.has("brand")); + assertFalse(propsNode.has("size")); + assertFalse(propsNode.has("count")); + assertFalse(propsNode.has("colors")); + } + } + + @Nested + @DisplayName("UNSET Operator Tests") + class UnsetOperatorTests { + + @ParameterizedTest(name = "{0}: UNSET top-level column and nested JSONB field") + @ArgumentsSource(AllStoresProvider.class) + void testUnsetTopLevelAndNestedFields(String storeName) throws Exception { + String docId = generateDocId("unset"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + List updates = + List.of( + // Top-level: sets column to NULL + SubDocumentUpdate.builder() + .subDocument("item") + .operator(UpdateOperator.UNSET) + .build(), + // Nested JSONB: removes key from JSON object + 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); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + // Verify top-level column is NULL/missing + JsonNode itemNode = resultJson.get("item"); + assertTrue(itemNode == null || itemNode.isNull()); + + // Verify nested JSONB key is removed, but other keys preserved + assertFalse(resultJson.get("props").has("brand")); + assertEquals("M", resultJson.get("props").get("size").asText()); + } + } + + @Nested + @DisplayName("ADD Operator Tests") + class AddOperatorTests { + + @ParameterizedTest(name = "{0}: ADD to all numeric types") + @ArgumentsSource(AllStoresProvider.class) + void testAddAllNumericTypes(String storeName) throws Exception { + String docId = generateDocId("add-numeric"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + List updates = + List.of( + // Top-level INT: 100 + 5 = 105 + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(5)) + .build(), + // Top-level INT (negative): 50 + (-15) = 35 + SubDocumentUpdate.builder() + .subDocument("quantity") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(-15)) + .build(), + // Top-level BIGINT: 1000000000000 + 500 = 1000000000500 + SubDocumentUpdate.builder() + .subDocument("big_number") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(500L)) + .build(), + // Top-level REAL: 3.5 + 1.0 = 4.5 + SubDocumentUpdate.builder() + .subDocument("rating") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(1.0f)) + .build(), + // Top-level DOUBLE: 50.0 + 2.5 = 52.5 + SubDocumentUpdate.builder() + .subDocument("weight") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(2.5)) + .build(), + // Nested JSONB: 200 + 50 = 250 + SubDocumentUpdate.builder() + .subDocument("sales.total") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(50)) + .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()); + assertEquals(105, resultJson.get("price").asInt()); + assertEquals(35, resultJson.get("quantity").asInt()); + assertEquals(1000000000500L, resultJson.get("big_number").asLong()); + assertEquals(4.5, resultJson.get("rating").asDouble(), 0.01); + assertEquals(52.5, resultJson.get("weight").asDouble(), 0.01); + assertEquals(250, resultJson.get("sales").get("total").asInt()); + // Other fields preserved + assertEquals(10, resultJson.get("sales").get("count").asInt()); + } + + @ParameterizedTest(name = "{0}: ADD on non-numeric field (TEXT column)") + @ArgumentsSource(AllStoresProvider.class) + void testAddOnNonNumericField(String storeName) throws Exception { + String docId = generateDocId("add-non-numeric"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // Try to ADD to 'item' which is a TEXT field + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("item") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(10)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + assertThrows(Exception.class, () -> collection.update(query, updates, options)); + } + } + + @Nested + @DisplayName("APPEND_TO_LIST Operator Tests") + class AppendToListOperatorTests { + + @ParameterizedTest(name = "{0}: APPEND_TO_LIST for top-level and nested arrays") + @ArgumentsSource(AllStoresProvider.class) + void testAppendToListAllCases(String storeName) throws Exception { + String docId = generateDocId("append"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + List updates = + List.of( + // Top-level array: append to existing tags + SubDocumentUpdate.builder() + .subDocument("tags") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"newTag1", "newTag2"})) + .build(), + // Nested JSONB array: append to existing props.colors + SubDocumentUpdate.builder() + .subDocument("props.colors") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"green", "yellow"})) + .build(), + // Nested JSONB: append to non-existent array (creates it) + SubDocumentUpdate.builder() + .subDocument("sales.regions") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"US", "EU"})) + .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 top-level array append + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(4, tagsNode.size()); + assertEquals("tag1", tagsNode.get(0).asText()); + assertEquals("tag2", tagsNode.get(1).asText()); + assertEquals("newTag1", tagsNode.get(2).asText()); + assertEquals("newTag2", tagsNode.get(3).asText()); + + // Verify nested JSONB array append + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(4, colorsNode.size()); + assertEquals("red", colorsNode.get(0).asText()); + assertEquals("blue", colorsNode.get(1).asText()); + assertEquals("green", colorsNode.get(2).asText()); + assertEquals("yellow", colorsNode.get(3).asText()); + + // Verify non-existent array was created + JsonNode regionsNode = resultJson.get("sales").get("regions"); + assertNotNull(regionsNode); + assertTrue(regionsNode.isArray()); + assertEquals(2, regionsNode.size()); + assertEquals("US", regionsNode.get(0).asText()); + assertEquals("EU", regionsNode.get(1).asText()); + + // Verify other fields preserved + assertEquals("TestBrand", resultJson.get("props").get("brand").asText()); + assertEquals(200, resultJson.get("sales").get("total").asInt()); + } + + @ParameterizedTest(name = "{0}: APPEND_TO_LIST on non-array field (TEXT column)") + @ArgumentsSource(AllStoresProvider.class) + void testAppendToListOnNonArrayField(String storeName) throws Exception { + String docId = generateDocId("append-non-array"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // Try to APPEND_TO_LIST to 'item' which is a TEXT field + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("item") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"value1", "value2"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + assertThrows(Exception.class, () -> collection.update(query, updates, options)); + } + + @ParameterizedTest(name = "{0}: APPEND_TO_LIST on non-array field (INTEGER column)") + @ArgumentsSource(AllStoresProvider.class) + void testAppendToListOnIntegerField(String storeName) throws Exception { + String docId = generateDocId("append-integer"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // Try to APPEND_TO_LIST to 'price' which is an INTEGER field + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new Integer[] {100, 200})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + assertThrows(Exception.class, () -> collection.update(query, updates, options)); + } + } + + @Nested + @DisplayName("ADD_TO_LIST_IF_ABSENT Operator Tests") + class AddToListIfAbsentOperatorTests { + + @ParameterizedTest(name = "{0}: ADD_TO_LIST_IF_ABSENT for top-level and nested arrays") + @ArgumentsSource(AllStoresProvider.class) + void testAddToListIfAbsentAllCases(String storeName) throws Exception { + String docId = generateDocId("addifabsent"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + List updates = + List.of( + // Top-level: 'tag1' exists, 'newTag' is new → adds only 'newTag' + SubDocumentUpdate.builder() + .subDocument("tags") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) + .subDocumentValue(SubDocumentValue.of(new String[] {"tag1", "newTag"})) + .build(), + // Nested JSONB: 'red' exists, 'green' is new → adds only 'green' + SubDocumentUpdate.builder() + .subDocument("props.colors") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) + .subDocumentValue(SubDocumentValue.of(new String[] {"red", "green"})) + .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 top-level: original 2 + 1 new unique = 3 (order not guaranteed) + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(3, tagsNode.size()); + Set tagValues = new HashSet<>(); + tagsNode.forEach(n -> tagValues.add(n.asText())); + assertTrue(tagValues.contains("tag1")); + assertTrue(tagValues.contains("tag2")); + assertTrue(tagValues.contains("newTag")); + + // Verify nested JSONB: original 2 + 1 new unique = 3 (order not guaranteed) + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(3, colorsNode.size()); + Set colorValues = new HashSet<>(); + colorsNode.forEach(n -> colorValues.add(n.asText())); + assertTrue(colorValues.contains("red")); + assertTrue(colorValues.contains("blue")); + assertTrue(colorValues.contains("green")); + } + + @ParameterizedTest(name = "{0}: ADD_TO_LIST_IF_ABSENT on non-array field (TEXT column)") + @ArgumentsSource(AllStoresProvider.class) + void testAddToListIfAbsentOnNonArrayField(String storeName) throws Exception { + String docId = generateDocId("addifabsent-non-array"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // Try to ADD_TO_LIST_IF_ABSENT to 'item' which is a TEXT field + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("item") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) + .subDocumentValue(SubDocumentValue.of(new String[] {"value1", "value2"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + assertThrows(Exception.class, () -> collection.update(query, updates, options)); + } + } + + @Nested + @DisplayName("REMOVE_ALL_FROM_LIST Operator Tests") + class RemoveAllFromListOperatorTests { + + @ParameterizedTest(name = "{0}: REMOVE_ALL_FROM_LIST for top-level and nested arrays") + @ArgumentsSource(AllStoresProvider.class) + void testRemoveAllFromListAllCases(String storeName) throws Exception { + String docId = generateDocId("remove"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + List updates = + List.of( + // Top-level: remove 'tag1' → leaves tag2 + SubDocumentUpdate.builder() + .subDocument("tags") + .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"tag1"})) + .build(), + // Nested JSONB: remove 'red' → leaves blue + 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); + + assertTrue(result.isPresent()); + JsonNode resultJson = OBJECT_MAPPER.readTree(result.get().toJson()); + + // Verify top-level: tag1 removed, tag2 remains + JsonNode tagsNode = resultJson.get("tags"); + assertTrue(tagsNode.isArray()); + assertEquals(1, tagsNode.size()); + assertEquals("tag2", tagsNode.get(0).asText()); + + // Verify nested JSONB: red removed, blue remains + JsonNode colorsNode = resultJson.get("props").get("colors"); + assertTrue(colorsNode.isArray()); + assertEquals(1, colorsNode.size()); + assertEquals("blue", colorsNode.get(0).asText()); + + // Verify numbers unchanged (no-op since we didn't update it) + JsonNode numbersNode = resultJson.get("numbers"); + assertTrue(numbersNode.isArray()); + assertEquals(3, numbersNode.size()); + } + + @ParameterizedTest(name = "{0}: REMOVE_ALL_FROM_LIST on non-array field (TEXT column)") + @ArgumentsSource(AllStoresProvider.class) + void testRemoveAllFromListOnNonArrayField(String storeName) throws Exception { + String docId = generateDocId("remove-non-array"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // Try to REMOVE_ALL_FROM_LIST from 'item' which is a TEXT field + List updates = + List.of( + SubDocumentUpdate.builder() + .subDocument("item") + .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"value1"})) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + assertThrows(Exception.class, () -> collection.update(query, updates, options)); + } + } + + @Nested + class AllOperatorTests { + + @ParameterizedTest + @ArgumentsSource(AllStoresProvider.class) + void testMultipleUpdatesOnSameFieldThrowsException(String storeName) throws IOException { + String docId = generateDocId("multiple-updates-on-same-field"); + insertTestDocument(docId); + + Collection collection = getCollection(storeName); + Query query = buildQueryById(docId); + + // top-level primitives + List topLevelPrimitiveUpdates = + List.of( + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(5)) + .build(), + SubDocumentUpdate.builder() + .subDocument("price") + .operator(UpdateOperator.ADD) + .subDocumentValue(SubDocumentValue.of(-15)) + .build()); + + UpdateOptions options = + UpdateOptions.builder().returnDocumentType(ReturnDocumentType.AFTER_UPDATE).build(); + + // Since there are multiple updates on the same field, it should throw an exception + assertThrows( + Exception.class, () -> collection.update(query, topLevelPrimitiveUpdates, options)); + + // top-level arrays + List topLevelArrayUpdates = + List.of( + SubDocumentUpdate.builder() + .subDocument("tags") + .operator(UpdateOperator.APPEND_TO_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"tag1", "tag2"})) + .build(), + SubDocumentUpdate.builder() + .subDocument("tags") + .operator(UpdateOperator.REMOVE_ALL_FROM_LIST) + .subDocumentValue(SubDocumentValue.of(new String[] {"tag2"})) + .build()); + + assertThrows( + Exception.class, () -> collection.update(query, topLevelArrayUpdates, options)); + + // nested array updates + List nestedArrayUpdates = + List.of( + SubDocumentUpdate.builder() + .subDocument("sales.regions") + .operator(UpdateOperator.SET) + .subDocumentValue(SubDocumentValue.of(new String[] {"US", "EU", "APAC"})) + .build(), + SubDocumentUpdate.builder() + .subDocument("sales.regions") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) + .subDocumentValue(SubDocumentValue.of(new String[] {"EMEA"})) + .build()); + + assertThrows(Exception.class, () -> collection.update(query, nestedArrayUpdates, options)); + + // nested primitives + List nestedPrimitiveUpdates = + List.of( + SubDocumentUpdate.builder() + .subDocument("props.brand") + .operator(UpdateOperator.SET) + .subDocumentValue(SubDocumentValue.of("NewBrand")) + .build(), + SubDocumentUpdate.builder() + .subDocument("props.brand") + .operator(UpdateOperator.ADD_TO_LIST_IF_ABSENT) + .subDocumentValue(SubDocumentValue.of("NewBrand2")) + .build()); + + assertThrows( + Exception.class, () -> collection.update(query, nestedPrimitiveUpdates, options)); + } + } + } +} diff --git a/document-store/src/integrationTest/resources/expected/add_all_numeric_types_expected.json b/document-store/src/integrationTest/resources/expected/add_all_numeric_types_expected.json new file mode 100644 index 000000000..56c24e273 --- /dev/null +++ b/document-store/src/integrationTest/resources/expected/add_all_numeric_types_expected.json @@ -0,0 +1,12 @@ +{ + "item": "NumericTestItem", + "price": 105, + "quantity": 35, + "big_number": 1000000000500, + "rating": 4.5, + "weight": 52.5, + "sales": { + "total": 250, + "count": 10 + } +} diff --git a/document-store/src/integrationTest/resources/expected/set_all_field_types_expected.json b/document-store/src/integrationTest/resources/expected/set_all_field_types_expected.json new file mode 100644 index 000000000..3e621ca23 --- /dev/null +++ b/document-store/src/integrationTest/resources/expected/set_all_field_types_expected.json @@ -0,0 +1,62 @@ +{ + "id": "1", + "item": "UpdatedItem", + "price": 999, + "quantity": 50, + "date": "2014-03-01 13:30:00.0", + "in_stock": false, + "tags": [ + "tag4", + "tag5", + "tag6" + ], + "categoryTags": [ + "Hygiene", + "PersonalCare" + ], + "props": { + "size": "M", + "brand": "NewBrand", + "colors": [ + "Blue", + "Green" + ], + "seller": { + "name": "Metro Chemicals Pvt. Ltd.", + "address": { + "city": "Mumbai", + "pincode": 400004 + } + }, + "source-loc": [ + "warehouse-A", + "store-1" + ], + "product-code": "SOAP-DET-001" + }, + "sales": { + "regions": [ + "US", + "EU", + "APAC" + ] + }, + "numbers": [ + 10, + 20, + 30 + ], + "scores": [ + 1.1, + 2.2, + 3.3 + ], + "flags": [ + true, + false, + true + ], + "big_number": 9999999999, + "rating": 4.5, + "weight": 123.456 +} diff --git a/document-store/src/integrationTest/resources/schema/flat_collection_test_schema.sql b/document-store/src/integrationTest/resources/schema/flat_collection_test_schema.sql new file mode 100644 index 000000000..83aac0856 --- /dev/null +++ b/document-store/src/integrationTest/resources/schema/flat_collection_test_schema.sql @@ -0,0 +1,19 @@ +CREATE TABLE "%s" ( + "id" TEXT PRIMARY KEY, + "item" TEXT, + "price" INTEGER, + "quantity" INTEGER, + "date" TIMESTAMPTZ, + "in_stock" BOOLEAN, + "tags" TEXT[], + "categoryTags" TEXT[], + "props" JSONB, + "sales" JSONB, + "numbers" INTEGER[], + "scores" DOUBLE PRECISION[], + "flags" BOOLEAN[], + "big_number" BIGINT, + "rating" REAL, + "created_date" DATE, + "weight" DOUBLE PRECISION +); 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 c7423e6c0..c2e1dd37e 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 @@ -805,7 +805,7 @@ private void executeUpdate( UpdateOperator operator = update.getOperator(); Params.Builder paramsBuilder = Params.newBuilder(); - PostgresUpdateOperationParser unifiedParser = UPDATE_PARSER_MAP.get(operator); + PostgresUpdateOperationParser parser = UPDATE_PARSER_MAP.get(operator); String fragment; @@ -817,10 +817,11 @@ private void executeUpdate( .update(update) .paramsBuilder(paramsBuilder) .columnType(colMeta.getPostgresType()) + .isArray(colMeta.isArray()) .build(); - fragment = unifiedParser.parseNonJsonbField(input); + fragment = parser.parseNonJsonbField(input); } else { - // parseInternal() returns just the value expression + // this handles nested jsonb fields UpdateParserInput jsonbInput = UpdateParserInput.builder() .baseField(String.format("\"%s\"", columnName)) @@ -829,11 +830,19 @@ private void executeUpdate( .paramsBuilder(paramsBuilder) .columnType(colMeta.getPostgresType()) .build(); - String valueExpr = unifiedParser.parseInternal(jsonbInput); + String valueExpr = parser.parseInternal(jsonbInput); fragment = String.format("\"%s\" = %s", columnName, valueExpr); } - // Transfer params from builder to our list - params.addAll(paramsBuilder.build().getObjectParams().values()); + for (Object paramValue : paramsBuilder.build().getObjectParams().values()) { + if (isTopLevel && colMeta.isArray() && paramValue != null) { + Object[] arrayValues = (Object[]) paramValue; + Array sqlArray = + connection.createArrayOf(colMeta.getPostgresType().getSqlType(), arrayValues); + params.add(sqlArray); + } else { + params.add(paramValue); + } + } setFragments.add(fragment); } @@ -852,11 +861,9 @@ private void executeUpdate( try (PreparedStatement ps = connection.prepareStatement(sql)) { int idx = 1; - // Add SET clause params for (Object param : params) { ps.setObject(idx++, param); } - // Add WHERE clause params for (Object param : filterParams.getObjectParams().values()) { ps.setObject(idx++, param); } @@ -1079,12 +1086,12 @@ private boolean createOrReplaceWithRetry(Key key, Document document, boolean isR .collect(Collectors.toList()); String sql = buildCreateOrReplaceSql(allColumns, docColumns, quotedPkColumn); - LOGGER.debug("Upsert SQL: {}", sql); + LOGGER.debug("CreateOrReplace SQL: {}", sql); - return executeUpsert(sql, parsed); + return executeUpsertReturningIsInsert(sql, parsed); } catch (PSQLException e) { - return handlePSQLExceptionForUpsert(e, key, document, tableName, isRetry); + return handlePSQLExceptionForCreateOrReplace(e, key, document, tableName, isRetry); } catch (SQLException e) { LOGGER.error("SQLException in createOrReplace. key: {} content: {}", key, document, e); throw new IOException(e); @@ -1233,6 +1240,27 @@ private String buildCreateOrReplaceSql( } private boolean executeUpsert(String sql, TypedDocument parsed) throws SQLException { + try (Connection conn = client.getPooledConnection(); + PreparedStatement ps = conn.prepareStatement(sql)) { + int index = 1; + for (String column : parsed.getColumns()) { + setParameter( + conn, + ps, + index++, + parsed.getValue(column), + parsed.getType(column), + parsed.isArray(column)); + } + try (ResultSet rs = ps.executeQuery()) { + return rs.next(); + } + } + } + + /** Returns true if INSERT, false if UPDATE. */ + private boolean executeUpsertReturningIsInsert(String sql, TypedDocument parsed) + throws SQLException { try (Connection conn = client.getPooledConnection(); PreparedStatement ps = conn.prepareStatement(sql)) { int index = 1; @@ -1247,8 +1275,6 @@ private boolean executeUpsert(String sql, TypedDocument parsed) throws SQLExcept } try (ResultSet rs = ps.executeQuery()) { if (rs.next()) { - // is_insert is true if xmax = 0 (new row), false if updated. This helps us differentiate - // b/w creates/upserts return rs.getBoolean("is_insert"); } } diff --git a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java index ffcc283c0..686228aab 100644 --- a/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java +++ b/document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java @@ -1378,6 +1378,14 @@ private void addColumnToJsonNode( } break; + case "float4": + case "real": + float floatValue = resultSet.getFloat(columnIndex); + if (!resultSet.wasNull()) { + jsonNode.put(columnName, floatValue); + } + break; + case "float8": case "double": double doubleValue = resultSet.getDouble(columnIndex); 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 57fcbc430..ab35ce4a7 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 @@ -9,6 +9,14 @@ public class PostgresAddToListIfAbsentParser implements PostgresUpdateOperationP @Override public String parseNonJsonbField(final UpdateParserInput input) { + if (!input.isArray()) { + throw new IllegalArgumentException( + String.format( + "ADD_TO_LIST_IF_ABSENT operator can only be applied to array columns. " + + "Column '%s' is not an array type.", + input.getBaseField())); + } + final SubDocumentValue value = input.getUpdate().getSubDocumentValue(); // Extract array values directly for top-level array columns 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 5c07f00fa..80440ef1c 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 @@ -8,6 +8,14 @@ public class PostgresAppendToListParser implements PostgresUpdateOperationParser @Override public String parseNonJsonbField(final UpdateParserInput input) { + if (!input.isArray()) { + throw new IllegalArgumentException( + String.format( + "APPEND_TO_LIST operator can only be applied to array columns. " + + "Column '%s' is not an array type.", + input.getBaseField())); + } + final SubDocumentValue value = input.getUpdate().getSubDocumentValue(); // Extract array values directly for top-level array columns 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 eded52341..73930e125 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 @@ -16,6 +16,14 @@ public class PostgresRemoveAllFromListParser implements PostgresUpdateOperationP @Override public String parseNonJsonbField(final UpdateParserInput input) { + if (!input.isArray()) { + throw new IllegalArgumentException( + String.format( + "REMOVE_ALL_FROM_LIST operator can only be applied to array columns. " + + "Column '%s' is not an array type.", + input.getBaseField())); + } + final PostgresSubDocumentArrayGetter subDocArrayGetter = new PostgresSubDocumentArrayGetter(); final SubDocumentArray array = input.getUpdate().getSubDocumentValue().accept(subDocArrayGetter); 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 d3763a60e..f753f1710 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 @@ -8,10 +8,13 @@ import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate; import org.hypertrace.core.documentstore.postgres.Params; import org.hypertrace.core.documentstore.postgres.Params.Builder; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; +import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentArrayGetter; import org.hypertrace.core.documentstore.postgres.subdoc.PostgresSubDocumentValueParser; @AllArgsConstructor public class PostgresSetValueParser implements PostgresUpdateOperationParser { + private final PostgresUpdateOperationParser leafParser; private final int leafNodePathSize; @@ -22,13 +25,30 @@ public PostgresSetValueParser() { @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()); + if (input.isArray()) { + // For array columns, extract as Object[] and add as single param + Object[] values = + input + .getUpdate() + .getSubDocumentValue() + .accept(new PostgresSubDocumentArrayGetter()) + .values(); + input.getParamsBuilder().addObjectParam(values); + return String.format("\"%s\" = ?", input.getBaseField()); + } else { + // For scalar columns, use value parser which returns proper expression with type cast + String valueExpr = + input + .getUpdate() + .getSubDocumentValue() + .accept(new PostgresSubDocumentValueParser(input.getParamsBuilder())); + // For JSONB columns, use the returned expression (e.g., "?::jsonb" for nested documents) + // For other columns, use plain "?" + if (input.getColumnType() == PostgresDataType.JSONB) { + return String.format("\"%s\" = %s", input.getBaseField(), valueExpr); + } + return String.format("\"%s\" = ?", input.getBaseField()); + } } @Override 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 a82c3d911..d81a1869d 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 @@ -3,11 +3,24 @@ import static org.hypertrace.core.documentstore.model.subdoc.SubDocument.PATH_SEPARATOR; import static org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.formatSubDocPath; +import org.hypertrace.core.documentstore.postgres.query.v1.parser.filter.nonjson.field.PostgresDataType; + public class PostgresUnsetPathParser implements PostgresUpdateOperationParser { @Override public String parseNonJsonbField(final UpdateParserInput input) { - return String.format("\"%s\" = NULL", input.getBaseField()); + String baseField = input.getBaseField(); + + if (input.isArray()) { + // Array columns → empty array + return String.format("\"%s\" = '{}'", baseField); + } else if (input.getColumnType() == PostgresDataType.JSONB) { + // JSONB columns → empty object + return String.format("\"%s\" = '{}'::jsonb", baseField); + } else { + // Other columns → NULL + return String.format("\"%s\" = NULL", baseField); + } } @Override 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 249491004..35e3efca1 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 @@ -33,5 +33,6 @@ class UpdateParserInput { Params.Builder paramsBuilder; // only for flat collections PostgresDataType columnType; + boolean isArray; } } diff --git a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java index c59b2afd7..622996c94 100644 --- a/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java +++ b/document-store/src/test/java/org/hypertrace/core/documentstore/postgres/FlatPostgresCollectionTest.java @@ -265,7 +265,6 @@ void testUpsertRetriesOnUndefinedColumn() throws Exception { PSQLException psqlException = createPSQLException(PSQLState.UNDEFINED_COLUMN); when(mockPreparedStatement.executeQuery()).thenThrow(psqlException).thenReturn(mockResultSet); when(mockResultSet.next()).thenReturn(true); - when(mockResultSet.getBoolean("is_insert")).thenReturn(true); doNothing().when(mockSchemaRegistry).invalidate(COLLECTION_NAME); @@ -289,13 +288,13 @@ void testUpsertRetriesOnDatatypeMismatch() throws Exception { PSQLException psqlException = createPSQLException(PSQLState.DATATYPE_MISMATCH); when(mockPreparedStatement.executeQuery()).thenThrow(psqlException).thenReturn(mockResultSet); when(mockResultSet.next()).thenReturn(true); - when(mockResultSet.getBoolean("is_insert")).thenReturn(false); doNothing().when(mockSchemaRegistry).invalidate(COLLECTION_NAME); boolean result = flatPostgresCollection.upsert(key, document); - assertFalse(result); + // upsert always returns true if it succeeds + assertTrue(result); verify(mockSchemaRegistry, times(1)).invalidate(COLLECTION_NAME); verify(mockPreparedStatement, times(2)).executeQuery(); }