[POC] Introduce new stack structure

packages/backend-ai/src/conversation/factory.ts

@@ -16,6 +16,7 @@ import {
 import path from 'path';
 import { CallerDirectoryExtractor } from '@aws-amplify/platform-core';
 import { AiModel } from '@aws-amplify/data-schema-types';
+import { Stack } from 'aws-cdk-lib';
 
 class ConversationHandlerFunctionGenerator
   implements ConstructContainerEntryGenerator
@@ -72,11 +73,14 @@ class DefaultConversationHandlerFunctionFactory
     private readonly callerStack: string | undefined
   ) {}
 
-  getInstance = ({
-    constructContainer,
-    outputStorageStrategy,
-    resourceNameValidator,
-  }: ConstructFactoryGetInstanceProps): ConversationHandlerFunction => {
+  getInstance = (
+    {
+      constructContainer,
+      outputStorageStrategy,
+      resourceNameValidator,
+    }: ConstructFactoryGetInstanceProps,
+    stack?: Stack
+  ): ConversationHandlerFunction => {
     resourceNameValidator?.validate(this.props.name);
     if (!this.generator) {
       const props = { ...this.props };
@@ -87,7 +91,8 @@ class DefaultConversationHandlerFunctionFactory
       );
     }
     return constructContainer.getOrCompute(
-      this.generator
+      this.generator,
+      stack
     ) as ConversationHandlerFunction;
   };
 

packages/backend-auth/src/factory.ts

@@ -98,7 +98,8 @@ export class AmplifyAuthFactory implements ConstructFactory<BackendAuth> {
    * Get a singleton instance of AmplifyAuth
    */
   getInstance = (
-    getInstanceProps: ConstructFactoryGetInstanceProps
+    getInstanceProps: ConstructFactoryGetInstanceProps,
+    stack?: Stack
   ): BackendAuth => {
     const { constructContainer, importPathVerifier, resourceNameValidator } =
       getInstanceProps;
@@ -111,9 +112,16 @@ export class AmplifyAuthFactory implements ConstructFactory<BackendAuth> {
       resourceNameValidator?.validate(this.props.name);
     }
     if (!this.generator) {
-      this.generator = new AmplifyAuthGenerator(this.props, getInstanceProps);
+      this.generator = new AmplifyAuthGenerator(
+        this.props,
+        getInstanceProps,
+        stack
+      );
     }
-    return constructContainer.getOrCompute(this.generator) as BackendAuth;
+    return constructContainer.getOrCompute(
+      this.generator,
+      stack
+    ) as BackendAuth;
   };
 }
 
