Allow schedules to update search attributes

temporalio · Feb 8, 2025 · d8ee220 · d8ee220
1 parent ba05b76
commit d8ee220
Show file tree

Hide file tree

Showing 3 changed files with 129 additions and 3 deletions.
diff --git a/packages/client/src/schedule-client.ts b/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
+          ? {
+              indexedFields: encodeUnifiedSearchAttributes(opts.searchAttributes, opts.typedSearchAttributes),
+            }
+          : undefined,
     };
     try {
       return await this.workflowService.updateSchedule(req);

diff --git a/packages/client/src/schedule-types.ts b/packages/client/src/schedule-types.ts
@@ -147,9 +147,8 @@ export type CompiledScheduleOptions = Replace<
 /**
  * The specification of an updated Schedule, as expected by {@link ScheduleHandle.update}.
  */
-// TODO(thomas): support search attributes on update (ts issue #1475)
 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'>;

diff --git a/packages/test/src/test-schedules.ts b/packages/test/src/test-schedules.ts
@@ -11,10 +11,11 @@ import {
   ScheduleSummary,
   ScheduleUpdateOptions,
   SearchAttributes,
+  ScheduleDescription,
 } from '@temporalio/client';
 import { msToNumber } from '@temporalio/common/lib/time';
 import { SearchAttributeType, TypedSearchAttributes } from '@temporalio/common';
-import { registerDefaultCustomSearchAttributes, RUN_INTEGRATION_TESTS } from './helpers';
+import { registerDefaultCustomSearchAttributes, RUN_INTEGRATION_TESTS, waitUntil } from './helpers';
 
 export interface Context {
   client: Client;
@@ -749,4 +750,124 @@ 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.getAttributes().find(([k]) => k.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: [TypedSearchAttributes.createAttribute('CustomIntField', SearchAttributeType.INT, 1)],
+    });
+
+    // Check the search attributes are part of the schedule description.
+    const desc = await handle.describe();
+    t.deepEqual(desc.searchAttributes, {
+      CustomKeywordField: ['keyword-one'],
+      CustomIntField: [1],
+    });
+    t.deepEqual(
+      desc.typedSearchAttributes,
+      new TypedSearchAttributes([
+        TypedSearchAttributes.createAttribute('CustomIntField', SearchAttributeType.INT, 1),
+        TypedSearchAttributes.createAttribute('CustomKeywordField', SearchAttributeType.KEYWORD, '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: [
+          TypedSearchAttributes.createAttribute('CustomIntField', SearchAttributeType.INT, 2),
+          // Add a new typed search attribute.
+          TypedSearchAttributes.createAttribute('CustomTextField', SearchAttributeType.TEXT, 'new-text'),
+        ],
+      }));
+
+      let desc = await waitForAttributeChange(handle, 'CustomTextField', true);
+      t.deepEqual(desc.searchAttributes, {
+        CustomKeywordField: ['keyword-two'],
+        CustomIntField: [2],
+        CustomDoubleField: [1.5],
+        CustomTextField: ['new-text'],
+      });
+      t.deepEqual(
+        desc.typedSearchAttributes,
+        new TypedSearchAttributes([
+          TypedSearchAttributes.createAttribute('CustomIntField', SearchAttributeType.INT, 2),
+          TypedSearchAttributes.createAttribute('CustomKeywordField', SearchAttributeType.KEYWORD, 'keyword-two'),
+          TypedSearchAttributes.createAttribute('CustomTextField', SearchAttributeType.TEXT, 'new-text'),
+          TypedSearchAttributes.createAttribute('CustomDoubleField', SearchAttributeType.DOUBLE, 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: [TypedSearchAttributes.createAttribute('CustomIntField', SearchAttributeType.INT, 3)],
+      }));
+
+      desc = await waitForAttributeChange(handle, 'CustomTextField', false);
+      t.deepEqual(desc.searchAttributes, {
+        CustomKeywordField: ['keyword-three'],
+        CustomIntField: [3],
+      });
+      t.deepEqual(
+        desc.typedSearchAttributes,
+        new TypedSearchAttributes([
+          TypedSearchAttributes.createAttribute('CustomIntField', SearchAttributeType.INT, 3),
+          TypedSearchAttributes.createAttribute('CustomKeywordField', SearchAttributeType.KEYWORD, 'keyword-three'),
+        ])
+      );
+
+      // Remove all search attributes.
+      await handle.update((desc) => ({
+        ...desc,
+        searchAttributes: {},
+        typedSearchAttributes: [],
+      }));
+
+      desc = await waitForAttributeChange(handle, 'CustomIntField', false);
+      t.deepEqual(desc.searchAttributes, {});
+      t.deepEqual(desc.typedSearchAttributes, new TypedSearchAttributes([]));
+    } finally {
+      await handle.delete();
+    }
+  });
 }