feat(workflow): add upsert memo command (#1321)

Co-authored-by: Tharun Rajendran <tharunrajendran@Tharuns-MacBook-Pro.local>
Co-authored-by: James Watkins-Harvey <james.watkinsharvey@temporal.io>

Author: 3 people

packages/test/src/helpers.ts

@@ -5,7 +5,7 @@ import ava, { TestFn } from 'ava';
 import * as grpc from '@grpc/grpc-js';
 import asyncRetry from 'async-retry';
 import { v4 as uuid4 } from 'uuid';
-import { inWorkflowContext } from '@temporalio/workflow';
+import { inWorkflowContext, WorkflowInfo } from '@temporalio/workflow';
 import { Payload, PayloadCodec } from '@temporalio/common';
 import { Worker as RealWorker, WorkerOptions } from '@temporalio/worker';
 import * as worker from '@temporalio/worker';
@@ -16,6 +16,7 @@ import {
   TestWorkflowEnvironment as RealTestWorkflowEnvironment,
   TimeSkippingTestWorkflowEnvironmentOptions,
 } from '@temporalio/testing';
+import { LoggerSinksInternal as DefaultLoggerSinks } from '@temporalio/workflow/lib/logs';
 
 export function u8(s: string): Uint8Array {
   // TextEncoder requires lib "dom"
@@ -260,3 +261,18 @@ export async function getRandomPort(fn = (_port: number) => Promise.resolve()):
     });
   });
 }
+
+export function asSdkLoggerSink(
+  fn: (info: WorkflowInfo, message: string, attrs?: Record<string, unknown>) => Promise<void>,
+  opts?: Omit<worker.InjectedSinkFunction<any>, 'fn'>
+): worker.InjectedSinks<DefaultLoggerSinks> {
+  return {
+    __temporal_logger: {
+      trace: { fn, ...opts },
+      debug: { fn, ...opts },
+      info: { fn, ...opts },
+      warn: { fn, ...opts },
+      error: { fn, ...opts },
+    },
+  };
+}

packages/test/src/test-integration-workflows.ts

@@ -15,7 +15,7 @@ import { activityStartedSignal } from './workflows/definitions';
 import * as workflows from './workflows';
 import { Context, helpers, makeTestFunction } from './helpers-integration';
 import { overrideSdkInternalFlag } from './mock-internal-flags';
-import { RUN_TIME_SKIPPING_TESTS } from './helpers';
+import { RUN_TIME_SKIPPING_TESTS, asSdkLoggerSink } from './helpers';
 
 const test = makeTestFunction({ workflowsPath: __filename, workflowInterceptorModules: [__filename] });
 
@@ -947,6 +947,95 @@ if (RUN_TIME_SKIPPING_TESTS) {
   });
 }
 