@@ -124,6 +132,7 @@ class AmplifyAuthGenerator implements ConstructContainerEntryGenerator {
   constructor(
     private readonly props: AmplifyAuthProps,
     private readonly getInstanceProps: ConstructFactoryGetInstanceProps,
+    private readonly stack?: Stack,
     private readonly authAccessBuilder = _authAccessBuilder,
     private readonly authAccessPolicyArbiterFactory = new AuthAccessPolicyArbiterFactory()
   ) {
@@ -169,7 +178,8 @@ class AmplifyAuthGenerator implements ConstructContainerEntryGenerator {
       ([triggerEvent, handlerFactory]) => {
         (authConstruct.resources.userPool as UserPool).addTrigger(
           UserPoolOperation.of(triggerEvent),
-          handlerFactory.getInstance(this.getInstanceProps).resources.lambda
+          handlerFactory.getInstance(this.getInstanceProps, this.stack)
+            .resources.lambda
         );
       }
     );

packages/backend-data/src/factory.ts

@@ -38,7 +38,7 @@ import {
   CDKContextKey,
   TagName,
 } from '@aws-amplify/platform-core';
-import { Aspects, IAspect, Tags } from 'aws-cdk-lib';
+import { Aspects, IAspect, Stack, Tags } from 'aws-cdk-lib';
 import { convertJsResolverDefinition } from './convert_js_resolvers.js';
 import { AppSyncPolicyGenerator } from './app_sync_policy_generator.js';
 import {
@@ -77,7 +77,10 @@ export class DataFactory implements ConstructFactory<AmplifyData> {
   /**
    * Gets an instance of the Data construct
    */
-  getInstance = (props: ConstructFactoryGetInstanceProps): AmplifyData => {
+  getInstance = (
+    props: ConstructFactoryGetInstanceProps,
+    stack?: Stack
+  ): AmplifyData => {
     const {
       constructContainer,
       outputStorageStrategy,
@@ -106,7 +109,10 @@ export class DataFactory implements ConstructFactory<AmplifyData> {
         outputStorageStrategy
       );
     }
-    return constructContainer.getOrCompute(this.generator) as AmplifyData;
+    return constructContainer.getOrCompute(
+      this.generator,
+      stack
+    ) as AmplifyData;
   };
 }
 

packages/backend-function/src/factory.ts

@@ -177,18 +177,24 @@ class FunctionFactory implements ConstructFactory<AmplifyFunction> {
   /**
    * Creates an instance of AmplifyFunction within the provided Amplify context
    */
-  getInstance = ({
-    constructContainer,
-    outputStorageStrategy,
-    resourceNameValidator,
-  }: ConstructFactoryGetInstanceProps): AmplifyFunction => {
+  getInstance = (
+    {
+      constructContainer,
+      outputStorageStrategy,
+      resourceNameValidator,
+    }: ConstructFactoryGetInstanceProps,
+    stack?: Stack
+  ): AmplifyFunction => {
     if (!this.generator) {
       this.generator = new FunctionGenerator(
         this.hydrateDefaults(resourceNameValidator),
         outputStorageStrategy
       );
     }
-    return constructContainer.getOrCompute(this.generator) as AmplifyFunction;
+    return constructContainer.getOrCompute(
+      this.generator,
+      stack
+    ) as AmplifyFunction;
   };
 
   private hydrateDefaults = (

packages/backend-storage/src/factory.ts

@@ -32,7 +32,8 @@ export class AmplifyStorageFactory
    * Get a singleton instance of the Bucket
    */
   getInstance = (
-    getInstanceProps: ConstructFactoryGetInstanceProps
+    getInstanceProps: ConstructFactoryGetInstanceProps,
+    stack?: Stack
   ): AmplifyStorage => {
     const { constructContainer, importPathVerifier, resourceNameValidator } =
       getInstanceProps;
@@ -50,7 +51,8 @@ export class AmplifyStorageFactory
       );
     }
     const amplifyStorage = constructContainer.getOrCompute(
-      this.generator
+      this.generator,
+      stack
     ) as AmplifyStorage;
 
     /*

packages/backend/API.md

@@ -34,6 +34,16 @@ import { Stack } from 'aws-cdk-lib';
 
 export { a }
 
+// @public (undocumented)
+export type AmplifyStackBase = {
+    readonly stack: Stack;
+};
+
+// @public (undocumented)
+export type AmplifyStackResources = AmplifyStackBase & {
+    [K in keyof DefineStackProps]: Omit<ReturnType<DefineStackProps[K]['getInstance']>, keyof ResourceAccessAcceptorFactory>;
+};
+
 export { AuthCfnResources }
 
 export { AuthResources }
@@ -82,6 +92,14 @@ export { defineData }
 
 export { defineFunction }
 
+// @public
+export const defineStack: (name: string, constructFactories: DefineStackProps) => ConstructFactory<ResourceProvider<AmplifyStackResources>>;
+
+// @public (undocumented)
+export type DefineStackProps = Record<string, ConstructFactory<ResourceProvider & Partial<ResourceAccessAcceptorFactory<never>>>> & {
+    [K in keyof AmplifyStackBase]?: never;
+};
+
 export { defineStorage }
 
 export { FunctionResources }

packages/backend/src/engine/singleton_construct_container.ts

@@ -9,6 +9,7 @@ import { getBackendIdentifier } from '../backend_identifier.js';
 import { DefaultBackendSecretResolver } from './backend-secret/backend_secret_resolver.js';
 import { BackendIdScopedSsmEnvironmentEntriesGenerator } from './backend_id_scoped_ssm_environment_entries_generator.js';
 import { BackendIdScopedStableBackendIdentifiers } from '../backend_id_scoped_stable_backend_identifiers.js';
+import { Stack } from 'aws-cdk-lib';
 
 /**
  * Serves as a DI container and shared state store for initializing Amplify constructs
@@ -33,10 +34,13 @@ export class SingletonConstructContainer implements ConstructContainer {
    * Otherwise, the generator is called and the value is cached and returned
    */
   getOrCompute = (
-    generator: ConstructContainerEntryGenerator
+    generator: ConstructContainerEntryGenerator,
+    scope?: Stack
   ): ResourceProvider => {
     if (!this.providerCache.has(generator)) {
-      const scope = this.stackResolver.getStackFor(generator.resourceGroupName);
+      if (!scope) {
+        scope = this.stackResolver.getStackFor(generator.resourceGroupName);
+      }
       const backendId = getBackendIdentifier(scope);
       const ssmEnvironmentEntriesGenerator =
         new BackendIdScopedSsmEnvironmentEntriesGenerator(scope, backendId);

packages/backend/src/index.ts

@@ -24,3 +24,6 @@ export { defineStorage } from '@aws-amplify/backend-storage';
 
 // function
 export { defineFunction } from '@aws-amplify/backend-function';
+
+// stack
+export * from './stack_factory.js';

packages/backend/src/stack_factory.ts

@@ -0,0 +1,126 @@
+import {
+  ConstructContainerEntryGenerator,
+  ConstructFactory,
+  ConstructFactoryGetInstanceProps,
+  GenerateContainerEntryProps,
+  ResourceAccessAcceptorFactory,
+  ResourceProvider,
+} from '@aws-amplify/plugin-types';
+import { Stack } from 'aws-cdk-lib';
+
+export type AmplifyStackBase = {
+  readonly stack: Stack;
+};
+
+export type DefineStackProps = Record<
+  string,
+  ConstructFactory<
+    ResourceProvider & Partial<ResourceAccessAcceptorFactory<never>>
+  >
+> & { [K in keyof AmplifyStackBase]?: never };
+
+export type AmplifyStackResources<T extends DefineStackProps> =
+  AmplifyStackBase & {
+    [K in keyof T]: Omit<
+      ReturnType<T[K]['getInstance']>,
+      keyof ResourceAccessAcceptorFactory
+    >;
+  };
+
+class StackGenerator<
+  T extends Record<string, ConstructFactory<ResourceProvider>>
+> implements ConstructContainerEntryGenerator<AmplifyStackResources<T>>
+{
+  readonly resourceGroupName: string;
+
+  constructor(
+    name: string,
+    readonly getInstanceProps: ConstructFactoryGetInstanceProps,
+    readonly constructFactories: T
+  ) {
+    this.resourceGroupName = name;
+  }
+
+  generateContainerEntry = ({
+    scope,
+  }: GenerateContainerEntryProps): ResourceProvider<
+    AmplifyStackResources<T>
+  > => {
+    // register providers but don't actually execute anything yet
+    Object.values(this.constructFactories).forEach((factory) => {
+      if (typeof factory.provides === 'string') {
+        this.getInstanceProps.constructContainer.registerConstructFactory(
+          factory.provides,
+          factory
+        );
+      }
+    });
+    // now invoke all the factories and collect the constructs into nestedResources
+    const nestedResources: {
+      [K in keyof Record<
+        string,
+        ConstructFactory<ResourceProvider>
+      >]: ReturnType<
+        Record<string, ConstructFactory<ResourceProvider>>[K]['getInstance']
+      >;
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    } = {} as any;
+    Object.entries(this.constructFactories).forEach(
+      ([resourceName, constructFactory]) => {
+        // The type inference on this.resources is not happy about this assignment because it doesn't know the exact type of .getInstance()
+        // However, the assignment is okay because we are iterating over the entries of constructFactories and assigning the resource name to the corresponding instance
+        nestedResources[
+          resourceName as keyof Record<
+            string,
+            ConstructFactory<ResourceProvider>
+          >
+        ] = constructFactory.getInstance(
+          this.getInstanceProps,
+          scope as Stack
+          // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        ) as any;
+      }
+    );
+    return {
+      resources: {
+        ...nestedResources,
+        stack: scope as Stack,
+      } as AmplifyStackResources<T>,
+    };
+  };
+}
+
+class StackFactory<T extends DefineStackProps>
+  implements ConstructFactory<ResourceProvider<AmplifyStackResources<T>>>
+{
+  private generator: ConstructContainerEntryGenerator;
+
+  constructor(
+    private readonly name: string,
+    private readonly constructFactories: T
+  ) {}
+
+  getInstance(
+    props: ConstructFactoryGetInstanceProps
+  ): ResourceProvider<AmplifyStackResources<T>> {
+    if (!this.generator) {
+      this.generator = new StackGenerator(
+        this.name,
+        props,
+        this.constructFactories
+      );
+    }
+    return props.constructContainer.getOrCompute(
+      this.generator
+    ) as ResourceProvider<AmplifyStackResources<T>>;
+  }
+}
+
+/**
+ * TODO
+ */
+export const defineStack = <T extends DefineStackProps>(
+  name: string,
+  constructFactories: T
+): ConstructFactory<ResourceProvider<AmplifyStackResources<T>>> =>
+  new StackFactory(name, constructFactories);

packages/plugin-types/API.md

@@ -110,7 +110,7 @@ export type BranchName = string;
 
 // @public
 export type ConstructContainer = {
-    getOrCompute: (generator: ConstructContainerEntryGenerator) => ResourceProvider;
+    getOrCompute: (generator: ConstructContainerEntryGenerator, stack?: Stack) => ResourceProvider;
     registerConstructFactory: (token: string, provider: ConstructFactory) => void;
     getConstructFactory: <T extends ResourceProvider>(token: string) => ConstructFactory<T> | undefined;
 };
@@ -124,7 +124,7 @@ export type ConstructContainerEntryGenerator<T extends object = object> = {
 // @public
 export type ConstructFactory<T extends ResourceProvider = ResourceProvider> = {
     readonly provides?: string;
-    getInstance: (props: ConstructFactoryGetInstanceProps) => T;
+    getInstance: (props: ConstructFactoryGetInstanceProps, stack?: Stack) => T;
 };
 
 // @public (undocumented)

packages/plugin-types/src/construct_container.ts

@@ -4,6 +4,7 @@ import { BackendSecretResolver } from './backend_secret_resolver.js';
 import { ResourceProvider } from './resource_provider.js';
 import { SsmEnvironmentEntriesGenerator } from './ssm_environment_entries_generator.js';
 import { StableBackendIdentifiers } from './stable_backend_identifiers.js';
+import { Stack } from 'aws-cdk-lib';
 /**
  * Initializes a CDK Construct in a given scope
  */
@@ -34,7 +35,8 @@ export type GenerateContainerEntryProps = {
  */
 export type ConstructContainer = {
   getOrCompute: (
-    generator: ConstructContainerEntryGenerator
+    generator: ConstructContainerEntryGenerator,
+    stack?: Stack
   ) => ResourceProvider;
   registerConstructFactory: (token: string, provider: ConstructFactory) => void;
   getConstructFactory: <T extends ResourceProvider>(