Schedule update search attributes (#1625)

Author: THardy98

packages/client/src/schedule-client.ts

@@ -307,6 +307,12 @@ export class ScheduleClient extends BaseClient {
       },
       identity: this.options.identity,
       requestId: uuid4(),
+      searchAttributes:
+        opts.searchAttributes || opts.typedSearchAttributes // eslint-disable-line deprecation/deprecation
+          ? {
+              indexedFields: encodeUnifiedSearchAttributes(opts.searchAttributes, opts.typedSearchAttributes), // eslint-disable-line deprecation/deprecation
+            }
+          : undefined,
     };
     try {
       return await this.workflowService.updateSchedule(req);

packages/client/src/schedule-types.ts

@@ -142,7 +142,7 @@ export type CompiledScheduleOptions = Replace<
  * The specification of an updated Schedule, as expected by {@link ScheduleHandle.update}.
  */
 export type ScheduleUpdateOptions<A extends ScheduleOptionsAction = ScheduleOptionsAction> = Replace<
-  Omit<ScheduleOptions, 'scheduleId' | 'memo' | 'searchAttributes' | 'typedSearchAttributes'>,
+  Omit<ScheduleOptions, 'scheduleId' | 'memo'>,
   {
     action: A;
     state: Omit<ScheduleOptions['state'], 'triggerImmediately' | 'backfill'>;

packages/test/src/test-schedules.ts

@@ -10,10 +10,16 @@ import {
   ScheduleHandle,
   ScheduleSummary,
   ScheduleUpdateOptions,
+  ScheduleDescription,
 } from '@temporalio/client';
 import { msToNumber } from '@temporalio/common/lib/time';
-import { SearchAttributes, TypedSearchAttributes } from '@temporalio/common';
-import { registerDefaultCustomSearchAttributes, RUN_INTEGRATION_TESTS } from './helpers';
+import {
+  SearchAttributeType,
+  SearchAttributes,
+  TypedSearchAttributes,
+  defineSearchAttributeKey,
+} from '@temporalio/common';
+import { registerDefaultCustomSearchAttributes, RUN_INTEGRATION_TESTS, waitUntil } from './helpers';
 import { defaultSAKeys } from './helpers-integration';
 
 export interface Context {
@@ -751,4 +757,128 @@ if (RUN_INTEGRATION_TESTS) {
       await handle.delete();
     }
   });
+
+  test.serial('Can update search attributes of a schedule', async (t) => {
+    const { client } = t.context;
+    const scheduleId = `can-update-search-attributes-of-schedule-${randomUUID()}`;
+
+    // Helper to wait for search attribute changes to propagate.
+    const waitForAttributeChange = async (
+      handle: ScheduleHandle,
+      attributeName: string,
+      shouldExist: boolean
+    ): Promise<ScheduleDescription> => {
+      await waitUntil(async () => {
+        const desc = await handle.describe();
+        const exists =
+          desc.typedSearchAttributes.getAll().find((pair) => pair.key.name === attributeName) !== undefined;
+        return exists === shouldExist;
+      }, 300);
+      return await handle.describe();
+    };
+
+    // Create a schedule with search attributes.
+    const handle = await client.schedule.create({
+      scheduleId,
+      spec: {
+        calendars: [{ hour: { start: 2, end: 7, step: 1 } }],
+      },
+      action: {
+        type: 'startWorkflow',
+        workflowType: dummyWorkflow,
+        taskQueue,
+      },
+      searchAttributes: {
+        CustomKeywordField: ['keyword-one'],
+      },
+      typedSearchAttributes: [{ key: defineSearchAttributeKey('CustomIntField', SearchAttributeType.INT), value: 1 }],
+    });
+
+    // Check the search attributes are part of the schedule description.
+    const desc = await handle.describe();
+    // eslint-disable-next-line deprecation/deprecation
+    t.deepEqual(desc.searchAttributes, {
+      CustomKeywordField: ['keyword-one'],
+      CustomIntField: [1],
+    });
+    t.deepEqual(
+      desc.typedSearchAttributes,
+      new TypedSearchAttributes([
+        { key: defineSearchAttributeKey('CustomIntField', SearchAttributeType.INT), value: 1 },
+        { key: defineSearchAttributeKey('CustomKeywordField', SearchAttributeType.KEYWORD), value: 'keyword-one' },
+      ])
+    );
+
+    // Perform a series of updates to schedule's search attributes.
+    try {
+      // Update existing search attributes, add new ones.
+      await handle.update((desc) => ({
+        ...desc,
+        searchAttributes: {
+          CustomKeywordField: ['keyword-two'],
+          // Add a new search attribute.
+          CustomDoubleField: [1.5],
+        },
+        typedSearchAttributes: [
+          { key: defineSearchAttributeKey('CustomIntField', SearchAttributeType.INT), value: 2 },
+          // Add a new typed search attribute.
+          { key: defineSearchAttributeKey('CustomTextField', SearchAttributeType.TEXT), value: 'new-text' },
+        ],
+      }));
+
+      let desc = await waitForAttributeChange(handle, 'CustomTextField', true);
+      // eslint-disable-next-line deprecation/deprecation
+      t.deepEqual(desc.searchAttributes, {
+        CustomKeywordField: ['keyword-two'],
+        CustomIntField: [2],
+        CustomDoubleField: [1.5],
+        CustomTextField: ['new-text'],
+      });
+      t.deepEqual(
+        desc.typedSearchAttributes,
+        new TypedSearchAttributes([
+          { key: defineSearchAttributeKey('CustomIntField', SearchAttributeType.INT), value: 2 },
+          { key: defineSearchAttributeKey('CustomKeywordField', SearchAttributeType.KEYWORD), value: 'keyword-two' },
+          { key: defineSearchAttributeKey('CustomTextField', SearchAttributeType.TEXT), value: 'new-text' },
+          { key: defineSearchAttributeKey('CustomDoubleField', SearchAttributeType.DOUBLE), value: 1.5 },
+        ])
+      );
+
+      // Update and remove some search attributes. We remove a search attribute by omitting an existing key from the update.
+      await handle.update((desc) => ({
+        ...desc,
+        searchAttributes: {
+          CustomKeywordField: ['keyword-three'],
+        },
+        typedSearchAttributes: [{ key: defineSearchAttributeKey('CustomIntField', SearchAttributeType.INT), value: 3 }],
+      }));
+
+      desc = await waitForAttributeChange(handle, 'CustomTextField', false);
+      // eslint-disable-next-line deprecation/deprecation
+      t.deepEqual(desc.searchAttributes, {
+        CustomKeywordField: ['keyword-three'],
+        CustomIntField: [3],
+      });
+      t.deepEqual(
+        desc.typedSearchAttributes,
+        new TypedSearchAttributes([
+          { key: defineSearchAttributeKey('CustomIntField', SearchAttributeType.INT), value: 3 },
+          { key: defineSearchAttributeKey('CustomKeywordField', SearchAttributeType.KEYWORD), value: 'keyword-three' },
+        ])
+      );
+
+      // Remove all search attributes.
+      await handle.update((desc) => ({
+        ...desc,
+        searchAttributes: {},
+        typedSearchAttributes: [],
+      }));
+
+      desc = await waitForAttributeChange(handle, 'CustomIntField', false);
+      t.deepEqual(desc.searchAttributes, {}); // eslint-disable-line deprecation/deprecation
+      t.deepEqual(desc.typedSearchAttributes, new TypedSearchAttributes([]));
+    } finally {
+      await handle.delete();
+    }
+  });
 }