1 files changed, 74 insertions, 6 deletions
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/SqlUnitModel.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/SqlUnitModel.java
index bb6fa14..568bbe7 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/SqlUnitModel.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/SqlUnitModel.java
@@ -2,17 +2,41 @@ package io.trygvis.container.compiler;
 
 import io.trygvis.container.compiler.model.TypeRef;
 import io.trygvis.persistence.EntityMirror;
+import io.trygvis.persistence.SequenceMirror;
 
 import javax.lang.model.element.Element;
-import java.util.Collections;
 import java.util.Map;
 import java.util.TreeMap;
 
-import static java.util.Collections.addAll;
-
 public class SqlUnitModel {
     private Map<TypeRef, EntityMirror> entities = new TreeMap<>();
-    private Map<EntityMirror, Element> elements = new TreeMap<>();
+    private Map<EntityMirror, Element> entityElements = new TreeMap<>();
+    private Map<String, SequenceMirror> sequences = new TreeMap<>();
+    private Map<SequenceMirror, Element> sequenceElements = new TreeMap<>();
+    private String packageName;
+
+    public String getPackageName() {
+        if (packageName == null) {
+            throw new InternalErrorException("packageName is not set yet.");
+        }
+
+        return packageName;
+    }
+
+    public void setPackageName(String packageName) {
+        this.packageName = packageName;
+    }
+
+    // -----------------------------------------------------------------------
+    // Entity Mirrors
+    // -----------------------------------------------------------------------
+
+    public SqlUnitModel add(EntityMirror entity, Element element) {
+        entities.put(entity.type, entity);
+        entityElements.put(entity, element);
+
+        return this;
+    }
 
     public SqlUnitModel add(EntityMirror... entities) {
         for (EntityMirror entity : entities) {
@@ -34,7 +58,51 @@ public class SqlUnitModel {
         return entities;
     }
 
-    public Element elementForEntity(EntityMirror entity) {
-        return elements.get(entity);
+    public Element element(EntityMirror entity) {
+        return entityElements.get(entity);
+    }
+
+    // -----------------------------------------------------------------------
+    // Sequences
+    // -----------------------------------------------------------------------
+
+    public SqlUnitModel add(SequenceMirror sequence, Element element) {
+        sequences.put(sequence.name, sequence);
+        sequenceElements.put(sequence, element);
+        return this;
+    }
+
+    public SqlUnitModel add(SequenceMirror... sequenceMirrors) {
+        for (SequenceMirror sequenceMirror : sequenceMirrors) {
+            this.sequences.put(sequenceMirror.name, sequenceMirror);
+        }
+        return this;
+    }
+
+    public SequenceMirror getDefaultSequence() {
+        switch (sequences.size()) {
+            case 1:
+                return sequences.values().iterator().next();
+            case 0:
+                throw new CompilerException("No sequenced defined");
+            default:
+                throw new CompilerException("More than one sequence defined, the field must specify which sequence to use with @SequenceGenerator");
+        }
+    }
+
+    public SequenceMirror getSequence(String name) {
+        SequenceMirror sequence = sequences.get(name);
+        if (sequence == null) {
+            throw new CompilerException("Could not find sequence: " + name);
+        }
+        return sequence;
+    }
+
+    public Map<String, SequenceMirror> getSequences() {
+        return sequences;
+    }
+
+    public Element element(SequenceMirror sequence) {
+        return sequenceElements.get(sequence);
     }
 }