Add API to count workflows (#2518)

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/client/WorkflowClient.java

@@ -240,6 +240,15 @@ WorkflowStub newUntypedWorkflowStub(
    */
   Stream<WorkflowExecutionMetadata> listExecutions(@Nullable String query);
 
+  /**
+   * Count workflow executions using the Visibility API.
+   *
+   * @param query Temporal Visibility query, for syntax see <a
+   *     href="https://docs.temporal.io/visibility#list-filter">Visibility docs</a>
+   * @return count result object
+   */
+  WorkflowExecutionCount countWorkflows(@Nullable String query);
+
   /**
    * Streams history events for a workflow execution for the provided {@code workflowId}.
    *

temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java

@@ -241,6 +241,13 @@ public Stream<WorkflowExecutionMetadata> listExecutions(@Nullable String query)
     return listExecutions(query, null);
   }
 
+  @Override
+  public WorkflowExecutionCount countWorkflows(@Nullable String query) {
+    WorkflowClientCallsInterceptor.CountWorkflowsInput input =
+        new WorkflowClientCallsInterceptor.CountWorkflowsInput(query);
+    return workflowClientCallsInvoker.countWorkflows(input).getCount();
+  }
+
   Stream<WorkflowExecutionMetadata> listExecutions(
       @Nullable String query, @Nullable Integer pageSize) {
     ListWorkflowExecutionIterator iterator =

temporal-sdk/src/main/java/io/temporal/client/WorkflowExecutionCount.java

@@ -0,0 +1,53 @@
+package io.temporal.client;
+
+import io.temporal.api.workflowservice.v1.CountWorkflowExecutionsResponse;
+import io.temporal.internal.common.SearchAttributesUtil;
+import java.util.List;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+
+/** Result of counting workflow executions. */
+public class WorkflowExecutionCount {
+  /** Individual aggregation group record. */
+  public static class AggregationGroup {
+    private final List<List<?>> groupValues;
+    private final long count;
+
+    AggregationGroup(List<io.temporal.api.common.v1.Payload> groupValues, long count) {
+      this.groupValues =
+          groupValues.stream().map(SearchAttributesUtil::decode).collect(Collectors.toList());
+      this.count = count;
+    }
+
+    /** Values of the group. */
+    public List<List<?>> getGroupValues() {
+      return groupValues;
+    }
+
+    /** Count of workflows in the group. */
+    public long getCount() {
+      return count;
+    }
+  }
+
+  private final long count;
+  private final List<AggregationGroup> groups;
+
+  public WorkflowExecutionCount(@Nonnull CountWorkflowExecutionsResponse response) {
+    this.count = response.getCount();
+    this.groups =
+        response.getGroupsList().stream()
+            .map(g -> new AggregationGroup(g.getGroupValuesList(), g.getCount()))
+            .collect(Collectors.toList());
+  }
+
+  /** Total number of workflows matching the request. */
+  public long getCount() {
+    return count;
+  }
+
+  /** Aggregation groups returned by the service. */
+  public List<AggregationGroup> getGroups() {
+    return groups;
+  }
+}

temporal-sdk/src/main/java/io/temporal/common/interceptors/WorkflowClientCallsInterceptor.java

@@ -75,6 +75,8 @@ public interface WorkflowClientCallsInterceptor {
 
   DescribeWorkflowOutput describe(DescribeWorkflowInput input);
 
+  CountWorkflowOutput countWorkflows(CountWorkflowsInput input);
+
   final class WorkflowStartInput {
     private final String workflowId;
     private final String workflowType;
@@ -602,4 +604,29 @@ public WorkflowExecutionDescription getDescription() {
       return description;
     }
   }
+
+  final class CountWorkflowsInput {
+    private final String query;
+
+    public CountWorkflowsInput(@Nullable String query) {
+      this.query = query;
+    }
+
+    @Nullable
+    public String getQuery() {
+      return query;
+    }
+  }
+
+  final class CountWorkflowOutput {
+    private final WorkflowExecutionCount count;
+
+    public CountWorkflowOutput(WorkflowExecutionCount count) {
+      this.count = count;
+    }
+
+    public WorkflowExecutionCount getCount() {
+      return count;
+    }
+  }
 }

temporal-sdk/src/main/java/io/temporal/common/interceptors/WorkflowClientCallsInterceptorBase.java

@@ -72,4 +72,9 @@ public TerminateOutput terminate(TerminateInput input) {
   public DescribeWorkflowOutput describe(DescribeWorkflowInput input) {
     return next.describe(input);
   }
+
+  @Override
+  public CountWorkflowOutput countWorkflows(CountWorkflowsInput input) {
+    return next.countWorkflows(input);
+  }
 }

temporal-sdk/src/main/java/io/temporal/internal/client/RootWorkflowClientInvoker.java

@@ -676,6 +676,17 @@ public DescribeWorkflowOutput describe(DescribeWorkflowInput input) {
         new WorkflowExecutionDescription(response, dataConverterWithWorkflowContext));
   }
 
+  @Override
+  public CountWorkflowOutput countWorkflows(CountWorkflowsInput input) {
+    CountWorkflowExecutionsRequest.Builder req =
+        CountWorkflowExecutionsRequest.newBuilder().setNamespace(clientOptions.getNamespace());
+    if (input.getQuery() != null) {
+      req.setQuery(input.getQuery());
+    }
+    CountWorkflowExecutionsResponse resp = genericClient.countWorkflowExecutions(req.build());
+    return new CountWorkflowOutput(new WorkflowExecutionCount(resp));
+  }
+
   private static <R> R convertResultPayloads(
       Optional<Payloads> resultValue,
       Class<R> resultClass,

temporal-sdk/src/main/java/io/temporal/internal/client/external/GenericWorkflowClient.java

@@ -44,6 +44,8 @@ CompletableFuture<GetWorkflowExecutionHistoryResponse> getWorkflowExecutionHisto
   CompletableFuture<ListWorkflowExecutionsResponse> listWorkflowExecutionsAsync(
       ListWorkflowExecutionsRequest listRequest);
 
+  CountWorkflowExecutionsResponse countWorkflowExecutions(CountWorkflowExecutionsRequest request);
+
   CreateScheduleResponse createSchedule(CreateScheduleRequest request);
 
   CompletableFuture<ListSchedulesResponse> listSchedulesAsync(ListSchedulesRequest request);

temporal-sdk/src/main/java/io/temporal/internal/client/external/GenericWorkflowClientImpl.java

@@ -215,6 +215,18 @@ public CompletableFuture<ListWorkflowExecutionsResponse> listWorkflowExecutionsA
         grpcRetryerOptions);
   }
 
