feat: add selection support in older api (#115)

* feat: add selection support in older api * fixed the bulk array update internal query * fixed closable iterator issue * refactored common method, and fixed handling of outer column * refactor: clean up and removed duplicate logic * adds warn logs for debugging * fix for insight dto * adds unit and integration test for selection clause * converting warn log to debug log * Update document-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresQueryParser.java Co-authored-by: Suresh Prakash <[email protected]> * addressed comments for reusing ID field, and rededunt condition Co-authored-by: Suresh Prakash <[email protected]>
hypertrace · Aug 12, 2022 · 86fbc5d · 86fbc5d
1 parent 0a35e0e
commit 86fbc5d
Show file tree

Hide file tree

Showing 5 changed files with 169 additions and 250 deletions.
diff --git a/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreTest.java b/document-store/src/integrationTest/java/org/hypertrace/core/documentstore/DocStoreTest.java
@@ -19,6 +19,7 @@
 
 import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
@@ -2099,6 +2100,25 @@ public void testSearch(String dataStoreName) throws IOException {
         assertTrue(exception.getMessage().contains(expected));
       }
     }
+
+    // test for selections
+    {
+      Query query = new Query();
+      query.addSelection("entityId");
+      query.addSelection("entityType");
+      query.setFilter(new Filter(Filter.Op.EQ, "_id", key1.toString()));
+      Iterator<Document> results = collection.search(query);
+      List<Document> documents = new ArrayList<>();
+      while (results.hasNext()) {
+        documents.add(results.next());
+      }
+      Assertions.assertEquals(1, documents.size());
+      Map<String, String> result =
+          OBJECT_MAPPER.readValue(documents.get(0).toJson(), new TypeReference<>() {});
+      Assertions.assertEquals(2, result.size());
+      Assertions.assertEquals("SERVICE", result.get("entityType"));
+      Assertions.assertEquals("e3ffc6f0-fc92-3a9c-9fa0-26269184d1aa", result.get("entityId"));
+    }
   }
 
   @ParameterizedTest

diff --git a/...nt-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java b/...nt-store/src/main/java/org/hypertrace/core/documentstore/postgres/PostgresCollection.java
@@ -1,6 +1,7 @@
 package org.hypertrace.core.documentstore.postgres;
 
 import com.fasterxml.jackson.core.JsonParseException;
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ArrayNode;
@@ -324,17 +325,18 @@ public BulkUpdateResult bulkOperationOnArrayValue(BulkArrayValueUpdateRequest re
 
   @Override
   public CloseableIterator<Document> search(Query query) {
+    String selection = PostgresQueryParser.parseSelections(query.getSelections());
+    StringBuilder sqlBuilder =
+        new StringBuilder(String.format("SELECT %s FROM ", selection)).append(collectionName);
+
     String filters = null;
-    StringBuilder sqlBuilder = new StringBuilder("SELECT * FROM ").append(collectionName);
     Params.Builder paramsBuilder = Params.newBuilder();
 
     // If there is a filter in the query, parse it fully.
     if (query.getFilter() != null) {
       filters = PostgresQueryParser.parseFilter(query.getFilter(), paramsBuilder);
     }
 
-    LOGGER.debug("Sending query to PostgresSQL: {} : {}", collectionName, filters);
-
     if (filters != null) {
       sqlBuilder.append(" WHERE ").append(filters);
     }
@@ -354,13 +356,20 @@ public CloseableIterator<Document> search(Query query) {
       sqlBuilder.append(" OFFSET ").append(offset);
     }
 
+    String pgSqlQuery = sqlBuilder.toString();
     try {
       PreparedStatement preparedStatement =
-          buildPreparedStatement(sqlBuilder.toString(), paramsBuilder.build());
+          buildPreparedStatement(pgSqlQuery, paramsBuilder.build());
+      LOGGER.debug("Executing search query to PostgresSQL:{}", preparedStatement.toString());
       ResultSet resultSet = preparedStatement.executeQuery();
-      return new PostgresResultIterator(resultSet);
+      CloseableIterator closeableIterator =
+          query.getSelections().size() > 0
+              ? new PostgresResultIteratorWithMetaData(resultSet)
+              : new PostgresResultIterator(resultSet);
+      return closeableIterator;
     } catch (SQLException e) {
-      LOGGER.error("SQLException querying documents. query: {}", query, e);
+      LOGGER.error(
+          "SQLException in querying documents - query: {}, sqlQuery:{}", query, pgSqlQuery, e);
     }
 
     return EMPTY_ITERATOR;
@@ -739,8 +748,7 @@ private Optional<Long> getCreatedTime(Key key) throws IOException {
 
   private CloseableIterator<Document> searchDocsForKeys(Set<Key> keys) {
     List<String> keysAsStr = keys.stream().map(Key::toString).collect(Collectors.toList());
-    Query query =
-        new Query().withSelection("*").withFilter(new Filter(Filter.Op.IN, ID, keysAsStr));
+    Query query = new Query().withFilter(new Filter(Filter.Op.IN, ID, keysAsStr));
     return search(query);
   }
 
@@ -778,6 +786,7 @@ private CloseableIterator<Document> executeQueryV1(
     try {
       PreparedStatement preparedStatement =
           buildPreparedStatement(sqlQuery, queryParser.getParamsBuilder().build());
+      LOGGER.debug("Executing executeQueryV1 sqlQuery:{}", preparedStatement.toString());
       ResultSet resultSet = preparedStatement.executeQuery();
       CloseableIterator closeableIterator =
           query.getSelections().size() > 0
@@ -1075,11 +1084,7 @@ protected Document prepareDocument() throws SQLException, IOException {
       Map<String, Object> jsonNode = new HashMap();
       for (int i = 1; i <= columnCount; i++) {
         String columnName = resultSetMetaData.getColumnName(i);
-        int columnType = resultSetMetaData.getColumnType(i);
-        String columnValue =
-            columnType == Types.ARRAY
-                ? MAPPER.writeValueAsString(resultSet.getArray(i).getArray())
-                : resultSet.getString(i);
+        String columnValue = getColumnValue(resultSetMetaData, columnName, i);
         if (StringUtils.isNotEmpty(columnValue)) {
           JsonNode leafNodeValue = MAPPER.readTree(columnValue);
           if (PostgresUtils.isEncodedNestedField(columnName)) {
@@ -1093,6 +1098,24 @@ protected Document prepareDocument() throws SQLException, IOException {
       return new JSONDocument(MAPPER.writeValueAsString(jsonNode));
     }
 
+    private String getColumnValue(
+        ResultSetMetaData resultSetMetaData, String columnName, int columnIndex)
+        throws SQLException, JsonProcessingException {
+      int columnType = resultSetMetaData.getColumnType(columnIndex);
+      // check for array
+      if (columnType == Types.ARRAY) {
+        return MAPPER.writeValueAsString(resultSet.getArray(columnIndex).getArray());
+      }
+
+      // check for ID column
+      if (columnName.equals(ID)) {
+        return MAPPER.writeValueAsString(resultSet.getString(columnIndex));
+      }
+
+      // rest of the columns
+      return resultSet.getString(columnIndex);
+    }
+
     private void handleNestedField(
         String columnName, Map<String, Object> rootNode, JsonNode leafNodeValue) {
       List<String> keys = PostgresUtils.splitNestedField(columnName);