feat(sdk, compiler): displaying logs of tf-aws functions (#3622)

fixes: #1973
## Description
* Collecting and printing `log(...)` outputs when running a test on target tf-aws

## Checklist

- [x] Title matches [Winglang's style guide](https://www.winglang.io/contributing/start-here/pull_requests#how-are-pull-request-titles-formatted)
- [x] Description explains motivation and solution
- [ ] Tests added (always)
- [x] Docs updated (only required for features)
- [ ] Added `pr/e2e-full` label if this feature requires end-to-end testing

*By submitting this pull request, I confirm that my contribution is made under the terms of the [Wing Cloud Contribution License](https://github.com/winglang/wing/blob/main/CONTRIBUTION_LICENSE.md)*.

examples/tests/sdk_tests/function/invoke.w

@@ -2,15 +2,18 @@ bring cloud;
 bring util;
 
 let payload = "hello";
+log("log preflight");
 
 let f = new cloud.Function(inflight (input: str): str => {
+  log("log inside function\ncontains 2 lines");
   let target = util.tryEnv("WING_TARGET");
   assert(target?); // make sure WING_TARGET is defined in all environments
 
   return "${input}-response";
 });
 
 test "invoke" {
+  log("log inside test");
   let x = f.invoke("hello");
   assert(x == "hello-response");
 }

examples/tests/sdk_tests/function/logging.js

@@ -0,0 +1,3 @@
+exports.logging = function () {
+  console.log("hello world");
+};

examples/tests/sdk_tests/function/logging.w

@@ -0,0 +1,35 @@
+bring cloud;
+
+class Util {
+  extern "./logging.js" static inflight logging(): void;
+}
+
+let f1 = new cloud.Function(inflight (input: str): void => {
+  log("log inside f1");
+}) as "f1";
+
+let f2 = new cloud.Function(inflight (input: str): void => {
+  f1.invoke("");
+  log("log inside f2");
+  f1.invoke("");
+}) as "f2";
+
+/**
+should log: 
+
+hello world
+log inside f1
+log inside f2
+log inside f1
+hello world
+log inside f1
+log inside f2
+log inside f1
+*/
+
+test "logging" {
+  Util.logging();
+  f2.invoke("");
+  Util.logging();
+  f2.invoke("");
+}

libs/wingsdk/.projenrc.ts

@@ -86,6 +86,7 @@ const project = new cdk.JsiiProject({
     "@aws-sdk/[email protected]",
     "@aws-sdk/[email protected]",
     "@aws-sdk/[email protected]",
+    "@types/aws-lambda",
     // the following 2 deps are required by @aws-sdk/util-utf8-node
     "@aws-sdk/[email protected]",
     "@aws-sdk/[email protected]",
@@ -107,7 +108,6 @@ const project = new cdk.JsiiProject({
     `@cdktf/provider-aws@^15.0.0`, // only for testing Wing plugins
     "wing-api-checker@workspace:^",
     "bump-pack@workspace:^",
-    "@types/aws-lambda",
     "@types/fs-extra",
     "@types/mime-types",
     "@types/express",

libs/wingsdk/src/cloud/function.ts

@@ -2,6 +2,7 @@ import { mkdirSync, writeFileSync } from "fs";
 import { join } from "path";
 import { Construct } from "constructs";
 import { fqnForType } from "../constants";
+import { Code } from "../core";
 import { App } from "../core/app";
 import { CaseConventions, ResourceNames } from "../shared/resource-names";
 import { Duration } from "../std/duration";
@@ -83,11 +84,7 @@ export abstract class Function extends Resource implements IInflightHost {
     handler._registerBind(this, ["handle", "$inflight_init"]);
 
     const inflightClient = handler._toInflight();
-    const lines = new Array<string>();
-
-    lines.push("exports.handler = async function(event) {");
-    lines.push(`  return await (${inflightClient.text}).handle(event);`);
-    lines.push("};");
+    const lines = this._generateLines(inflightClient);
 
     const assetName = ResourceNames.generateName(this, {
       // Avoid characters that may cause path issues
@@ -110,6 +107,22 @@ export abstract class Function extends Resource implements IInflightHost {
     }
   }
 
+  /**
+   * Generates the code lines for the cloud function, can be overridden by the targets
+   * @param inflightClient inflight client code
+   * @returns cloud function code string
+   * @internal
+   */
+  protected _generateLines(inflightClient: Code): string[] {
+    const lines = new Array<string>();
+
+    lines.push("exports.handler = async function(event) {");
+    lines.push(`  return await (${inflightClient.text}).handle(event);`);
+    lines.push("};");
+
+    return lines;
+  }
+
   /**
    * Add an environment variable to the function.
    */

libs/wingsdk/src/shared-aws/function.inflight.ts

@@ -1,22 +1,83 @@
 import { InvokeCommand, LambdaClient } from "@aws-sdk/client-lambda";
 import { fromUtf8, toUtf8 } from "@aws-sdk/util-utf8-node";
+import { Context } from "aws-lambda";
 import { IFunctionClient } from "../cloud";
+import { Trace } from "../std";
 
 export class FunctionClient implements IFunctionClient {
   constructor(
     private readonly functionArn: string,
+    private readonly constructPath: string,
     private readonly lambdaClient = new LambdaClient({})
   ) {}
 
   /**
-   * Invoke the function, passing the given payload as an argument.
+   * Reading the function's logs,
+   * along with any logs of a function that was called by the parent function
+   *
+   * @param logGroupName function's context logGroupName
+   * @param logStreamName function's context logGroupName
+   * @param constructPath cdk's path to construct
+   * @returns a list of Traces
    */
-  public async invoke(payload: string): Promise<string> {
+
+  private readLogs(logs: Trace[]): Trace[] {
+    const logsCollector: Trace[] = [];
+
+    for (const log of logs) {
+      const invocationLog = log.data.message?.match(/Invoking .*:/gm);
+
+      if (invocationLog) {
+        const logData = log.data.message.split("\t") ?? [];
+        const parsedLogs: Trace[] =
+          JSON.parse(
+            Buffer.from(logData[logData.length - 1], "base64").toString(
+              "binary"
+            )
+          )?.logs ?? [];
+
+        logsCollector.push(...this.readLogs(parsedLogs));
+      } else {
+        logsCollector.push(log);
+      }
+    }
+
+    return logsCollector;
+  }
+
+  /**
+   * Verify the function's return payload
+   *
+   * @returns the function's return payload, if verified
+   */
+  private verify(value: { context?: Context; payload: string }): string {
+    if (typeof value.payload !== "string") {
+      throw new Error(
+        `function returned value of type ${typeof value.payload}, not string`
+      );
+    }
+    return value.payload;
+  }
+
+  /**
+   * Invoke the function
+   */
+  private async executeFunction(payload: string): Promise<{
+    context?: Context & {
+      logs: Trace[];
+    };
+    payload: string;
+  }> {
     const command = new InvokeCommand({
       FunctionName: this.functionArn,
       Payload: fromUtf8(JSON.stringify(payload)),
+      ClientContext: Buffer.from(
+        JSON.stringify({ constructPath: this.constructPath }),
+        "binary"
+      ).toString("base64"),
     });
     const response = await this.lambdaClient.send(command);
+
     if (response.FunctionError) {
       throw new Error(
         `Invoke failed with message: "${
@@ -25,14 +86,47 @@ export class FunctionClient implements IFunctionClient {
       );
     }
     if (!response.Payload) {
-      return "";
+      return { payload: "" };
     }
     const value = JSON.parse(toUtf8(response.Payload)) ?? "";
-    if (typeof value !== "string") {
-      throw new Error(
-        `function returned value of type ${typeof value}, not string`
-      );
-    }
     return value;
   }
+
+  /**
+   * Invoke the function, passing the given payload as an argument.
+   *  @returns the function returned payload only
+   */
+  public async invoke(payload: string): Promise<string> {
+    const value = await this.executeFunction(payload);
+    const functionName = value?.context?.functionName;
+
+    // kind of hacky, but this is the most convenient way to pass those arguments to the calling function
+    console.log(
+      `Invoking ${functionName}:\t${Buffer.from(
+        JSON.stringify({
+          logs: value?.context?.logs,
+        }),
+        "binary"
+      ).toString("base64")}`
+    );
+
+    return this.verify(value);
+  }
+
+  /**
+   * Invoke the function, passing the given payload as an argument.
+   *
+   * @returns the function returned payload and logs
+   */
+  public async invokeWithLogs(payload: string): Promise<[string, Trace[]]> {
+    const traces: Trace[] = [];
+
+    const value = await this.executeFunction(payload);
+
+    if (value.context?.logs) {
+      traces.push(...this.readLogs(value.context.logs));
+    }
+
+    return [this.verify(value), traces];
+  }
 }

libs/wingsdk/src/shared-aws/function.ts

@@ -1,5 +1,7 @@
 import { PolicyStatement } from "./types";
-import { IInflightHost } from "../std";
+import { Code } from "../core";
+import { IInflightHost, TraceType } from "../std";
+import { FUNCTION_TYPE } from "../target-sim/schema-resources";
 import { Function as TfAwsFunction } from "../target-tf-aws";
 
 /**
@@ -37,3 +39,41 @@ export class Function {
     return undefined;
   }
 }
+
+/**
+ * Generates the code lines for the cloud function,
+ * overridden by the aws targets to have the function context too,
+ * as well as collecting the logs as Traces and keeping them in context.logs for later use.
+ * Eventually, this enables us displaying user defined logs, called in aws lambdas,
+ * in the user's terminal while testing.
+ * @param inflightClient inflight client code
+ * @returns cloud function code string
+ * @internal
+ */
+export function _generateAwsFunctionLines(inflightClient: Code): string[] {
+  const lines = new Array<string>();
+
+  lines.push("exports.handler = async function(event, context) {");
+  lines.push(`
+const $originalLog = console.log.bind({});
+
+console.log = (...args) => {
+const logs = args.map(item => ({
+  data: { message: item },
+  sourceType: "${FUNCTION_TYPE}",
+  sourcePath: context?.clientContext?.constructPath,
+  type: "${TraceType.LOG}",
+  timestamp: new Date().toISOString()
+}));
+!context.logs ? context.logs = [...logs] : context.logs.push(...logs);
+
+$originalLog(args);
+};`);
+
+  lines.push(
+    `  return { payload: (await (${inflightClient.text}).handle(event)) ?? "", context };`
+  );
+  lines.push("};");
+
+  return lines;
+}

libs/wingsdk/src/shared-aws/test-runner.inflight.ts

@@ -1,5 +1,5 @@
 import { FunctionClient } from "./function.inflight";
-import { ITestRunnerClient, TestResult } from "../std";
+import { ITestRunnerClient, TestResult, Trace } from "../std";
 
 export class TestRunnerClient implements ITestRunnerClient {
   // A map from test names to their corresponding function ARNs.
@@ -24,11 +24,14 @@ export class TestRunnerClient implements ITestRunnerClient {
     if (!functionArn) {
       throw new Error(`No test found with path "${path}"`);
     }
-    const client = new FunctionClient(functionArn);
+    const client = new FunctionClient(functionArn, path);
+    let traces: Trace[] = [];
     let pass = false;
     let error: string | undefined;
+
     try {
-      await client.invoke("");
+      const [_, functionTraces] = await client.invokeWithLogs("");
+      traces.push(...functionTraces);
       pass = true;
     } catch (e) {
       error = (e as any).stack;
@@ -37,7 +40,7 @@ export class TestRunnerClient implements ITestRunnerClient {
       path,
       pass,
       error,
-      traces: [], // TODO: https://github.com/winglang/wing/issues/1973
+      traces,
     };
   }
 }