Allow setting headers and waitTime seconds (#213)

indeedeng · Jan 3, 2024 · 08a5d25 · 08a5d25
1 parent db7f7c7
commit 08a5d25
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git a/src/main/java/io/iworkflow/core/Client.java b/src/main/java/io/iworkflow/core/Client.java
@@ -48,6 +48,12 @@ public Client(final Registry registry, final ClientOptions clientOptions) {
         this.unregisteredClient = new UnregisteredClient(clientOptions);
     }
 
+    public Client(final Registry registry, final ClientOptions clientOptions, final UnregisteredClient unregisteredClient) {
+        this.registry = registry;
+        this.clientOptions = clientOptions;
+        this.unregisteredClient = unregisteredClient;
+    }
+
     public UnregisteredClient getUnregisteredClient() {
         return unregisteredClient;
     }

diff --git a/src/main/java/io/iworkflow/core/ClientOptions.java b/src/main/java/io/iworkflow/core/ClientOptions.java
@@ -2,6 +2,9 @@
 
 import org.immutables.value.Value;
 
+import java.util.Map;
+import java.util.Optional;
+
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 @Value.Immutable
@@ -12,6 +15,10 @@ public abstract class ClientOptions {
 
     public abstract ObjectEncoder getObjectEncoder();
 
+    public abstract Optional<Integer> getLongPollApiMaxWaitTimeSeconds();
+
+    public abstract Map<String,String> getRequestHeaders();
+
     @Value.Default
     public ServiceApiRetryConfig getServiceApiRetryConfig() {
         return ImmutableServiceApiRetryConfig.builder()

diff --git a/src/main/java/io/iworkflow/core/UnregisteredClient.java b/src/main/java/io/iworkflow/core/UnregisteredClient.java
@@ -60,9 +60,17 @@ public UnregisteredClient(final ClientOptions clientOptions) {
         );
         feignBuilder.errorDecoder(new ServerErrorRetryDecoder());
         apiClient.setFeignBuilder(feignBuilder);
+        feignBuilder.requestInterceptor(
+                (requestTemplate) -> clientOptions.getRequestHeaders().forEach(requestTemplate::header)
+        );
 
         this.defaultApi = apiClient.buildClient(DefaultApi.class);
     }
+
+    public UnregisteredClient(final ClientOptions clientOptions, final DefaultApi openApiClient) {
+        this.clientOptions = clientOptions;
+        this.defaultApi = openApiClient;
+    }
 
     /**
      * @param workflowType           required
@@ -323,6 +331,10 @@ private List<StateCompletionOutput> getWorkflowResults(
                 .workflowId(workflowId)
                 .workflowRunId(workflowRunId);
 
+        if(withWait && clientOptions.getLongPollApiMaxWaitTimeSeconds().isPresent()) {
+            request.waitTimeSeconds(clientOptions.getLongPollApiMaxWaitTimeSeconds().get());
+        }
+
         final WorkflowGetResponse workflowGetResponse;
         try {
             if (withWait) {
@@ -352,6 +364,11 @@ public <T> T waitForStateExecutionCompletion(
         final WorkflowWaitForStateCompletionRequest request = new WorkflowWaitForStateCompletionRequest()
                 .stateExecutionId(stateExecutionId)
                 .workflowId(workflowId);
+
+        if(clientOptions.getLongPollApiMaxWaitTimeSeconds().isPresent()) {
+            request.waitTimeSeconds(clientOptions.getLongPollApiMaxWaitTimeSeconds().get());
+        }
+
         final WorkflowWaitForStateCompletionResponse response;
         try {
             response = defaultApi.apiV1WorkflowWaitForStateCompletionPost(request);