feat(Analytics) Allow dateRangeField to be configurable for timeSerie…

…s chart (datahub-project#11366)
daha · Sep 20, 2024 · c5e5925 · c5e5925
1 parent f6dde2b
commit c5e5925
Showing 1 changed file with 44 additions and 6 deletions.
diff --git a/...l-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsService.java b/...l-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsService.java
@@ -78,7 +78,8 @@ public List<NamedLine> getTimeseriesChart(
       Optional<String> dimension, // Length 1 for now
       Map<String, List<String>> filters,
       Map<String, List<String>> mustNotFilters,
-      Optional<String> uniqueOn) {
+      Optional<String> uniqueOn,
+      String dateRangeField) {
 
     log.debug(
         String.format(
@@ -87,11 +88,11 @@ public List<NamedLine> getTimeseriesChart(
             + String.format("filters: %s, uniqueOn: %s", filters, uniqueOn));
 
     AggregationBuilder filteredAgg =
-        getFilteredAggregation(filters, mustNotFilters, Optional.of(dateRange));
+        getFilteredAggregation(filters, mustNotFilters, Optional.of(dateRange), dateRangeField);
 
     AggregationBuilder dateHistogram =
         AggregationBuilders.dateHistogram(DATE_HISTOGRAM)
-            .field("timestamp")
+            .field(dateRangeField)
             .calendarInterval(new DateHistogramInterval(granularity.name().toLowerCase()));
     uniqueOn.ifPresent(s -> dateHistogram.subAggregation(getUniqueQuery(s)));
 
@@ -128,6 +129,25 @@ public List<NamedLine> getTimeseriesChart(
     }
   }
 
+  public List<NamedLine> getTimeseriesChart(
+      String indexName,
+      DateRange dateRange,
+      DateInterval granularity,
+      Optional<String> dimension, // Length 1 for now
+      Map<String, List<String>> filters,
+      Map<String, List<String>> mustNotFilters,
+      Optional<String> uniqueOn) {
+    return getTimeseriesChart(
+        indexName,
+        dateRange,
+        granularity,
+        dimension,
+        filters,
+        mustNotFilters,
+        uniqueOn,
+        "timestamp");
+  }
+
   private int extractCount(MultiBucketsAggregation.Bucket bucket, boolean didUnique) {
     return didUnique
         ? (int) bucket.getAggregations().<Cardinality>get(UNIQUE).getValue()
@@ -323,20 +343,38 @@ private Filter executeAndExtract(SearchRequest searchRequest) {
     }
   }
 
+  // Make dateRangeField as customizable
   private AggregationBuilder getFilteredAggregation(
       Map<String, List<String>> mustFilters,
       Map<String, List<String>> mustNotFilters,
-      Optional<DateRange> dateRange) {
+      Optional<DateRange> dateRange,
+      String dateRangeField) {
     BoolQueryBuilder filteredQuery = QueryBuilders.boolQuery();
     mustFilters.forEach((key, values) -> filteredQuery.must(QueryBuilders.termsQuery(key, values)));
     mustNotFilters.forEach(
         (key, values) -> filteredQuery.mustNot(QueryBuilders.termsQuery(key, values)));
-    dateRange.ifPresent(range -> filteredQuery.must(dateRangeQuery(range)));
+    dateRange.ifPresent(range -> filteredQuery.must(dateRangeQuery(range, dateRangeField)));
     return AggregationBuilders.filter(FILTERED, filteredQuery);
   }
 
+  private AggregationBuilder getFilteredAggregation(
+      Map<String, List<String>> mustFilters,
+      Map<String, List<String>> mustNotFilters,
+      Optional<DateRange> dateRange) {
+    // Use timestamp as dateRangeField
+    return getFilteredAggregation(mustFilters, mustNotFilters, dateRange, "timestamp");
+  }
+
   private QueryBuilder dateRangeQuery(DateRange dateRange) {
-    return QueryBuilders.rangeQuery("timestamp").gte(dateRange.getStart()).lt(dateRange.getEnd());
+    // Use timestamp as dateRangeField
+    return dateRangeQuery(dateRange, "timestamp");
+  }
+
+  // Make dateRangeField as customizable
+  private QueryBuilder dateRangeQuery(DateRange dateRange, String dateRangeField) {
+    return QueryBuilders.rangeQuery(dateRangeField)
+        .gte(dateRange.getStart())
+        .lt(dateRange.getEnd());
   }
 
   private AggregationBuilder getUniqueQuery(String uniqueOn) {