fix(appbuilder): improve open handler logic

roger-zhangg · web-flow · commit 18fe3db40ab9 · 2024-10-23T16:57:58.000-07:00
## Problem
open handler version may search and open the wrong handler.

## Solution
Improve the logic to identify handler in ruby/python
Improve the logic search in limited location in java/node/dotnet
diff --git a/packages/core/src/awsService/appBuilder/utils.ts b/packages/core/src/awsService/appBuilder/utils.ts
@@ -17,6 +17,8 @@ import { samDeployUrl } from '../../shared/constants'
 import path from 'path'
 import fs from '../../shared/fs/fs'
 import { getLogger } from '../../shared/logger/logger'
+import { RuntimeFamily, getFamily } from '../../lambda/models/samLambdaRuntime'
+import { showMessage } from '../../shared/utilities/messages'
 const localize = nls.loadMessageBundle()
 
 export async function runOpenTemplate(arg?: TreeNode) {
@@ -28,36 +30,134 @@ export async function runOpenTemplate(arg?: TreeNode) {
     await vscode.window.showTextDocument(document)
 }
 
-export async function runOpenHandler(arg: ResourceNode) {
+/**
+ * Find and open the lambda handler with given ResoruceNode
+ * If not found, a NoHandlerFound error will be raised
+ * @param arg ResourceNode
+ */
+export async function runOpenHandler(arg: ResourceNode): Promise<void> {
     const folderUri = path.dirname(arg.resource.location.fsPath)
-    let handler: string | undefined
-    let extension = '*'
-    if (arg.resource.resource.Runtime?.includes('java')) {
-        handler = arg.resource.resource.Handler?.split('::')[0]
-        if (handler?.includes('.')) {
-            handler = handler.split('.')[1]
-        }
-        extension = 'java'
-    } else if (arg.resource.resource.Runtime?.includes('dotnet')) {
-        handler = arg.resource.resource.Handler?.split('::')[1]
-        if (handler?.includes('.')) {
-            handler = handler.split('.')[1]
-        }
-        extension = 'cs'
-    } else {
-        handler = arg.resource.resource.Handler?.split('.')[0]
+    if (!arg.resource.resource.CodeUri) {
+        throw new ToolkitError('No CodeUri provided in template, cannot open handler', { code: 'NoCodeUriProvided' })
     }
-    const handlerFile = (
-        await vscode.workspace.findFiles(
-            new vscode.RelativePattern(folderUri, `**/${handler}.${extension}`),
-            new vscode.RelativePattern(folderUri, '.aws-sam')
-        )
-    )[0]
+
+    if (!arg.resource.resource.Handler) {
+        throw new ToolkitError('No Handler provided in template, cannot open handler', { code: 'NoHandlerProvided' })
+    }
+
+    if (!arg.resource.resource.Runtime) {
+        throw new ToolkitError('No Runtime provided in template, cannot open handler', { code: 'NoRuntimeProvided' })
+    }
+
+    const handlerFile = await getLambdaHandlerFile(
+        vscode.Uri.file(folderUri),
+        arg.resource.resource.CodeUri,
+        arg.resource.resource.Handler,
+        arg.resource.resource.Runtime
+    )
     if (!handlerFile) {
-        throw new ToolkitError(`No handler file found with name "${handler}"`, { code: 'NoHandlerFound' })
+        throw new ToolkitError(`No handler file found with name "${arg.resource.resource.Handler}"`, {
+            code: 'NoHandlerFound',
+        })
     }
-    const document = await vscode.workspace.openTextDocument(handlerFile)
-    await vscode.window.showTextDocument(document)
+    await vscode.workspace.openTextDocument(handlerFile).then(async (doc) => await vscode.window.showTextDocument(doc))
+}
+
+// create a set to store all supported runtime in the following function
+const supportedRuntimeForHandler = new Set<RuntimeFamily>([
+    RuntimeFamily.Ruby,
+    RuntimeFamily.Python,
+    RuntimeFamily.NodeJS,
+    RuntimeFamily.DotNet,
+    RuntimeFamily.Java,
+])
+
+/**
+ * Get the actual Lambda handler file, in vscode.Uri format, from the template
+ * file and handler name. If not found, return undefined.
+ *
+ * @param folderUri The root folder for sam project
+ * @param codeUri codeUri prop in sam template
+ * @param handler handler prop in sam template
+ * @param runtime runtime prop in sam template
+ * @returns
+ */
+export async function getLambdaHandlerFile(
+    folderUri: vscode.Uri,
+    codeUri: string,
+    handler: string,
+    runtime: string
+): Promise<vscode.Uri | undefined> {
+    const family = getFamily(runtime)
+    if (!supportedRuntimeForHandler.has(family)) {
+        throw new ToolkitError(`Runtime ${runtime} is not supported for open handler button`, {
+            code: 'RuntimeNotSupported',
+        })
+    }
+
+    const handlerParts = handler.split('.')
+    // sample: app.lambda_handler -> app.rb
+    if (family === RuntimeFamily.Ruby) {
+        // Ruby supports namespace/class handlers as well, but the path is
+        // guaranteed to be slash-delimited so we can assume the first part is
+        // the path
+        return vscode.Uri.joinPath(folderUri, codeUri, handlerParts.slice(0, handlerParts.length - 1).join('/') + '.rb')
+    }
+
+    // sample:app.lambda_handler -> app.py
+    if (family === RuntimeFamily.Python) {
+        // Otherwise (currently Node.js and Python) handle dot-delimited paths
+        return vscode.Uri.joinPath(folderUri, codeUri, handlerParts.slice(0, handlerParts.length - 1).join('/') + '.py')
+    }
+
+    // sample: app.handler -> app.mjs/app.js
+    // More likely to be mjs if NODEJS version>=18, now searching for both
+    if (family === RuntimeFamily.NodeJS) {
+        const handlerName = handlerParts.slice(0, handlerParts.length - 1).join('/')
+        const handlerPath = path.dirname(handlerName)
+        const handlerFile = path.basename(handlerName)
+        const pattern = new vscode.RelativePattern(
+            vscode.Uri.joinPath(folderUri, codeUri, handlerPath),
+            `${handlerFile}.{js,mjs}`
+        )
+        return searchHandlerFile(folderUri, pattern)
+    }
+    // search directly under Code uri for Dotnet and java
+    // sample: ImageResize::ImageResize.Function::FunctionHandler -> Function.cs
+    if (family === RuntimeFamily.DotNet) {
+        const handlerName = path.basename(handler.split('::')[1].replaceAll('.', '/'))
+        const pattern = new vscode.RelativePattern(vscode.Uri.joinPath(folderUri, codeUri), `${handlerName}.cs`)
+        return searchHandlerFile(folderUri, pattern)
+    }
+
+    // sample: resizer.App::handleRequest -> App.java
+    if (family === RuntimeFamily.Java) {
+        const handlerName = handler.split('::')[0].replaceAll('.', '/')
+        const pattern = new vscode.RelativePattern(vscode.Uri.joinPath(folderUri, codeUri), `**/${handlerName}.java`)
+        return searchHandlerFile(folderUri, pattern)
+    }
+}
+
+/**
+    Searches for a handler file in the given pattern and returns the first match.
+    If no match is found, returns undefined.
+*/
+export async function searchHandlerFile(
+    folderUri: vscode.Uri,
+    pattern: vscode.RelativePattern
+): Promise<vscode.Uri | undefined> {
+    const handlerFile = await vscode.workspace.findFiles(pattern, new vscode.RelativePattern(folderUri, '.aws-sam'))
+    if (handlerFile.length === 0) {
+        return undefined
+    }
+    if (handlerFile.length > 1) {
+        getLogger().warn(`Multiple handler files found with name "${path.basename(handlerFile[0].fsPath)}"`)
+        void showMessage('warn', `Multiple handler files found with name "${path.basename(handlerFile[0].fsPath)}"`)
+    }
+    if (await fs.exists(handlerFile[0])) {
+        return handlerFile[0]
+    }
+    return undefined
 }
 
 async function promptUserForTemplate() {
diff --git a/packages/core/src/lambda/models/samLambdaRuntime.ts b/packages/core/src/lambda/models/samLambdaRuntime.ts
@@ -22,6 +22,7 @@ export enum RuntimeFamily {
     DotNet,
     Go,
     Java,
+    Ruby,
 }
 
 export type RuntimePackageType = 'Image' | 'Zip'
@@ -65,6 +66,7 @@ export const javaRuntimes: ImmutableSet<Runtime> = ImmutableSet<Runtime>([
     'java21',
 ])
 export const dotNetRuntimes: ImmutableSet<Runtime> = ImmutableSet<Runtime>(['dotnet6', 'dotnet8'])
+export const rubyRuntimes: ImmutableSet<Runtime> = ImmutableSet<Runtime>(['ruby3.2', 'ruby3.3'])
 
 /**
  * Deprecated runtimes can be found at https://docs.aws.amazon.com/lambda/latest/dg/runtime-support-policy.html
@@ -83,13 +85,16 @@ export const deprecatedRuntimes: ImmutableSet<Runtime> = ImmutableSet<Runtime>([
     'nodejs8.10',
     'nodejs10.x',
     'nodejs12.x',
+    'ruby2.5',
+    'ruby2.7',
 ])
 const defaultRuntimes = ImmutableMap<RuntimeFamily, Runtime>([
     [RuntimeFamily.NodeJS, 'nodejs20.x'],
     [RuntimeFamily.Python, 'python3.12'],
     [RuntimeFamily.DotNet, 'dotnet6'],
     [RuntimeFamily.Go, 'go1.x'],
     [RuntimeFamily.Java, 'java17'],
+    [RuntimeFamily.Ruby, 'ruby3.3'],
 ])
 
 export const samZipLambdaRuntimes: ImmutableSet<Runtime> = ImmutableSet.union([
@@ -163,6 +168,8 @@ export function getFamily(runtime: string): RuntimeFamily {
         return RuntimeFamily.Go
     } else if (javaRuntimes.has(runtime)) {
         return RuntimeFamily.Java
+    } else if (rubyRuntimes.has(runtime)) {
+        return RuntimeFamily.Ruby
     }
     return RuntimeFamily.Unknown
 }
@@ -212,6 +219,10 @@ export function getRuntimeFamily(langId: string): RuntimeFamily {
             return RuntimeFamily.Python
         case 'go':
             return RuntimeFamily.Go
+        case 'java':
+            return RuntimeFamily.Java
+        case 'ruby':
+            return RuntimeFamily.Ruby
         default:
             return RuntimeFamily.Unknown
     }
diff --git a/packages/core/src/test/awsService/appBuilder/utils.test.ts b/packages/core/src/test/awsService/appBuilder/utils.test.ts
@@ -33,33 +33,81 @@ const scenarios: TestScenario[] = [
         handler: 'resizer.App::handleRequest',
         codeUri: 'ResizerFunction',
         fileLocation: 'ResizerFunction/src/main/java/resizer/App.java',
-        fileInfo: 'test',
+        fileInfo: 'testjava',
         regex: /App.java/g,
     },
     {
         runtime: 'dotnet6',
         handler: 'ImageResize::ImageResize.Function::FunctionHandler',
         codeUri: 'ImageResize/',
         fileLocation: 'ImageResize/Function.cs',
-        fileInfo: 'test',
+        fileInfo: 'testdotnet',
         regex: /Function.cs/g,
     },
     {
         runtime: 'python3.9',
         handler: 'app.lambda_handler',
         codeUri: 'hello_world/',
         fileLocation: 'hello_world/app.py',
-        fileInfo: 'test',
+        fileInfo: 'testpython',
         regex: /app.py/g,
     },
     {
         runtime: 'nodejs18.x',
         handler: 'app.handler',
         codeUri: 'src/',
         fileLocation: 'src/app.js',
-        fileInfo: 'test',
+        fileInfo: 'testnode',
         regex: /app.js/g,
     },
+    {
+        runtime: 'ruby3.2',
+        handler: 'app.lambda_handler',
+        codeUri: 'hello_world/',
+        fileLocation: 'hello_world/app.rb',
+        fileInfo: 'testruby',
+        regex: /app.rb/g,
+    },
+    {
+        runtime: 'java21',
+        handler: 'resizer.App::handleRequest',
+        codeUri: 'ResizerFunction',
+        fileLocation: 'ResizerFunction/src/foo/bar/main/java/resizer/App.java',
+        fileInfo: 'testjava2',
+        regex: /App.java/g,
+    },
+    {
+        runtime: 'dotnet8',
+        handler: 'ImageResize::ImageResize.Function::FunctionHandler',
+        codeUri: 'ImageResize/src/test',
+        fileLocation: 'ImageResize/src/test/Function.cs',
+        fileInfo: 'testdotnet2',
+        regex: /Function.cs/g,
+    },
+    {
+        runtime: 'python3.12',
+        handler: 'app.foo.bar.lambda_handler',
+        codeUri: 'hello_world/test123',
+        fileLocation: 'hello_world/test123/app/foo/bar.py',
+        fileInfo: 'testpython2',
+        regex: /bar.py/g,
+    },
+    {
+        runtime: 'nodejs20.x',
+        handler: 'app.foo.bar.handler',
+        codeUri: 'src/test123',
+        fileLocation: 'src/test123/app/foo/bar.js',
+        fileInfo: 'testnode2',
+        regex: /bar.js/g,
+    },
+    {
+        runtime: 'ruby3.3',
+        handler: 'app/foo/bar.lambda_handler',
+        codeUri: 'hello_world/test456',
+        fileLocation: 'hello_world/test456/app/foo/bar.rb',
+        fileInfo: 'testruby2',
+        regex: /bar.rb/g,
+    },
 ]
 
 describe('AppBuilder Utils', function () {
@@ -165,6 +213,43 @@ describe('AppBuilder Utils', function () {
                     await assertTextEditorContains(scenario.fileInfo)
                 })
             }
+
+            it(`should warn for multiple java handler found`, async function () {
+                const rNode = new ResourceNode(
+                    {
+                        samTemplateUri: vscode.Uri.file(path.join(tempFolder, 'template.yaml')),
+                        workspaceFolder: workspace,
+                        projectRoot: vscode.Uri.file(tempFolder),
+                    },
+                    {
+                        Id: 'MyFunction',
+                        Type: SERVERLESS_FUNCTION_TYPE,
+                        Runtime: 'java21',
+                        Handler: 'resizer.App::handleRequest',
+                        CodeUri: 'ResizerFunction',
+                    }
+                )
+                // When 2 java handler with right name under code URI
+                await fs.mkdir(
+                    path.join(tempFolder, ...path.dirname('ResizerFunction/src/main/java/resizer/App.java').split('/'))
+                )
+                await fs.writeFile(
+                    path.join(tempFolder, ...'ResizerFunction/src/main/java/resizer/App.java'.split('/')),
+                    'testjava'
+                )
+                await fs.mkdir(
+                    path.join(tempFolder, ...path.dirname('ResizerFunction/src/main/java/resizer2/App.java').split('/'))
+                )
+                await fs.writeFile(
+                    path.join(tempFolder, ...'ResizerFunction/src/main/java/resizer2/App.java'.split('/')),
+                    'testjava'
+                )
+                // Then should warn
+                getTestWindow().onDidShowMessage((msg) =>
+                    assert(msg.assertWarn('Multiple handler files found with name App.java"'))
+                )
+                await runOpenHandler(rNode)
+            })
         }),
         describe('open template', function () {
             let sandbox: sinon.SinonSandbox
