Merge pull request #1258 from aws/normj/annotations-fix-delete

Fix issue of syncing CF template when all LambdaFunctions are removed.

Author: normj

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/CloudFormationTemplateFinder.cs

@@ -34,12 +34,28 @@ public string DetermineProjectRootDirectory(string sourceFilePath)
         }
 
         public string FindCloudFormationTemplate(string projectRootDirectory)
+        {
+            var templateAbsolutePath = DetermineCloudFormationTemplatePath(projectRootDirectory);
+
+            if (!_fileManager.Exists(templateAbsolutePath))
+                _fileManager.Create(templateAbsolutePath).Close();
+            
+            return templateAbsolutePath;
+        }
+
+        public bool DoesCloudFormationTemplateExist(string projectRootDirectory)
+        {
+            var templateAbsolutePath = DetermineCloudFormationTemplatePath(projectRootDirectory);
+            return _fileManager.Exists(templateAbsolutePath);
+        }
+
+        private string DetermineCloudFormationTemplatePath(string projectRootDirectory)
         {
             if (!_directoryManager.Exists(projectRootDirectory))
                 throw new DirectoryNotFoundException("Failed to find the project root directory");
 
             var templateAbsolutePath = string.Empty;
-            
+
             var defaultConfigFile = _directoryManager.GetFiles(projectRootDirectory, "aws-lambda-tools-defaults.json", SearchOption.AllDirectories)
                 .FirstOrDefault();
 
@@ -51,10 +67,7 @@ public string FindCloudFormationTemplate(string projectRootDirectory)
             // set the template path inside the project root directory. 
             if (string.IsNullOrEmpty(templateAbsolutePath))
                 templateAbsolutePath = Path.Combine(projectRootDirectory, "serverless.template");
-                
-            if (!_fileManager.Exists(templateAbsolutePath))
-                _fileManager.Create(templateAbsolutePath).Close();
-            
+
             return templateAbsolutePath;
         }
 

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/Generator.cs

@@ -46,8 +46,9 @@ public void Execute(GeneratorExecutionContext context)
                     return;
                 }
 
-                // If there are no Lambda methods, return early
-                if (!receiver.LambdaMethods.Any())
+                // If no project directory was detected then skip the generator.
+                // This is most likely to happen when the project is empty and doesn't have any classes in it yet.
+                if(string.IsNullOrEmpty(receiver.ProjectDirectory))
                 {
                     return;
                 }
@@ -69,7 +70,6 @@ public void Execute(GeneratorExecutionContext context)
                 var annotationReport = new AnnotationReport();
 
                 var templateFinder = new CloudFormationTemplateFinder(_fileManager, _directoryManager);
-                var projectRootDirectory = string.Empty;
 
                 foreach (var lambdaMethod in receiver.LambdaMethods)
                 {
@@ -112,15 +112,19 @@ public void Execute(GeneratorExecutionContext context)
                     diagnosticReporter.Report(Diagnostic.Create(DiagnosticDescriptors.CodeGeneration, Location.None, $"{model.GeneratedMethod.ContainingType.Name}.g.cs", sourceText));
 
                     annotationReport.LambdaFunctions.Add(model);
+                }
 
-                    if (string.IsNullOrEmpty(projectRootDirectory))
-                        projectRootDirectory = templateFinder.DetermineProjectRootDirectory(lambdaMethod.SyntaxTree.FilePath);
+                // Run the CloudFormation sync if any LambdaMethods exists. Also run if no LambdaMethods exists but there is a
+                // CloudFormation template in case orphaned functions in the template need to be removed.
+                // Both checks are required because if there is no template but there are LambdaMethods the CF template the template will be created.
+                if (receiver.LambdaMethods.Any() || templateFinder.DoesCloudFormationTemplateExist(receiver.ProjectDirectory))
+                {
+                    annotationReport.CloudFormationTemplatePath = templateFinder.FindCloudFormationTemplate(receiver.ProjectDirectory);
+                    annotationReport.ProjectRootDirectory = receiver.ProjectDirectory;
+                    var cloudFormationJsonWriter = new CloudFormationJsonWriter(_fileManager, _directoryManager, _jsonWriter, diagnosticReporter);
+                    cloudFormationJsonWriter.ApplyReport(annotationReport);
                 }
 
-                annotationReport.CloudFormationTemplatePath = templateFinder.FindCloudFormationTemplate(projectRootDirectory);
-                annotationReport.ProjectRootDirectory = projectRootDirectory;
-                var cloudFormationJsonWriter = new CloudFormationJsonWriter(_fileManager, _directoryManager,_jsonWriter, diagnosticReporter);
-                cloudFormationJsonWriter.ApplyReport(annotationReport);
             }
             catch (Exception e)
             {
@@ -135,7 +139,7 @@ public void Execute(GeneratorExecutionContext context)
         public void Initialize(GeneratorInitializationContext context)
         {
             // Register a syntax receiver that will be created for each generation pass
-            context.RegisterForSyntaxNotifications(() => new SyntaxReceiver());
+            context.RegisterForSyntaxNotifications(() => new SyntaxReceiver(_fileManager, _directoryManager));
         }
     }
 }

Libraries/src/Amazon.Lambda.Annotations.SourceGenerator/SyntaxReceiver.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using Amazon.Lambda.Annotations.SourceGenerator.FileIO;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 
@@ -12,8 +13,25 @@ internal class SyntaxReceiver : ISyntaxContextReceiver
 
         public List<ClassDeclarationSyntax> StartupClasses { get; private set; } = new List<ClassDeclarationSyntax>();
 
+        public string ProjectDirectory { get; private set; }
+
+        private readonly IFileManager _fileManager;
+        private readonly IDirectoryManager _directoryManager;
+
+        public SyntaxReceiver(IFileManager fileManager, IDirectoryManager directoryManager)
+        {
+            _fileManager = fileManager;
+            _directoryManager = directoryManager;
+        }
+
         public void OnVisitSyntaxNode(GeneratorSyntaxContext context)
         {
+            if(this.ProjectDirectory == null && context.Node is ClassDeclarationSyntax)
+            {
+                var templateFinder = new CloudFormationTemplateFinder(_fileManager, _directoryManager);
+                this.ProjectDirectory = templateFinder.DetermineProjectRootDirectory(context.Node.SyntaxTree.FilePath);
+            }
+
             // any method with at least one attribute is a candidate of function generation
             if (context.Node is MethodDeclarationSyntax methodDeclarationSyntax && methodDeclarationSyntax.AttributeLists.Count > 0)
             {