1 files changed, 21 insertions, 6 deletions
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/MyProcessor.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/MyProcessor.java
index 735d688..43ba49e 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/MyProcessor.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/MyProcessor.java
@@ -2,6 +2,7 @@ package io.trygvis.container.compiler;
 
 import io.trygvis.container.log.Log;
 import io.trygvis.persistence.SqlEntity;
+import io.trygvis.persistence.SqlEntitySet;
 import org.springframework.transaction.annotation.Transactional;
 
 import javax.annotation.processing.Completion;
@@ -12,8 +13,10 @@ import javax.lang.model.SourceVersion;
 import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.Element;
 import javax.lang.model.element.ExecutableElement;
+import javax.lang.model.element.PackageElement;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.type.DeclaredType;
+import javax.lang.model.util.ElementFilter;
 import javax.lang.model.util.Elements;
 import javax.lang.model.util.Types;
 import javax.persistence.Entity;
@@ -23,6 +26,7 @@ import java.util.Set;
 import static java.util.Arrays.asList;
 import static java.util.Collections.emptyList;
 import static java.util.Collections.emptySet;
+import static javax.lang.model.util.ElementFilter.typesIn;
 
 public class MyProcessor implements Processor {
 
@@ -57,13 +61,14 @@ public class MyProcessor implements Processor {
         return new HashSet<>(asList(
                 Transactional.class.getName(),
                 Log.class.getName(),
-                Entity.class.getName()));
+                Entity.class.getName(),
+                SqlEntitySet.class.getName()));
     }
 
     @Override
     public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
         try {
-            return work(annotations, roundEnv);
+            return work(roundEnv);
         } catch (RuntimeException e) {
             throw e;
         } catch (Exception e) {
@@ -71,18 +76,25 @@ public class MyProcessor implements Processor {
         }
     }
 
-    public boolean work(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) throws Exception {
+    boolean first = true;
+    public boolean work(RoundEnvironment roundEnv) throws Exception {
+        if (!first) {
+            return false;
+        }
+        first = false;
+
         TypeElement tx = elements.getTypeElement(Transactional.class.getCanonicalName());
         TypeElement log = elements.getTypeElement(Log.class.getCanonicalName());
         TypeElement entity = elements.getTypeElement(Entity.class.getCanonicalName());
 
         EntityHandler entityHandler = new EntityHandler(processingEnv);
 
-        Set<? extends Element> sqlEntities = roundEnv.getElementsAnnotatedWith(SqlEntity.class);
-        entityHandler.phase1(sqlEntities);
+        Set<TypeElement> sqlEntities = typesIn(roundEnv.getElementsAnnotatedWith(SqlEntity.class));
+        Set<PackageElement> packages = ElementFilter.packagesIn(roundEnv.getElementsAnnotatedWith(SqlEntitySet.class));
+        entityHandler.phase1(sqlEntities, packages);
 
         for (Element element : roundEnv.getRootElements()) {
-            System.out.println("Processing: " + element);
+            System.out.println("Processing: " + element.asType());
             for (AnnotationMirror annotationMirror : element.getAnnotationMirrors()) {
                 DeclaredType annotationType = annotationMirror.getAnnotationType();
 
@@ -97,6 +109,9 @@ public class MyProcessor implements Processor {
                 }
             }
         }
+
+        entityHandler.phase3();
+
         return true;
     }
 }