feat: add strict stages mapping validation (#741)

Author: davidfrigolet

core/flamingock-core-api/src/main/java/io/flamingock/api/annotations/EnableFlamingock.java

@@ -209,4 +209,12 @@
      * @see SetupType
      */
     SetupType setup() default SetupType.DEFAULT;
+
+    /**
+     * If true, the annotation processor will validate that all code-based changes
+     * (classes annotated with @Change) are mapped to some stage. When unmapped changes
+     * are found and this flag is true, a RuntimeException is thrown. Default is false
+     * (only a warning is emitted).
+     */
+    boolean strictStageMapping() default false;
 }

core/flamingock-processor/src/main/java/io/flamingock/core/processor/FlamingockAnnotationProcessor.java

@@ -51,6 +51,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.BiPredicate;
 import java.util.stream.Collectors;
 
 /**
@@ -209,6 +210,9 @@ public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment
                 noLegacyChangesByPackage,
                 flamingockAnnotation
         );
+
+        validateAllChangesAreMappedToStages(legacyChanges, noLegacyChangesByPackage, pipeline, flamingockAnnotation.strictStageMapping());
+
         Serializer serializer = new Serializer(processingEnv, logger);
         String setup = flamingockAnnotation.setup().toString();
         String configFile = flamingockAnnotation.configFile();
@@ -721,4 +725,92 @@ private void validateConfiguration(EnableFlamingock pipelineAnnotation, boolean
         }
     }
 
+    /**
+     * Validates that code-based preview changes (non-legacy) are attached to some stage.
+     * It checks package names from the changes map against the pipeline stages' sourcesPackage
+     * (including the SYSTEM stage). If any stage does not represent any package, a RuntimeException is thrown.
+     *
+     * @param legacyCodedChanges legacy code changes (kept for completeness; not used for package mapping)
+     * @param noLegacyCodedChangesByPackage non-legacy changes grouped by source package
+     * @param pipeline built preview pipeline
+     * @param strictStageMapping if true, unmapped changes will throw an exception; if false, will log a warning
+     */
+    private void validateAllChangesAreMappedToStages(List<CodePreviewChange> legacyCodedChanges,
+                                                     Map<String, List<CodePreviewChange>> noLegacyCodedChangesByPackage,
+                                                     PreviewPipeline pipeline,
+                                                     Boolean strictStageMapping) {
+        if (noLegacyCodedChangesByPackage == null || noLegacyCodedChangesByPackage.isEmpty()) {
+            return;
+        }
+
+        // Helper to test if a change package is covered by a stage package (exact or parent)
+        BiPredicate<String, String> covers = (stagePkg, changePkg) -> {
+            if (stagePkg == null || changePkg == null) return false;
+            return changePkg.equals(stagePkg) || changePkg.startsWith(stagePkg + ".");
+        };
+
+        List<String> unmapped = new ArrayList<>();
+
+        for (String pkg : noLegacyCodedChangesByPackage.keySet()) {
+            if (pkg == null) {
+                continue;
+            }
+            boolean matched = false;
+
+            // Check system stage
+            PreviewStage system = pipeline.getSystemStage();
+            if (system != null) {
+                String sysPkg = system.getSourcesPackage();
+                if (covers.test(sysPkg, pkg)) {
+                    matched = true;
+                } else if (system.getTasks() != null) {
+                    for (io.flamingock.internal.common.core.preview.AbstractPreviewTask task : system.getTasks()) {
+                        if (task instanceof CodePreviewChange) {
+                            String taskPkg = ((CodePreviewChange) task).getSourcePackage();
+                            if (covers.test(taskPkg, pkg)) {
+                                matched = true;
+                                break;
+                            }
+                        }
+                    }
+                }
+            }
+
+            // Check regular stages
+            if (!matched && pipeline.getStages() != null) {
+                for (PreviewStage stage : pipeline.getStages()) {
+                    String stagePkg = stage.getSourcesPackage();
+                    if (covers.test(stagePkg, pkg)) {
+                        matched = true;
+                        break;
+                    }
+                    if (stage.getTasks() != null) {
+                        for (io.flamingock.internal.common.core.preview.AbstractPreviewTask task : stage.getTasks()) {
+                            if (task instanceof CodePreviewChange) {
+                                String taskPkg = ((CodePreviewChange) task).getSourcePackage();
+                                if (covers.test(taskPkg, pkg)) {
+                                    matched = true;
+                                    break;
+                                }
+                            }
+                        }
+                        if (matched) break;
+                    }
+                }
+            }
+
+            if (!matched) {
+                unmapped.add(pkg);
+            }
+        }
+
+        if (!unmapped.isEmpty()) {
+            String message = "Changes are not mapped to any stage: " + String.join(", ", unmapped);
+            if (Boolean.TRUE.equals(strictStageMapping)) {
+                throw new RuntimeException(message);
+            } else {
+                logger.warn(message);
+            }
+        }
+    }
 }

