Skip to content

Commit 13e591e

Browse files
suddendustsuddendust
committed
Added test cases
1 parent 7f7defe commit 13e591e

4 files changed

Lines changed: 314 additions & 2 deletions

File tree

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

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,12 +53,14 @@
5353
import static org.hypertrace.core.documentstore.utils.Utils.assertDocsAndSizeEqualWithoutOrder;
5454
import static org.hypertrace.core.documentstore.utils.Utils.convertJsonToMap;
5555
import static org.hypertrace.core.documentstore.utils.Utils.readFileFromResource;
56+
import static org.junit.Assert.assertNotNull;
5657
import static org.junit.jupiter.api.Assertions.assertEquals;
5758
import static org.junit.jupiter.api.Assertions.assertFalse;
5859
import static org.junit.jupiter.api.Assertions.assertNotEquals;
5960
import static org.junit.jupiter.api.Assertions.assertThrows;
6061
import static org.junit.jupiter.api.Assertions.assertTrue;
6162

63+
import com.fasterxml.jackson.databind.JsonNode;
6264
import com.fasterxml.jackson.databind.ObjectMapper;
6365
import com.typesafe.config.Config;
6466
import com.typesafe.config.ConfigFactory;
@@ -97,6 +99,7 @@
9799
import org.hypertrace.core.documentstore.model.options.UpdateOptions.MissingDocumentStrategy;
98100
import org.hypertrace.core.documentstore.model.subdoc.SubDocumentUpdate;
99101
import org.hypertrace.core.documentstore.model.subdoc.SubDocumentValue;
102+
import org.hypertrace.core.documentstore.postgres.PostgresDatastore;
100103
import org.hypertrace.core.documentstore.query.Aggregation;
101104
import org.hypertrace.core.documentstore.query.Filter;
102105
import org.hypertrace.core.documentstore.query.Pagination;
@@ -109,6 +112,7 @@
109112
import org.hypertrace.core.documentstore.utils.Utils;
110113
import org.junit.jupiter.api.AfterAll;
111114
import org.junit.jupiter.api.BeforeAll;
115+
import org.junit.jupiter.api.Disabled;
112116
import org.junit.jupiter.api.Nested;
113117
import org.junit.jupiter.api.extension.ExtensionContext;
114118
import org.junit.jupiter.params.ParameterizedTest;
@@ -126,6 +130,9 @@ public class DocStoreQueryV1Test {
126130

127131
private static final String COLLECTION_NAME = "myTest";
128132
private static final String UPDATABLE_COLLECTION_NAME = "updatable_collection";
133+
private static final String FLAT_COLLECTION_NAME = "myTestFlat";
134+
private static final String PG_FLAT_COLLECTION_INSERT_STATMENTS_FILE_LOC =
135+
"query/pg_flat_collection_insert.json";
129136

130137
private static Map<String, Datastore> datastoreMap;
131138

@@ -175,12 +182,79 @@ public static void init() throws IOException {
175182
createCollectionData("query/collection_data.json", COLLECTION_NAME);
176183
}
177184

185+
private static void createFlatCollectionSchema(
186+
Datastore postgresDatastore, String collectionName) {
187+
String createTableSQL =
188+
String.format(
189+
"CREATE TABLE \"%s\" ("
190+
+ "\"_id\" INTEGER PRIMARY KEY,"
191+
+ "\"item\" TEXT,"
192+
+ "\"price\" INTEGER,"
193+
+ "\"quantity\" INTEGER,"
194+
+ "\"date\" TIMESTAMPTZ,"
195+
+ "\"props\" JSONB,"
196+
+ "\"sales\" JSONB"
197+
+ ");",
198+
collectionName);
199+
200+
// Access PostgresDatastore directly to get client connection
201+
PostgresDatastore pgDatastore = (PostgresDatastore) postgresDatastore;
202+
203+
try {
204+
205+
// Execute the CREATE TABLE SQL directly using the public getPostgresClient method
206+
try (java.sql.Connection connection = pgDatastore.getPostgresClient();
207+
java.sql.PreparedStatement statement = connection.prepareStatement(createTableSQL)) {
208+
statement.execute();
209+
System.out.println("Created flat collection table: " + collectionName);
210+
}
211+
} catch (Exception e) {
212+
System.err.println("Failed to create flat collection schema: " + e.getMessage());
213+
e.printStackTrace();
214+
}
215+
}
216+
217+
private static void executeInsertStatements(PostgresDatastore pgDatastore) {
218+
try {
219+
// Read JSON file from resources
220+
String jsonContent =
221+
readFileFromResource(PG_FLAT_COLLECTION_INSERT_STATMENTS_FILE_LOC).orElseThrow();
222+
ObjectMapper mapper = new ObjectMapper();
223+
JsonNode rootNode = mapper.readTree(jsonContent);
224+
225+
// Extract statements array
226+
JsonNode statementsNode = rootNode.get("statements");
227+
if (statementsNode == null || !statementsNode.isArray()) {
228+
throw new RuntimeException("Invalid JSON format: 'statements' array not found");
229+
}
230+
231+
try (java.sql.Connection connection = pgDatastore.getPostgresClient()) {
232+
for (com.fasterxml.jackson.databind.JsonNode statementNode : statementsNode) {
233+
String statement = statementNode.asText().trim();
234+
if (!statement.isEmpty()) {
235+
try (java.sql.PreparedStatement preparedStatement =
236+
connection.prepareStatement(statement)) {
237+
preparedStatement.executeUpdate();
238+
}
239+
}
240+
}
241+
}
242+
} catch (Exception e) {
243+
System.err.println("Failed to execute INSERT statements: " + e.getMessage());
244+
}
245+
}
246+
178247
private static void createCollectionData(final String resourcePath, final String collectionName)
179248
throws IOException {
180249
final Map<Key, Document> documents = Utils.buildDocumentsFromResource(resourcePath);
181250
datastoreMap.forEach(
182251
(k, v) -> {
183252
v.deleteCollection(collectionName);
253+
// for Postgres, we also create the flat collection
254+
if (v instanceof PostgresDatastore) {
255+
createFlatCollectionSchema(v, FLAT_COLLECTION_NAME);
256+
executeInsertStatements((PostgresDatastore) v);
257+
}
184258
v.createCollection(collectionName, null);
185259
Collection collection = v.getCollection(collectionName);
186260
collection.bulkUpsert(documents);
@@ -3674,4 +3748,150 @@ private static void testCountApi(
36743748
final long expectedSize = convertJsonToMap(fileContent).size();
36753749
assertEquals(expectedSize, actualSize);
36763750
}
3751+
3752+
@ParameterizedTest
3753+
@ArgumentsSource(PostgresProvider.class)
3754+
void testFlatPostgresCollectionFindAll(String dataStoreName) throws IOException {
3755+
Datastore datastore = datastoreMap.get(dataStoreName);
3756+
Collection flatCollection =
3757+
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
3758+
3759+
// Test basic query to retrieve all documents
3760+
Query query = Query.builder().build();
3761+
CloseableIterator<Document> iterator = flatCollection.find(query);
3762+
3763+
// Count documents
3764+
long count = 0;
3765+
while (iterator.hasNext()) {
3766+
Document doc = iterator.next();
3767+
count++;
3768+
// Verify document has content (basic validation)
3769+
assertNotNull(doc);
3770+
assertNotNull(doc.toJson());
3771+
assertTrue(doc.toJson().length() > 0);
3772+
assertEquals(DocumentType.FLAT, doc.getDocumentType());
3773+
}
3774+
iterator.close();
3775+
3776+
// Should have 8 documents from the INSERT statements
3777+
assertEquals(8, count);
3778+
}
3779+
3780+
@ParameterizedTest
3781+
@ArgumentsSource(PostgresProvider.class)
3782+
void testFlatPostgresCollectionFilterByItem(String dataStoreName) throws IOException {
3783+
Datastore datastore = datastoreMap.get(dataStoreName);
3784+
Collection flatCollection =
3785+
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
3786+
3787+
// Test filtering by item
3788+
Query itemQuery =
3789+
Query.builder()
3790+
.setFilter(
3791+
RelationalExpression.of(
3792+
IdentifierExpression.of("item"), EQ, ConstantExpression.of("Soap")))
3793+
.build();
3794+
3795+
CloseableIterator<Document> soapIterator = flatCollection.find(itemQuery);
3796+
long soapCount = 0;
3797+
while (soapIterator.hasNext()) {
3798+
Document doc = soapIterator.next();
3799+
// Verify it's a soap document by checking JSON contains "Soap"
3800+
assertTrue(doc.toJson().contains("\"Soap\""));
3801+
soapCount++;
3802+
}
3803+
soapIterator.close();
3804+
3805+
// Should have 3 soap documents (IDs 1, 5, 8)
3806+
assertEquals(3, soapCount);
3807+
}
3808+
3809+
@ParameterizedTest
3810+
@ArgumentsSource(PostgresProvider.class)
3811+
void testFlatPostgresCollectionCount(String dataStoreName) throws IOException {
3812+
Datastore datastore = datastoreMap.get(dataStoreName);
3813+
Collection flatCollection =
3814+
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
3815+
3816+
// Test count method - all documents
3817+
long totalCount = flatCollection.count(Query.builder().build());
3818+
assertEquals(8, totalCount);
3819+
3820+
// Test count with filter - soap documents only
3821+
Query soapQuery =
3822+
Query.builder()
3823+
.setFilter(
3824+
RelationalExpression.of(
3825+
IdentifierExpression.of("item"), EQ, ConstantExpression.of("Soap")))
3826+
.build();
3827+
long soapCountQuery = flatCollection.count(soapQuery);
3828+
assertEquals(3, soapCountQuery);
3829+
}
3830+
3831+
/**
3832+
* This test is disabled for now because flat collections do not support search on nested queries
3833+
* in JSONB fields (ex. props.brand)
3834+
*
3835+
* @param dataStoreName the datastore name, in this case, Postgres
3836+
* @throws IOException
3837+
*/
3838+
@ParameterizedTest
3839+
@ArgumentsSource(PostgresProvider.class)
3840+
@Disabled
3841+
void testFlatPostgresCollectionNestedFieldQuery(String dataStoreName) throws IOException {
3842+
Datastore datastore = datastoreMap.get(dataStoreName);
3843+
Collection flatCollection =
3844+
datastore.getCollectionForType(FLAT_COLLECTION_NAME, DocumentType.FLAT);
3845+
3846+
// Test querying nested field in props JSONB column - filter by brand
3847+
Query brandQuery =
3848+
Query.builder()
3849+
.setFilter(
3850+
RelationalExpression.of(
3851+
IdentifierExpression.of("props.brand"), EQ, ConstantExpression.of("Dettol")))
3852+
.build();
3853+
3854+
CloseableIterator<Document> brandIterator = flatCollection.find(brandQuery);
3855+
long brandCount = 0;
3856+
while (brandIterator.hasNext()) {
3857+
Document doc = brandIterator.next();
3858+
// Verify it contains the expected brand
3859+
assertTrue(doc.toJson().contains("\"Dettol\""));
3860+
brandCount++;
3861+
}
3862+
brandIterator.close();
3863+
3864+
// Should have 1 Dettol document (ID 1)
3865+
assertEquals(1, brandCount);
3866+
3867+
// Test querying deeply nested field - filter by seller city
3868+
Query cityQuery =
3869+
Query.builder()
3870+
.setFilter(
3871+
RelationalExpression.of(
3872+
IdentifierExpression.of("props.seller.address.city"),
3873+
EQ,
3874+
ConstantExpression.of("Mumbai")))
3875+
.build();
3876+
3877+
CloseableIterator<Document> cityIterator = flatCollection.find(cityQuery);
3878+
long cityCount = 0;
3879+
while (cityIterator.hasNext()) {
3880+
Document doc = cityIterator.next();
3881+
// Verify it contains Mumbai
3882+
assertTrue(doc.toJson().contains("\"Mumbai\""));
3883+
cityCount++;
3884+
}
3885+
cityIterator.close();
3886+
3887+
// Should have 2 Mumbai documents (IDs 1, 3)
3888+
assertEquals(2, cityCount);
3889+
3890+
// Test count with nested field filter
3891+
long brandCountQuery = flatCollection.count(brandQuery);
3892+
assertEquals(1, brandCountQuery);
3893+
3894+
long cityCountQuery = flatCollection.count(cityQuery);
3895+
assertEquals(2, cityCountQuery);
3896+
}
36773897
}

‎document-store/src/integrationTest/resources/query/pg_flat_collection_insert.json‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
{
2+
"statements": [
3+
"INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"props\", \"sales\"\n) VALUES (\n1, 'Soap', 10, 2, '2014-03-01T08:00:00Z',\n'{\"colors\": [\"Blue\", \"Green\"], \"brand\": \"Dettol\", \"size\": \"M\", \"seller\": {\"name\": \"Metro Chemicals Pvt. Ltd.\", \"address\": {\"city\": \"Mumbai\", \"pincode\": 400004}}}',\nNULL\n)",
4+
"INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"props\", \"sales\"\n) VALUES (\n2, 'Mirror', 20, 1, '2014-03-01T09:00:00Z',\nNULL,\nNULL\n)",
5+
"INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"props\", \"sales\"\n) VALUES (\n3, 'Shampoo', 5, 10, '2014-03-15T09:00:00Z',\n'{\"colors\": [\"Black\"], \"brand\": \"Sunsilk\", \"size\": \"L\", \"seller\": {\"name\": \"Metro Chemicals Pvt. Ltd.\", \"address\": {\"city\": \"Mumbai\", \"pincode\": 400004}}}',\nNULL\n)",
6+
"INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"props\", \"sales\"\n) VALUES (\n4, 'Shampoo', 5, 20, '2014-04-04T11:21:39.736Z',\nNULL,\nNULL\n)",
7+
"INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"props\", \"sales\"\n) VALUES (\n5, 'Soap', 20, 5, '2014-04-04T21:23:13.331Z',\n'{\"colors\": [\"Orange\", \"Blue\"], \"brand\": \"Lifebuoy\", \"size\": \"S\", \"seller\": {\"name\": \"Hans and Co.\", \"address\": {\"city\": \"Kolkata\", \"pincode\": 700007}}}',\nNULL\n)",
8+
"INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"props\", \"sales\"\n) VALUES (\n6, 'Comb', 7.5, 5, '2015-06-04T05:08:13Z',\nNULL,\nNULL\n)",
9+
"INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"props\", \"sales\"\n) VALUES (\n7, 'Comb', 7.5, 10, '2015-09-10T08:43:00Z',\n'{\"colors\": [], \"seller\": {\"name\": \"Go Go Plastics\", \"address\": {\"city\": \"Kolkata\", \"pincode\": 700007}}}',\nNULL\n)",
10+
"INSERT INTO \"myTestFlat\" (\n\"_id\", \"item\", \"price\", \"quantity\", \"date\", \"props\", \"sales\"\n) VALUES (\n8, 'Soap', 10, 5, '2016-02-06T20:20:13Z',\nNULL,\nNULL\n)"
11+
]
12+
}

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import org.hypertrace.core.documentstore.Document;
55
import org.hypertrace.core.documentstore.model.options.QueryOptions;
66
import org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser;
7+
import org.hypertrace.core.documentstore.postgres.query.v1.transformer.FlatPostgresFieldTransformer;
78
import org.hypertrace.core.documentstore.query.Query;
89
import org.slf4j.Logger;
910
import org.slf4j.LoggerFactory;
@@ -12,6 +13,7 @@
1213
* PostgreSQL collection implementation for flat documents. All fields are stored as top-level
1314
* PostgreSQL columns.
1415
*/
16+
// todo: Throw unsupported op exception for all write methods
1517
public class FlatPostgresCollection extends PostgresCollection {
1618

1719
private static final Logger LOGGER = LoggerFactory.getLogger(FlatPostgresCollection.class);
@@ -37,11 +39,19 @@ public CloseableIterator<Document> find(
3739
@Override
3840
public long count(
3941
org.hypertrace.core.documentstore.query.Query query, QueryOptions queryOptions) {
40-
PostgresQueryParser queryParser = new PostgresQueryParser(tableIdentifier, query);
42+
PostgresQueryParser queryParser =
43+
new PostgresQueryParser(
44+
tableIdentifier,
45+
query,
46+
new org.hypertrace.core.documentstore.postgres.query.v1.transformer
47+
.FlatPostgresFieldTransformer());
4148
return countWithParser(query, queryParser);
4249
}
4350

4451
private PostgresQueryParser createParser(Query query) {
45-
return new PostgresQueryParser(tableIdentifier, PostgresQueryExecutor.transformAndLog(query));
52+
return new PostgresQueryParser(
53+
tableIdentifier,
54+
PostgresQueryExecutor.transformAndLog(query),
55+
new FlatPostgresFieldTransformer());
4656
}
4757
}

0 commit comments

Comments
 (0)