Skip to content

Commit c509cab

Browse files
suddendustsuddendust
committed
Refactor
1 parent 545439a commit c509cab

4 files changed

Lines changed: 191 additions & 234 deletions

File tree

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

Lines changed: 66 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -3,23 +3,33 @@
33
import com.fasterxml.jackson.databind.ObjectMapper;
44
import com.fasterxml.jackson.databind.node.ObjectNode;
55
import com.typesafe.config.ConfigFactory;
6+
import java.io.BufferedReader;
67
import java.io.IOException;
8+
import java.io.InputStream;
9+
import java.io.InputStreamReader;
10+
import java.nio.charset.StandardCharsets;
711
import java.sql.Connection;
812
import java.sql.PreparedStatement;
913
import java.util.HashMap;
1014
import java.util.Map;
15+
import java.util.stream.Stream;
1116
import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
1217
import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
1318
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
1419
import org.hypertrace.core.documentstore.expression.operators.RelationalOperator;
20+
import org.hypertrace.core.documentstore.model.options.MissingColumnStrategy;
1521
import org.hypertrace.core.documentstore.postgres.PostgresDatastore;
1622
import org.hypertrace.core.documentstore.query.Query;
23+
import org.junit.jupiter.api.extension.ExtensionContext;
24+
import org.junit.jupiter.params.provider.Arguments;
25+
import org.junit.jupiter.params.provider.ArgumentsProvider;
1726
import org.slf4j.Logger;
1827
import org.slf4j.LoggerFactory;
1928
import org.testcontainers.containers.GenericContainer;
2029
import org.testcontainers.containers.wait.strategy.Wait;
2130
import org.testcontainers.utility.DockerImageName;
2231

