Merge pull request #1 from piaverous/allow-typescript-handler

Feat: Allow typescript handler option

README.md

@@ -39,6 +39,7 @@ Most options are set under `custom.warmup` in the `serverless.yaml` file.
 
 * **folderName** (default `_warmup`)
 * **cleanFolder** (default `true`)
+* **tsHandler** (default `false`)
 * **name** (default `${service}-${stage}-warmup-plugin`)
 * **role** (default to role in the provider)
 * **tags** (default to serverless default tags)
@@ -62,6 +63,7 @@ custom:
   warmup:
     enabled: true # Whether to warm up functions by default or not
     folderName: '_warmup' # Name of the folder created for the generated warmup 
+    tsHandler: true # Whether the handler function for warmup should be generated in TS or not
     cleanFolder: false
     memorySize: 256
     name: 'make-them-pop'

src/index.js

@@ -61,6 +61,7 @@ class WarmUp {
       this.functionsToWarmup,
       this.resolvedOptions.region,
       this.warmupOpts.pathFile,
+      this.warmupOpts.tsHandler,
     );
     WarmUp.addWarmUpFunctionToService(this.serverless.service, this.warmupOpts);
   }
@@ -140,6 +141,7 @@ class WarmUp {
     const config = (typeof possibleConfig === 'object') ? possibleConfig : {};
     const folderName = (typeof config.folderName === 'string') ? config.folderName : '_warmup';
     const pathFolder = path.join(this.serverless.config.servicePath, folderName);
+    const tsHandler = (typeof config.tsHandler === 'boolean') ? config.tsHandler : defaultOpts.tsHandler;
 
     /* eslint-disable no-nested-ternary */
     // Keep backwards compatibility for now
@@ -152,7 +154,8 @@ class WarmUp {
     return {
       folderName,
       pathFolder,
-      pathFile: `${pathFolder}/index.js`,
+      tsHandler,
+      pathFile: `${pathFolder}/index.${tsHandler ? 'ts' : 'js'}`,
       pathHandler: `${folderName}/index.warmUp`,
       cleanFolder: (typeof config.cleanFolder === 'boolean') ? config.cleanFolder : defaultOpts.cleanFolder,
       name: (typeof config.name === 'string') ? config.name : defaultOpts.name,
@@ -230,6 +233,7 @@ class WarmUp {
     const globalDefaultOpts = {
       folderName: '_warmup',
       cleanFolder: true,
+      tsHandler: false,
       memorySize: 128,
       name: `${service.service}-${stage}-warmup-plugin`,
       events: [{ schedule: 'rate(5 minutes)' }],
@@ -299,14 +303,26 @@ class WarmUp {
    *
    * @return {Promise}
    * */
-  async createWarmUpFunctionArtifact(functions, region, pathFile) {
+  async createWarmUpFunctionArtifact(functions, region, pathFile, tsHandler) {
     /** Log warmup start */
     this.serverless.cli.log(`WarmUp: setting ${functions.length} lambdas to be warm`);
 
     /** Log functions being warmed up */
     functions.forEach(func => this.serverless.cli.log(`WarmUp: ${func.name}`));
 
-    const warmUpFunction = `"use strict";
+    const warmUpFunction = tsHandler
+      ? WarmUp.getTSWarmupFunction(region, functions)
+      : WarmUp.getJSWarmupFunction(region, functions);
+
+    /** Write warm up file */
+    return fs.outputFile(pathFile, warmUpFunction);
+  }
+
+  /**
+   * @description Write the handler in javascript
+   * */
+  static getJSWarmupFunction(region, functions) {
+    return `"use strict";
 
 /** Generated by Serverless WarmUp Plugin at ${new Date().toISOString()} */
 const aws = require("aws-sdk");
@@ -353,9 +369,59 @@ module.exports.warmUp = async (event, context) => {
 
   console.log(\`Warm Up Finished with \${invokes.filter(r => !r).length} invoke errors\`);
 }`;
+  }
 
-    /** Write warm up file */
-    return fs.outputFile(pathFile, warmUpFunction);
+  /**
+   * @description Write the handler in typescript
+   * */
+  static getTSWarmupFunction(region, functions) {
+    return `"use strict";
+
+/** Generated by Serverless WarmUp Plugin at ${new Date().toISOString()} */
+import aws from "aws-sdk";
+aws.config.region = "${region}";
+const lambda = new aws.Lambda();
+const functions = ${JSON.stringify(functions)};
+
+export const warmUp = async (event: any, context) => {
+  console.log("Warm Up Start");
+
+  const invokes = await Promise.all(functions.map(async (func) => {
+    let concurrency;
+    const functionConcurrency = process.env["WARMUP_CONCURRENCY_" + func.name.toUpperCase().replace(/-/g, '_')]
+
+    if (functionConcurrency) {
+      concurrency = parseInt(functionConcurrency);
+      console.log(\`Warming up function: \${func.name} with concurrency: \${concurrency} (from function-specific environment variable)\`);
+    } else if (process.env.WARMUP_CONCURRENCY) {
+      concurrency = parseInt(process.env.WARMUP_CONCURRENCY);
+      console.log(\`Warming up function: \${func.name} with concurrency: \${concurrency} (from global environment variable)\`);
+    } else {
+      concurrency = func.config.concurrency;
+      console.log(\`Warming up function: \${func.name} with concurrency: \${concurrency}\`);
+    }
+
+    const params = {
+      ClientContext: Buffer.from(\`{"custom":\${func.config.payload}}\`).toString('base64'),
+      FunctionName: func.name,
+      InvocationType: "RequestResponse",
+      LogType: "None",
+      Qualifier: process.env.SERVERLESS_ALIAS || "$LATEST",
+      Payload: func.config.payload
+    };
+
+    try {
+      await Promise.all(Array(concurrency).fill(0).map(async () => await lambda.invoke(params).promise()));
+      console.log(\`Warm Up Invoke Success: \${func.name}\`);
+      return true;
+    } catch (e) {
+      console.log(\`Warm Up Invoke Error: \${func.name}\`, e);
+      return false;
+    }
+  }));
+
+  console.log(\`Warm Up Finished with \${invokes.filter(r => !r).length} invoke errors\`);
+}`;
   }
 
   /**