+export async function upsertAndReadMemo(memo: Record<string, unknown>): Promise<Record<string, unknown> | undefined> {
+  workflow.upsertMemo(memo);
+  return workflow.workflowInfo().memo;
+}
+
+test('Workflow can upsert memo', async (t) => {
+  const { createWorker, startWorkflow } = helpers(t);
+  const worker = await createWorker();
+  await worker.runUntil(async () => {
+    const handle = await startWorkflow(upsertAndReadMemo, {
+      memo: {
+        alpha: 'bar1',
+        bravo: 'bar3',
+        charlie: { delta: 'bar2', echo: 12 },
+        foxtrot: 'bar4',
+      },
+      args: [
+        {
+          alpha: 'bar11',
+          bravo: null,
+          charlie: { echo: 34, golf: 'bar5' },
+          hotel: 'bar6',
+        },
+      ],
+    });
+    const result = await handle.result();
+    t.deepEqual(result, {
+      alpha: 'bar11',
+      charlie: { echo: 34, golf: 'bar5' },
+      foxtrot: 'bar4',
+      hotel: 'bar6',
+    });
+    const { memo } = await handle.describe();
+    t.deepEqual(memo, {
+      alpha: 'bar11',
+      charlie: { echo: 34, golf: 'bar5' },
+      foxtrot: 'bar4',
+      hotel: 'bar6',
+    });
+  });
+});
+
+test('Sink functions contains upserted memo', async (t) => {
+  const { createWorker, executeWorkflow } = helpers(t);
+  const recordedMessages = Array<{ message: string; memo: Record<string, unknown> | undefined }>();
+  const sinks = asSdkLoggerSink(async (info, message, _attrs) => {
+    recordedMessages.push({
+      message,
+      memo: info.memo,
+    });
+  });
+  const worker = await createWorker({ sinks });
+  await worker.runUntil(async () => {
+    await executeWorkflow(upsertAndReadMemo, {
+      memo: {
+        note1: 'aaa',
+        note2: 'bbb',
+        note4: 'eee',
+      },
+      args: [
+        {
+          note2: 'ccc',
+          note3: 'ddd',
+          note4: null,
+        },
+      ],
+    });
+  });
+
+  t.deepEqual(recordedMessages, [
+    {
+      message: 'Workflow started',
+      memo: {
+        note1: 'aaa',
+        note2: 'bbb',
+        note4: 'eee',
+      },
+    },
+    {
+      message: 'Workflow completed',
+      memo: {
+        note1: 'aaa',
+        note2: 'ccc',
+        note3: 'ddd',
+      },
+    },
+  ]);
+});
+
 export const interceptors: workflow.WorkflowInterceptorsFactory = () => {
   const interceptorsFactoryFunc = module.exports[`${workflow.workflowInfo().workflowType}Interceptors`];
   if (typeof interceptorsFactoryFunc === 'function') {

packages/test/src/test-sinks.ts

@@ -2,19 +2,11 @@
 import test from 'ava';
 import { v4 as uuid4 } from 'uuid';
 import { Connection, WorkflowClient } from '@temporalio/client';
-import {
-  DefaultLogger,
-  InjectedSinks,
-  Runtime,
-  InjectedSinkFunction,
-  WorkerOptions,
-  LogEntry,
-} from '@temporalio/worker';
-import { LoggerSinksInternal as DefaultLoggerSinks } from '@temporalio/workflow/lib/logs';
+import { DefaultLogger, InjectedSinks, Runtime, WorkerOptions, LogEntry } from '@temporalio/worker';
 import { SearchAttributes, WorkflowInfo } from '@temporalio/workflow';
 import { UnsafeWorkflowInfo } from '@temporalio/workflow/src/interfaces';
 import { SdkComponent } from '@temporalio/common';
-import { RUN_INTEGRATION_TESTS, Worker, registerDefaultCustomSearchAttributes } from './helpers';
+import { RUN_INTEGRATION_TESTS, Worker, asSdkLoggerSink, registerDefaultCustomSearchAttributes } from './helpers';
 import { defaultOptions } from './mock-native-worker';
 import * as workflows from './workflows';
 
@@ -27,21 +19,6 @@ class DependencyError extends Error {
   }
 }
 
-function asSdkLoggerSink(
-  fn: (info: WorkflowInfo, message: string, attrs?: Record<string, unknown>) => Promise<void>,
-  opts?: Omit<InjectedSinkFunction<any>, 'fn'>
-): InjectedSinks<DefaultLoggerSinks> {
-  return {
-    __temporal_logger: {
-      trace: { fn, ...opts },
-      debug: { fn, ...opts },
-      info: { fn, ...opts },
-      warn: { fn, ...opts },
-      error: { fn, ...opts },
-    },
-  };
-}
-
 if (RUN_INTEGRATION_TESTS) {
   const recordedLogs: { [workflowId: string]: LogEntry[] } = {};
 

packages/test/src/workflows/index.ts

@@ -96,3 +96,4 @@ export * from './upsert-and-read-search-attributes';
 export * from './url-whatwg';
 export * from './wait-on-user';
 export * from './workflow-cancellation-scenarios';
+export * from './upsert-and-read-memo';

packages/test/src/workflows/upsert-and-read-memo.ts

@@ -0,0 +1,6 @@
+import { upsertMemo, workflowInfo } from '@temporalio/workflow';
+
+export async function upsertAndReadMemo(memo: Record<string, unknown>): Promise<Record<string, unknown> | undefined> {
+  upsertMemo(memo);
+  return workflowInfo().memo;
+}

packages/workflow/src/workflow.ts

@@ -1366,6 +1366,74 @@ export function upsertSearchAttributes(searchAttributes: SearchAttributes): void
   });
 }
 
+/**
+ * Updates this Workflow's Memos by merging the provided `memo` with existing
+ * Memos (as returned by `workflowInfo().memo`).
+ *
+ * New memo is merged by replacing properties of the same name _at the first
+ * level only_. Setting a property to value `undefined` or `null` clears that
+ * key from the Memo.
+ *
+ * For example:
+ *
+ * ```ts
+ * upsertMemo({
+ *   key1: value,
+ *   key3: { subkey1: value }
+ *   key4: value,
+ * });
+ * upsertMemo({
+ *   key2: value
+ *   key3: { subkey2: value }
+ *   key4: undefined,
+ * });
+ * ```
+ *
+ * would result in the Workflow having these Memo:
+ *
+ * ```ts
+ * {
+ *   key1: value,
+ *   key2: value,
+ *   key3: { subkey2: value }  // Note this object was completely replaced
+ *   // Note that key4 was completely removed
+ * }
+ * ```
+ *
+ * @param memo The Record to merge.
+ */
+export function upsertMemo(memo: Record<string, unknown>): void {
+  const activator = assertInWorkflowContext('Workflow.upsertMemo(...) may only be used from a Workflow Execution.');
+
+  if (memo == null) {
+    throw new Error('memo must be a non-null Record');
+  }
+
+  activator.pushCommand({
+    modifyWorkflowProperties: {
+      upsertedMemo: {
+        fields: mapToPayloads(
+          activator.payloadConverter,
+          // Convert null to undefined
+          Object.fromEntries(Object.entries(memo).map(([k, v]) => [k, v ?? undefined]))
+        ),
+      },
+    },
+  });
+
+  activator.mutateWorkflowInfo((info: WorkflowInfo): WorkflowInfo => {
+    return {
+      ...info,
+      memo: Object.fromEntries(
+        Object.entries({
+          ...info.memo,
+          ...memo,
+        }).filter(([_, v]) => v != null)
+      ),
+    };
+  });
+}
+
 export const stackTraceQuery = defineQuery<string>('__stack_trace');
 export const enhancedStackTraceQuery = defineQuery<EnhancedStackTrace>('__enhanced_stack_trace');
 export const workflowMetadataQuery = defineQuery<temporal.api.sdk.v1.IWorkflowMetadata>('__temporal_workflow_metadata');