core/flamingock-processor/src/test/java/io/flamingock/core/processor/PipelinePreProcessorTest.java

@@ -31,6 +31,8 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Collections;
@@ -391,6 +393,47 @@ void shouldThrowErrorForMultipleSystemStagesInYaml() throws Exception {
         }
     }
 
+    @Test
+    @DisplayName("Should throw error when changes exist that are not mapped to any stage and mapping is strict")
+    void shouldThrowErrorForUnmappedChanges() throws Exception {
+        FlamingockAnnotationProcessor processor = new FlamingockAnnotationProcessor();
+
+        // Prepare changes map with one mapped and one unmapped package (lists can be empty for this test)
+        Map<String, List<io.flamingock.internal.common.core.preview.CodePreviewChange>> changesByPackage = new HashMap<>();
+        changesByPackage.put("com.example.migrations", Collections.emptyList());
+        changesByPackage.put("com.example.unmapped", Collections.emptyList()); // Unmapped
+
+        // Provide a non-empty changes list for the mapped stage so the builder doesn't throw
+        io.flamingock.internal.common.core.preview.CodePreviewChange mockChange = new io.flamingock.internal.common.core.preview.CodePreviewChange();
+
+        // Build a pipeline with a single stage covering com.example.migrations
+        PreviewStage mappedStage = PreviewStage.defaultBuilder(StageType.DEFAULT)
+                .setName("migrations")
+                .setSourcesPackage("com.example.migrations")
+                .setSourcesRoots(Collections.emptyList())
+                .setResourcesRoot("src/main/resources")
+                .setChanges(Collections.singletonList(mockChange)) // <- non-empty list
+                .build();
+
+        PreviewPipeline pipeline = new PreviewPipeline(Collections.singletonList(mappedStage));
+
+        // Invoke private validator by reflection and assert it throws (InvocationTargetException wraps the RuntimeException)
+        Method validator = FlamingockAnnotationProcessor.class.getDeclaredMethod(
+                "validateAllChangesAreMappedToStages", List.class, Map.class, PreviewPipeline.class, Boolean.class);
+        validator.setAccessible(true);
+
+        InvocationTargetException ex = assertThrows(InvocationTargetException.class, () ->
+                validator.invoke(processor, Collections.emptyList(), changesByPackage, pipeline, true));
+
+        Throwable cause = ex.getCause();
+        assertNotNull(cause, "Validator should throw a RuntimeException cause");
+        assertInstanceOf(RuntimeException.class, cause, "Cause should be RuntimeException");
+        assertTrue(cause.getMessage().contains("not mapped to any stage"), "Should have error about unmapped changes");
+        assertTrue(cause.getMessage().contains("com.example.unmapped"), "Should mention the unmapped package");
+    }
+
+
+
     // Helper methods using reflection to test the internal pipeline building logic
     private PreviewPipeline buildPipelineFromAnnotation(FlamingockAnnotationProcessor processor, EnableFlamingock annotation, Map<String, List<AbstractPreviewTask>> changes) throws Exception {
         // Set up minimal processor state
@@ -585,6 +628,12 @@ public EnableFlamingock build() {
                 @Override public Stage[] stages() { return stages; }
                 @Override public String configFile() { return configFile; }
                 @Override public io.flamingock.api.SetupType setup() { return io.flamingock.api.SetupType.DEFAULT; }
+
+                @Override
+                public boolean strictStageMapping() {
+                    return true;
+                }
+
                 @Override public Class<? extends java.lang.annotation.Annotation> annotationType() { return EnableFlamingock.class; }
             };
         }