[POC] Use single stack for top level functions

packages/backend/API.md

@@ -70,7 +70,12 @@ export { ConstructFactoryGetInstanceProps }
 export { defineAuth }
 
 // @public
-export const defineBackend: <T extends DefineBackendProps>(constructFactories: T) => Backend<T>;
+export const defineBackend: <T extends DefineBackendProps>(constructFactories: T, options?: DefineBackendOptions) => Backend<T>;
+
+// @public (undocumented)
+export type DefineBackendOptions = {
+    useSingleStack?: boolean;
+};
 
 // @public (undocumented)
 export type DefineBackendProps = Record<string, ConstructFactory<ResourceProvider & Partial<ResourceAccessAcceptorFactory<never>>>> & {

packages/backend/src/backend.ts

@@ -33,3 +33,8 @@ export type Backend<T extends DefineBackendProps> = BackendBase & {
     keyof ResourceAccessAcceptorFactory
   >;
 };
+
+export type DefineBackendOptions = {
+  // TODO This is not the best name given that there are still nested stacks in the system.
+  useSingleStack?: boolean;
+};

packages/backend/src/backend_factory.test.ts

@@ -57,6 +57,7 @@ void describe('Backend', () => {
       {
         testConstructFactory,
       },
+      undefined,
       rootStack
     );
 
@@ -96,6 +97,7 @@ void describe('Backend', () => {
       {
         testConstructFactory,
       },
+      undefined,
       rootStack
     );
 
@@ -140,6 +142,7 @@ void describe('Backend', () => {
       {
         testConstructFactory,
       },
+      undefined,
       rootStack
     );
     assert.equal(
@@ -149,7 +152,7 @@ void describe('Backend', () => {
   });
 
   void it('stores attribution metadata in root stack', () => {
-    new BackendFactory({}, rootStack);
+    new BackendFactory({}, undefined, rootStack);
     const rootStackTemplate = Template.fromStack(rootStack);
     assert.equal(
       JSON.parse(rootStackTemplate.toJSON().Description).stackType,
@@ -158,27 +161,27 @@ void describe('Backend', () => {
   });
 
   void it('registers branch linker for branch deployments', () => {
-    new BackendFactory({}, rootStack);
+    new BackendFactory({}, undefined, rootStack);
     const rootStackTemplate = Template.fromStack(rootStack);
     rootStackTemplate.resourceCountIs('Custom::AmplifyBranchLinkerResource', 1);
   });
 
   void it('does not register branch linker for sandbox deployments', () => {
     const rootStack = createStackAndSetContext('sandbox');
-    new BackendFactory({}, rootStack);
+    new BackendFactory({}, undefined, rootStack);
     const rootStackTemplate = Template.fromStack(rootStack);
     rootStackTemplate.resourceCountIs('Custom::AmplifyBranchLinkerResource', 0);
   });
 
   void describe('createStack', () => {
     void it('returns nested stack', () => {
-      const backend = new BackendFactory({}, rootStack);
+      const backend = new BackendFactory({}, undefined, rootStack);
       const testStack = backend.createStack('testStack');
       assert.strictEqual(rootStack.node.findChild('testStack'), testStack);
     });
 
     void it('throws if stack has already been created with specified name', () => {
-      const backend = new BackendFactory({}, rootStack);
+      const backend = new BackendFactory({}, undefined, rootStack);
       backend.createStack('testStack');
       assert.throws(() => backend.createStack('testStack'), {
         message: 'Custom stack named testStack has already been created',
@@ -188,7 +191,7 @@ void describe('Backend', () => {
 
   void it('can add custom output', () => {
     const rootStack = createStackAndSetContext('sandbox');
-    const backend = new BackendFactory({}, rootStack);
+    const backend = new BackendFactory({}, undefined, rootStack);
     const clientConfigPartial: DeepPartialAmplifyGeneratedConfigs<ClientConfig> =
       {
         version: '1.1',

packages/backend/src/backend_factory.ts

@@ -6,6 +6,7 @@ import {
 import { Stack } from 'aws-cdk-lib';
 import {
   NestedStackResolver,
+  SingleStackResolver,
   StackResolver,
 } from './engine/nested_stack_resolver.js';
 import { SingletonConstructContainer } from './engine/singleton_construct_container.js';
@@ -18,7 +19,11 @@ import { createDefaultStack } from './default_stack_factory.js';
 import { getBackendIdentifier } from './backend_identifier.js';
 import { platformOutputKey } from '@aws-amplify/backend-output-schemas';
 import { fileURLToPath } from 'node:url';
-import { Backend, DefineBackendProps } from './backend.js';
+import {
+  Backend,
+  DefineBackendOptions,
+  DefineBackendProps,
+} from './backend.js';
 import { AmplifyBranchLinkerConstruct } from './engine/branch-linker/branch_linker_construct.js';
 import {
   ClientConfig,
@@ -55,16 +60,27 @@ export class BackendFactory<
    * Initialize an Amplify backend with the given construct factories and in the given CDK App.
    * If no CDK App is specified a new one is created
    */
-  constructor(constructFactories: T, stack: Stack = createDefaultStack()) {
+  constructor(
+    constructFactories: T,
+    options?: DefineBackendOptions,
+    stack: Stack = createDefaultStack()
+  ) {
     new AttributionMetadataStorage().storeAttributionMetadata(
       stack,
       rootStackTypeIdentifier,
       fileURLToPath(new URL('../package.json', import.meta.url))
     );
-    this.stackResolver = new NestedStackResolver(
-      stack,
-      new AttributionMetadataStorage()
-    );
+    if (options?.useSingleStack) {
+      this.stackResolver = new SingleStackResolver(
+        stack,
+        new AttributionMetadataStorage()
+      );
+    } else {
+      this.stackResolver = new NestedStackResolver(
+        stack,
+        new AttributionMetadataStorage()
+      );
+    }
 
     const constructContainer = new SingletonConstructContainer(
       this.stackResolver
@@ -148,11 +164,13 @@ export class BackendFactory<
 /**
  * Creates a new Amplify backend instance and returns it
  * @param constructFactories - list of backend factories such as those created by `defineAuth` or `defineData`
+ * @param options - optional options
  */
 export const defineBackend = <T extends DefineBackendProps>(
-  constructFactories: T
+  constructFactories: T,
+  options?: DefineBackendOptions
 ): Backend<T> => {
-  const backend = new BackendFactory(constructFactories);
+  const backend = new BackendFactory(constructFactories, options);
   return {
     ...backend.resources,
     createStack: backend.createStack,

packages/backend/src/engine/nested_stack_resolver.ts

@@ -55,3 +55,48 @@ export class NestedStackResolver implements StackResolver {
     return this.stacks[resourceGroupName];
   };
 }
+
+/**
+ * TODO
+ */
+export class SingleStackResolver implements StackResolver {
+  private readonly stacks: Record<string, Stack> = {};
+
+  /**
+   * Initialize with a root stack
+   */
+  constructor(
+    private readonly rootStack: Stack,
+    private readonly attributionMetadataStorage: AttributionMetadataStorage
+  ) {}
+
+  // TODO de-duplicate this logic
+  createCustomStack = (name: string): Stack => {
+    if (this.stacks[name]) {
+      throw new Error(`Custom stack named ${name} has already been created`);
+    }
+    const stack = this.getStackForInternal(name);
+    // this is safe even if stack is cached from an earlier invocation because storeAttributionMetadata is a noop if the stack description already exists
+    this.attributionMetadataStorage.storeAttributionMetadata(
+      stack,
+      `custom`,
+      fileURLToPath(new URL('../../package.json', import.meta.url))
+    );
+    return stack;
+  };
+
+  getStackFor = (): Stack => {
+    return this.rootStack;
+  };
+
+  // TODO this is a hack
+  private getStackForInternal = (resourceGroupName: string): Stack => {
+    if (!this.stacks[resourceGroupName]) {
+      this.stacks[resourceGroupName] = new NestedStack(
+        this.rootStack,
+        resourceGroupName
+      );
+    }
+    return this.stacks[resourceGroupName];
+  };
+}

test-projects/function-stack-1/amplify/backend.ts

@@ -9,6 +9,8 @@ const backend = defineBackend({
   auth,
   data,
   myApiFunction,
+}, {
+  useSingleStack: true
 });
 
 const eventSource = new DynamoEventSource(

test-projects/function-stack-2/amplify/backend.ts

@@ -8,6 +8,8 @@ const backend = defineBackend({
   auth,
   data,
   myApiFunction,
+},{
+  useSingleStack: true,
 });
 
 backend.myApiFunction.resources.lambda.addToRolePolicy(