Support filtering query based on close status (#394)

Author: andrewjdawson2016

src/main/java/com/uber/cadence/client/WorkflowStub.java

@@ -17,7 +17,9 @@
 
 package com.uber.cadence.client;
 
+import com.uber.cadence.QueryRejectCondition;
 import com.uber.cadence.WorkflowExecution;
+import com.uber.cadence.internal.common.QueryResponse;
 import java.lang.reflect.Type;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
@@ -124,6 +126,19 @@ <R> CompletableFuture<R> getResultAsync(
 
   <R> R query(String queryType, Class<R> resultClass, Type resultType, Object... args);
 
+  <R> QueryResponse<R> query(
+      String queryType,
+      Class<R> resultClass,
+      QueryRejectCondition queryRejectCondition,
+      Object... args);
+
+  <R> QueryResponse<R> query(
+      String queryType,
+      Class<R> resultClass,
+      Type resultType,
+      QueryRejectCondition queryRejectCondition,
+      Object... args);
+
   /** Request cancellation. */
   void cancel();
 

src/main/java/com/uber/cadence/internal/common/QueryResponse.java

@@ -0,0 +1,48 @@
+/*
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Modifications copyright (C) 2017 Uber Technologies, Inc.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+
+package com.uber.cadence.internal.common;
+
+import com.uber.cadence.QueryRejected;
+
+public final class QueryResponse<T> {
+  private final QueryRejected queryRejected;
+  private final T result;
+
+  public QueryResponse(QueryRejected queryRejected, T result) {
+    this.queryRejected = queryRejected;
+    this.result = result;
+  }
+
+  /**
+   * Returns the value of the QueryRejected property for this object.
+   *
+   * @return The value of the QueryRejected property for this object.
+   */
+  public QueryRejected getQueryRejected() {
+    return queryRejected;
+  }
+
+  /**
+   * Returns the value of the Result property for this object.
+   *
+   * @return The value of the Result property for this object.
+   */
+  public T getResult() {
+    return result;
+  }
+}

src/main/java/com/uber/cadence/internal/external/GenericWorkflowClientExternal.java

@@ -17,6 +17,7 @@
 
 package com.uber.cadence.internal.external;
 
+import com.uber.cadence.QueryWorkflowResponse;
 import com.uber.cadence.WorkflowExecution;
 import com.uber.cadence.WorkflowExecutionAlreadyStartedError;
 import com.uber.cadence.internal.common.SignalWithStartWorkflowExecutionParameters;
@@ -38,7 +39,7 @@ WorkflowExecution signalWithStartWorkflowExecution(
 
   void requestCancelWorkflowExecution(WorkflowExecution execution);
 
-  byte[] queryWorkflow(QueryWorkflowParameters queryParameters);
+  QueryWorkflowResponse queryWorkflow(QueryWorkflowParameters queryParameters);
 
   void terminateWorkflowExecution(TerminateWorkflowExecutionParameters terminateParameters);
 

src/main/java/com/uber/cadence/internal/external/GenericWorkflowClientExternalImpl.java

@@ -268,7 +268,7 @@ public void requestCancelWorkflowExecution(WorkflowExecution execution) {
   }
 
   @Override
-  public byte[] queryWorkflow(QueryWorkflowParameters queryParameters) {
+  public QueryWorkflowResponse queryWorkflow(QueryWorkflowParameters queryParameters) {
     QueryWorkflowRequest request = new QueryWorkflowRequest();
     request.setDomain(domain);
     WorkflowExecution execution = new WorkflowExecution();
@@ -278,12 +278,13 @@ public byte[] queryWorkflow(QueryWorkflowParameters queryParameters) {
     query.setQueryArgs(queryParameters.getInput());
     query.setQueryType(queryParameters.getQueryType());
     request.setQuery(query);
+    request.setQueryRejectCondition(queryParameters.getQueryRejectCondition());
     try {
       QueryWorkflowResponse response =
           Retryer.retryWithResult(
               Retryer.DEFAULT_SERVICE_OPERATION_RETRY_OPTIONS,
               () -> service.QueryWorkflow(request));
-      return response.getQueryResult();
+      return response;
     } catch (TException e) {
       throw CheckedExceptionWrapper.wrap(e);
     }

src/main/java/com/uber/cadence/internal/replay/QueryWorkflowParameters.java

@@ -17,6 +17,7 @@
 
 package com.uber.cadence.internal.replay;
 
+import com.uber.cadence.QueryRejectCondition;
 import java.nio.charset.StandardCharsets;
 
 public class QueryWorkflowParameters implements Cloneable {
@@ -29,6 +30,8 @@ public class QueryWorkflowParameters implements Cloneable {
 
   private String workflowId;
 
+  private QueryRejectCondition queryRejectCondition;
+
   public QueryWorkflowParameters() {}
 
   public byte[] getInput() {
@@ -83,6 +86,20 @@ public QueryWorkflowParameters withWorkflowId(String workflowId) {
     return this;
   }
 
+  public QueryRejectCondition getQueryRejectCondition() {
+    return queryRejectCondition;
+  }
+
+  public void setQueryRejectCondition(QueryRejectCondition queryRejectCondition) {
+    this.queryRejectCondition = queryRejectCondition;
+  }
+
+  public QueryWorkflowParameters withQueryRejectCondition(
+      QueryRejectCondition queryRejectCondition) {
+    this.queryRejectCondition = queryRejectCondition;
+    return this;
+  }
+
   @Override
   public String toString() {
     StringBuilder sb = new StringBuilder();
@@ -91,6 +108,7 @@ public String toString() {
     sb.append("Input: " + new String(input, 0, 512, StandardCharsets.UTF_8) + ", ");
     sb.append("WorkflowId: " + workflowId + ", ");
     sb.append("RunId: " + runId + ", ");
+    sb.append("QueryRejectCondition: " + queryRejectCondition + ", ");
     sb.append("}");
     return sb.toString();
   }
@@ -101,6 +119,7 @@ public QueryWorkflowParameters copy() {
     result.setRunId(runId);
     result.setQueryType(queryType);
     result.setWorkflowId(workflowId);
+    result.setQueryRejectCondition(queryRejectCondition);
     return result;
   }
 }

src/main/java/com/uber/cadence/internal/sync/TestWorkflowEnvironmentInternal.java

@@ -50,6 +50,7 @@
 import com.uber.cadence.PollForDecisionTaskRequest;
 import com.uber.cadence.PollForDecisionTaskResponse;
 import com.uber.cadence.QueryFailedError;
+import com.uber.cadence.QueryRejectCondition;
 import com.uber.cadence.QueryWorkflowRequest;
 import com.uber.cadence.QueryWorkflowResponse;
 import com.uber.cadence.RecordActivityTaskHeartbeatByIDRequest;
@@ -87,6 +88,7 @@
 import com.uber.cadence.client.WorkflowClientOptions;
 import com.uber.cadence.client.WorkflowOptions;
 import com.uber.cadence.client.WorkflowStub;
+import com.uber.cadence.internal.common.QueryResponse;
 import com.uber.cadence.internal.testservice.TestWorkflowService;
 import com.uber.cadence.serviceclient.IWorkflowService;
 import com.uber.cadence.testing.TestEnvironmentOptions;
@@ -901,6 +903,25 @@ public <R> R query(String queryType, Class<R> resultClass, Type resultType, Obje
         return next.query(queryType, resultClass, resultType, args);
       }
 
+      @Override
+      public <R> QueryResponse<R> query(
+          String queryType,
+          Class<R> resultClass,
+          QueryRejectCondition queryRejectCondition,
+          Object... args) {
+        return next.query(queryType, resultClass, queryRejectCondition, args);
+      }
+
+      @Override
+      public <R> QueryResponse<R> query(
+          String queryType,
+          Class<R> resultClass,
+          Type resultType,
+          QueryRejectCondition queryRejectCondition,
+          Object... args) {
+        return next.query(queryType, resultClass, resultType, queryRejectCondition, args);
+      }
+
       @Override
       public void cancel() {
         next.cancel();

src/main/java/com/uber/cadence/internal/sync/WorkflowStubImpl.java

@@ -20,6 +20,8 @@
 import com.uber.cadence.EntityNotExistsError;
 import com.uber.cadence.InternalServiceError;
 import com.uber.cadence.QueryFailedError;
+import com.uber.cadence.QueryRejectCondition;
+import com.uber.cadence.QueryWorkflowResponse;
 import com.uber.cadence.WorkflowExecution;
 import com.uber.cadence.WorkflowExecutionAlreadyStartedError;
 import com.uber.cadence.WorkflowType;
@@ -35,6 +37,7 @@
 import com.uber.cadence.converter.DataConverterException;
 import com.uber.cadence.converter.JsonDataConverter;
 import com.uber.cadence.internal.common.CheckedExceptionWrapper;
+import com.uber.cadence.internal.common.QueryResponse;
 import com.uber.cadence.internal.common.SignalWithStartWorkflowExecutionParameters;
 import com.uber.cadence.internal.common.StartWorkflowExecutionParameters;
 import com.uber.cadence.internal.common.WorkflowExecutionFailedException;
@@ -343,14 +346,40 @@ public <R> R query(String queryType, Class<R> resultClass, Object... args) {
 
   @Override
   public <R> R query(String queryType, Class<R> resultClass, Type resultType, Object... args) {
+    return query(queryType, resultClass, resultType, null, args).getResult();
+  }
+
+  @Override
+  public <R> QueryResponse<R> query(
+      String queryType,
+      Class<R> resultClass,
+      QueryRejectCondition queryRejectCondition,
+      Object... args) {
+    return query(queryType, resultClass, resultClass, queryRejectCondition, args);
+  }
+
+  @Override
+  public <R> QueryResponse<R> query(
+      String queryType,
+      Class<R> resultClass,
+      Type resultType,
+      QueryRejectCondition queryRejectCondition,
+      Object... args) {
     checkStarted();
     QueryWorkflowParameters p = new QueryWorkflowParameters();
     p.setInput(dataConverter.toData(args));
     p.setQueryType(queryType);
     p.setWorkflowId(execution.get().getWorkflowId());
+    p.setQueryRejectCondition(queryRejectCondition);
     try {
-      byte[] result = genericClient.queryWorkflow(p);
-      return dataConverter.fromData(result, resultClass, resultType);
+      QueryWorkflowResponse result = genericClient.queryWorkflow(p);
+      if (result.queryRejected == null) {
+        return new QueryResponse<>(
+            null, dataConverter.fromData(result.getQueryResult(), resultClass, resultType));
+      } else {
+        return new QueryResponse<>(result.getQueryRejected(), null);
+      }
+
     } catch (RuntimeException e) {
       Exception unwrapped = CheckedExceptionWrapper.unwrap(e);
       if (unwrapped instanceof EntityNotExistsError) {

src/main/java/com/uber/cadence/internal/testservice/TestWorkflowMutableStateImpl.java

@@ -1391,6 +1391,20 @@ public void requestCancelWorkflowExecution(RequestCancelWorkflowExecutionRequest
   public QueryWorkflowResponse query(QueryWorkflowRequest queryRequest) throws TException {
     QueryId queryId = new QueryId(executionId);
 
+    Optional<WorkflowExecutionCloseStatus> optCloseStatus = getCloseStatus();
+    if (optCloseStatus.isPresent() && queryRequest.getQueryRejectCondition() != null) {
+      WorkflowExecutionCloseStatus closeStatus = optCloseStatus.get();
+      boolean rejectNotOpen =
+          queryRequest.getQueryRejectCondition() == QueryRejectCondition.NOT_OPEN;
+      boolean rejectNotCompletedCleanly =
+          queryRequest.getQueryRejectCondition() == QueryRejectCondition.NOT_COMPLETED_CLEANLY
+              && closeStatus != WorkflowExecutionCloseStatus.COMPLETED;
+      if (rejectNotOpen || rejectNotCompletedCleanly) {
+        return new QueryWorkflowResponse()
+            .setQueryRejected(new QueryRejected().setCloseStatus(closeStatus));
+      }
+    }
+
     PollForDecisionTaskResponse task =
         new PollForDecisionTaskResponse()
             .setTaskToken(queryId.toBytes())

src/main/thrift/shared.thrift

@@ -288,6 +288,13 @@ enum EncodingType {
   ThriftRW,
 }
 
+enum QueryRejectCondition {
+  // NOT_OPEN indicates that query should be rejected if workflow is not open
+  NOT_OPEN
+  // NOT_COMPLETED_CLEANLY indicates that query should be rejected if workflow did not complete cleanly
+  NOT_COMPLETED_CLEANLY
+}
+
 struct DataBlob {
   10: optional EncodingType EncodingType
   20: optional binary Data
@@ -1305,10 +1312,17 @@ struct QueryWorkflowRequest {
   10: optional string domain
   20: optional WorkflowExecution execution
   30: optional WorkflowQuery query
+  // QueryRejectCondition can used to reject the query if workflow state does not satisify condition
+  40: optional QueryRejectCondition queryRejectCondition
+}
+
+struct QueryRejected {
+  10: optional WorkflowExecutionCloseStatus closeStatus
 }
 
 struct QueryWorkflowResponse {
   10: optional binary queryResult
+  20: optional QueryRejected queryRejected
 }
 
 struct WorkflowQuery {

src/test/java/com/uber/cadence/workflow/WorkflowTest.java

@@ -30,10 +30,12 @@
 import com.uber.cadence.GetWorkflowExecutionHistoryResponse;
 import com.uber.cadence.HistoryEvent;
 import com.uber.cadence.Memo;
+import com.uber.cadence.QueryRejectCondition;
 import com.uber.cadence.SearchAttributes;
 import com.uber.cadence.SignalExternalWorkflowExecutionFailedCause;
 import com.uber.cadence.TimeoutType;
 import com.uber.cadence.WorkflowExecution;
+import com.uber.cadence.WorkflowExecutionCloseStatus;
 import com.uber.cadence.WorkflowIdReusePolicy;
 import com.uber.cadence.activity.Activity;
 import com.uber.cadence.activity.ActivityMethod;
@@ -57,6 +59,7 @@
 import com.uber.cadence.common.MethodRetry;
 import com.uber.cadence.common.RetryOptions;
 import com.uber.cadence.converter.JsonDataConverter;
+import com.uber.cadence.internal.common.QueryResponse;
 import com.uber.cadence.internal.common.WorkflowExecutionUtils;
 import com.uber.cadence.internal.sync.DeterministicRunnerTest;
 import com.uber.cadence.internal.worker.PollerOptions;
@@ -2453,6 +2456,12 @@ public void testSignalUntyped() {
         });
     execution.set(client.start());
     assertEquals("Hello World!", client.getResult(String.class));
+    assertEquals("World!", client.query("QueryableWorkflow::getState", String.class));
+    QueryResponse<String> queryResponse =
+        client.query("QueryableWorkflow::getState", String.class, QueryRejectCondition.NOT_OPEN);
+    assertNull(queryResponse.getResult());
+    assertEquals(
+        WorkflowExecutionCloseStatus.COMPLETED, queryResponse.getQueryRejected().closeStatus);
   }
 
   static final AtomicInteger decisionCount = new AtomicInteger();