(aws-lambda / assets): Ability to use token in Code.fromAsset, or stage lambda code lazily

[Chriscbr]

Describe the feature

Provide lazy strings or inputs to Code.fromAsset:

import { Lazy, Stack, StackProps } from "aws-cdk-lib";
import * as lambda from "aws-cdk-lib/aws-lambda";
import { Construct } from "constructs";
import { join } from "path";

export class CdkTestStack extends Stack {
  constructor(scope: Construct, id: string, props?: StackProps) {
    super(scope, id, props);

    let path = Lazy.string({ produce: () => join(__dirname, "lmb") });
    new lambda.Function(this, "CdkTestFunction1", {
      code: lambda.Code.fromAsset(path),
      handler: "hello.handler",
      runtime: lambda.Runtime.NODEJS_18_X,
    });
  }
}

Use Case

In some situations, I might not have all of the lambda code available for my CDK app available before I run cdk synth or cdk deploy. For example, in Winglang (language that uses CDK constructs under the hood) we're trying to generate some lambda code dynamically based on the graph of resources in the construct tree. In this case, I won't know the exact contents of my lambda code until the entire tree is constructed (but we can assume the contents will be known before app.synth() is called).

Proposed Solution

No response

Other Information

I tried extending the abstract Code class and implementing the bind() method, but looking at the existing implementation of AssetCode here, I couldn't figure out a way to instantiate s3_assets.Asset lazily. Even if I use the bundling option from the assets module described here, it seems like the bundling command gets invoked while the tree is being created (in other words, in the middle of the evaluation of new MyCdkStack()), not afterwards.

Acknowledgements

• I may be able to implement this feature request
• This feature might incur a breaking change

CDK version used

2.92.0

Environment details (OS name and version, etc.)

macOS

[pahud]

Makes sense. Thank you for the use case sharing.

[martinjlowm]

I've hacked together a temporary solution to this (although it doesn't operate directly on fromAsset). As mentioned in #28717, my wish was to incorporate lazy builds with Nix and I've come to a point where I'm almost using the CDK's API as it was meant to.

The setup consists of:

1. a custom App-wide synthesizer
2. an aspect to traverse assets
3. a single build command once the tree has been traversed once

It looks like the following,

Aspect (bundler)

export class NixBundler implements IAspect {
  built = false;

  assets: LazyS3Asset[] = [];

  public visit(node: IConstruct): void {
    if (node instanceof CfnFunction) {
      const code = node.code as CfnFunction.CodeProperty;
      if (Tokenization.isResolvable(code.s3Bucket) && Tokenization.isResolvable(code.s3Key)) {
        // We have to store the asset reference as a property on each key,
        // because the asset construct is an intermediate and is not directly
        // exposed in the visitor
        this.assets.push((code.s3Bucket as unknown as LazyS3AssetProperty).asset);
      }
    }
  }

  build() {
    if (this.built) {
      return;
    }

    if (!this.assets.length) {
      return;
    }

    // https://github.com/NixOS/nix/issues/9042 - let's assume there are no duplicate entries
    const result = execSync(`nix build ${this.assets.map((asset) => asset.installable).join(' ')} --impure --print-out-paths`);

    const outputs = result.toString().trim().split('\n');
    for (const [asset, output] of zip(this.assets, outputs)) {
      if (!(asset && output)) {
        continue;
      }

      asset.sourcePath = output;
    }

    this.built = true;
  }
}

LazyAssetCode

export class LazyS3AssetProperty implements IResolvable {
  private _property?: string;

  public constructor(public readonly asset: LazyS3Asset) {
  }

  creationStack: string[] = [];
  typeHint?: ResolutionTypeHint | undefined;

  public resolve() {
    return this._property;
  }

  public resolveTo(value: string) {
    this._property = value;
  }

  public toString(): string {
    return Token.asString(this);
  }
}

export class LazyS3Asset {
  public bucketName = new LazyS3AssetProperty(this);
  public objectKey = new LazyS3AssetProperty(this);

  set sourcePath(v: string) {
    const asset = this.fn(v);

    // Not super happy with this - we're essentially manually resolving the tokens
    this.bucketName.resolveTo(asset.s3BucketName);
    this.objectKey.resolveTo(asset.s3ObjectKey);
  }

  public static from(installable: string, fn: (sourcePath: string) => s3_assets.Asset) {
    return new LazyS3Asset(installable, fn);
  }

  private constructor(public readonly installable: string, private readonly fn: (sourcePath: string) => s3_assets.Asset) {
  }
}

/**
 * (Nix) Lambda code from a local directory.
 */
export class LazyAssetCode extends Code {
  private asset?: s3_assets.Asset;
  public sourcePath?: string;

  constructor(private readonly installable: string, private readonly options: s3_assets.AssetOptions = { }) {
    super();
  }

  public bind(scope: Construct): CodeConfig {
    const s3Location = LazyS3Asset.from(this.installable, (sourcePath: string) => {
      this.asset = new s3_assets.Asset(scope, 'Code', {
        path: sourcePath,
        deployTime: true,
      });

      return this.asset;
    });

    return {
      s3Location: {
        bucketName: s3Location.bucketName as unknown as string,
        objectKey: s3Location.objectKey as unknown as string,
      },
    };
  }

  public bindToResource(resource: CfnResource, options: ResourceBindOptions = { }) {
    if (!this.asset) {
      return;
    }

    const resourceProperty = options.resourceProperty || 'Code';

    // https://github.com/aws/aws-cdk/issues/1432
    this.asset.addResourceMetadata(resource, resourceProperty);
  }
}

Wrapped App

class DefaultStackBundlerSynthesizer extends DefaultStackSynthesizer {
  constructor(public readonly bundler: NixBundler) {
    super();
  }

  synthesize(session: ISynthesisSession): void {
    // Collects the asset-constructs
    this.bundler.build();
    super.synthesize(session);
  }
}

export class CustomApp extends App {
  constructor(public readonly bundler: NixBundler = new NixBundler()) {
    super({ defaultStackSynthesizer: new DefaultStackBundlerSynthesizer(bundler)});

    // Traverses the tree
    Aspects.of(this).add(this.bundler);
  }
}

Usage

const func = new Function(this, 'entry', {
  runtime: lambda.Runtime.NODEJS_18_X,
  handler: 'hello.handler',
  code: new LazyAssetCode('nixpkgs#hello'),
});

With this, Nix is invoked with all the installables once the tree has been traversed and blocks until all the assets are available. It either utilizes the binary cache (best case) or forwards the work to an arbitrary number of builders as is configured by the system.