feat(toolkit-lib)!: parameterize context storage (#534)

Previously, context was always being written to `cdk.context.json` in
some directory, even for `fromAssemblyBuilder()` or
`fromAssemblyDirectory()` which should probably both not write context
anywhere.

In this PR, we replace the context parameter to the toolkit with the
concept of a `contextStore`, which has two operations: `read()` and
`update()`.

There are simple file and in-memory implementations for simple cases, as
well as the more complex CDK app case which is implemented as
`CdkAppMultiContext`. That last one is the default for `fromCdkApp()`.

The complex one is implemented in terms of the prior
`Context/ContextBag/Settings` classes, which are probably overly complex
for the job this class is doing but keeping the old classes is an easy
way to make sure that the behavior stays compatible.

BREAKING CHANGES: `context` is now called `contextStore`. Message
`CDK_ASSEMBLY_I0042` no longer knows about the filename where changes
get written.

---
By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache-2.0 license

---------

Co-authored-by: Momo Kornher <kornherm@amazon.co.uk>

Author: rix0rrr

packages/@aws-cdk/toolkit-lib/docs/message-registry.md

@@ -135,7 +135,7 @@ Please let us know by [opening an issue](https://github.com/aws/aws-cdk-cli/issu
 | `CDK_TOOLKIT_I0101` | A notice that is marked as informational | `info` | n/a |
 | `CDK_ASSEMBLY_I0010` | Generic environment preparation debug messages | `debug` | n/a |
 | `CDK_ASSEMBLY_W0010` | Emitted if the found framework version does not support context overflow | `warn` | n/a |
-| `CDK_ASSEMBLY_I0042` | Writing updated context | `debug` | {@link UpdatedContext} |
+| `CDK_ASSEMBLY_I0042` | Writing context updates | `debug` | {@link UpdatedContext} |
 | `CDK_ASSEMBLY_I0240` | Context lookup was stopped as no further progress was made.  | `debug` | {@link MissingContext} |
 | `CDK_ASSEMBLY_I0241` | Fetching missing context. This is an iterative message that may appear multiple times with different missing keys. | `debug` | {@link MissingContext} |
 | `CDK_ASSEMBLY_I1000` | Cloud assembly output starts | `debug` | n/a |

packages/@aws-cdk/toolkit-lib/lib/api/cloud-assembly/context-store.ts

@@ -0,0 +1,191 @@
+import { promises as fs } from 'fs';
+import * as os from 'os';
+import * as path from 'path';
+import { ToolkitError } from '../../toolkit/toolkit-error';
+import type { ContextBag } from '../context';
+import { Context } from '../context';
+import { Settings, TRANSIENT_CONTEXT_KEY } from '../settings';
+
+/**
+ * A storage place for context used in synthesis
+ */
+export interface IContextStore {
+  /**
+   * Read the context from the context store, plus all updates we have made so far.
+   */
+  read(): Promise<Record<string, unknown>>;
+
+  /**
+   * Commit the given updates to the context store
+   *
+   * `undefined` is used as a value to indicate that the key needs to be removed.
+   *
+   * If a context value is an object that is a superset of `{ [TRANSIENT_CONTEXT_KEY]: true }`
+   * it *should* be returned by subsequent `read()` operations on this object,
+   * but it *should not* be persisted to permanent storage.
+   *
+   * You can use the `persistableContext()` function to filter a context dictionary
+   * down to remove all values that shouldn't be persisted.
+   */
+  update(updates: Record<string, unknown>): Promise<void>;
+}
+
+/**
+ * A context store as used by a CDK app.
+ *
+ * Will source context from the following locations:
+ *
+ * - Any context values passed to the constructor (expected
+ *   to come from the command line, treated as ephemeral).
+ * - The `context` key in `<appDirectory>/cdk.json`.
+ * - `<appDirectory>/cdk.context.json`.
+ * - The `context` key in `~/.cdk.json`.
+ *
+ * Updates will be written to `<appDirectory>/cdk.context.json`.
+ */
+export class CdkAppMultiContext implements IContextStore {
+  private _context?: Context;
+  private configContextFile: string;
+  private projectContextFile: string;
+  private userConfigFile: string;
+
+  constructor(appDirectory: string, private readonly commandlineContext?: Record<string, unknown>) {
+    this.configContextFile = path.join(appDirectory, 'cdk.json');
+    this.projectContextFile = path.join(appDirectory, 'cdk.context.json');
+    this.userConfigFile = path.join(os.homedir() ?? '/tmp', '.cdk.json');
+  }
+
+  public async read(): Promise<Record<string, unknown>> {
+    const context = await this.asyncInitialize();
+    return context.all;
+  }
+
+  public async update(updates: Record<string, unknown>): Promise<void> {
+    const context = await this.asyncInitialize();
+    for (const [key, value] of Object.entries(updates)) {
+      context.set(key, value);
+    }
+
+    await context.save(this.projectContextFile);
+  }
+
+  /**
+   * Initialize the `Context` object
+   *
+   * This code all exists to reuse code that's already there, to minimize
+   * the chances of the new code behaving subtly differently than the
+   * old code.
+   *
+   * It might be most of this is unnecessary now...
+   */
+  private async asyncInitialize(): Promise<Context> {
+    if (this._context) {
+      return this._context;
+    }
+
+    const contextSources: ContextBag[] = [
+      { bag: new Settings(this.commandlineContext, true) },
+      {
+        fileName: this.configContextFile,
+        bag: (await settingsFromFile(this.configContextFile)).subSettings(['context']).makeReadOnly(),
+      },
+      {
+        fileName: this.projectContextFile,
+        bag: await settingsFromFile(this.projectContextFile),
+      },
+      {
+        fileName: this.userConfigFile,
+        bag: (await settingsFromFile(this.userConfigFile)).subSettings(['context']).makeReadOnly(),
+      },
+    ];
+
+    this._context = new Context(...contextSources);
+    return this._context;
+  }
+}
+
+/**
+ * On-disk context stored in a single file
+ */
+export class FileContext implements IContextStore {
+  private _cache?: Record<string, unknown>;
+
+  constructor(private readonly fileName: string) {
+  }
+
+  public async read(): Promise<Record<string, unknown>> {
+    if (!this._cache) {
+      try {
+        this._cache = JSON.parse(await fs.readFile(this.fileName, 'utf-8'));
+      } catch (e: any) {
+        if (e.code === 'ENOENT') {
+          this._cache = {};
+        } else {
+          throw e;
+        }
+      }
+    }
+    if (!this._cache || typeof this._cache !== 'object') {
+      throw new ToolkitError(`${this.fileName} must contain an object, got: ${JSON.stringify(this._cache)}`);
+    }
+    return this._cache;
+  }
+
+  public async update(updates: Record<string, unknown>): Promise<void> {
+    this._cache = {
+      ...await this.read(),
+      ...updates,
+    };
+
+    const persistable = persistableContext(this._cache);
+    await fs.writeFile(this.fileName, JSON.stringify(persistable, undefined, 2), 'utf-8');
+  }
+}
+
+/**
+ * An in-memory context store
+ */
+export class MemoryContext implements IContextStore {
+  private context: Record<string, unknown> = {};
+
+  constructor(initialContext?: Record<string, unknown>) {
+    this.context = { ...initialContext };
+  }
+
+  public read(): Promise<Record<string, unknown>> {
+    return Promise.resolve(this.context);
+  }
+
+  public update(updates: Record<string, unknown>): Promise<void> {
+    this.context = {
+      ...this.context,
+      ...updates,
+    };
+
+    return Promise.resolve();
+  }
+}
+
+/**
+ * Filter the given context, leaving only entries that should be persisted
+ */
+export function persistableContext(context: Record<string, unknown>): Record<string, unknown> {
+  return Object.fromEntries(Object.entries(context)
+    .filter(([_, value]) => !isTransientValue(value)));
+}
+
+function isTransientValue(x: unknown) {
+  return x && typeof x === 'object' && (x as any)[TRANSIENT_CONTEXT_KEY];
+}
+
+async function settingsFromFile(filename: string): Promise<Settings> {
+  try {
+    const data = JSON.parse(await fs.readFile(filename, 'utf-8'));
+    return new Settings(data);
+  } catch (e: any) {
+    if (e.code === 'ENOENT') {
+      return new Settings();
+    }
+    throw e;
+  }
+}

packages/@aws-cdk/toolkit-lib/lib/api/cloud-assembly/index.ts

@@ -3,3 +3,4 @@ export type { StackSelector } from './stack-selector';
 export * from './cached-source';
 export * from './source-builder';
 export * from './types';
+export * from './context-store';

packages/@aws-cdk/toolkit-lib/lib/api/cloud-assembly/private/context-aware-source.ts

@@ -2,9 +2,9 @@ import type { MissingContext } from '@aws-cdk/cloud-assembly-schema';
 import * as contextproviders from '../../../context-providers';
 import type { ToolkitServices } from '../../../toolkit/private';
 import { ToolkitError } from '../../../toolkit/toolkit-error';
-import { PROJECT_CONTEXT, type Context } from '../../context';
 import type { IoHelper } from '../../io/private';
 import { IO } from '../../io/private';
+import type { IContextStore } from '../context-store';
 import type { ICloudAssemblySource, IReadableCloudAssembly } from '../types';
 
 export interface ContextAwareCloudAssemblyProps {
@@ -15,16 +15,9 @@ export interface ContextAwareCloudAssemblyProps {
   readonly services: ToolkitServices;
 
   /**
-   * Application context
+   * Location to read and write context
    */
-  readonly context: Context;
-
-  /**
-   * The file used to store application context in (relative to cwd).
-   *
-   * @default "cdk.context.json"
-   */
-  readonly contextFile?: string;
+  readonly contextStore: IContextStore;
 
   /**
    * Enable context lookups.
@@ -55,14 +48,12 @@ export interface ContextAwareCloudAssemblyProps {
  */
 export class ContextAwareCloudAssemblySource implements ICloudAssemblySource {
   private canLookup: boolean;
-  private context: Context;
-  private contextFile: string;
+  private context: IContextStore;
   private ioHelper: IoHelper;
 
   constructor(private readonly source: ICloudAssemblySource, private readonly props: ContextAwareCloudAssemblyProps) {
     this.canLookup = props.lookups ?? true;
-    this.context = props.context;
-    this.contextFile = props.contextFile ?? PROJECT_CONTEXT; // @todo new feature not needed right now
+    this.context = props.contextStore;
     this.ioHelper = props.services.ioHelper;
   }
 
@@ -100,20 +91,17 @@ export class ContextAwareCloudAssemblySource implements ICloudAssemblySource {
 
         if (tryLookup) {
           await this.ioHelper.notify(IO.CDK_ASSEMBLY_I0241.msg('Some context information is missing. Fetching...', { missingKeys }));
-          await contextproviders.provideContextValues(
+          const contextUpdates = await contextproviders.provideContextValues(
             assembly.manifest.missing,
-            this.context,
             this.props.services.sdkProvider,
             this.props.services.pluginHost,
             this.ioHelper,
           );
 
-          // Cache the new context to disk
-          await this.ioHelper.notify(IO.CDK_ASSEMBLY_I0042.msg(`Writing updated context to ${this.contextFile}...`, {
-            contextFile: this.contextFile,
-            context: this.context.all,
+          await this.ioHelper.notify(IO.CDK_ASSEMBLY_I0042.msg('Writing context updates...', {
+            context: contextUpdates,
           }));
-          await this.context.save(this.contextFile);
+          await this.context.update(contextUpdates);
 
           // Execute again. Unlock the assembly here so that the producer can acquire
           // a read lock on the directory again.