From 93b5907f99eb87650c7569b296d1e6f976c3d4bc Mon Sep 17 00:00:00 2001
From: david-leifker <114954101+david-leifker@users.noreply.github.com>
Date: Fri, 22 Mar 2024 21:47:00 -0500
Subject: [PATCH] fix(usage-stats): usage-stats error handling and filter
 (#10105)

---
 .../resolvers/dataset/DatasetUsageStatsResolver.java       | 7 +++++--
 .../graphql/types/usage/UsageQueryResultMapper.java        | 5 +++++
 .../java/com/linkedin/metadata/search/utils/ESUtils.java   | 3 ++-
 .../com/linkedin/metadata/resources/usage/UsageStats.java  | 6 +++---
 4 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetUsageStatsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetUsageStatsResolver.java
index 371911913c..03d08024eb 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetUsageStatsResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/dataset/DatasetUsageStatsResolver.java
@@ -8,6 +8,7 @@ import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.Entity;
 import com.linkedin.datahub.graphql.generated.UsageQueryResult;
 import com.linkedin.datahub.graphql.types.usage.UsageQueryResultMapper;
+import com.linkedin.metadata.utils.metrics.MetricUtils;
 import com.linkedin.usage.UsageClient;
 import com.linkedin.usage.UsageTimeRange;
 import graphql.schema.DataFetcher;
@@ -45,9 +46,11 @@ public class DatasetUsageStatsResolver implements DataFetcher<CompletableFuture<
                 usageClient.getUsageStats(resourceUrn.toString(), range);
             return UsageQueryResultMapper.map(usageQueryResult);
           } catch (Exception e) {
-            throw new RuntimeException(
-                String.format("Failed to load Usage Stats for resource %s", resourceUrn), e);
+            log.error(String.format("Failed to load Usage Stats for resource %s", resourceUrn), e);
+            MetricUtils.counter(this.getClass(), "usage_stats_dropped").inc();
           }
+
+          return UsageQueryResultMapper.EMPTY;
         });
   }
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/usage/UsageQueryResultMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/usage/UsageQueryResultMapper.java
index 444605cd99..b2758adc9b 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/usage/UsageQueryResultMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/usage/UsageQueryResultMapper.java
@@ -1,13 +1,18 @@
 package com.linkedin.datahub.graphql.types.usage;
 
 import com.linkedin.datahub.graphql.generated.UsageQueryResult;
+import com.linkedin.datahub.graphql.generated.UsageQueryResultAggregations;
 import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import java.util.List;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 
 public class UsageQueryResultMapper
     implements ModelMapper<com.linkedin.usage.UsageQueryResult, UsageQueryResult> {
 
+  public static final UsageQueryResult EMPTY =
+      new UsageQueryResult(List.of(), new UsageQueryResultAggregations(0, List.of(), List.of(), 0));
+
   public static final UsageQueryResultMapper INSTANCE = new UsageQueryResultMapper();
 
   public static UsageQueryResult map(
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java
index 6c45072164..3263773ca0 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java
@@ -188,7 +188,7 @@ public class ESUtils {
         .forEach(
             criterion -> {
               if (Set.of(Condition.EXISTS, Condition.IS_NULL).contains(criterion.getCondition())
-                  || !criterion.getValue().trim().isEmpty()
+                  || (criterion.hasValue() && !criterion.getValue().trim().isEmpty())
                   || criterion.hasValues()) {
                 if (!criterion.isNegated()) {
                   // `filter` instead of `must` (enables caching and bypasses scoring)
@@ -646,6 +646,7 @@ public class ESUtils {
    * <p>For all new code, we should be using the new 'values' field for performing multi-match. This
    * is simply retained for backwards compatibility of the search API.
    */
+  @Deprecated
   private static QueryBuilder buildEqualsFromCriterionWithValue(
       @Nonnull final String fieldName,
       @Nonnull final Criterion criterion,
diff --git a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/usage/UsageStats.java b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/usage/UsageStats.java
index 554b6e909e..f80dfa5ce0 100644
--- a/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/usage/UsageStats.java
+++ b/metadata-service/restli-servlet-impl/src/main/java/com/linkedin/metadata/resources/usage/UsageStats.java
@@ -386,14 +386,14 @@ public class UsageStats extends SimpleResourceTemplate<UsageAggregation> {
           Filter filter = new Filter();
           ArrayList<Criterion> criteria = new ArrayList<>();
           Criterion hasUrnCriterion =
-              new Criterion().setField("urn").setCondition(Condition.EQUAL).setValue(resource);
+              new Criterion().setField("urn").setCondition(Condition.EQUAL).setValues(new StringArray(resource));
           criteria.add(hasUrnCriterion);
           if (startTime != null) {
             Criterion startTimeCriterion =
                 new Criterion()
                     .setField(ES_FIELD_TIMESTAMP)
                     .setCondition(Condition.GREATER_THAN_OR_EQUAL_TO)
-                    .setValue(startTime.toString());
+                    .setValues(new StringArray(startTime.toString()));
             criteria.add(startTimeCriterion);
           }
           if (endTime != null) {
@@ -401,7 +401,7 @@ public class UsageStats extends SimpleResourceTemplate<UsageAggregation> {
                 new Criterion()
                     .setField(ES_FIELD_TIMESTAMP)
                     .setCondition(Condition.LESS_THAN_OR_EQUAL_TO)
-                    .setValue(endTime.toString());
+                    .setValues(new StringArray(endTime.toString()));
             criteria.add(endTimeCriterion);
           }