1 files changed, 33 insertions, 0 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
new file mode 100644
index 0000000..20d5f3b
--- /dev/null
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/SqlUnitModel.java
@@ -0,0 +1,33 @@
+package io.trygvis.container.compiler;
+
+import io.trygvis.container.compiler.model.TypeRef;
+import io.trygvis.persistence.EntityMirror;
+
+import javax.lang.model.element.Element;
+import java.util.Map;
+import java.util.TreeMap;
+
+public class SqlUnitModel {
+    private Map<TypeRef, EntityMirror> entities = new TreeMap<>();
+    private Map<EntityMirror, Element> elements = new TreeMap<>();
+
+    public void add(EntityMirror entity) {
+        entities.put(entity.type, entity);
+    }
+
+    public EntityMirror get(TypeRef type) {
+        EntityMirror entity = entities.get(type);
+        if (entity == null) {
+            throw new InternalErrorException("Could not find entity for " + type.fqName);
+        }
+        return entity;
+    }
+
+    public Map<TypeRef, EntityMirror> getEntities() {
+        return entities;
+    }
+
+    public Element elementForEntity(EntityMirror entity) {
+        return elements.get(entity);
+    }
+}