Add search attributes to dev server (#1624)

THardy98 · web-flow · commit a0470c98bc2d · 2025-03-19T10:26:44.000-07:00
diff --git a/packages/core-bridge/ts/index.ts b/packages/core-bridge/ts/index.ts
@@ -1,6 +1,7 @@
-import { LogLevel, Duration } from '@temporalio/common';
+import { LogLevel, Duration, SearchAttributeType } from '@temporalio/common';
 import type { TLSConfig, ProxyConfig, HttpConnectProxyConfig } from '@temporalio/common/lib/internal-non-workflow';
 import { WorkerTuner } from './worker-tuner';
+import { SearchAttributeKey } from '@temporalio/common/src/search-attributes';
 
 export {
   WorkerTuner,
@@ -511,6 +512,10 @@ export interface DevServerConfig {
    * be supported in the future.
    */
   extraArgs?: string[];
+  /**
+   * Search attributes to be registered with the dev server.
+   */
+  searchAttributes?: SearchAttributeKey<SearchAttributeType>[];
 }
 
 /**
diff --git a/packages/test/src/helpers-integration-multi-codec.ts b/packages/test/src/helpers-integration-multi-codec.ts
@@ -10,7 +10,7 @@ import {
   createLocalTestEnvironment,
   makeConfigurableEnvironmentTestFn,
 } from './helpers-integration';
-import { ByteSkewerPayloadCodec, registerDefaultCustomSearchAttributes, Worker } from './helpers';
+import { ByteSkewerPayloadCodec, Worker } from './helpers';
 
 // Note: re-export shared workflows (or long workflows)
 export * from './workflows';
@@ -43,7 +43,6 @@ export function makeTestFn(makeBundle: () => Promise<WorkflowBundle>): TestFn<Te
           const env = await createLocalTestEnvironment({
             client: { dataConverter },
           });
-          await registerDefaultCustomSearchAttributes(env.connection);
 
           configs.push({
             loadedDataConverter,
diff --git a/packages/test/src/helpers-integration.ts b/packages/test/src/helpers-integration.ts
@@ -27,14 +27,9 @@ import {
 } from '@temporalio/worker';
 import * as workflow from '@temporalio/workflow';
 import { temporal } from '@temporalio/proto';
+import { defineSearchAttributeKey, SearchAttributeType } from '@temporalio/common/lib/search-attributes';
 import { ConnectionInjectorInterceptor } from './activities/interceptors';
-import {
-  Worker,
-  TestWorkflowEnvironment,
-  test as anyTest,
-  bundlerOptions,
-  registerDefaultCustomSearchAttributes,
-} from './helpers';
+import { Worker, TestWorkflowEnvironment, test as anyTest, bundlerOptions } from './helpers';
 
 export interface Context {
   env: TestWorkflowEnvironment;
@@ -89,12 +84,22 @@ export async function createTestWorkflowBundle({
   });
 }
 
+export const defaultSAKeys = {
+  CustomIntField: defineSearchAttributeKey('CustomIntField', SearchAttributeType.INT),
+  CustomBoolField: defineSearchAttributeKey('CustomBoolField', SearchAttributeType.BOOL),
+  CustomKeywordField: defineSearchAttributeKey('CustomKeywordField', SearchAttributeType.KEYWORD),
+  CustomTextField: defineSearchAttributeKey('CustomTextField', SearchAttributeType.TEXT),
+  CustomDatetimeField: defineSearchAttributeKey('CustomDatetimeField', SearchAttributeType.DATETIME),
+  CustomDoubleField: defineSearchAttributeKey('CustomDoubleField', SearchAttributeType.DOUBLE),
+};
+
 export async function createLocalTestEnvironment(
   opts?: LocalTestWorkflowEnvironmentOptions
 ): Promise<TestWorkflowEnvironment> {
   return await TestWorkflowEnvironment.createLocal({
     ...(opts || {}), // Use provided options or default to an empty object
     server: {
+      searchAttributes: Object.values(defaultSAKeys),
       ...(opts?.server || {}), // Use provided server options or default to an empty object
       extraArgs: [
         ...defaultDynamicConfigOptions.flatMap((opt) => ['--dynamic-config-value', opt]),
@@ -130,7 +135,6 @@ export function makeTestFunction(opts: {
     recordedLogs: opts.recordedLogs,
     createTestContext: async (_t: ExecutionContext): Promise<Context> => {
       const env = await createLocalTestEnvironment(opts.workflowEnvironmentOpts);
-      await registerDefaultCustomSearchAttributes(env.connection);
       return {
         workflowBundle: await createTestWorkflowBundle({
           workflowsPath: opts.workflowsPath,
diff --git a/packages/test/src/test-integration-workflows.ts b/packages/test/src/test-integration-workflows.ts
@@ -10,11 +10,19 @@ import { CancelReason } from '@temporalio/worker/lib/activity';
 import * as workflow from '@temporalio/workflow';
 import { defineQuery, defineSignal } from '@temporalio/workflow';
 import { SdkFlags } from '@temporalio/workflow/lib/flags';
-import { ActivityCancellationType, ApplicationFailure, WorkflowExecutionAlreadyStartedError } from '@temporalio/common';
+import {
+  ActivityCancellationType,
+  ApplicationFailure,
+  defineSearchAttributeKey,
+  SearchAttributePair,
+  SearchAttributeType,
+  TypedSearchAttributes,
+  WorkflowExecutionAlreadyStartedError,
+} from '@temporalio/common';
 import { signalSchedulingWorkflow } from './activities/helpers';
 import { activityStartedSignal } from './workflows/definitions';
 import * as workflows from './workflows';
-import { Context, helpers, makeTestFunction } from './helpers-integration';
+import { Context, createLocalTestEnvironment, helpers, makeTestFunction } from './helpers-integration';
 import { overrideSdkInternalFlag } from './mock-internal-flags';
 import { asSdkLoggerSink, loadHistory, RUN_TIME_SKIPPING_TESTS } from './helpers';
 
@@ -1303,3 +1311,29 @@ test('Count workflow executions', async (t) => {
     ],
   });
 });
+
+test('can register search attributes to dev server', async (t) => {
+  const key = defineSearchAttributeKey('new-search-attr', SearchAttributeType.INT);
+  const newSearchAttribute: SearchAttributePair = { key, value: 12 };
+
+  // Create new test environment with search attribute registered.
+  const env = await createLocalTestEnvironment({
+    server: {
+      searchAttributes: [key],
+    },
+  });
+
+  const newClient = env.client;
+  // Expect workflow with search attribute to start without error.
+  const handle = await newClient.workflow.start(completableWorkflow, {
+    args: [true],
+    workflowId: randomUUID(),
+    taskQueue: 'new-env-tq',
+    typedSearchAttributes: [newSearchAttribute],
+  });
+  // Expect workflow description to have search attribute.
+  const desc = await handle.describe();
+  t.deepEqual(desc.typedSearchAttributes, new TypedSearchAttributes([newSearchAttribute]));
+  t.deepEqual(desc.searchAttributes, { 'new-search-attr': [12] }); // eslint-disable-line deprecation/deprecation
+  await env.teardown();
+});
diff --git a/packages/test/src/test-schedules.ts b/packages/test/src/test-schedules.ts
@@ -12,8 +12,9 @@ import {
   ScheduleUpdateOptions,
 } from '@temporalio/client';
 import { msToNumber } from '@temporalio/common/lib/time';
-import { SearchAttributes, SearchAttributeType, TypedSearchAttributes } from '@temporalio/common';
+import { SearchAttributes, TypedSearchAttributes } from '@temporalio/common';
 import { registerDefaultCustomSearchAttributes, RUN_INTEGRATION_TESTS } from './helpers';
+import { defaultSAKeys } from './helpers-integration';
 
 export interface Context {
   client: Client;
@@ -168,9 +169,7 @@ if (RUN_INTEGRATION_TESTS) {
         searchAttributes: {
           CustomKeywordField: ['test-value2'],
         },
-        typedSearchAttributes: new TypedSearchAttributes([
-          { key: { name: 'CustomIntField', type: SearchAttributeType.INT }, value: 42 },
-        ]),
+        typedSearchAttributes: new TypedSearchAttributes([{ key: defaultSAKeys.CustomIntField, value: 42 }]),
       },
     });
 
@@ -188,8 +187,8 @@ if (RUN_INTEGRATION_TESTS) {
       t.deepEqual(
         describedSchedule.action.typedSearchAttributes,
         new TypedSearchAttributes([
-          { key: { name: 'CustomIntField', type: SearchAttributeType.INT }, value: 42 },
-          { key: { name: 'CustomKeywordField', type: SearchAttributeType.KEYWORD }, value: 'test-value2' },
+          { key: defaultSAKeys.CustomIntField, value: 42 },
+          { key: defaultSAKeys.CustomKeywordField, value: 'test-value2' },
         ])
       );
     } finally {
@@ -216,9 +215,7 @@ if (RUN_INTEGRATION_TESTS) {
         searchAttributes: {
           CustomKeywordField: ['test-value2'],
         },
-        typedSearchAttributes: new TypedSearchAttributes([
-          { key: { name: 'CustomIntField', type: SearchAttributeType.INT }, value: 42 },
-        ]),
+        typedSearchAttributes: new TypedSearchAttributes([{ key: defaultSAKeys.CustomIntField, value: 42 }]),
       },
     });
 
@@ -237,8 +234,8 @@ if (RUN_INTEGRATION_TESTS) {
       t.deepEqual(
         describedSchedule.action.typedSearchAttributes,
         new TypedSearchAttributes([
-          { key: { name: 'CustomIntField', type: SearchAttributeType.INT }, value: 42 },
-          { key: { name: 'CustomKeywordField', type: SearchAttributeType.KEYWORD }, value: 'test-value2' },
+          { key: defaultSAKeys.CustomIntField, value: 42 },
+          { key: defaultSAKeys.CustomKeywordField, value: 'test-value2' },
         ])
       );
     } finally {
@@ -351,9 +348,7 @@ if (RUN_INTEGRATION_TESTS) {
         searchAttributes: {
           CustomKeywordField: ['test-value2'],
         },
-        typedSearchAttributes: new TypedSearchAttributes([
-          { key: { name: 'CustomIntField', type: SearchAttributeType.INT }, value: 42 },
-        ]),
+        typedSearchAttributes: new TypedSearchAttributes([{ key: defaultSAKeys.CustomIntField, value: 42 }]),
       },
     });
 
@@ -598,9 +593,7 @@ if (RUN_INTEGRATION_TESTS) {
             taskQueue,
           },
           searchAttributes,
-          typedSearchAttributes: new TypedSearchAttributes([
-            { key: { name: 'CustomIntField', type: SearchAttributeType.INT }, value: 42 },
-          ]),
+          typedSearchAttributes: new TypedSearchAttributes([{ key: defaultSAKeys.CustomIntField, value: 42 }]),
         })
       );
     }
diff --git a/packages/test/src/test-typed-search-attributes.ts b/packages/test/src/test-typed-search-attributes.ts
@@ -28,6 +28,7 @@ const test = makeTestFunction({
   workflowEnvironmentOpts: {
     server: {
       namespace: 'test-typed-search-attributes',
+      searchAttributes: [],
     },
   },
 });
diff --git a/packages/testing/src/index.ts b/packages/testing/src/index.ts
@@ -27,6 +27,7 @@ import {
   Logger,
   defaultFailureConverter,
   defaultPayloadConverter,
+  TypedSearchAttributes,
 } from '@temporalio/common';
 import { msToNumber, msToTs, tsToMs } from '@temporalio/common/lib/time';
 import { ActivityInterceptorsFactory, DefaultLogger, NativeConnection, Runtime } from '@temporalio/worker';
@@ -302,6 +303,18 @@ export class TestWorkflowEnvironment {
   ): Promise<TestWorkflowEnvironment> {
     const { supportsTimeSkipping, namespace, ...rest } = opts;
     const optsWithDefaults = addDefaults(filterNullAndUndefined(rest));
+
+    // Add search attributes to CLI server arguments
+    if ('searchAttributes' in optsWithDefaults.server && optsWithDefaults.server.searchAttributes) {
+      let newArgs: string[] = [];
+      for (const { name, type } of optsWithDefaults.server.searchAttributes) {
+        newArgs.push('--search-attribute');
+        newArgs.push(`${name}=${TypedSearchAttributes.toMetadataType(type)}`);
+      }
+      newArgs = newArgs.concat(optsWithDefaults.server.extraArgs ?? []);
+      optsWithDefaults.server.extraArgs = newArgs;
+    }
+
     const server = await Runtime.instance().createEphemeralServer(optsWithDefaults.server);
     const address = getEphemeralServerTarget(server);
 
