Skip to content

Commit 431e7c0

Browse files
kotharironakkotharironak
authored
feat: add support for selections of new query api (#90)
* feat: add support for selection type expression * adds integration test and fixes for it * uses fieldname as alias if it is not provided explicitly * fixed the test cases after fixing alias issue * initial logic of handling nested field * Handle the nested fields for selection * adds test cases for nested fields * addressed type related comment for precast * addressed rest of comments, and refactored a bit
1 parent 851a0f4 commit 431e7c0

14 files changed

+548
-36
lines changed

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

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package org.hypertrace.core.documentstore;
22

3+
import static org.hypertrace.core.documentstore.expression.operators.FunctionOperator.MULTIPLY;
34
import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.AND;
45
import static org.hypertrace.core.documentstore.expression.operators.LogicalOperator.OR;
56
import static org.hypertrace.core.documentstore.expression.operators.RelationalOperator.EQ;
@@ -37,6 +38,7 @@
3738
import org.bson.codecs.configuration.CodecConfigurationException;
3839
import org.hypertrace.core.documentstore.Filter.Op;
3940
import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
41+
import org.hypertrace.core.documentstore.expression.impl.FunctionExpression;
4042
import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
4143
import org.hypertrace.core.documentstore.expression.impl.LogicalExpression;
4244
import org.hypertrace.core.documentstore.expression.impl.RelationalExpression;
@@ -1567,6 +1569,78 @@ public void testQueryV1ForFilterWithLogicalExpressionAndOr(String dataStoreName)
15671569
dataStoreName, resultDocs, 6, "mongo/filter_with_logical_and_or_operator.json");
15681570
}
15691571

