MINOR - Improve Data Quality Results List from Search

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityTimeSeriesRepository.java

@@ -2,6 +2,7 @@
 
 import static org.openmetadata.schema.type.EventType.ENTITY_UPDATED;
 import static org.openmetadata.schema.type.Include.ALL;
+import static org.openmetadata.service.Entity.getEntityFields;
 
 import java.beans.IntrospectionException;
 import java.io.IOException;
@@ -11,9 +12,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 import javax.json.Json;
-import javax.json.JsonArray;
 import javax.json.JsonObject;
 import javax.json.JsonObjectBuilder;
 import javax.json.JsonPatch;
@@ -46,6 +47,7 @@ public abstract class EntityTimeSeriesRepository<T extends EntityTimeSeriesInter
   protected final String entityType;
   protected final Class<T> entityClass;
   protected final CollectionDAO daoCollection;
+  protected final Set<String> allowedFields;
 
   public EntityTimeSeriesRepository(
       String collectionPath,
@@ -58,34 +60,47 @@ public EntityTimeSeriesRepository(
     this.entityType = entityType;
     this.searchRepository = Entity.getSearchRepository();
     this.daoCollection = Entity.getCollectionDAO();
+    this.allowedFields = getEntityFields(entityClass);
     Entity.registerEntity(entityClass, entityType, this);
   }
 
   @Transaction
-  public T createNewRecord(T recordEntity, String extension, String recordFQN) {
-    recordEntity.setId(UUID.randomUUID());
-    storeInternal(recordEntity, recordFQN, extension);
-    storeRelationshipInternal(recordEntity);
-    postCreate(recordEntity);
-    return recordEntity;
+  public T createNewRecord(T recordEntity, String recordFQN) {
+    return createNewRecord(recordEntity, recordFQN, null);
   }
 
-  public T createNewRecord(T recordEntity, String recordFQN) {
+  @Transaction
+  public T createNewRecord(T recordEntity, String recordFQN, String extension) {
     recordEntity.setId(UUID.randomUUID());
-    storeInternal(recordEntity, recordFQN);
+    if (extension != null) {
+      storeInternal(recordEntity, recordFQN, extension);
+    } else {
+      storeInternal(recordEntity, recordFQN);
+    }
     storeRelationshipInternal(recordEntity);
     postCreate(recordEntity);
     return recordEntity;
   }
 
   @Transaction
   protected void storeInternal(T recordEntity, String recordFQN) {
-    timeSeriesDao.insert(recordFQN, entityType, JsonUtils.pojoToJson(recordEntity));
+    storeInternal(recordEntity, recordFQN, null);
   }
 
   @Transaction
   protected void storeInternal(T recordEntity, String recordFQN, String extension) {
-    timeSeriesDao.insert(recordFQN, extension, entityType, JsonUtils.pojoToJson(recordEntity));
+    if (extension != null) {
+      timeSeriesDao.insert(recordFQN, extension, entityType, JsonUtils.pojoToJson(recordEntity));
+    } else {
+      timeSeriesDao.insert(recordFQN, entityType, JsonUtils.pojoToJson(recordEntity));
+    }
+  }
+
+  public final EntityUtil.Fields getFields(String fields) {
+    if ("*".equals(fields)) {
+      return new EntityUtil.Fields(allowedFields, String.join(",", allowedFields));
+    }
+    return new EntityUtil.Fields(allowedFields, fields);
   }
 
   protected void storeRelationshipInternal(T recordEntity) {
@@ -124,7 +139,6 @@ protected void setUpdatedFields(T updated, String user) {
   protected void validatePatchFields(T updated, T original) {
     // Nothing to do in the default implementation
   }
-  ;
 
   @Transaction
   public final void addRelationship(
@@ -214,37 +228,36 @@ public ResultList<T> listWithOffset(
       boolean latest,
       boolean skipErrors) {
     int total = timeSeriesDao.listCount(filter, startTs, endTs, latest);
-    List<T> entityList = new ArrayList<>();
-    List<EntityError> errors = null;
-    int offsetInt = getOffset(offset);
-    String afterOffset = getAfterOffset(offsetInt, limitParam, total);
-    String beforeOffset = getBeforeOffset(offsetInt, limitParam);
-
-    if (limitParam > 0) {
-      List<String> jsons =
-          timeSeriesDao.listWithOffset(filter, limitParam, offsetInt, startTs, endTs, latest);
-      Map<String, List<?>> entityListMap = getEntityList(jsons, skipErrors);
-      entityList = (List<T>) entityListMap.get("entityList");
-      if (skipErrors) {
-        errors = (List<EntityError>) entityListMap.get("errors");
-      }
-      return getResultList(entityList, beforeOffset, afterOffset, total, errors);
-    } else {
-      return getResultList(entityList, null, null, total);
-    }
+    return ListWithOffsetInternal(
+        offset, filter, limitParam, startTs, endTs, latest, skipErrors, total);
   }
 
   public ResultList<T> listWithOffset(
       String offset, ListFilter filter, int limitParam, boolean skipErrors) {
     int total = timeSeriesDao.listCount(filter);
+    return ListWithOffsetInternal(offset, filter, limitParam, null, null, false, skipErrors, total);
+  }
+
+  private ResultList<T> ListWithOffsetInternal(
+      String offset,
+      ListFilter filter,
+      int limitParam,
+      Long startTs,
+      Long endTs,
+      boolean latest,
+      boolean skipErrors,
+      int total) {
     List<T> entityList = new ArrayList<>();
     List<EntityError> errors = null;
 
     int offsetInt = getOffset(offset);
     String afterOffset = getAfterOffset(offsetInt, limitParam, total);
     String beforeOffset = getBeforeOffset(offsetInt, limitParam);
     if (limitParam > 0) {
-      List<String> jsons = timeSeriesDao.listWithOffset(filter, limitParam, offsetInt);
+      List<String> jsons =
+          (startTs != null && endTs != null)
+              ? timeSeriesDao.listWithOffset(filter, limitParam, offsetInt, startTs, endTs, latest)
+              : timeSeriesDao.listWithOffset(filter, limitParam, offsetInt);
       Map<String, List<?>> entityListMap = getEntityList(jsons, skipErrors);
       entityList = (List<T>) entityListMap.get("entityList");
       if (skipErrors) {
@@ -397,21 +410,40 @@ public ResultList<T> listLatestFromSearch(
     List<T> entityList = new ArrayList<>();
     setIncludeSearchFields(searchListFilter);
     setExcludeSearchFields(searchListFilter);
+    String aggregationPath = "$.sterms#byTerms.buckets";
     String aggregationStr =
-        "{\"aggregations\": {\"byTerms\": {\"terms\": {\"field\": \"%s\", \"size\":100},\"aggs\": {\"latest\": "
-            + "{\"top_hits\": {\"size\": 1, \"sort_field\":\"timestamp\",\"sort_order\":\"desc\"}}}}}}";
+        "{\"aggregations\":{\"byTerms\":{\"terms\": {\"field\":\"%s\",\"size\":100},\"aggs\":{\"latest\":"
+            + "{\"top_hits\":{\"size\":1,\"sort_field\":\"timestamp\",\"sort_order\":\"desc\"}}}}}}";
     aggregationStr = String.format(aggregationStr, groupBy);
     JsonObject aggregation = JsonUtils.readJson(aggregationStr).asJsonObject();
     JsonObject jsonObjResults =
         searchRepository.aggregate(q, entityType, aggregation, searchListFilter);
-    List<JsonObject> jsonTestCaseResults = parseListLatestAggregation(jsonObjResults);
 
-    for (JsonObject json : jsonTestCaseResults) {
-      T entity = setFieldsInternal(JsonUtils.readOrConvertValue(json, entityClass), fields);
-      setInheritedFields(entity);
-      clearFieldsInternal(entity, fields);
-      entityList.add(entity);
-    }
+    Optional<List> jsonObjects =
+        JsonUtils.readJsonAtPath(jsonObjResults.toString(), aggregationPath, List.class);
+    jsonObjects.ifPresent(
+        jsonObjectList -> {
+          for (Map<String, String> json : (List<Map<String, String>>) jsonObjectList) {
+            String bucketAggregationPath = "top_hits#latest.hits.hits";
+            Optional<List> hits =
+                JsonUtils.readJsonAtPath(
+                    JsonUtils.pojoToJson(json), bucketAggregationPath, List.class);
+            hits.ifPresent(
+                hitList -> {
+                  for (Map<String, String> hit : (List<Map<String, String>>) hitList) {
+                    JsonObject source = getSourceDocument(JsonUtils.pojoToJson(hit));
+                    T entity =
+                        setFieldsInternal(
+                            JsonUtils.readOrConvertValue(source, entityClass), fields);
+                    if (entity != null) {
+                      setInheritedFields(entity);
+                      clearFieldsInternal(entity, fields);
+                      entityList.add(entity);
+                    }
+                  }
+                });
+          }
+        });
     return new ResultList<>(entityList, null, null, entityList.size());
   }
 
@@ -447,48 +479,27 @@ protected List<String> getExcludeSearchFields() {
     return new ArrayList<>();
   }
 
-  private List<JsonObject> parseListLatestAggregation(JsonObject jsonObjResults) {
-    JsonObject jsonByTerms = jsonObjResults.getJsonObject("sterms#byTerms");
-    List<JsonObject> jsonTestCaseResults = new ArrayList<>();
+  private JsonObject getSourceDocument(String hit) {
     List<String> includeSearchFields = getIncludeSearchFields();
     List<String> excludeSearchFields = getExcludeSearchFields();
-    Optional.ofNullable(jsonByTerms)
-        .map(jbt -> jbt.getJsonArray("buckets"))
-        .ifPresent(
-            termsBucket -> {
-              for (JsonValue bucket : termsBucket) {
-                JsonObject hitsBucket = bucket.asJsonObject().getJsonObject("top_hits#latest");
-                if (hitsBucket != null) {
-                  JsonObject hitsTwo = hitsBucket.getJsonObject("hits");
-                  if (hitsTwo != null) {
-                    JsonArray hits = hitsTwo.getJsonArray("hits");
-                    if (hits != null) {
-                      for (JsonValue hit : hits) {
-                        JsonObject source = hit.asJsonObject().getJsonObject("_source");
-                        // Aggregation results will return all fields by default, so we need to
-                        // filter out the fields
-                        // that are not included in the search fields
-                        if (source != null
-                            && (!CommonUtil.nullOrEmpty(includeSearchFields)
-                                || !CommonUtil.nullOrEmpty(excludeSearchFields))) {
-                          JsonObjectBuilder sourceCopy = Json.createObjectBuilder();
-                          for (Map.Entry<String, JsonValue> entry : source.entrySet()) {
-                            if (includeSearchFields.contains(entry.getKey())
-                                || (CommonUtil.nullOrEmpty(includeSearchFields)
-                                    && !excludeSearchFields.contains(entry.getKey()))) {
-                              sourceCopy.add(entry.getKey(), entry.getValue());
-                            }
-                          }
-                          jsonTestCaseResults.add(sourceCopy.build());
-                        } else {
-                          if (source != null) jsonTestCaseResults.add(source);
-                        }
-                      }
-                    }
-                  }
-                }
-              }
-            });
-    return jsonTestCaseResults;
+    JsonObject hitJson = JsonUtils.readJson(hit).asJsonObject();
+    JsonObject source = hitJson.asJsonObject().getJsonObject("_source");
+    // Aggregation results will return all fields by default,
+    // so we need to filter out the fields that are not included
+    // in the search fields
+    if (source != null
+        && (!CommonUtil.nullOrEmpty(includeSearchFields)
+            || !CommonUtil.nullOrEmpty(excludeSearchFields))) {
+      JsonObjectBuilder sourceCopy = Json.createObjectBuilder();
+      for (Map.Entry<String, JsonValue> entry : source.entrySet()) {
+        if (includeSearchFields.contains(entry.getKey())
+            || (CommonUtil.nullOrEmpty(includeSearchFields)
+                && !excludeSearchFields.contains(entry.getKey()))) {
+          sourceCopy.add(entry.getKey(), entry.getValue());
+        }
+      }
+      return sourceCopy.build();
+    }
+    return source;
   }
 }

openmetadata-service/src/main/java/org/openmetadata/service/resources/dqtests/TestCaseResultResource.java

@@ -222,10 +222,41 @@ public ResultList<TestCaseResult> listTestCaseResultsFromSearch(
       @Parameter(
               description =
                   "Get the latest test case result for each test case -- requires `testSuiteId`. Offset and limit are ignored",
-              schema = @Schema(type = "boolean", example = "true|false"))
+              schema =
+                  @Schema(
+                      type = "boolean",
+                      example = "false",
+                      allowableValues = {"true", "false"}))
           @QueryParam("latest")
           @DefaultValue("false")
           String latest,
+      @Parameter(
+              description = "Filter for test case result by type (e.g. column, table, all)",
+              schema =
+                  @Schema(
+                      type = "string",
+                      example = "all",
+                      allowableValues = {"column", "table", "all"}))
+          @QueryParam("testCaseType")
+          @DefaultValue("all")
+          String type,
+      @Parameter(
+              description =
+                  "Filter for test case by data quality dimension (e.g. OpenMetadata, dbt, etc.)",
+              schema =
+                  @Schema(
+                      type = "string",
+                      allowableValues = {
+                        "Completeness",
+                        "Accuracy",
+                        "Consistency",
+                        "Validity",
+                        "Uniqueness",
+                        "Integrity",
+                        "SQL"
+                      }))
+          @QueryParam("dataQualityDimension")
+          String dataQualityDimension,
       @Parameter(
               description = "search query term to use in list",
               schema = @Schema(type = "string"))
@@ -235,7 +266,7 @@ public ResultList<TestCaseResult> listTestCaseResultsFromSearch(
     if (latest.equals("true") && testSuiteId == null) {
       throw new IllegalArgumentException("latest=true requires testSuiteId");
     }
-    EntityUtil.Fields fields = new EntityUtil.Fields(Set.of(""), fieldParams);
+    EntityUtil.Fields fields = repository.getFields(fieldParams);
     SearchListFilter searchListFilter = new SearchListFilter();
     Optional.ofNullable(startTimestamp)
         .ifPresent(ts -> searchListFilter.addQueryParam("startTimestamp", ts.toString()));
@@ -247,6 +278,9 @@ public ResultList<TestCaseResult> listTestCaseResultsFromSearch(
         .ifPresent(tcf -> searchListFilter.addQueryParam("testCaseFQN", tcf));
     Optional.ofNullable(testSuiteId)
         .ifPresent(tsi -> searchListFilter.addQueryParam("testSuiteId", tsi));
+    Optional.ofNullable(type).ifPresent(t -> searchListFilter.addQueryParam("testCaseType", t));
+    Optional.ofNullable(dataQualityDimension)
+        .ifPresent(dqd -> searchListFilter.addQueryParam("dataQualityDimension", dqd));
 
     ResourceContextInterface resourceContextInterface = getResourceContext(testCaseFQN);
     // Override OperationContext to change the entity to table
@@ -259,7 +293,7 @@ public ResultList<TestCaseResult> listTestCaseResultsFromSearch(
           securityContext,
           fields,
           searchListFilter,
-          "testSuites.id",
+          "testCaseFQN.keyword",
           q,
           operationContext,
           resourceContextInterface);