chore(examples): update api-websocket-lambda-dynamodb example

examples/api-websocket-lambda-dynamodb/index.ts

@@ -1,26 +1,21 @@
 import * as pulumicdk from '@pulumi/cdk';
 import { AssetCode, Function, Runtime } from 'aws-cdk-lib/aws-lambda';
-import { CfnApi, CfnDeployment, CfnIntegration, CfnRoute, CfnStage } from 'aws-cdk-lib/aws-apigatewayv2';
-import { App, Duration, RemovalPolicy, Stack, StackProps } from 'aws-cdk-lib';
-import { Effect, PolicyStatement, Role, ServicePrincipal } from 'aws-cdk-lib/aws-iam';
+import { WebSocketApi, WebSocketStage } from 'aws-cdk-lib/aws-apigatewayv2';
+import { Duration, RemovalPolicy } from 'aws-cdk-lib';
 import { AttributeType, Table } from 'aws-cdk-lib/aws-dynamodb';
-import { Construct } from 'constructs';
-
-const config = {
-    stage: 'dev',
-    region: 'ap-southeast-2',
-    account_id: '',
-};
+import { WebSocketLambdaIntegration } from 'aws-cdk-lib/aws-apigatewayv2-integrations';
+import { Output } from '@pulumi/pulumi';
+import { LogGroup, RetentionDays } from 'aws-cdk-lib/aws-logs';
 
 class ChatAppStack extends pulumicdk.Stack {
+    public readonly url: Output<string>;
+    public readonly table: Output<string>;
     constructor(id: string) {
         super(id);
 
         // initialise api
         const name = id + '-api';
-        const api = new CfnApi(this, name, {
-            name: 'ChatAppApi',
-            protocolType: 'WEBSOCKET',
+        const api = new WebSocketApi(this, name, {
             routeSelectionExpression: '$request.body.action',
         });
         const table = new Table(this, `${name}-table`, {
@@ -33,10 +28,15 @@ class ChatAppStack extends pulumicdk.Stack {
             removalPolicy: RemovalPolicy.DESTROY,
         });
 
+        const logs = new LogGroup(this, 'websocket-lambda-logs', {
+            removalPolicy: RemovalPolicy.DESTROY,
+            retention: RetentionDays.ONE_DAY,
+        });
         const connectFunc = new Function(this, 'connect-lambda', {
+            logGroup: logs,
             code: new AssetCode('./onconnect'),
             handler: 'app.handler',
-            runtime: Runtime.NODEJS_16_X,
+            runtime: Runtime.NODEJS_LATEST,
             timeout: Duration.seconds(300),
             memorySize: 256,
             environment: {
@@ -47,9 +47,10 @@ class ChatAppStack extends pulumicdk.Stack {
         table.grantReadWriteData(connectFunc);
 
         const disconnectFunc = new Function(this, 'disconnect-lambda', {
+            logGroup: logs,
             code: new AssetCode('./ondisconnect'),
             handler: 'app.handler',
-            runtime: Runtime.NODEJS_16_X,
+            runtime: Runtime.NODEJS_LATEST,
             timeout: Duration.seconds(300),
             memorySize: 256,
             environment: {
@@ -60,113 +61,43 @@ class ChatAppStack extends pulumicdk.Stack {
         table.grantReadWriteData(disconnectFunc);
 
         const messageFunc = new Function(this, 'message-lambda', {
+            logGroup: logs,
             code: new AssetCode('./sendmessage'),
             handler: 'app.handler',
-            runtime: Runtime.NODEJS_16_X,
+            runtime: Runtime.NODEJS_LATEST,
             timeout: Duration.seconds(300),
             memorySize: 256,
-            initialPolicy: [
-                new PolicyStatement({
-                    actions: ['execute-api:ManageConnections'],
-                    resources: [
-                        'arn:aws:execute-api:' + config['region'] + ':' + config['account_id'] + ':' + api.ref + '/*',
-                    ],
-                    effect: Effect.ALLOW,
-                }),
-            ],
             environment: {
                 TABLE_NAME: table.tableName,
             },
         });
+        api.grantManageConnections(messageFunc);
 
         table.grantReadWriteData(messageFunc);
 
-        // access role for the socket api to access the socket lambda
-        const policy = new PolicyStatement({
-            effect: Effect.ALLOW,
-            resources: [connectFunc.functionArn, disconnectFunc.functionArn, messageFunc.functionArn],
-            actions: ['lambda:InvokeFunction'],
-        });
-
-        const role = new Role(this, `${name}-iam-role`, {
-            assumedBy: new ServicePrincipal('apigateway.amazonaws.com'),
-        });
-        role.addToPolicy(policy);
-
-        // lambda integration
-        const connectIntegration = new CfnIntegration(this, 'connect-lambda-integration', {
-            apiId: api.ref,
-            integrationType: 'AWS_PROXY',
-            integrationUri:
-                'arn:aws:apigateway:' +
-                config['region'] +
-                ':lambda:path/2015-03-31/functions/' +
-                connectFunc.functionArn +
-                '/invocations',
-            credentialsArn: role.roleArn,
-        });
-        const disconnectIntegration = new CfnIntegration(this, 'disconnect-lambda-integration', {
-            apiId: api.ref,
-            integrationType: 'AWS_PROXY',
-            integrationUri:
-                'arn:aws:apigateway:' +
-                config['region'] +
-                ':lambda:path/2015-03-31/functions/' +
-                disconnectFunc.functionArn +
-                '/invocations',
-            credentialsArn: role.roleArn,
+        api.addRoute('$connect', {
+            integration: new WebSocketLambdaIntegration('connect-lambda', connectFunc),
         });
-        const messageIntegration = new CfnIntegration(this, 'message-lambda-integration', {
-            apiId: api.ref,
-            integrationType: 'AWS_PROXY',
-            integrationUri:
-                'arn:aws:apigateway:' +
-                config['region'] +
-                ':lambda:path/2015-03-31/functions/' +
-                messageFunc.functionArn +
-                '/invocations',
-            credentialsArn: role.roleArn,
+        api.addRoute('$disconnect', {
+            integration: new WebSocketLambdaIntegration('disconnect-lambda', disconnectFunc),
         });
-
-        const connectRoute = new CfnRoute(this, 'connect-route', {
-            apiId: api.ref,
-            routeKey: '$connect',
-            authorizationType: 'NONE',
-            target: 'integrations/' + connectIntegration.ref,
-        });
-
-        const disconnectRoute = new CfnRoute(this, 'disconnect-route', {
-            apiId: api.ref,
-            routeKey: '$disconnect',
-            authorizationType: 'NONE',
-            target: 'integrations/' + disconnectIntegration.ref,
-        });
-
-        const messageRoute = new CfnRoute(this, 'message-route', {
-            apiId: api.ref,
-            routeKey: 'sendmessage',
-            authorizationType: 'NONE',
-            target: 'integrations/' + messageIntegration.ref,
-        });
-
-        const deployment = new CfnDeployment(this, `${name}-deployment`, {
-            apiId: api.ref,
+        api.addRoute('sendmessage', {
+            integration: new WebSocketLambdaIntegration('message-lambda', messageFunc),
         });
 
-        new CfnStage(this, `${name}-stage`, {
-            apiId: api.ref,
+        const stage = new WebSocketStage(this, `${name}-stage`, {
             autoDeploy: true,
-            deploymentId: deployment.ref,
             stageName: 'dev',
+            webSocketApi: api,
         });
 
-        deployment.node.addDependency(connectRoute);
-        deployment.node.addDependency(disconnectRoute);
-        deployment.node.addDependency(messageRoute);
+        this.table = this.asOutput(table.tableName);
+        this.url = this.asOutput(stage.url);
 
         this.synth();
     }
 }
 
 const stack = new ChatAppStack('chat-app');
-export default stack.outputs;
+export const url = stack.url;
+export const table = stack.table;

examples/api-websocket-lambda-dynamodb/onconnect/app.js

@@ -1,23 +1,27 @@
 // Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: MIT-0
 
-const AWS = require('aws-sdk');
+const { DynamoDBDocumentClient, PutCommand } = require('@aws-sdk/lib-dynamodb');
+const { DynamoDBClient } = require('@aws-sdk/client-dynamodb');
 
-const ddb = new AWS.DynamoDB.DocumentClient({ apiVersion: '2012-08-10', region: process.env.AWS_REGION });
+const client = new DynamoDBClient({});
+const ddb = DynamoDBDocumentClient.from(client);
 
 exports.handler = async event => {
-  const putParams = {
+  console.log(`event: ${JSON.stringify(event)}`);
+  const command = new PutCommand({
     TableName: process.env.TABLE_NAME,
     Item: {
       connectionId: event.requestContext.connectionId
     }
-  };
+  });
 
   try {
-    await ddb.put(putParams).promise();
+    await ddb.send(command)
   } catch (err) {
     return { statusCode: 500, body: 'Failed to connect: ' + JSON.stringify(err) };
   }
 
+  console.log("Successfully connected");
   return { statusCode: 200, body: 'Connected.' };
 };

examples/api-websocket-lambda-dynamodb/ondisconnect/app.js

@@ -7,23 +7,27 @@
 // $disconnect is a best-effort event. 
 // API Gateway will try its best to deliver the $disconnect event to your integration, but it cannot guarantee delivery.
 
-const AWS = require('aws-sdk');
+const { DeleteCommand, DynamoDBDocumentClient } = require('@aws-sdk/lib-dynamodb');
+const { DynamoDBClient } = require('@aws-sdk/client-dynamodb');
 
-const ddb = new AWS.DynamoDB.DocumentClient({ apiVersion: '2012-08-10', region: process.env.AWS_REGION });
+const client = new DynamoDBClient({});
+const ddb = DynamoDBDocumentClient.from(client);
 
 exports.handler = async event => {
-  const deleteParams = {
+  console.log(`event: ${JSON.stringify(event)}`);
+  const command = new DeleteCommand({
     TableName: process.env.TABLE_NAME,
     Key: {
       connectionId: event.requestContext.connectionId
     }
-  };
+  });
 
   try {
-    await ddb.delete(deleteParams).promise();
+    await ddb.send(command);
   } catch (err) {
     return { statusCode: 500, body: 'Failed to disconnect: ' + JSON.stringify(err) };
   }
 
+  console.log("Successfully disconnected");
   return { statusCode: 200, body: 'Disconnected.' };
 };

examples/api-websocket-lambda-dynamodb/package.json

@@ -4,6 +4,7 @@
         "@types/node": "^10.0.0"
     },
     "dependencies": {
+        "@aws-sdk/client-apigatewaymanagementapi": "^3.632.0",
         "@aws-sdk/client-dynamodb": "^3.632.0",
         "@aws-sdk/lib-dynamodb": "^3.632.0",
         "@pulumi/aws": "^4.6.0",

examples/api-websocket-lambda-dynamodb/sendmessage/app.js

@@ -1,35 +1,48 @@
 // Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: MIT-0
 
-const AWS = require('aws-sdk');
+const { DeleteCommand, DynamoDBDocumentClient, ScanCommand } = require('@aws-sdk/lib-dynamodb');
+const { DynamoDBClient } = require('@aws-sdk/client-dynamodb');
+const { ApiGatewayManagementApiClient, PostToConnectionCommand } = require('@aws-sdk/client-apigatewaymanagementapi');
 
-const ddb = new AWS.DynamoDB.DocumentClient({ apiVersion: '2012-08-10', region: process.env.AWS_REGION });
+const client = new DynamoDBClient({});
+const ddb = DynamoDBDocumentClient.from(client);
+let apigwManagementApi;
 
 const { TABLE_NAME } = process.env;
 
 exports.handler = async event => {
+  console.log(`event: ${JSON.stringify(event)}`);
   let connectionData;
+  const command = new ScanCommand({
+    TableName: TABLE_NAME,
+    ProjectionExpression: 'connectionId',
+  });
 
   try {
-    connectionData = await ddb.scan({ TableName: TABLE_NAME, ProjectionExpression: 'connectionId' }).promise();
+    connectionData = await ddb.send(command);
   } catch (e) {
     return { statusCode: 500, body: e.stack };
   }
 
-  const apigwManagementApi = new AWS.ApiGatewayManagementApi({
-    apiVersion: '2018-11-29',
-    endpoint: event.requestContext.domainName + '/' + event.requestContext.stage
-  });
+  const endpoint = event.requestContext.domainName + '/' + event.requestContext.stage;
+  if (!apigwManagementApi || apigwManagementApi?.config.endpoint !== endpoint) {
+    apigwManagementApi = new ApiGatewayManagementApiClient({
+      apiVersion: '2018-11-29',
+      endpoint,
+    });
+  }
 
   const postData = JSON.parse(event.body).data;
 
   const postCalls = connectionData.Items.map(async ({ connectionId }) => {
+    console.log(`Found connection ${connectionId}`);
     try {
-      await apigwManagementApi.postToConnection({ ConnectionId: connectionId, Data: postData }).promise();
+      await apigwManagementApi.send(new PostToConnectionCommand({ ConnectionId: connectionId, Data: postData }));
     } catch (e) {
       if (e.statusCode === 410) {
         console.log(`Found stale connection, deleting ${connectionId}`);
-        await ddb.delete({ TableName: TABLE_NAME, Key: { connectionId } }).promise();
+        await ddb.send(new DeleteCommand({ TableName: TABLE_NAME, Key: { connectionId } }));
       } else {
         throw e;
       }
@@ -42,5 +55,6 @@ exports.handler = async event => {
     return { statusCode: 500, body: e.stack };
   }
 
+  console.log("Successfully sent message");
   return { statusCode: 200, body: 'Data sent.' };
 };