wip

14 files changed, 289 insertions, 80 deletions

diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/CompilerException.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/CompilerException.java
index a3801ab..e9fcb87 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/CompilerException.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/CompilerException.java
@@ -5,6 +5,10 @@ import javax.lang.model.element.Element;
 public class CompilerException extends RuntimeException {
     public final Element element;
 
+    public CompilerException(String message) {
+        this(null, message);
+    }
+
     public CompilerException(Element element, String message) {
         super(message);
         this.element = element;
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/EntityHandler.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/EntityHandler.java
index cf21887..c623af7 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/EntityHandler.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/EntityHandler.java
@@ -4,7 +4,7 @@ import io.trygvis.container.compiler.model.ClassG;
 import io.trygvis.container.compiler.model.TypeRef;
 import io.trygvis.persistence.EntityMirror;
 import io.trygvis.persistence.FieldMirror;
-import io.trygvis.persistence.GeneratorSupport;
+import io.trygvis.persistence.GeneratorConfiguration;
 import io.trygvis.persistence.SqlEntity;
 import io.trygvis.persistence.TypeHandler;
 
@@ -18,6 +18,7 @@ import javax.lang.model.element.VariableElement;
 import javax.lang.model.type.TypeMirror;
 import javax.persistence.Id;
 import javax.tools.JavaFileObject;
+import java.io.IOException;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.List;
@@ -30,8 +31,8 @@ import static java.lang.Character.isUpperCase;
 import static javax.lang.model.util.ElementFilter.fieldsIn;
 
 public class EntityHandler extends AbstractHandler {
-    private GeneratorSupport generatorSupport = new GeneratorSupport();
-    private List<EntityMirror> entities = new ArrayList<>();
+    private GeneratorConfiguration generatorConfiguration = new GeneratorConfiguration();
+    private SqlUnitModel sqlUnit = new SqlUnitModel();
     private PackageElement packageElement;
 
     public EntityHandler(ProcessingEnvironment processingEnv) {
@@ -48,7 +49,7 @@ public class EntityHandler extends AbstractHandler {
 
                         TypeHandler typeHandler = (TypeHandler) typeHandlerClass.newInstance();
                         String type = entity.asType().toString();
-                        generatorSupport.addTypeHandler(new TypeRef(type, type), typeHandler);
+                        generatorConfiguration.addTypeHandler(new TypeRef(type, type), typeHandler);
                         System.out.println("Loaded TypeHandler for " + type + " through " + typeHandlerClass.getCanonicalName());
                         break;
                 }
@@ -58,7 +59,7 @@ public class EntityHandler extends AbstractHandler {
 
         System.out.println("packages = " + packages);
         if (packages.size() == 0) {
-            throw new CompilerException(null, "There has to be exactly one @SqlEntitySet annotated package.");
+            throw new CompilerException("There has to be exactly one @SqlEntitySet annotated package.");
         }
         packageElement = packages.iterator().next();
 
@@ -78,12 +79,12 @@ public class EntityHandler extends AbstractHandler {
         throw new CompilerException(type, "Could not find annotation " + c.getSimpleName());
     }
 
-    public void processEntity(TypeElement element) throws Exception {
-        EntityMirror entityMirror = new EntityMirror(generatorSupport, new TypeRef(types.getDeclaredType(element)),
+    public void recordEntity(TypeElement element) throws Exception {
+        EntityMirror entityMirror = new EntityMirror(generatorConfiguration, new TypeRef(types.getDeclaredType(element)),
                 sqlName(element.getSimpleName().toString()));
 
         for (VariableElement f : fieldsIn(elements.getAllMembers(element))) {
-            entityMirror.add(fromElement(generatorSupport, f));
+            entityMirror.add(fromElement(generatorConfiguration, f));
         }
 
         List<FieldMirror> idFields = new ArrayList<>();
@@ -100,44 +101,47 @@ public class EntityHandler extends AbstractHandler {
             throw new CompilerException(element, "This implementation only support a single @Id annotated field.");
         }
 
-        String p = elements.getPackageOf(element).getQualifiedName().toString();
+        sqlUnit.add(entityMirror);
+    }
 
+    public void generate(EntityMirror entityMirror) throws IOException {
         ClassG g = new ClassG(entityMirror.daoType);
-        g.addPublicStaticFinalField(String.class, "createTableSql").value(toJavaString(entityMirror.createTableSql()));
+        g.addPublicStaticFinalField(String.class, "createTableSql").value(toJavaString(entityMirror.createTableSql(sqlUnit)));
         g.addPublicStaticFinalField(String.class, "dropTableSql").value(toJavaString(entityMirror.dropTableSql()));
         g.addPublicStaticFinalField(String.class, "insertIntoSql").value(toJavaString(entityMirror.insertIntoSql()));
         g.addPublicStaticFinalField(String.class, "deleteFromSql").value(toJavaString(entityMirror.deleteFromSql()));
-        entityMirror.insertInto(g);
+        entityMirror.insertInto(sqlUnit, g);
+        entityMirror.delete(g);
+        entityMirror.deleteById(g);
 
-        JavaFileObject sourceFile = processingEnv.getFiler().createSourceFile(entityMirror.daoType.fqName, element);
+        JavaFileObject sourceFile = processingEnv.getFiler().createSourceFile(entityMirror.daoType.fqName, sqlUnit.elementForEntity(entityMirror));
         try (PrintWriter w = new PrintWriter(sourceFile.openWriter())) {
             g.write(w);
         }
-
-        entities.add(entityMirror);
     }
 
-    public FieldMirror fromElement(GeneratorSupport generatorSupport, VariableElement var) {
-        TypeRef type = new TypeRef(var.asType().toString());
+    public FieldMirror fromElement(GeneratorConfiguration generatorConfiguration, VariableElement var) {
+        TypeRef type = new TypeRef(var.asType());
 //        System.out.print("element = ");
 //        elements.printElements(new PrintWriter(System.out), var);
         String javaName = var.getSimpleName().toString();
         String sqlName = sqlName(javaName);
         boolean notNull = false;
-        FieldMirror field;
+        boolean unique = false;
         boolean id = isId(var);
-        if (generatorSupport.isPrimitive(type)) {
-            field = new PrimitiveFieldMirror(type, javaName, sqlName, id, notNull);
-        } else if (generatorSupport.hasTypeHandler(type)) {
-            if (id) {
-                throw new CompilerException(var, "A @Id field has to be a primitive or embedded.");
-            }
+        boolean primitive = generatorConfiguration.isPrimitive(type);
+        if (id && !primitive) {
+            throw new CompilerException(var, "A @Id field has to be a primitive or embedded.");
+        }
 
-            field = new FieldMirror.ReferenceFieldMirror(type, javaName, sqlName, notNull);
+        FieldMirror field;
+        if (primitive) {
+            field = new PrimitiveFieldMirror(type, javaName, sqlName, id, notNull, unique);
+        } else if (generatorConfiguration.hasTypeHandler(type)) {
+            throw new CompilerException(var, "Missing type handler for type: " + type.fqName);
         } else {
-            throw new CompilerException(var, "Missing type handler for type: " + type);
+            field = new FieldMirror.ReferenceFieldMirror(type, javaName, sqlName, notNull, unique);
         }
-//        System.out.println("field = " + field);
         return field;
     }
 
@@ -145,12 +149,18 @@ public class EntityHandler extends AbstractHandler {
         return var.getAnnotation(Id.class) != null;
     }
 
-    // TODO: Rename DAO to Session.
     public void phase3() throws Exception {
+        for (EntityMirror entity : sqlUnit.getEntities().values()) {
+            generate(entity);
+        }
+        generateSession();
+    }
+
+    private void generateSession() throws IOException {
         String p = packageElement.getQualifiedName().toString();
 
         // TODO: Support a name prefix from @SqlEntitySet
-        TypeRef type = new TypeRef(p + ".Daos");
+        TypeRef type = new TypeRef(p + ".Session");
 
         ClassG g = new ClassG(type);
 
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/InternalErrorException.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/InternalErrorException.java
new file mode 100644
index 0000000..dfb6746
--- /dev/null
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/InternalErrorException.java
@@ -0,0 +1,9 @@
+package io.trygvis.container.compiler;
+
+public class InternalErrorException extends RuntimeException {
+
+    public InternalErrorException(String message) {
+        super(message);
+    }
+
+}
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 713b948..297b9d8 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
@@ -77,9 +77,15 @@ public class MyProcessor implements Processor {
             System.out.println("claimed = " + claimed);
             return claimed;
         } catch (CompilerException e) {
+            e.printStackTrace(System.err);
             Messager messager = processingEnv.getMessager();
             messager.printMessage(Diagnostic.Kind.ERROR, e.getMessage(), e.element);
             return true;
+        } catch (InternalErrorException e) {
+            e.printStackTrace(System.err);
+            Messager messager = processingEnv.getMessager();
+            messager.printMessage(Diagnostic.Kind.ERROR, e.getMessage());
+            return true;
         } catch (RuntimeException e) {
             throw e;
         } catch (Exception e) {
@@ -116,7 +122,7 @@ public class MyProcessor implements Processor {
                     new LogHandler(processingEnv).processLog((TypeElement) element);
                 }
                 if (types.isSameType(entity.asType(), annotationType)) {
-                    entityHandler.processEntity((TypeElement) element);
+                    entityHandler.recordEntity((TypeElement) element);
                 }
             }
         }
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/NotImplementedException.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/NotImplementedException.java
new file mode 100644
index 0000000..a1941b6
--- /dev/null
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/NotImplementedException.java
@@ -0,0 +1,7 @@
+package io.trygvis.container.compiler;
+
+public class NotImplementedException extends InternalErrorException {
+    public NotImplementedException() {
+        super("Not implemented");
+    }
+}
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);
+    }
+}
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/persistence/EntityMirror.java b/container-compiler-plugin/src/main/java/io/trygvis/persistence/EntityMirror.java
index 2a4952f..8e1602c 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/persistence/EntityMirror.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/persistence/EntityMirror.java
@@ -1,5 +1,7 @@
 package io.trygvis.persistence;
 
+import io.trygvis.container.compiler.NotImplementedException;
+import io.trygvis.container.compiler.SqlUnitModel;
 import io.trygvis.container.compiler.model.ClassG;
 import io.trygvis.container.compiler.model.ParameterRef;
 import io.trygvis.container.compiler.model.TypeRef;
@@ -10,19 +12,21 @@ import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
 
+import static io.trygvis.persistence.FieldMirror.PrimitiveFieldMirror;
+import static io.trygvis.persistence.FieldMirror.ReferenceFieldMirror;
 import static java.lang.String.format;
 import static org.apache.commons.lang.StringUtils.join;
 
-public class EntityMirror {
-    public final GeneratorSupport generatorSupport;
+public class EntityMirror implements Comparable<EntityMirror> {
+    public final GeneratorConfiguration generatorConfiguration;
     public final List<FieldMirror> fields = new ArrayList<>();
     public final List<FieldMirror> idFields = new ArrayList<>();
     public final TypeRef type;
     public final String tableName;
     public final TypeRef daoType;
 
-    public EntityMirror(GeneratorSupport generatorSupport, TypeRef type, String tableName) {
-        this.generatorSupport = generatorSupport;
+    public EntityMirror(GeneratorConfiguration generatorConfiguration, TypeRef type, String tableName) {
+        this.generatorConfiguration = generatorConfiguration;
         this.type = type;
         this.tableName = tableName;
 
@@ -31,18 +35,48 @@ public class EntityMirror {
 
     public void add(FieldMirror field) {
         fields.add(field);
-        if(field.id) {
+        if (field.id) {
             idFields.add(field);
         }
     }
 
-    public String createTableSql() {
+    public FieldMirror getIdField() {
+        return idFields.get(0);
+    }
+
+    public String createTableSql(SqlUnitModel unit) {
         List<String> columns = new ArrayList<>();
         for (FieldMirror field : fields) {
-            TypeHandler typeHandler = generatorSupport.typeHandler(field.type);
-            String s = "  " + field.sqlName + " " + typeHandler.sqlType(field);
-            if(field.id) {
-                s += " PRIMARY KEY";
+            String s;
+            if (field instanceof PrimitiveFieldMirror) {
+                TypeHandler typeHandler = generatorConfiguration.typeHandler(field.type);
+                s = "  " + field.sqlName + " " + typeHandler.sqlType(field);
+                if (field.id) {
+                    s += " PRIMARY KEY";
+                } else if (field.notNull) {
+                    s += " NOT NULL";
+                } else if (field.unique) {
+                    s += " UNIQUE";
+                }
+            } else if (field instanceof ReferenceFieldMirror) {
+                ReferenceFieldMirror ref = (ReferenceFieldMirror) field;
+                EntityMirror referenced = unit.get(ref.type);
+                if(referenced.idFields.size() == 1) {
+                    FieldMirror idField = referenced.idFields.get(0);
+                    TypeHandler typeHandler = generatorConfiguration.typeHandler(idField.type);
+                    s = "  " + field.sqlName + " " + typeHandler.sqlType(field);
+                    s += " REFERENCES " + referenced.tableName + "(" + idField.sqlName + ")";
+                    if (field.notNull) {
+                        s += " NOT NULL";
+                    } else if (field.unique) {
+                        s += " UNIQUE";
+                    }
+                }
+                else {
+                    throw new NotImplementedException();
+                }
+            } else {
+                throw new RuntimeException("Unknown field type: " + field.getClass());
             }
             columns.add(s);
         }
@@ -77,7 +111,7 @@ public class EntityMirror {
         return "DELETE FROM " + tableName + " WHERE " + join(ss, " AND ") + ";";
     }
 
-    public void insertInto(ClassG g) {
+    public void insertInto(SqlUnitModel unit, ClassG g) {
         TypeRef conType = g.addImport(Connection.class);
         TypeRef psType = g.addImport(PreparedStatement.class);
         TypeRef objectType = g.addImport(type);
@@ -89,8 +123,16 @@ public class EntityMirror {
         body.add("try(" + psType.name + " stmt = " + con.name + ".prepareStatement(insertIntoSql)) {");
         for (int i = 0; i < fields.size(); i++) {
             FieldMirror field = fields.get(i);
-            TypeHandler typeHandler = generatorSupport.typeHandler(field.type);
-            body.add("    stmt." + typeHandler.resultSetSetter(i + 1, "o", field) + ";");
+            if(field instanceof PrimitiveFieldMirror) {
+                TypeHandler typeHandler = generatorConfiguration.typeHandler(field.type);
+                body.add("    stmt." + typeHandler.setter(i + 1, "o." + field.javaName) + ";");
+            } else if (field instanceof ReferenceFieldMirror) {
+                ReferenceFieldMirror ref = (ReferenceFieldMirror) field;
+                EntityMirror referenced = unit.get(ref.type);
+                FieldMirror idField = referenced.getIdField();
+                TypeHandler typeHandler = generatorConfiguration.typeHandler(idField.type);
+                body.add("    stmt." + typeHandler.setter(i + 1, "o." + field.javaName + "." + idField.javaName) + ";");
+            }
         }
         body.add("    stmt.executeUpdate();");
         body.add("}");
@@ -98,4 +140,67 @@ public class EntityMirror {
         g.addStaticMethod(body, TypeRef.VOID, "insertInto", con, object).
                 exception(g.addImport(SQLException.class));
     }
+
+    public void delete(ClassG g) {
+        TypeRef conType = g.addImport(Connection.class);
+        TypeRef objectType = g.addImport(type);
+        ParameterRef con = new ParameterRef(conType, "con");
+        ParameterRef o = new ParameterRef(objectType, "o");
+
+        List<String> arguments = new ArrayList<>();
+        arguments.add(con.name);
+        for (FieldMirror field : idFields) {
+            arguments.add(o.name + "." + field.javaName);
+        }
+        List<String> body = new ArrayList<>();
+        body.add("deleteById(" + join(arguments, ", ") + ");");
+
+        g.addStaticMethod(body, TypeRef.VOID, "delete", con, o).
+                exception(g.addImport(SQLException.class));
+    }
+
+    public void deleteById(ClassG g) {
+        TypeRef conType = g.addImport(Connection.class);
+        TypeRef psType = g.addImport(PreparedStatement.class);
+        ParameterRef con = new ParameterRef(conType, "con");
+
+        ParameterRef[] parameters = new ParameterRef[idFields.size() + 1];
+        parameters[0] = con;
+
+        List<String> body = new ArrayList<>();
+
+        body.add("try(" + psType.name + " stmt = " + con.name + ".prepareStatement(deleteFromSql)) {");
+        for (int i = 0; i < idFields.size(); i++) {
+            FieldMirror field = idFields.get(i);
+            TypeHandler typeHandler = generatorConfiguration.typeHandler(field.type);
+            ParameterRef p = new ParameterRef(field.type, field.javaName);
+            parameters[i + 1] = p;
+            body.add("    stmt." + typeHandler.setter(i + 1, field.javaName) + ";");
+        }
+        body.add("    stmt.executeUpdate();");
+        body.add("}");
+
+        g.addStaticMethod(body, TypeRef.VOID, "deleteById", parameters).
+                exception(g.addImport(SQLException.class));
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (!(o instanceof EntityMirror)) return false;
+
+        EntityMirror that = (EntityMirror) o;
+
+        return type.equals(that.type);
+    }
+
+    @Override
+    public int hashCode() {
+        return type.hashCode();
+    }
+
+    @Override
+    public int compareTo(EntityMirror o) {
+        return type.compareTo(o.type);
+    }
 }
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/persistence/FieldMirror.java b/container-compiler-plugin/src/main/java/io/trygvis/persistence/FieldMirror.java
index bf0a0c3..f92cc2c 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/persistence/FieldMirror.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/persistence/FieldMirror.java
@@ -15,24 +15,26 @@ public abstract class FieldMirror {
     public final String sqlName;
     public final boolean id;
     public final boolean notNull;
+    public final boolean unique;
 
     public enum FieldType {
         PRIMITIVE,
         REFERENCE,
     }
 
-    protected FieldMirror(FieldType fieldType, TypeRef type, String javaName, String sqlName, boolean id, boolean notNull) {
+    protected FieldMirror(FieldType fieldType, TypeRef type, String javaName, String sqlName, boolean id, boolean notNull, boolean unique) {
         this.fieldType = fieldType;
         this.type = type;
         this.javaName = javaName;
         this.sqlName = sqlName;
         this.id = id;
         this.notNull = notNull;
+        this.unique = unique;
     }
 
     public static class PrimitiveFieldMirror extends FieldMirror {
-        public PrimitiveFieldMirror(TypeRef typeRef, String javaName, String sqlName, boolean id, boolean notNull) {
-            super(PRIMITIVE, typeRef, javaName, sqlName, id, notNull);
+        public PrimitiveFieldMirror(TypeRef typeRef, String javaName, String sqlName, boolean id, boolean notNull, boolean unique) {
+            super(PRIMITIVE, typeRef, javaName, sqlName, id, notNull, unique);
         }
 
         @Override
@@ -42,13 +44,14 @@ public abstract class FieldMirror {
                     ", javaName='" + javaName + '\'' +
                     ", sqlName='" + sqlName + '\'' +
                     ", notNull=" + notNull +
+                    ", unique=" + unique +
                     '}';
         }
     }
 
-    public static class  ReferenceFieldMirror extends FieldMirror {
-        public ReferenceFieldMirror(TypeRef typeRef, String javaName, String sqlName, boolean notNull) {
-            super(REFERENCE, typeRef, javaName, sqlName, false, notNull);
+    public static class ReferenceFieldMirror extends FieldMirror {
+        public ReferenceFieldMirror(TypeRef typeRef, String javaName, String sqlName, boolean notNull, boolean unique) {
+            super(REFERENCE, typeRef, javaName, sqlName, false, notNull, unique);
         }
 
         @Override
@@ -58,6 +61,7 @@ public abstract class FieldMirror {
                     ", javaName='" + javaName + '\'' +
                     ", sqlName='" + sqlName + '\'' +
                     ", notNull=" + notNull +
+                    ", unique=" + unique +
                     '}';
         }
     }
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/persistence/GeneratorSupport.java b/container-compiler-plugin/src/main/java/io/trygvis/persistence/GeneratorConfiguration.java
index 0563691..37d8146 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/persistence/GeneratorSupport.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/persistence/GeneratorConfiguration.java
@@ -1,12 +1,13 @@
 package io.trygvis.persistence;
 
+import io.trygvis.container.compiler.CompilerException;
 import io.trygvis.container.compiler.model.TypeRef;
 
 import java.util.Date;
 import java.util.Map;
 import java.util.TreeMap;
 
-public class GeneratorSupport {
+public class GeneratorConfiguration {
 
     private final Map<TypeRef, TypeHandler> primitiveTypeHandlers = new TreeMap<>();
     private final Map<TypeRef, TypeHandler> typeHandlers = new TreeMap<>();
@@ -26,8 +27,9 @@ public class GeneratorSupport {
 
     public TypeHandler typeHandler(TypeRef type) {
         TypeHandler typeHandler = typeHandlers.get(type);
-        if (typeHandler == null)
-            throw new RuntimeException("Unsupported field type: " + type);
+        if (typeHandler == null) {
+            throw new CompilerException("Unsupported field type: " + type.fqName);
+        }
         return typeHandler;
     }
 
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/persistence/TypeHandler.java b/container-compiler-plugin/src/main/java/io/trygvis/persistence/TypeHandler.java
index ff2ad0a..3aacd7a 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/persistence/TypeHandler.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/persistence/TypeHandler.java
@@ -1,14 +1,14 @@
 package io.trygvis.persistence;
 
 public interface TypeHandler {
-    String resultSetSetter(int i, String o, FieldMirror field);
+    String setter(int i, String ref);
 
     String sqlType(FieldMirror field);
 
     public static class IntTypeHandler implements TypeHandler {
         @Override
-        public String resultSetSetter(int i, String o, FieldMirror field) {
-            return "setInt(" + i + ", " + o + "." + field.javaName + ")";
+        public String setter(int i, String ref) {
+            return "setInt(" + i + ", " + ref + ")";
         }
 
         @Override
@@ -19,8 +19,8 @@ public interface TypeHandler {
 
     public static class LongTypeHandler implements TypeHandler {
         @Override
-        public String resultSetSetter(int i, String o, FieldMirror field) {
-            return "setLong(" + i + ", " + o + "." + field.javaName + ")";
+        public String setter(int i, String ref) {
+            return "setLong(" + i + ", " + ref + ")";
         }
 
         @Override
@@ -31,8 +31,8 @@ public interface TypeHandler {
 
     public static class StringTypeHandler implements TypeHandler {
         @Override
-        public String resultSetSetter(int i, String o, FieldMirror field) {
-            return "setString(" + i + ", " + o + "." + field.javaName + ")";
+        public String setter(int i, String ref) {
+            return "setString(" + i + ", " + ref + ")";
         }
 
         @Override
@@ -43,8 +43,8 @@ public interface TypeHandler {
 
     public static class DateTypeHandler implements TypeHandler {
         @Override
-        public String resultSetSetter(int i, String o, FieldMirror field) {
-            return "setTimestamp(" + i + ", new java.sql.Timestamp(" + o + "." + field.javaName + ".getTime()))";
+        public String setter(int i, String ref) {
+            return "setTimestamp(" + i + ", new java.sql.Timestamp(" + ref + ".getTime()))";
         }
 
         @Override
diff --git a/container-compiler-plugin/src/test/java/io/trygvis/container/compiler/ProcessorTest.java b/container-compiler-plugin/src/test/java/io/trygvis/container/compiler/ProcessorTest.java
index bac8209..b8d022d 100644
--- a/container-compiler-plugin/src/test/java/io/trygvis/container/compiler/ProcessorTest.java
+++ b/container-compiler-plugin/src/test/java/io/trygvis/container/compiler/ProcessorTest.java
@@ -38,21 +38,27 @@ public class ProcessorTest {
 
         task.setProcessors(asList(new MyProcessor()));
 
-        Boolean result = task.call();
+        boolean result = task.call();
+        if (!result) {
+            for (Diagnostic<? extends JavaFileObject> diagnostic : collector.getDiagnostics()) {
+                JavaFileObject source = diagnostic.getSource();
+                String error = "";
+                if (source != null) {
+                    error += source.toUri().getPath();
+                }
+                error +=  ":" + diagnostic.getLineNumber() + ":" + diagnostic.getColumnNumber();
+                System.out.println(error + ":" + diagnostic.getMessage(Locale.ENGLISH));
+            }
+        }
 
-        assertThat(fileManager.codes.keySet()).containsOnly(
-                "io.trygvis.persistence.test.Daos",
-                "io.trygvis.persistence.test.PersonDao");
         for (Map.Entry<String, String> entry : fileManager.codes.entrySet()) {
             System.out.println("=== " + entry.getKey());
             System.out.println(entry.getValue());
         }
 
-        for (Diagnostic<? extends JavaFileObject> diagnostic : collector.getDiagnostics()) {
-            System.out.println("diagnostic.source = ->" + diagnostic.getSource().getName() + "<-");
-            System.out.println("diagnostic.message = " + diagnostic.getMessage(Locale.ENGLISH));
-        }
-
+        assertThat(fileManager.codes.keySet()).containsOnly(
+                "io.trygvis.persistence.test.Session",
+                "io.trygvis.persistence.test.PersonDao");
         assertThat(collector.getDiagnostics()).isEmpty();
         assertThat(result).isTrue();
 
@@ -62,7 +68,7 @@ public class ProcessorTest {
     private JavaSourceFromString loadCode(String className) throws IOException {
         String path = "/" + className.replace('.', '/') + ".java";
         URL resource = getClass().getResource(path);
-        if(resource == null) {
+        if (resource == null) {
             throw new RuntimeException("Could not load code for: " + path);
         }
         return new JavaSourceFromString(className, IOUtils.toString(resource, UTF_8));
diff --git a/container-compiler-plugin/src/test/java/io/trygvis/persistence/EntityMirrorTest.java b/container-compiler-plugin/src/test/java/io/trygvis/persistence/EntityMirrorTest.java
index 33db82b..9d90eef 100644
--- a/container-compiler-plugin/src/test/java/io/trygvis/persistence/EntityMirrorTest.java
+++ b/container-compiler-plugin/src/test/java/io/trygvis/persistence/EntityMirrorTest.java
@@ -1,5 +1,6 @@
 package io.trygvis.persistence;
 
+import io.trygvis.container.compiler.SqlUnitModel;
 import io.trygvis.container.compiler.model.TypeRef;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
@@ -10,8 +11,11 @@ import static org.testng.Assert.assertEquals;
 
 public class EntityMirrorTest {
 
-    static FieldMirror idLong = new PrimitiveFieldMirror(new TypeRef(Long.class), "id", "id", true, false);
-    static FieldMirror idString = new PrimitiveFieldMirror(new TypeRef(String.class), "id", "id", true, false);
+    static TypeRef entityType = new TypeRef("Wat");
+    static FieldMirror idLong = new PrimitiveFieldMirror(new TypeRef(Long.class), "id", "id", true, false, true);
+    static FieldMirror idString = new PrimitiveFieldMirror(new TypeRef(String.class), "id", "id", true, false, true);
+    static FieldMirror name = new PrimitiveFieldMirror(new TypeRef(String.class), "name", "name", false, false, false);
+    static FieldMirror ref = new FieldMirror.ReferenceFieldMirror(entityType, "parent", "parent", false, false);
 
     @DataProvider(name = "insertIntoSql", parallel = true)
     public static Object[][] insertIntoProvider() {
@@ -30,18 +34,37 @@ public class EntityMirrorTest {
                         "CREATE TABLE my_table(%n" +
                                 "  id BIGINT PRIMARY KEY%n" +
                                 ");")},
+                new Object[]{
+                        new FieldMirror[]{idLong, name},
+                        "INSERT INTO my_table(id, name) VALUES(?, ?);",
+                        "DELETE FROM my_table WHERE id=?;", format(
+                        "CREATE TABLE my_table(%n" +
+                                "  id BIGINT PRIMARY KEY,%n" +
+                                "  name VARCHAR(1000)%n" +
+                                ");")},
+                new Object[]{
+                        new FieldMirror[]{idLong, ref},
+                        "INSERT INTO my_table(id, parent) VALUES(?, ?);",
+                        "DELETE FROM my_table WHERE id=?;", format(
+                        "CREATE TABLE my_table(%n" +
+                                "  id BIGINT PRIMARY KEY,%n" +
+                                "  parent BIGINT REFERENCES my_table(id)%n" +
+                                ");")},
         };
     }
 
     @Test(dataProvider = "insertIntoSql")
     public void testInsertIntoSql(FieldMirror[] fields, String insert, String delete, String create) throws Exception {
-        EntityMirror myTable = new EntityMirror(new GeneratorSupport(), new TypeRef("Wat"), "my_table");
+        EntityMirror myTable = new EntityMirror(new GeneratorConfiguration(), entityType, "my_table");
         for (FieldMirror field : fields) {
             myTable.add(field);
         }
 
+        SqlUnitModel unit = new SqlUnitModel();
+        unit.add(myTable);
+
         assertEquals(myTable.insertIntoSql(), insert);
         assertEquals(myTable.deleteFromSql(), delete);
-        assertEquals(myTable.createTableSql(), create);
+        assertEquals(myTable.createTableSql(unit), create);
     }
 }
diff --git a/container-compiler-plugin/src/test/java/io/trygvis/persistence/test/PersonTypeHandler.java b/container-compiler-plugin/src/test/java/io/trygvis/persistence/test/PersonTypeHandler.java
index fa803da..a872271 100644
--- a/container-compiler-plugin/src/test/java/io/trygvis/persistence/test/PersonTypeHandler.java
+++ b/container-compiler-plugin/src/test/java/io/trygvis/persistence/test/PersonTypeHandler.java
@@ -1,11 +1,10 @@
 package io.trygvis.persistence.test;
 
-import io.trygvis.persistence.FieldMirror;
-import io.trygvis.persistence.TypeHandler;
-
-public class PersonTypeHandler implements TypeHandler {
+@SuppressWarnings("UnusedDeclaration")
+public class PersonTypeHandler /*implements TypeHandler*/ {
+    /*
     @Override
-    public String resultSetSetter(int i, String o, FieldMirror field) {
+    public String setter(int i, String o, FieldMirror field) {
         return "setLong(" + i + ", o." + field.javaName + ".id)";
     }
 
@@ -13,4 +12,5 @@ public class PersonTypeHandler implements TypeHandler {
     public String sqlType(FieldMirror field) {
         return "LONG";
     }
+    */
 }
diff --git a/container-compiler-plugin/src/test/resources/io/trygvis/persistence/test/Person.java b/container-compiler-plugin/src/test/resources/io/trygvis/persistence/test/Person.java
index 131691f..3f1be35 100644
--- a/container-compiler-plugin/src/test/resources/io/trygvis/persistence/test/Person.java
+++ b/container-compiler-plugin/src/test/resources/io/trygvis/persistence/test/Person.java
@@ -8,7 +8,7 @@ import javax.persistence.ManyToOne;
 import java.util.Date;
 
 @Entity
-@SqlEntity("io.trygvis.persistence.test.PersonTypeHandler")
+//@SqlEntity("io.trygvis.persistence.test.PersonTypeHandler")
 public class Person {
     @Id
     public Long id;