Rearrange lambda quick editing to avoid download on activation

Author: reedham-aws

packages/core/src/lambda/activation.ts

@@ -17,7 +17,7 @@ import { ExtContext } from '../shared/extensions'
 import { invokeRemoteLambda } from './vue/remoteInvoke/invokeLambda'
 import { registerSamDebugInvokeVueCommand, registerSamInvokeVueCommand } from './vue/configEditor/samInvokeBackend'
 import { Commands } from '../shared/vscode/commands2'
-import { DefaultLambdaClient, getFunctionWithCredentials } from '../shared/clients/lambdaClient'
+import { DefaultLambdaClient } from '../shared/clients/lambdaClient'
 import { copyLambdaUrl } from './commands/copyLambdaUrl'
 import { ResourceNode } from '../awsService/appBuilder/explorer/nodes/resourceNode'
 import { isTreeNode, TreeNode } from '../shared/treeview/resourceTreeDataProvider'
@@ -29,47 +29,125 @@ import { ToolkitError, isError } from '../shared/errors'
 import { LogStreamFilterResponse } from '../awsService/cloudWatchLogs/wizard/liveTailLogStreamSubmenu'
 import { tempDirPath } from '../shared/filesystemUtilities'
 import fs from '../shared/fs/fs'
-import { deployFromTemp, editLambda, getReadme, openLambdaFolderForEdit } from './commands/editLambda'
-import { getTempLocation } from './utils'
+import {
+    confirmOutdatedChanges,
+    deleteFilesInFolder,
+    deployFromTemp,
+    getReadme,
+    openLambdaFolderForEdit,
+    watchForUpdates,
+} from './commands/editLambda'
+import { compareCodeSha, getFunctionInfo, getTempLocation, setFunctionInfo } from './utils'
 import { registerLambdaUriHandler } from './uriHandlers'
+import globals from '../shared/extensionGlobals'
 
 const localize = nls.loadMessageBundle()
 