1572+
@ParameterizedTest
1573+
@MethodSource("databaseContextProvider")
1574+
public void testQueryV1ForSelectionExpression(String dataStoreName) throws IOException {
1575+
Map<Key, Document> documents = createDocumentsFromResource("mongo/collection_data.json");
1576+
Datastore datastore = datastoreMap.get(dataStoreName);
1577+
Collection collection = datastore.getCollection(COLLECTION_NAME);
1578+
1579+
// add docs
1580+
boolean result = collection.bulkUpsert(documents);
1581+
Assertions.assertTrue(result);
1582+
1583+
// query docs
1584+
org.hypertrace.core.documentstore.query.Query query =
1585+
org.hypertrace.core.documentstore.query.Query.builder()
1586+
.setFilter(
1587+
RelationalExpression.of(
1588+
IdentifierExpression.of("price"), EQ, ConstantExpression.of(10)))
1589+
.addSelection(IdentifierExpression.of("item"))
1590+
.addSelection(IdentifierExpression.of("props.brand"))
1591+
.addSelection(IdentifierExpression.of("props.seller.name"))
1592+
.addSelection(
1593+
FunctionExpression.builder()
1594+
.operand(IdentifierExpression.of("price"))
1595+
.operator(MULTIPLY)
1596+
.operand(IdentifierExpression.of("quantity"))
1597+
.build(),
1598+
"total")
1599+
.build();
1600+
1601+
Iterator<Document> resultDocs = collection.aggregate(query);
1602+
assertSizeAndDocsEqual(
1603+
dataStoreName, resultDocs, 2, "mongo/test_selection_expression_result.json");
1604+
}
1605+
1606+
@ParameterizedTest
1607+
@MethodSource("databaseContextProvider")
1608+
public void testQueryV1FunctionalSelectionExpressionWithNestedFieldWithAlias(String dataStoreName)
1609+
throws IOException {
1610+
Map<Key, Document> documents = createDocumentsFromResource("mongo/collection_data.json");
1611+
Datastore datastore = datastoreMap.get(dataStoreName);
1612+
Collection collection = datastore.getCollection(COLLECTION_NAME);
1613+
1614+
// add docs
1615+
boolean result = collection.bulkUpsert(documents);
1616+
Assertions.assertTrue(result);
1617+
1618+
// query docs
1619+
org.hypertrace.core.documentstore.query.Query query =
1620+
org.hypertrace.core.documentstore.query.Query.builder()
1621+
.setFilter(
1622+
RelationalExpression.of(
1623+
IdentifierExpression.of("price"), EQ, ConstantExpression.of(10)))
1624+
.addSelection(IdentifierExpression.of("item"))
1625+
.addSelection(IdentifierExpression.of("props.brand"), "props_brand")
1626+
.addSelection(IdentifierExpression.of("props.seller.name"), "props_seller_name")
1627+
.addSelection(
1628+
FunctionExpression.builder()
1629+
.operand(IdentifierExpression.of("price"))
1630+
.operator(MULTIPLY)
1631+
.operand(IdentifierExpression.of("quantity"))
1632+
.build(),
1633+
"total")
1634+
.build();
1635+
1636+
Iterator<Document> resultDocs = collection.aggregate(query);
1637+
assertSizeAndDocsEqual(
1638+
dataStoreName,
1639+
resultDocs,
1640+
2,
1641+
"mongo/test_selection_expression_nested_fields_alias_result.json");
1642+
}
1643+
15701644
private Map<String, List<CreateUpdateTestThread>> executeCreateUpdateThreads(
15711645
Collection collection, Operation operation, int numThreads, SingleValueKey documentKey) {
15721646
List<CreateUpdateTestThread> threads = new ArrayList<CreateUpdateTestThread>();

‎document-store/src/integrationTest/resources/mongo/test_selection_expression_nested_fields_alias_result.json‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
[
2+
{
3+
"item": "Soap",
4+
"total": 20,
5+
"props_seller_name": "Metro Chemicals Pvt. Ltd.",
6+
"props_brand": "Dettol"
7+
},
8+
{
9+
"item": "Soap",
10+
"total": 50
11+
}
12+
]

‎document-store/src/integrationTest/resources/mongo/test_selection_expression_result.json‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
[
2+
{
3+
"item": "Soap",
4+
"total": 20,
5+
"props": {
6+
"brand": "Dettol",
7+
"seller": {
8+
"name": "Metro Chemicals Pvt. Ltd."
9+
}
10+
}
11+
},
12+
{
13+
"item": "Soap",
14+
"total": 50
15+
}
16+
]

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

Lines changed: 72 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
package org.hypertrace.core.documentstore.postgres;
22

3-
import com.fasterxml.jackson.core.JsonProcessingException;
3+
import com.fasterxml.jackson.databind.JsonNode;
44
import com.fasterxml.jackson.databind.ObjectMapper;
55
import com.fasterxml.jackson.databind.node.ObjectNode;
66
import com.google.common.annotations.VisibleForTesting;
@@ -9,17 +9,20 @@
99
import java.sql.Connection;
1010
import java.sql.PreparedStatement;
1111
import java.sql.ResultSet;
12+
import java.sql.ResultSetMetaData;
1213
import java.sql.SQLException;
1314
import java.sql.Statement;
1415
import java.sql.Timestamp;
1516
import java.util.Arrays;
17+
import java.util.HashMap;
1618
import java.util.HashSet;
1719
import java.util.List;
1820
import java.util.Map;
1921
import java.util.NoSuchElementException;
2022
import java.util.Set;
2123
import java.util.stream.Collectors;
2224
import lombok.SneakyThrows;
25+
import org.apache.commons.lang3.StringUtils;
2326
import org.hypertrace.core.documentstore.BulkArrayValueUpdateRequest;
2427
import org.hypertrace.core.documentstore.BulkDeleteResult;
2528
import org.hypertrace.core.documentstore.BulkUpdateRequest;
@@ -33,6 +36,7 @@
3336
import org.hypertrace.core.documentstore.Key;
3437
import org.hypertrace.core.documentstore.Query;
3538
import org.hypertrace.core.documentstore.UpdateResult;
39+
import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils;
3640
import org.slf4j.Logger;
3741
import org.slf4j.LoggerFactory;
3842

@@ -281,31 +285,40 @@ public CloseableIterator<Document> search(Query query) {
281285
@Override
282286
public CloseableIterator<Document> find(
283287
final org.hypertrace.core.documentstore.query.Query query) {
284-
throw new UnsupportedOperationException();
288+
return executeQueryV1(query);
285289
}
286290

287291
@Override
288292
public CloseableIterator<Document> aggregate(
289293
final org.hypertrace.core.documentstore.query.Query query) {
294+
return executeQueryV1(query);
295+
}
296+
297+
@Override
298+
public long count(org.hypertrace.core.documentstore.query.Query query) {
299+
throw new UnsupportedOperationException();
300+
}
301+
302+
private CloseableIterator<Document> executeQueryV1(
303+
final org.hypertrace.core.documentstore.query.Query query) {
290304
org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser queryParser =
291305
new org.hypertrace.core.documentstore.postgres.query.v1.PostgresQueryParser(collectionName);
292306
String sqlQuery = queryParser.parse(query);
293307
try {
294308
PreparedStatement preparedStatement =
295309
buildPreparedStatement(sqlQuery, queryParser.getParamsBuilder().build());
296310
ResultSet resultSet = preparedStatement.executeQuery();
297-
return new PostgresResultIterator(resultSet);
311+
CloseableIterator closeableIterator =
312+
query.getSelections().size() > 0
313+
? new PostgresResultIteratorWithMetaData(resultSet)
314+
: new PostgresResultIterator(resultSet);
315+
return closeableIterator;
298316
} catch (SQLException e) {
299317
LOGGER.error("SQLException querying documents. query: {}", query, e);
300318
}
301319
return EMPTY_ITERATOR;
302320
}
303321

304-
@Override
305-
public long count(org.hypertrace.core.documentstore.query.Query query) {
306-
throw new UnsupportedOperationException();
307-
}
308-
309322
@VisibleForTesting
310323
protected PreparedStatement buildPreparedStatement(String sqlQuery, Params params)
311324
throws SQLException, RuntimeException {
@@ -614,10 +627,10 @@ public void drop() {
614627

615628
static class PostgresResultIterator implements CloseableIterator {
616629

617-
private final ObjectMapper MAPPER = new ObjectMapper();
618-
private ResultSet resultSet;
619-
private boolean cursorMovedForward = false;
620-
private boolean hasNext = false;
630+
protected final ObjectMapper MAPPER = new ObjectMapper();
631+
protected ResultSet resultSet;
632+
protected boolean cursorMovedForward = false;
633+
protected boolean hasNext = false;
621634

622635
public PostgresResultIterator(ResultSet resultSet) {
623636
this.resultSet = resultSet;
@@ -651,7 +664,7 @@ public Document next() {
651664
}
652665
}
653666

654-
private Document prepareDocument() throws SQLException, JsonProcessingException, IOException {
667+
protected Document prepareDocument() throws SQLException, IOException {
655668
String documentString = resultSet.getString(DOCUMENT);
656669
ObjectNode jsonNode = (ObjectNode) MAPPER.readTree(documentString);
657670
jsonNode.remove(DOCUMENT_ID);
@@ -664,25 +677,59 @@ private Document prepareDocument() throws SQLException, JsonProcessingException,
664677
return new JSONDocument(MAPPER.writeValueAsString(jsonNode));
665678
}
666679

667-
private String prepareNumericBlock(String fieldName, Object value) {
668-
if (value instanceof Number) {
669-
String fmt = "case jsonb_typeof(%s)\n" + "WHEN ‘number’ THEN (%s)::numeric > ?\n" + "end";
670-
} else if (value instanceof Boolean) {
671-
String fmtBoolean =
672-
"case jsonb_typeof(<field>)\n"
673-
+ "WHEN 'boolean' THEN (<field>)::boolean > ?\n"
674-
+ "end";
675-
}
676-
return null;
677-
}
678-
679680
@SneakyThrows
680681
@Override
681682
public void close() {
682683
resultSet.close();
683684
}
684685
}
685686

687+
static class PostgresResultIteratorWithMetaData extends PostgresResultIterator {
688+
689+
public PostgresResultIteratorWithMetaData(ResultSet resultSet) {
690+
super(resultSet);
691+
}
692+
693+
@Override
694+
protected Document prepareDocument() throws SQLException, IOException {
695+
ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
696+
int columnCount = resultSetMetaData.getColumnCount();
697+
Map<String, Object> jsonNode = new HashMap();
698+
for (int i = 1; i <= columnCount; i++) {
699+
String columnName = resultSetMetaData.getColumnName(i);
700+
String columnValue = resultSet.getString(i);
701+
if (StringUtils.isNotEmpty(columnValue)) {
702+
JsonNode leafNodeValue = MAPPER.readTree(columnValue);
703+
if (PostgresUtils.isEncodedNestedField(columnName)) {
704+
handleNestedField(
705+
PostgresUtils.decodeAliasForNestedField(columnName), jsonNode, leafNodeValue);
706+
} else {
707+
jsonNode.put(columnName, leafNodeValue);
708+
}
709+
}
710+
}
711+
return new JSONDocument(MAPPER.writeValueAsString(jsonNode));
712+
}
713+
714+
private void handleNestedField(
715+
String columnName, Map<String, Object> rootNode, JsonNode leafNodeValue) {
716+
List<String> keys = PostgresUtils.splitNestedField(columnName);
717+
// find the leaf node or create one for adding property value
718+
Map<String, Object> curNode = rootNode;
719+
for (int l = 0; l < keys.size() - 1; l++) {
720+
String key = keys.get(l);
721+
Map<String, Object> node = (Map<String, Object>) curNode.get(key);
722+
if (node == null) {
723+
node = new HashMap<>();
724+
curNode.put(key, node);
725+
}
726+
curNode = node;
727+
}
728+
String leafKey = keys.get(keys.size() - 1);
729+
curNode.put(leafKey, leafNodeValue);
730+
}
731+
}
732+
686733
private static CloseableIterator<Document> createEmptyIterator() {
687734
return new CloseableIterator<>() {
688735
@Override

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/PostgresQueryParser.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@
77
import org.hypertrace.core.documentstore.postgres.Params;
88
import org.hypertrace.core.documentstore.postgres.Params.Builder;
99
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresFilterTypeExpressionVisitor;
10+
import org.hypertrace.core.documentstore.postgres.query.v1.vistors.PostgresSelectTypeExpressionVisitor;
1011
import org.hypertrace.core.documentstore.query.Pagination;
1112
import org.hypertrace.core.documentstore.query.Query;
13+
import org.hypertrace.core.documentstore.query.SelectionSpec;
1214
import org.hypertrace.core.documentstore.query.SortingSpec;
1315

1416
public class PostgresQueryParser {
@@ -30,6 +32,12 @@ public String parse(Query query) {
3032
StringBuilder sqlBuilder = new StringBuilder(String.format("SELECT * FROM %s", collection));
3133
paramsBuilder = Params.newBuilder();
3234

35+
// selection clause
36+
Optional<String> selectionClause = parseSelection(query.getSelections());
37+
if (selectionClause.isPresent()) {
38+
sqlBuilder = new StringBuilder();
39+
sqlBuilder.append(String.format("SELECT %s FROM %s", selectionClause.get(), collection));
40+
}
3341
// where clause
3442
Optional<String> whereFilter = parseFilter(query.getFilter());
3543
if (whereFilter.isPresent()) {
@@ -63,6 +71,10 @@ public String parse(Query query) {
6371
return sqlBuilder.toString();
6472
}
6573

74+
private Optional<String> parseSelection(List<SelectionSpec> selectionSpecs) {
75+
return Optional.ofNullable(PostgresSelectTypeExpressionVisitor.getSelections(selectionSpecs));
76+
}
77+
6678
private Optional<String> parseFilter(Optional<FilterTypeExpression> filterTypeExpression) {
6779
return filterTypeExpression.map(
6880
expression -> expression.accept(new PostgresFilterTypeExpressionVisitor(this)));

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresConstantExpressionVisitor.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,15 @@
11
package org.hypertrace.core.documentstore.postgres.query.v1.vistors;
22

3+
import lombok.NoArgsConstructor;
34
import org.hypertrace.core.documentstore.expression.impl.ConstantExpression;
45

6+
@NoArgsConstructor
57
public class PostgresConstantExpressionVisitor extends PostgresSelectTypeExpressionVisitor {
68

9+
public PostgresConstantExpressionVisitor(PostgresSelectTypeExpressionVisitor baseVisitor) {
10+
super(baseVisitor);
11+
}
12+
713
@Override
814
public Object visit(final ConstantExpression expression) {
915
return expression.getValue();

‎document-store/src/main/java/org/hypertrace/core/documentstore/postgres/query/v1/vistors/PostgresDataAccessorIdentifierExpressionVisitor.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
package org.hypertrace.core.documentstore.postgres.query.v1.vistors;
2+
3+
import lombok.NoArgsConstructor;
4+
import org.hypertrace.core.documentstore.expression.impl.IdentifierExpression;
5+
import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils;
6+
import org.hypertrace.core.documentstore.postgres.utils.PostgresUtils.Type;
7+
8+
@NoArgsConstructor
9+
public class PostgresDataAccessorIdentifierExpressionVisitor
10+
extends PostgresSelectTypeExpressionVisitor {
11+
12+
private Type type = Type.NUMERIC;
13+
14+
public PostgresDataAccessorIdentifierExpressionVisitor(Type type) {
15+
this.type = type;
16+
}
17+
18+
public PostgresDataAccessorIdentifierExpressionVisitor(
19+
PostgresSelectTypeExpressionVisitor baseVisitor) {
20+
super(baseVisitor);
21+
}
22+
23+
public PostgresDataAccessorIdentifierExpressionVisitor(
24+
PostgresSelectTypeExpressionVisitor baseVisitor, Type type) {
25+
super(baseVisitor);
26+
this.type = type;
27+
}
28+
29+
@Override
30+
public String visit(final IdentifierExpression expression) {
31+
String dataAccessor = PostgresUtils.prepareFieldDataAccessorExpr(expression.getName());
32+
if (type.equals(Type.NUMERIC)) {
33+
return PostgresUtils.prepareCast(dataAccessor, 1);
34+
} else if (type.equals(Type.BOOLEAN)) {
35+
return PostgresUtils.prepareCast(dataAccessor, true);
36+
}
37+
return PostgresUtils.prepareCast(dataAccessor, "");
38+
}
39+
}

0 commit comments

Comments
 (0)