+  @Override
+  public CountWorkflowExecutionsResponse countWorkflowExecutions(
+      CountWorkflowExecutionsRequest request) {
+    return grpcRetryer.retryWithResult(
+        () ->
+            service
+                .blockingStub()
+                .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope)
+                .countWorkflowExecutions(request),
+        grpcRetryerOptions);
+  }
+
   @Override
   public CreateScheduleResponse createSchedule(CreateScheduleRequest request) {
     return grpcRetryer.retryWithResult(

temporal-sdk/src/main/java/io/temporal/internal/common/SearchAttributesUtil.java

@@ -120,6 +120,19 @@ public static <T> List<T> decode(
     return (List<T>) data;
   }
 
+  @Nullable
+  public static List<?> decode(Payload payload) {
+    List<?> data = converter.decode(payload);
+    if (data.size() == 0) {
+      // User code should observe the empty collection as non-existent search attribute, because
+      // it's effectively the same.
+      // We use an empty collection for "unset". See:
+      // https://github.com/temporalio/temporal/issues/561
+      return null;
+    }
+    return data;
+  }
+
   @SuppressWarnings("unchecked")
   @Nullable
   public static <T> List<T> decodeAsType(

temporal-sdk/src/test/java/io/temporal/client/CountWorkflowsTest.java

@@ -0,0 +1,42 @@
+package io.temporal.client;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assume.assumeTrue;
+
+import io.temporal.testing.internal.SDKTestWorkflowRule;
+import io.temporal.workflow.shared.TestWorkflows;
+import org.junit.Rule;
+import org.junit.Test;
+
+public class CountWorkflowsTest {
+  @Rule
+  public SDKTestWorkflowRule testWorkflowRule =
+      SDKTestWorkflowRule.newBuilder()
+          .setWorkflowTypes(TestWorkflows.DoNothingNoArgsWorkflow.class)
+          .build();
+
+  @Test
+  public void countWorkflowExecutions_returnsAllExecutions() throws InterruptedException {
+    assumeTrue(
+        "Test Server doesn't support countWorkflowExecutions endpoint yet",
+        SDKTestWorkflowRule.useExternalService);
+
+    final int EXECUTIONS_COUNT = 5;
+
+    for (int i = 0; i < EXECUTIONS_COUNT; i++) {
+      WorkflowStub.fromTyped(testWorkflowRule.newWorkflowStub(TestWorkflows.NoArgsWorkflow.class))
+          .start();
+    }
+
+    // Visibility API may be eventual consistent
+    Thread.sleep(4_000);
+
+    String queryString =
+        "TaskQueue='" + testWorkflowRule.getTaskQueue() + "' GROUP BY ExecutionStatus";
+    WorkflowExecutionCount count = testWorkflowRule.getWorkflowClient().countWorkflows(queryString);
+    assertEquals(EXECUTIONS_COUNT, count.getCount());
+    assertEquals(1, count.getGroups().size());
+    assertEquals(5, count.getGroups().get(0).getCount());
+    assertEquals("Completed", count.getGroups().get(0).getGroupValues().get(0).get(0));
+  }
+}