-/**
- * Activates Lambda components.
- */
-export async function activate(context: ExtContext): Promise<void> {
-    try {
-        if (vscode.workspace.workspaceFolders) {
-            for (const workspaceFolder of vscode.workspace.workspaceFolders) {
-                // Making the comparison case insensitive because Windows can have `C\` or `c\`
-                const workspacePath = workspaceFolder.uri.fsPath.toLowerCase()
-                const tempPath = path.join(tempDirPath, 'lambda').toLowerCase()
-                if (workspacePath.startsWith(tempPath)) {
-                    const name = path.basename(workspaceFolder.uri.fsPath)
-                    const region = path.basename(path.dirname(workspaceFolder.uri.fsPath))
-                    const getFunctionOutput = await getFunctionWithCredentials(region, name)
-                    const configuration = getFunctionOutput.Configuration
-                    await editLambda(
-                        {
-                            name,
-                            region,
-                            // Configuration as any due to the difference in types between sdkV2 and sdkV3
-                            configuration: configuration as any,
-                        },
-                        true
+async function openReadme() {
+    const readmeUri = vscode.Uri.file(await getReadme())
+    // We only want to do it if there's not a readme already
+    const isPreviewOpen = vscode.window.tabGroups.all.some((group) =>
+        group.tabs.some((tab) => tab.label.includes('README'))
+    )
+    if (!isPreviewOpen) {
+        await vscode.commands.executeCommand('markdown.showPreviewToSide', readmeUri)
+    }
+}
+
+async function quickEditActivation() {
+    if (vscode.workspace.workspaceFolders) {
+        for (const workspaceFolder of vscode.workspace.workspaceFolders) {
+            // Making the comparison case insensitive because Windows can have `C\` or `c\`
+            const workspacePath = workspaceFolder.uri.fsPath.toLowerCase()
+            const tempPath = path.join(tempDirPath, 'lambda').toLowerCase()
+            if (workspacePath.includes(tempPath)) {
+                const name = path.basename(workspaceFolder.uri.fsPath)
+                const region = path.basename(path.dirname(workspaceFolder.uri.fsPath))
+
+                const lambda = { name, region, configuration: undefined }
+
+                watchForUpdates(lambda, vscode.Uri.file(workspacePath))
+
+                await openReadme()
+
+                // Open handler function
+                try {
+                    const handler = await getFunctionInfo(lambda, 'handlerFile')
+                    const lambdaLocation = path.join(workspacePath, handler)
+                    await openLambdaFile(lambdaLocation, vscode.ViewColumn.One)
+                } catch (e) {
+                    void vscode.window.showWarningMessage(
+                        localize('AWS.lambda.openFile.failure', `Failed to determine handler location: ${e}`)
                     )
+                }
 
-                    const readmeUri = vscode.Uri.file(await getReadme())
-                    await vscode.commands.executeCommand('markdown.showPreviewToSide', readmeUri)
+                // Check if there are changes that need overwritten
+                try {
+                    // Checking if there are changes that need to be overwritten
+                    const prompt = localize(
+                        'AWS.lambda.download.confirmOutdatedSync',
+                        'There are changes to your function in the cloud since you last edited locally, do you want to overwrite your local changes?'
+                    )
+
+                    // Adding delay to give the authentication time to catch up
+                    await new Promise((resolve) => globals.clock.setTimeout(resolve, 1000))
+
+                    const overwriteChanges = !(await compareCodeSha(lambda))
+                        ? await confirmOutdatedChanges(prompt)
+                        : false
+                    if (overwriteChanges) {
+                        // Close all open tabs from this workspace
+                        const workspaceUri = vscode.Uri.file(workspacePath)
+                        for (const tabGroup of vscode.window.tabGroups.all) {
+                            const tabsToClose = tabGroup.tabs.filter(
+                                (tab) =>
+                                    tab.input instanceof vscode.TabInputText &&
+                                    tab.input.uri.fsPath.startsWith(workspaceUri.fsPath)
+                            )
+                            if (tabsToClose.length > 0) {
+                                await vscode.window.tabGroups.close(tabsToClose)
+                            }
+                        }
+
+                        // Delete all files in the directory
+                        await deleteFilesInFolder(workspacePath)
+
+                        // Remove workspace folder
+                        const workspaceIndex = vscode.workspace.workspaceFolders?.findIndex(
+                            (folder) => folder.uri.fsPath.toLowerCase() === workspacePath
+                        )
+                        if (workspaceIndex !== undefined && workspaceIndex >= 0) {
+                            vscode.workspace.updateWorkspaceFolders(workspaceIndex, 1)
+                        }
+
+                        await setFunctionInfo(lambda, { undeployed: false })
+
+                        // Show message to user about next steps
+                        void vscode.window.showInformationMessage(
+                            localize(
+                                'AWS.lambda.refresh.complete',
+                                'Local workspace cleared. Navigate to the Toolkit explorer to get fresh code from the cloud.'
+                            )
+                        )
+                    }
+                } catch (e) {
+                    void vscode.window.showWarningMessage(
+                        localize(
+                            'AWS.lambda.pull.failure',
+                            `Failed to pull latest changes from the cloud, you can still edit locally: ${e}`
+                        )
+                    )
                 }
             }
         }
-    } catch (e) {
-        void vscode.window.showWarningMessage(
-            localize('AWS.lambda.open.failure', `Unable to edit Lambda Function locally: ${e}`)
-        )
     }
+}
+
+/**
+ * Activates Lambda components.
+ */
+export async function activate(context: ExtContext): Promise<void> {
+    await quickEditActivation()
 
     context.extensionContext.subscriptions.push(
         Commands.register('aws.deleteLambda', async (node: LambdaFunctionNode | TreeNode) => {

packages/core/src/lambda/commands/downloadLambda.ts

@@ -194,7 +194,7 @@ async function downloadAndUnzipLambda(
     }
 }
 
-export async function openLambdaFile(lambdaLocation: string): Promise<void> {
+export async function openLambdaFile(lambdaLocation: string, viewColumn?: vscode.ViewColumn): Promise<void> {
     if (!(await fs.exists(lambdaLocation))) {
         const warning = localize(
             'AWS.lambda.download.fileNotFound',
@@ -206,7 +206,7 @@ export async function openLambdaFile(lambdaLocation: string): Promise<void> {
         throw new Error()
     }
     const doc = await vscode.workspace.openTextDocument(vscode.Uri.file(lambdaLocation))
-    await vscode.window.showTextDocument(doc)
+    await vscode.window.showTextDocument(doc, viewColumn)
 }
 
 async function addLaunchConfigEntry(

packages/core/src/lambda/commands/editLambda.ts

@@ -22,6 +22,7 @@ import { LambdaFunctionNodeDecorationProvider } from '../explorer/lambdaFunction
 import path from 'path'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { ToolkitError } from '../../shared/errors'
+import { getFunctionWithCredentials } from '../../shared/clients/lambdaClient'
 
 const localize = nls.loadMessageBundle()
 
@@ -84,7 +85,7 @@ export async function promptForSync(lambda: LambdaFunction, projectUri: vscode.U
     }
 }
 
-async function confirmOutdatedChanges(prompt: string): Promise<boolean> {
+export async function confirmOutdatedChanges(prompt: string): Promise<boolean> {
     return await showConfirmationMessage({
         prompt,
         confirm: localize('AWS.lambda.upload.overwrite', 'Overwrite'),
@@ -128,28 +129,62 @@ export async function deployFromTemp(lambda: LambdaFunction, projectUri: vscode.
     })
 }
 
+export async function deleteFilesInFolder(location: string) {
+    const entries = await fs.readdir(location)
+    await Promise.all(
+        entries.map((entry) => fs.delete(path.join(location, entry[0]), { recursive: true, force: true }))
+    )
+}
+
 export async function editLambdaCommand(functionNode: LambdaFunctionNode) {
     const region = functionNode.regionCode
     const functionName = functionNode.configuration.FunctionName!
-    return await editLambda({ name: functionName, region, configuration: functionNode.configuration })
+    return await editLambda({ name: functionName, region, configuration: functionNode.configuration }, 'explorer')
 }
 
-export async function editLambda(lambda: LambdaFunction, onActivation?: boolean) {
+export async function overwriteChangesForEdit(lambda: LambdaFunction, downloadLocation: string) {
+    try {
+        // Clear directory contents instead of deleting to avoid Windows EBUSY errors
+        if (await fs.existsDir(downloadLocation)) {
+            await deleteFilesInFolder(downloadLocation)
+        } else {
+            await fs.mkdir(downloadLocation)
+        }
+
+        await downloadLambdaInLocation(lambda, 'local', downloadLocation)
+
+        // Watching for updates, then setting info, then removing the badges must be done in this order
+        // This is because the files creating can throw the watcher, which sometimes leads to changes being marked as undeployed
+        watchForUpdates(lambda, vscode.Uri.file(downloadLocation))
+
+        await setFunctionInfo(lambda, {
+            lastDeployed: globals.clock.Date.now(),
+            undeployed: false,
+            sha: lambda.configuration!.CodeSha256,
+            handlerFile: getLambdaDetails(lambda.configuration!).fileName,
+        })
+        await LambdaFunctionNodeDecorationProvider.getInstance().removeBadge(
+            vscode.Uri.file(downloadLocation),
+            vscode.Uri.from({ scheme: 'lambda', path: `${lambda.region}/${lambda.name}` })
+        )
+    } catch {
+        throw new ToolkitError('Failed to download Lambda function', { code: 'failedDownload' })
+    }
+}
+
+export async function editLambda(lambda: LambdaFunction, source?: 'workspace' | 'explorer') {
     return await telemetry.lambda_quickEditFunction.run(async () => {
-        telemetry.record({ source: onActivation ? 'workspace' : 'explorer' })
+        telemetry.record({ source })
         const downloadLocation = getTempLocation(lambda.name, lambda.region)
-        const downloadLocationName = vscode.workspace.asRelativePath(downloadLocation, true)
 
         // We don't want to do anything if the folder already exists as a workspace folder, it means it's already being edited
-        if (
-            vscode.workspace.workspaceFolders?.some((folder) => folder.uri.fsPath === downloadLocation && !onActivation)
-        ) {
+        if (vscode.workspace.workspaceFolders?.some((folder) => folder.uri.fsPath === downloadLocation)) {
             return downloadLocation
         }
 
         const prompt = localize(
             'AWS.lambda.download.confirmOutdatedSync',
-            'There are changes to your Function in the cloud since you last edited locally, do you want to overwrite your local changes?'
+            'There are changes to your function in the cloud since you last edited locally, do you want to overwrite your local changes?'
         )
 
         // We want to overwrite changes in the following cases:
@@ -169,37 +204,9 @@ export async function editLambda(lambda: LambdaFunction, onActivation?: boolean)
             (!(await compareCodeSha(lambda)) ? await confirmOutdatedChanges(prompt) : false)
 
         if (overwriteChanges) {
-            try {
-                // Clear directory contents instead of deleting to avoid Windows EBUSY errors
-                if (await fs.existsDir(downloadLocation)) {
-                    const entries = await fs.readdir(downloadLocation)
-                    await Promise.all(
-                        entries.map((entry) =>
-                            fs.delete(path.join(downloadLocation, entry[0]), { recursive: true, force: true })
-                        )
-                    )
-                } else {
-                    await fs.mkdir(downloadLocation)
-                }
-
-                await downloadLambdaInLocation(lambda, downloadLocationName, downloadLocation)
-
-                // Watching for updates, then setting info, then removing the badges must be done in this order
-                // This is because the files creating can throw the watcher, which sometimes leads to changes being marked as undeployed
-                watchForUpdates(lambda, vscode.Uri.file(downloadLocation))
-                await setFunctionInfo(lambda, {
-                    lastDeployed: globals.clock.Date.now(),
-                    undeployed: false,
-                    sha: lambda.configuration!.CodeSha256,
-                })
-                await LambdaFunctionNodeDecorationProvider.getInstance().removeBadge(
-                    vscode.Uri.file(downloadLocation),
-                    vscode.Uri.from({ scheme: 'lambda', path: `${lambda.region}/${lambda.name}` })
-                )
-            } catch {
-                throw new ToolkitError('Failed to download Lambda function', { code: 'failedDownload' })
-            }
-        } else {
+            await overwriteChangesForEdit(lambda, downloadLocation)
+        } else if (source === 'explorer') {
+            // If the source is the explorer, we want to open, otherwise we just wait to open in the workspace
             const lambdaLocation = path.join(downloadLocation, getLambdaDetails(lambda.configuration!).fileName)
             await openLambdaFile(lambdaLocation)
             watchForUpdates(lambda, vscode.Uri.file(downloadLocation))
@@ -212,22 +219,24 @@ export async function editLambda(lambda: LambdaFunction, onActivation?: boolean)
 export async function openLambdaFolderForEdit(name: string, region: string) {
     const downloadLocation = getTempLocation(name, region)
 
-    if (
-        vscode.workspace.workspaceFolders?.some((workspaceFolder) =>
-            workspaceFolder.uri.fsPath.toLowerCase().startsWith(downloadLocation.toLowerCase())
-        )
-    ) {
-        // If the folder already exists in the workspace, show that folder
-        await vscode.commands.executeCommand('workbench.action.focusSideBar')
-        await vscode.commands.executeCommand('workbench.view.explorer')
-    } else {
-        await fs.mkdir(downloadLocation)
-
-        await vscode.commands.executeCommand('vscode.openFolder', vscode.Uri.file(downloadLocation), {
-            newWindow: true,
-            noRecentEntry: true,
-        })
-    }
+    // Do all authentication work before opening workspace to avoid race condition
+    const getFunctionOutput = await getFunctionWithCredentials(region, name)
+    const configuration = getFunctionOutput.Configuration
+
+    // Download and set up Lambda code before opening workspace
+    await editLambda(
+        {
+            name,
+            region,
+            configuration: configuration as any,
+        },
+        'workspace'
+    )
+
+    await vscode.commands.executeCommand('vscode.openFolder', vscode.Uri.file(downloadLocation), {
+        newWindow: true,
+        noRecentEntry: true,
+    })
 }
 
 export async function getReadme(): Promise<string> {

packages/core/src/lambda/uriHandlers.ts

@@ -10,7 +10,6 @@ import { SearchParams } from '../shared/vscode/uriHandler'
 import { openLambdaFolderForEdit } from './commands/editLambda'
 import { showConfirmationMessage } from '../shared/utilities/messages'
 import globals from '../shared/extensionGlobals'
-import { getFunctionWithCredentials } from '../shared/clients/lambdaClient'
 import { telemetry } from '../shared/telemetry/telemetry'
 import { ToolkitError } from '../shared/errors'
 
@@ -20,9 +19,6 @@ export function registerLambdaUriHandler() {
     async function openFunctionHandler(params: ReturnType<typeof parseOpenParams>) {
         await telemetry.lambda_uriHandler.run(async () => {
             try {
-                // We just want to be able to get the function - if it fails we abort and throw the error
-                await getFunctionWithCredentials(params.region, params.functionName)
-
                 if (params.isCfn === 'true') {
                     const response = await showConfirmationMessage({
                         prompt: localize(