32+
/** Base class for write tests */
2333
public abstract class BaseWriteTest {
2434

2535
protected static final Logger LOGGER = LoggerFactory.getLogger(BaseWriteTest.class);
@@ -34,26 +44,32 @@ public abstract class BaseWriteTest {
3444
protected static GenericContainer<?> postgresContainer;
3545
protected static Datastore postgresDatastore;
3646

37-
protected static final String FLAT_COLLECTION_SCHEMA_SQL =
38-
"CREATE TABLE \"%s\" ("
39-
+ "\"id\" TEXT PRIMARY KEY,"
40-
+ "\"item\" TEXT,"
41-
+ "\"price\" INTEGER,"
42-
+ "\"quantity\" INTEGER,"
43-
+ "\"date\" TIMESTAMPTZ,"
44-
+ "\"in_stock\" BOOLEAN,"
45-
+ "\"tags\" TEXT[],"
46-
+ "\"categoryTags\" TEXT[],"
47-
+ "\"props\" JSONB,"
48-
+ "\"sales\" JSONB,"
49-
+ "\"numbers\" INTEGER[],"
50-
+ "\"scores\" DOUBLE PRECISION[],"
51-
+ "\"flags\" BOOLEAN[],"
52-
+ "\"big_number\" BIGINT,"
53-
+ "\"rating\" REAL,"
54-
+ "\"created_date\" DATE,"
55-
+ "\"weight\" DOUBLE PRECISION"
56-
+ ");";
47+
// Maps for multi-store tests
48+
protected static Map<String, Datastore> datastoreMap = new HashMap<>();
49+
protected static Map<String, Collection> collectionMap = new HashMap<>();
50+
51+
protected Collection getCollection(String storeName) {
52+
return collectionMap.get(storeName);
53+
}
54+
55+
private static final String FLAT_COLLECTION_SCHEMA_PATH =
56+
"schema/flat_collection_test_schema.sql";
57+
58+
protected static String loadFlatCollectionSchema() {
59+
try (InputStream is =
60+
BaseWriteTest.class.getClassLoader().getResourceAsStream(FLAT_COLLECTION_SCHEMA_PATH);
61+
BufferedReader reader =
62+
new BufferedReader(new InputStreamReader(is, StandardCharsets.UTF_8))) {
63+
StringBuilder sb = new StringBuilder();
64+
String line;
65+
while ((line = reader.readLine()) != null) {
66+
sb.append(line).append(" ");
67+
}
68+
return sb.toString().trim();
69+
} catch (Exception e) {
70+
throw new RuntimeException("Failed to load schema from " + FLAT_COLLECTION_SCHEMA_PATH, e);
71+
}
72+
}
5773

5874
protected static void initMongo() {
5975
mongoContainer =
@@ -106,7 +122,7 @@ protected static void shutdownPostgres() {
106122

107123
protected static void createFlatCollectionSchema(
108124
PostgresDatastore pgDatastore, String tableName) {
109-
String createTableSQL = String.format(FLAT_COLLECTION_SCHEMA_SQL, tableName);
125+
String createTableSQL = String.format(loadFlatCollectionSchema(), tableName);
110126

111127
try (Connection connection = pgDatastore.getPostgresClient();
112128
PreparedStatement statement = connection.prepareStatement(createTableSQL)) {
@@ -118,17 +134,6 @@ protected static void createFlatCollectionSchema(
118134
}
119135
}
120136

121-
protected static void clearTable(String tableName) {
122-
PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore;
123-
String deleteSQL = String.format("DELETE FROM \"%s\"", tableName);
124-
try (Connection connection = pgDatastore.getPostgresClient();
125-
PreparedStatement statement = connection.prepareStatement(deleteSQL)) {
126-
statement.executeUpdate();
127-
} catch (Exception e) {
128-
LOGGER.error("Failed to clear table {}: {}", tableName, e.getMessage(), e);
129-
}
130-
}
131-
132137
protected static String generateDocId(String prefix) {
133138
return prefix + "-" + System.currentTimeMillis() + "-" + (int) (Math.random() * 10000);
134139
}
@@ -147,7 +152,7 @@ protected Query buildQueryById(String docId) {
147152
.build();
148153
}
149154

150-
protected Document createTestDocument(String docId) throws IOException {
155+
protected Document createTestDocument(String docId) {
151156
Key key = new SingleValueKey(DEFAULT_TENANT, docId);
152157
String keyStr = key.toString();
153158

@@ -179,4 +184,32 @@ protected Document createTestDocument(String docId) throws IOException {
179184
protected Key createKey(String docId) {
180185
return new SingleValueKey(DEFAULT_TENANT, docId);
181186
}
187+
188+
protected static void clearTable(String tableName) {
189+
PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore;
190+
String deleteSQL = String.format("DELETE FROM \"%s\"", tableName);
191+
try (Connection connection = pgDatastore.getPostgresClient();
192+
PreparedStatement statement = connection.prepareStatement(deleteSQL)) {
193+
statement.executeUpdate();
194+
} catch (Exception e) {
195+
LOGGER.error("Failed to clear table {}: {}", tableName, e.getMessage(), e);
196+
}
197+
}
198+
199+
protected void insertTestDocument(String docId, Collection collection) throws IOException {
200+
Key key = createKey(docId);
201+
Document document = createTestDocument(docId);
202+
collection.upsert(key, document);
203+
}
204+
205+
/** Provides all MissingColumnStrategy values for parameterized tests */
206+
protected static class MissingColumnStrategyProvider implements ArgumentsProvider {
207+
@Override
208+
public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
209+
return Stream.of(
210+
Arguments.of(MissingColumnStrategy.SKIP),
211+
Arguments.of(MissingColumnStrategy.THROW),
212+
Arguments.of(MissingColumnStrategy.IGNORE_DOCUMENT));
213+
}
214+
}
182215
}

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

Lines changed: 20 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import java.util.Map;
2626
import java.util.Optional;
2727
import java.util.Set;
28-
import java.util.stream.Stream;
2928
import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
3029
import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
3130
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
@@ -47,10 +46,7 @@
4746
import org.junit.jupiter.api.DisplayName;
4847
import org.junit.jupiter.api.Nested;
4948
import org.junit.jupiter.api.Test;
50-
import org.junit.jupiter.api.extension.ExtensionContext;
5149
import org.junit.jupiter.params.ParameterizedTest;
52-
import org.junit.jupiter.params.provider.Arguments;
53-
import org.junit.jupiter.params.provider.ArgumentsProvider;
5450
import org.junit.jupiter.params.provider.ArgumentsSource;
5551
import org.testcontainers.junit.jupiter.Testcontainers;
5652

@@ -105,22 +101,10 @@ private static void executeInsertStatements() {
105101
@BeforeEach
106102
public void setupData() {
107103
// Clear and repopulate with initial data before each test
108-
clearTable();
104+
clearTable(FLAT_COLLECTION_NAME);
109105
executeInsertStatements();
110106
}
111107

112-
private static void clearTable() {
113-
PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore;
114-
String deleteSQL = String.format("DELETE FROM \"%s\"", FLAT_COLLECTION_NAME);
115-
try (Connection connection = pgDatastore.getPostgresClient();
116-
PreparedStatement statement = connection.prepareStatement(deleteSQL)) {
117-
statement.executeUpdate();
118-
LOGGER.info("Cleared table: {}", FLAT_COLLECTION_NAME);
119-
} catch (Exception e) {
120-
LOGGER.error("Failed to clear table: {}", e.getMessage(), e);
121-
}
122-
}
123-
124108
@AfterEach
125109
public void cleanup() {
126110
// Data is cleared in @BeforeEach, but cleanup here for safety
@@ -138,7 +122,7 @@ class UpsertTests {
138122
@Test
139123
@DisplayName("Should create new document when key doesn't exist and return true")
140124
void testUpsertNewDocument() throws Exception {
141-
String docId = getRandomDocId(4);
125+
String docId = generateDocId("test");
142126

143127
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
144128
objectNode.put("id", docId);
@@ -163,7 +147,7 @@ void testUpsertNewDocument() throws Exception {
163147
@Test
164148
@DisplayName("Should merge with existing document preserving unspecified fields")
165149
void testUpsertMergesWithExistingDocument() throws Exception {
166-
String docId = getRandomDocId(4);
150+
String docId = generateDocId("test");
167151

168152
// First, create a document with multiple fields
169153
ObjectNode initialNode = OBJECT_MAPPER.createObjectNode();
@@ -204,8 +188,8 @@ void testUpsertMergesWithExistingDocument() throws Exception {
204188
@Test
205189
@DisplayName("Upsert vs CreateOrReplace: upsert preserves, createOrReplace resets to default")
206190
void testUpsertVsCreateOrReplaceBehavior() throws Exception {
207-
String docId1 = getRandomDocId(4);
208-
String docId2 = getRandomDocId(4);
191+
String docId1 = generateDocId("test");
192+
String docId2 = generateDocId("test");
209193

210194
// Setup: Create two identical documents
211195
ObjectNode initialNode = OBJECT_MAPPER.createObjectNode();
@@ -263,7 +247,7 @@ void testUpsertVsCreateOrReplaceBehavior() throws Exception {
263247
@Test
264248
@DisplayName("Should skip unknown fields in upsert (default SKIP strategy)")
265249
void testUpsertSkipsUnknownFields() throws Exception {
266-
String docId = getRandomDocId(4);
250+
String docId = generateDocId("test");
267251

268252
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
269253
objectNode.put("id", docId);
@@ -308,7 +292,7 @@ class CreateTests {
308292
@DisplayName("Should create document with all supported data types")
309293
void testCreateWithAllDataTypes() throws Exception {
310294
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
311-
String docId = getRandomDocId(4);
295+
String docId = generateDocId("test");
312296

313297
objectNode.put("id", docId);
314298
objectNode.put("item", "Comprehensive Test Item");
@@ -411,7 +395,7 @@ void testCreateWithAllDataTypes() throws Exception {
411395
@DisplayName("Should throw DuplicateDocumentException when creating with existing key")
412396
void testCreateDuplicateDocument() throws Exception {
413397

414-
String docId = getRandomDocId(4);
398+
String docId = generateDocId("test");
415399
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
416400
objectNode.put("id", "dup-doc-200");
417401
objectNode.put("item", "First Item");
@@ -438,7 +422,7 @@ void testCreateDuplicateDocument() throws Exception {
438422
void testUnknownFieldsAsPerMissingColumnStrategy(MissingColumnStrategy missingColumnStrategy)
439423
throws Exception {
440424

441-
String docId = getRandomDocId(4);
425+
String docId = generateDocId("test");
442426

443427
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
444428
objectNode.put("id", docId);
@@ -487,7 +471,7 @@ void testEmptyMissingColumnStrategyConfigUsesDefault() throws Exception {
487471
Collection collectionWithEmptyStrategy = getFlatCollectionWithStrategy("");
488472

489473
// Test that it uses default SKIP strategy (unknown fields are skipped, not thrown)
490-
String docId = getRandomDocId(4);
474+
String docId = generateDocId("test");
491475
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
492476
objectNode.put("id", docId);
493477
objectNode.put("item", "Test Item");
@@ -508,7 +492,7 @@ void testEmptyMissingColumnStrategyConfigUsesDefault() throws Exception {
508492
void testInvalidMissingColumnStrategyConfigUsesDefault() throws Exception {
509493
Collection collectionWithInvalidStrategy = getFlatCollectionWithStrategy("INVALID_STRATEGY");
510494

511-
String docId = getRandomDocId(4);
495+
String docId = generateDocId("test");
512496
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
513497
objectNode.put("id", docId);
514498
objectNode.put("item", "Test Item");
@@ -627,7 +611,7 @@ void testCreateRefreshesSchemaOnUndefinedColumnError() throws Exception {
627611
void testUnparsableValuesAsPerMissingColStrategy(MissingColumnStrategy missingColumnStrategy)
628612
throws Exception {
629613

630-
String docId = getRandomDocId(4);
614+
String docId = generateDocId("test");
631615

632616
// Try to insert a string value into an integer column with wrong type
633617
// The unparseable column should be skipped, not throw an exception
@@ -681,11 +665,6 @@ void testUnparsableValuesAsPerMissingColStrategy(MissingColumnStrategy missingCo
681665
}
682666
}
683667

684-
private String getRandomDocId(int len) {
685-
return org.testcontainers.shaded.org.apache.commons.lang3.RandomStringUtils.random(
686-
len, true, false);
687-
}
688-
689668
private static Collection getFlatCollectionWithStrategy(String strategy) {
690669
String postgresConnectionUrl =
691670
String.format("jdbc:postgresql://localhost:%s/", postgresContainer.getMappedPort(5432));
@@ -720,19 +699,6 @@ interface ResultSetConsumer {
720699
void accept(ResultSet rs) throws Exception;
721700
}
722701

723-
static class MissingColumnStrategyProvider implements ArgumentsProvider {
724-
725-
@Override
726-
public Stream<? extends Arguments> provideArguments(ExtensionContext context) {
727-
return Stream.of(MissingColumnStrategy.values())
728-
.filter(
729-
strategy ->
730-
(strategy == MissingColumnStrategy.THROW)
731-
|| (strategy == MissingColumnStrategy.SKIP))
732-
.map(Arguments::of);
733-
}
734-
}
735-
736702
@Nested
737703
@DisplayName("CreateOrReplace Operations")
738704
class CreateOrReplaceTests {
@@ -742,7 +708,7 @@ class CreateOrReplaceTests {
742708
"Should create new document and return true. Cols not specified should be set of default NULL")
743709
void testCreateOrReplaceNewDocument() throws Exception {
744710

745-
String docId = getRandomDocId(4);
711+
String docId = generateDocId("test");
746712

747713
ObjectNode objectNode = OBJECT_MAPPER.createObjectNode();
748714
objectNode.put("id", "upsert-new-doc-100");
@@ -773,7 +739,7 @@ void testCreateOrReplaceNewDocument() throws Exception {
773739
@Test
774740
@DisplayName("Should replace existing document and return false")
775741
void testCreateOrReplaceExistingDocument() throws Exception {
776-
String docId = getRandomDocId(4);
742+
String docId = generateDocId("test");
777743
ObjectNode initialNode = OBJECT_MAPPER.createObjectNode();
778744
initialNode.put("id", docId);
779745
initialNode.put("item", "Original Item");
@@ -835,7 +801,7 @@ void testCreateOrReplaceSkipsUnknownFields() throws Exception {
835801
@Test
836802
@DisplayName("Should handle JSONB fields in createOrReplace")
837803
void testCreateOrReplaceWithJsonbField() throws Exception {
838-
String docId = getRandomDocId(4);
804+
String docId = generateDocId("test");
839805
ObjectNode initialNode = OBJECT_MAPPER.createObjectNode();
840806
initialNode.put("id", docId);
841807
initialNode.put("item", "Item with props");
@@ -2207,7 +2173,7 @@ class AddSubdocOperatorTests {
22072173
@Test
22082174
@DisplayName("Should ADD to all numeric types via bulkUpdate")
22092175
void testAddAllNumericTypes() throws Exception {
2210-
String docId = getRandomDocId(4);
2176+
String docId = generateDocId("test");
22112177
Key key = new SingleValueKey(DEFAULT_TENANT, docId);
22122178
ObjectNode node = OBJECT_MAPPER.createObjectNode();
22132179
node.put("item", "NumericTestItem");
@@ -2314,7 +2280,7 @@ void testAddAllNumericTypes() throws Exception {
23142280
@DisplayName("Should handle ADD on NULL column (treat as 0)")
23152281
void testAddOnNullColumn() throws Exception {
23162282
// Create a document with NULL numeric columns
2317-
String docId = getRandomDocId(4);
2283+
String docId = generateDocId("test");
23182284
Key key = new SingleValueKey(DEFAULT_TENANT, docId);
23192285
ObjectNode node = OBJECT_MAPPER.createObjectNode();
23202286
node.put("item", "NullPriceItem");
@@ -2418,7 +2384,7 @@ class AppendToListOperatorTests {
24182384
@Test
24192385
@DisplayName("Should APPEND_TO_LIST for top-level and nested arrays via bulkUpdate")
24202386
void testAppendToListAllCases() throws Exception {
2421-
String docId = getRandomDocId(4);
2387+
String docId = generateDocId("test");
24222388
Key key = new SingleValueKey(DEFAULT_TENANT, docId);
24232389
ObjectNode node = OBJECT_MAPPER.createObjectNode();
24242390
node.put("item", "AppendTestItem");
@@ -2512,7 +2478,7 @@ class AddToListIfAbsentOperatorTests {
25122478
@Test
25132479
@DisplayName("Should ADD_TO_LIST_IF_ABSENT for top-level and nested arrays via bulkUpdate")
25142480
void testAddToListIfAbsentAllCases() throws Exception {
2515-
String docId = getRandomDocId(4);
2481+
String docId = generateDocId("test");
25162482
Key key = new SingleValueKey(DEFAULT_TENANT, docId);
25172483
ObjectNode node = OBJECT_MAPPER.createObjectNode();
25182484
node.put("item", "AddIfAbsentTestItem");
@@ -2585,7 +2551,7 @@ class RemoveAllFromListOperatorTests {
25852551
@Test
25862552
@DisplayName("Should REMOVE_ALL_FROM_LIST for top-level and nested arrays via bulkUpdate")
25872553
void testRemoveAllFromListAllCases() throws Exception {
2588-
String docId = getRandomDocId(4);
2554+
String docId = generateDocId("test");
25892555
Key key = new SingleValueKey(DEFAULT_TENANT, docId);
25902556
ObjectNode node = OBJECT_MAPPER.createObjectNode();
25912557
node.put("item", "RemoveTestItem");

0 commit comments

Comments
 (0)