Ad support for updating schedule search attributes (#2168)

Support update SA on a schedule

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/client/schedules/ScheduleUpdate.java

@@ -20,12 +20,33 @@
 
 package io.temporal.client.schedules;
 
+import io.temporal.common.SearchAttributes;
+
 /** An update returned from a schedule updater. */
 public final class ScheduleUpdate {
   private final Schedule schedule;
+  private final SearchAttributes typedSearchAttributes;
 
+  /**
+   * Create a new ScheduleUpdate.
+   *
+   * @param schedule schedule to replace the existing schedule with
+   */
   public ScheduleUpdate(Schedule schedule) {
     this.schedule = schedule;
+    this.typedSearchAttributes = null;
+  }
+
+  /**
+   * Create a new ScheduleUpdate.
+   *
+   * @param schedule schedule to replace the existing schedule with
+   * @param typedSearchAttributes search attributes to replace the existing search attributes with.
+   *     Returning null will not update the search attributes.
+   */
+  public ScheduleUpdate(Schedule schedule, SearchAttributes typedSearchAttributes) {
+    this.schedule = schedule;
+    this.typedSearchAttributes = typedSearchAttributes;
   }
 
   /**
@@ -36,4 +57,13 @@ public ScheduleUpdate(Schedule schedule) {
   public Schedule getSchedule() {
     return schedule;
   }
+
+  /**
+   * Get the search attributes to update.
+   *
+   * @return search attributes to update
+   */
+  public SearchAttributes getTypedSearchAttributes() {
+    return typedSearchAttributes;
+  }
 }

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

@@ -26,6 +26,7 @@
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.temporal.api.common.v1.Memo;
+import io.temporal.api.common.v1.SearchAttributes;
 import io.temporal.api.schedule.v1.*;
 import io.temporal.api.workflowservice.v1.*;
 import io.temporal.client.ListScheduleListDescriptionIterator;
@@ -277,16 +278,23 @@ public void updateSchedule(UpdateScheduleInput input) {
       return;
     }
 
-    UpdateScheduleRequest request =
+    UpdateScheduleRequest.Builder request =
         UpdateScheduleRequest.newBuilder()
             .setNamespace(clientOptions.getNamespace())
             .setIdentity(clientOptions.getIdentity())
             .setScheduleId(input.getDescription().getId())
             .setRequestId(UUID.randomUUID().toString())
-            .setSchedule(scheduleRequestHeader.scheduleToProto(schedule.getSchedule()))
-            .build();
+            .setSchedule(scheduleRequestHeader.scheduleToProto(schedule.getSchedule()));
+    if (schedule.getTypedSearchAttributes() != null) {
+      SearchAttributes encodedSa =
+          SearchAttributesUtil.encodeTyped(schedule.getTypedSearchAttributes());
+      if (encodedSa == null) {
+        encodedSa = SearchAttributes.getDefaultInstance();
+      }
+      request.setSearchAttributes(encodedSa);
+    }
     try {
-      genericClient.updateSchedule(request);
+      genericClient.updateSchedule(request.build());
     } catch (Exception e) {
       throw new ScheduleException(e);
     }

temporal-sdk/src/test/java/io/temporal/client/schedules/ScheduleTest.java

@@ -29,6 +29,7 @@
 import io.temporal.common.SearchAttributes;
 import io.temporal.common.converter.EncodedValues;
 import io.temporal.common.interceptors.ScheduleClientInterceptor;
+import io.temporal.testUtils.Eventually;
 import io.temporal.testing.internal.SDKTestWorkflowRule;
 import io.temporal.workflow.shared.TestWorkflows;
 import java.time.Duration;
@@ -415,9 +416,12 @@ public void describeSchedules() {
   public void updateSchedules() {
     ScheduleClient client = createScheduleClient();
     // Create the schedule
+    String keywordSAValue = "keyword";
     ScheduleOptions options =
         ScheduleOptions.newBuilder()
             .setMemo(Collections.singletonMap("memokey2", "memoval2"))
+            .setTypedSearchAttributes(
+                SearchAttributes.newBuilder().set(CUSTOM_KEYWORD_SA, keywordSAValue).build())
             .build();
     String scheduleId = UUID.randomUUID().toString();
     Schedule schedule = createTestSchedule().build();
@@ -469,7 +473,7 @@ public void updateSchedules() {
                   .setAction(input.getDescription().getSchedule().getAction())
                   .setSpec(ScheduleSpec.newBuilder().build());
           builder.setState(ScheduleState.newBuilder().setPaused(true).build());
-          return new ScheduleUpdate(builder.build());
+          return new ScheduleUpdate(builder.build(), null);
         });
     description = handle.describe();
     //
@@ -481,6 +485,32 @@ public void updateSchedules() {
     //
     Assert.assertNotEquals(expectedUpdateTime, description.getInfo().getLastUpdatedAt());
     Assert.assertEquals(true, description.getSchedule().getState().isPaused());
+    Assert.assertEquals(1, description.getTypedSearchAttributes().size());
+    Assert.assertEquals(
+        keywordSAValue, description.getTypedSearchAttributes().get(CUSTOM_KEYWORD_SA));
+    // Update the schedule search attribute by clearing them
+    handle.update(
+        (ScheduleUpdateInput input) ->
+            new ScheduleUpdate(input.getDescription().getSchedule(), SearchAttributes.EMPTY));
+    Eventually.assertEventually(
+        Duration.ofSeconds(1),
+        () -> {
+          ScheduleDescription desc = handle.describe();
+          Assert.assertEquals(0, desc.getTypedSearchAttributes().size());
+        });
+    // Update the schedule search attribute by adding a new search attribute
+    handle.update(
+        (ScheduleUpdateInput input) ->
+            new ScheduleUpdate(
+                input.getDescription().getSchedule(),
+                SearchAttributes.newBuilder().set(CUSTOM_KEYWORD_SA, "newkeyword").build()));
+    Eventually.assertEventually(
+        Duration.ofSeconds(1),
+        () -> {
+          ScheduleDescription desc = handle.describe();
+          Assert.assertEquals(1, desc.getTypedSearchAttributes().size());
+          Assert.assertEquals("newkeyword", desc.getTypedSearchAttributes().get(CUSTOM_KEYWORD_SA));
+        });
     // Cleanup schedule
     handle.delete();
   }

temporal-sdk/src/test/java/io/temporal/testUtils/Eventually.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2022 Temporal Technologies, Inc. All Rights Reserved.
+ *
+ * Copyright (C) 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 material except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License 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 io.temporal.testUtils;
+
+import java.time.Duration;
+import java.time.Instant;
+
+public class Eventually {
+  public static void assertEventually(Duration timeout, Runnable command) {
+    final Instant start = Instant.now();
+    final Instant deadline = start.plus(timeout);
+
+    boolean failed;
+    do {
+      try {
+        command.run();
+        failed = false;
+      } catch (Throwable t) {
+        failed = true;
+        if (Instant.now().isBefore(deadline)) {
+          // Try again after a short nap
+          try {
+            Thread.sleep(100);
+          } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new RuntimeException(e);
+          }
+        } else {
+          throw t;
+        }
+      }
+    } while (failed);
+  }
+}