feat: add getNumber to AggregationResult (#1851) (#1861)

st4sik · web-flow · commit b9c2c3ff775f · 2025-05-29T11:25:59.000-07:00
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/AggregationResult.java
@@ -84,6 +84,26 @@ public Double getDouble(String alias) {
     }
   }
 
+  /**
+   * Returns a result value for the given alias.
+   *
+   * @param alias A custom alias provided in the query or an autogenerated alias in the form of
+   *     'property_\d'
+   * @return An aggregation result value as a {@link Number} for the given alias.
+   */
+  public Number getNumber(String alias) {
+    Value<?> value = properties.get(alias);
+    switch (value.getType()) {
+      case LONG:
+        return (Long) value.get();
+      case DOUBLE:
+        return (Double) value.get();
+      default:
+        throw new RuntimeException(
+            String.format("Unsupported type %s received for alias '%s'.", value.getType(), alias));
+    }
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/AggregationResultTest.java
@@ -77,5 +77,23 @@ public void shouldThrowRuntimeExceptionOnUnknownTypes() {
     RuntimeException e2 =
         assertThrows(RuntimeException.class, () -> aggregationResult.getDouble("qty_avg"));
     assertThat(e2.getMessage()).isEqualTo("Unsupported type BOOLEAN received for alias 'qty_avg'.");
+
+    RuntimeException e3 =
+        assertThrows(RuntimeException.class, () -> aggregationResult.getNumber("qty_avg"));
+    assertThat(e3.getMessage()).isEqualTo("Unsupported type BOOLEAN received for alias 'qty_avg'.");
+  }
+
+  @Test
+  public void shouldGetDoubleAggregatedResultValueAsNumber() {
+    AggregationResult aggregationResult =
+        new AggregationResult(ImmutableMap.of("qty_avg", DoubleValue.of(45.9322)));
+    assertThat(aggregationResult.getNumber("qty_avg")).isEqualTo(45.9322);
+  }
+
+  @Test
+  public void shouldGetLongAggregatedResultValueAsNumber() {
+    AggregationResult aggregationResult =
+        new AggregationResult(ImmutableMap.of("count", LongValue.of(50)));
+    assertThat(aggregationResult.getNumber("count")).isEqualTo(50);
   }
 }
