wip

9 files changed, 310 insertions, 149 deletions

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 19205b1..d9e13ac 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
@@ -21,6 +21,7 @@ import javax.persistence.Id;
 import javax.tools.JavaFileObject;
 import java.io.IOException;
 import java.io.PrintWriter;
+import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -108,10 +109,12 @@ public class EntityHandler extends AbstractHandler {
     }
 
     public void generate(EntityMirror entityMirror) throws IOException {
-        TypeRef stringType = new TypeRef(String.class);
-        TypeRef sqlEntityDescType = new TypeRef(SqlEntityDesc.class);
-
         ClassG g = new ClassG(PUBLIC, entityMirror.daoType);
+
+        TypeRef stringType = g.imports.add(String.class);
+        TypeRef sqlEntityDescType = g.imports.add(SqlEntityDesc.class);
+        TypeRef sqlExceptionType = g.imports.add(SQLException.class);
+
         g.addPublicStaticFinalField(stringType, "createTableSql").
                 value(toJavaString(entityMirror.createTableSql(sqlUnit)));
         g.addPublicStaticFinalField(stringType, "dropTableSql").
@@ -125,10 +128,12 @@ public class EntityHandler extends AbstractHandler {
                 toJavaString(entityMirror.defaultFields()) +
                 ")";
         g.addPublicStaticFinalField(sqlEntityDescType, "desc").value(desc);
-        entityMirror.insertInto(sqlUnit, g);
-        entityMirror.delete(g);
-        entityMirror.deleteById(g);
-        entityMirror.query(sqlUnit, g);
+        ClassG.InnerClassG typedQuery = g.addInnerClass(entityMirror.queryType(g.imports));
+        typedQuery.inner.addMethod(entityMirror.fromResultSet(sqlUnit, g.imports));
+        g.addMethod(entityMirror.insertInto(sqlUnit, g.imports));
+        g.addMethod(entityMirror.delete(g.imports));
+        g.addMethod(entityMirror.deleteById(g.imports));
+        g.addMethod(entityMirror.query(sqlUnit, g.imports));
 
         JavaFileObject sourceFile = processingEnv.getFiler().createSourceFile(entityMirror.daoType.fqName, sqlUnit.elementForEntity(entityMirror));
         try (PrintWriter w = new PrintWriter(sourceFile.openWriter())) {
@@ -183,7 +188,7 @@ public class EntityHandler extends AbstractHandler {
         ClassG g = new ClassG(PUBLIC, type);
 
 /*
-        TypeRef conType = g.addImport(Connection.class);
+        TypeRef conType = g.add(Connection.class);
         Parameters parameters = new Parameters();
         ParameterRef c = parameters.addParameter(conType, "c");
 
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
index dfb6746..0f76813 100644
--- 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
@@ -6,4 +6,7 @@ public class InternalErrorException extends RuntimeException {
         super(message);
     }
 
+    public InternalErrorException(Exception e) {
+        super(e);
+    }
 }
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 20d5f3b..bb6fa14 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
@@ -4,15 +4,22 @@ import io.trygvis.container.compiler.model.TypeRef;
 import io.trygvis.persistence.EntityMirror;
 
 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<>();
 
-    public void add(EntityMirror entity) {
-        entities.put(entity.type, entity);
+    public SqlUnitModel add(EntityMirror... entities) {
+        for (EntityMirror entity : entities) {
+            this.entities.put(entity.type, entity);
+        }
+
+        return this;
     }
 
     public EntityMirror get(TypeRef type) {
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java
index f944b1d..b5db075 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/TransactionalHandler.java
@@ -2,6 +2,7 @@ package io.trygvis.container.compiler;
 
 import io.trygvis.container.compiler.model.ClassG;
 import io.trygvis.container.compiler.model.FieldRef;
+import io.trygvis.container.compiler.model.MethodRef;
 import io.trygvis.container.compiler.model.Parameters;
 import io.trygvis.container.compiler.model.TypeRef;
 import io.trygvis.container.tx.PlatformTransactionManager;
@@ -77,13 +78,13 @@ public class TransactionalHandler extends AbstractHandler {
             TypeKind kind = returnTypeMirror.getKind();
 
             boolean isVoid = kind == TypeKind.VOID;
-            TypeRef returnType = isVoid ? TypeRef.VOID : g.addImport(returnTypeMirror);
+            TypeRef returnType = isVoid ? TypeRef.VOID : g.imports.add(returnTypeMirror);
 
             Parameters parameters = new Parameters();
             List<String> arguments = new ArrayList<>();
             for (VariableElement ve : ee.getParameters()) {
 //                    parameters.add("final " + ve.asType().toString() + " " + ve.getSimpleName().toString());
-                TypeRef k = g.addImport(ve.asType());
+                TypeRef k = g.imports.add(ve.asType());
                 parameters.addParameter(k, ve.getSimpleName().toString());
                 arguments.add(ve.getSimpleName().toString());
             }
@@ -107,7 +108,8 @@ public class TransactionalHandler extends AbstractHandler {
             body.add("            }");
             body.add("        });");
 
-            g.addMethod(body, returnType, ee.getSimpleName().toString(), parameters);
+            MethodRef m = new MethodRef(PUBLIC, returnType, ee.getSimpleName().toString(), parameters, body);
+            g.addMethod(m);
         }
 
         JavaFileObject sourceFile = processingEnv.getFiler().createSourceFile((p.length() == 0 ? "" : p + ".") + className, element);
@@ -122,11 +124,11 @@ public class TransactionalHandler extends AbstractHandler {
         Parameters parameters = new Parameters();
         List<String> goesToSuper = new ArrayList<>();
 
-        TypeRef k = g.addImport(PlatformTransactionManager.class);
+        TypeRef k = g.imports.add(PlatformTransactionManager.class);
         Parameters.ParameterRef transactionManager = parameters.addParameter(k, "transactionManager");
 
         for (VariableElement p : constructor.getParameters()) {
-            k = g.addImport(p.asType());
+            k = g.imports.add(p.asType());
             String name = p.getSimpleName().toString();
             parameters.addParameter(k, name);
             goesToSuper.add(name);
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java
index e71bbb7..c99feae 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/ClassG.java
@@ -1,13 +1,11 @@
 package io.trygvis.container.compiler.model;
 
-import io.trygvis.container.compiler.Utils;
 import org.apache.commons.lang.StringUtils;
 
 import javax.annotation.Generated;
 import javax.lang.model.type.TypeMirror;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
 import java.util.Set;
 import java.util.TreeSet;
@@ -23,18 +21,17 @@ public class ClassG {
     public final TypeRef type;
     private TypeRef extendsType;
     private List<TypeRef> implementsTypes = new ArrayList<>();
-    private final Set<TypeRef> imports = new TreeSet<>();
+    public final Imports imports = new Imports();
     private final Set<FieldRef> fields = new TreeSet<>();
     private final List<MethodRef> methods = new ArrayList<>();
     private final List<Constructor> constructors = new ArrayList<>();
     private final List<InnerClassG> innerClasses = new ArrayList<>();
 
-    public static class InnerClassG extends ClassG {
-        public final ClassG parent;
+    public static class InnerClassG {
+        public final ClassG inner;
 
-        public InnerClassG(ClassG parent, int modifiers, TypeRef type) {
-            super(modifiers, type);
-            this.parent = parent;
+        public InnerClassG(ClassG inner) {
+            this.inner = inner;
         }
     }
 
@@ -57,44 +54,15 @@ public class ClassG {
         return this;
     }
 
-    public TypeRef addImport(final TypeMirror type) {
-        return addImport(new TypeRef(type));
-    }
-
-    public TypeRef addImport(Class<?> c) {
-        return addImport(new TypeRef(c));
-    }
-
-    public TypeRef addImport(TypeRef klass) {
-        String fqName = klass.fqName;
-        String name = klass.toString();
-        for (TypeRef i : imports) {
-            if (i.fqName.equals(fqName)) {
-                return i;
-            }
-
-            // If we've already found an import with the same name, use the fq version
-            if (i.plainName.equals(name)) {
-                name = fqName;
-                break;
-            }
-        }
-
-        TypeRef ref = new TypeRef(fqName, name);
-        imports.add(ref);
-
-        return ref;
-    }
-
     public FieldRef addField(TypeMirror klass, String name) {
-        TypeRef type = addImport(klass);
+        TypeRef type = imports.add(klass);
         FieldRef ref = new FieldRef(PRIVATE | FINAL, type, name);
         fields.add(ref);
         return ref;
     }
 
     public FieldRef addField(int modifiers, TypeRef type, String name) {
-        TypeRef t = addImport(type);
+        TypeRef t = imports.add(type);
         FieldRef ref = new FieldRef(modifiers, t, name);
         fields.add(ref);
         return ref;
@@ -118,44 +86,43 @@ public class ClassG {
         return constructor;
     }
 
-    public MethodRef addMethod(List<String> body, TypeRef returnType, String name, Parameters parameters) {
-        MethodRef ref = new MethodRef(PUBLIC, returnType, name, parameters, body);
-        methods.add(ref);
-        return ref;
+    public ClassG addMethod(MethodRef methodRef) {
+        this.methods.add(methodRef);
+        return this;
     }
 
-    public MethodRef addStaticMethod(List<String> body, TypeRef returnType, String name, Parameters parameters) {
-        MethodRef ref = new MethodRef(PUBLIC | STATIC, returnType, name, parameters, body);
-        methods.add(ref);
-        return ref;
-    }
+//    public MethodRef addMethod(List<String> body, TypeRef returnType, String name, Parameters parameters) {
+//        MethodRef ref = new MethodRef(PUBLIC, returnType, name, parameters, body);
+//        methods.add(ref);
+//        return ref;
+//    }
+//
+//    public MethodRef addStaticMethod(List<String> body, TypeRef returnType, String name, Parameters parameters) {
+//        MethodRef ref = new MethodRef(PUBLIC | STATIC, returnType, name, parameters, body);
+//        methods.add(ref);
+//        return ref;
+//    }
 
-    public InnerClassG addInnerClass(int modifiers, TypeRef type) {
-        InnerClassG inner = new InnerClassG(this, modifiers, type);
-        innerClasses.add(inner);
-        return inner;
+    public InnerClassG addInnerClass(ClassG inner) {
+        InnerClassG i = new InnerClassG(inner);
+        innerClasses.add(i);
+        return i;
     }
 
     public final List<String> generate() {
-        TypeRef generatedType = addImport(Generated.class);
+        return generate(false);
+    }
+
+    public final List<String> generate(boolean isInner) {
+        TypeRef generatedType = imports.add(Generated.class);
 
         List<String> body = new ArrayList<>();
-        boolean isInner = this instanceof InnerClassG;
         if (!type.inUnnamedPackage() && !isInner) {
             body.add("package " + type.packageName() + ";");
             body.add("");
         }
         if (!isInner) {
-            // TODO: Add imports from inner classes
-            for (TypeRef i : getImports()) {
-                if (i.isPrimitive() || i.inUnnamedPackage()) {
-                    continue;
-                }
-                body.add("import " + i.fqName + ";");
-            }
-            if (!imports.isEmpty()) {
-                body.add("");
-            }
+            body.addAll(imports.generate());
         }
 
         String extendsString = extendsType == null ? "" : " extends " + extendsType;
@@ -173,7 +140,7 @@ public class ClassG {
 
         for (InnerClassG innerClass : innerClasses) {
             body.add("");
-            addAll(1, body, innerClass.generate());
+            addAll(1, body, innerClass.inner.generate(true));
         }
 
         for (Constructor constructor : constructors) {
@@ -197,14 +164,6 @@ public class ClassG {
         }
     }
 
-    protected Collection<? extends TypeRef> getImports() {
-        List<TypeRef> imports = new ArrayList<>(this.imports);
-        for (InnerClassG c : innerClasses) {
-            imports.addAll(c.getImports());
-        }
-        return imports;
-    }
-
     private List<String> write(MethodRef method) {
         List<String> body = new ArrayList<>();
         String returnString;
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/Imports.java b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/Imports.java
new file mode 100644
index 0000000..9ed6097
--- /dev/null
+++ b/container-compiler-plugin/src/main/java/io/trygvis/container/compiler/model/Imports.java
@@ -0,0 +1,54 @@
+package io.trygvis.container.compiler.model;
+
+import javax.lang.model.type.TypeMirror;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.TreeSet;
+
+public class Imports {
+    private final Set<TypeRef> imports = new TreeSet<>();
+
+    public TypeRef add(final TypeMirror type) {
+        return add(new TypeRef(type));
+    }
+
+    public TypeRef add(Class<?> c) {
+        return add(new TypeRef(c));
+    }
+
+    public TypeRef add(TypeRef klass) {
+        String fqName = klass.fqName;
+        String name = klass.toString();
+        for (TypeRef i : imports) {
+            if (i.fqName.equals(fqName)) {
+                return i;
+            }
+
+            // If we've already found an import with the same name, use the fq version
+            if (i.plainName.equals(name)) {
+                name = fqName;
+                break;
+            }
+        }
+
+        TypeRef ref = new TypeRef(fqName, name);
+        imports.add(ref);
+
+        return ref;
+    }
+
+    public List<String> generate() {
+        List<String> body = new ArrayList<>();
+        for (TypeRef i : imports) {
+            if (i.isPrimitive() || i.inUnnamedPackage()) {
+                continue;
+            }
+            body.add("import " + i.fqName + ";");
+        }
+        if (!imports.isEmpty()) {
+            body.add("");
+        }
+        return body;
+    }
+}
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 74c4829..e06855f 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
@@ -3,16 +3,18 @@ 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.Imports;
+import io.trygvis.container.compiler.model.MethodRef;
 import io.trygvis.container.compiler.model.Parameters;
 import io.trygvis.container.compiler.model.TypeRef;
 import io.trygvis.persistence.sql.AbstractTypedQuery;
-import io.trygvis.persistence.sql.SqlEntityDesc;
 
 import javax.persistence.TypedQuery;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Types;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -41,10 +43,13 @@ public class EntityMirror implements Comparable<EntityMirror> {
         this.daoType = new TypeRef(type.plainName + "Dao").args(type.args);
     }
 
-    public void add(FieldMirror field) {
-        fields.add(field);
-        if (field.id) {
-            idFields.add(field);
+    public void add(FieldMirror... fields) {
+        for (FieldMirror field : fields) {
+            this.fields.add(field);
+            if (field.id) {
+                this.idFields.add(field);
+            }
+
         }
     }
 
@@ -128,9 +133,11 @@ public class EntityMirror implements Comparable<EntityMirror> {
         return join(names, ", ");
     }
 
-    public void insertInto(SqlUnitModel unit, ClassG g) {
-        TypeRef conType = g.addImport(Connection.class);
-        TypeRef psType = g.addImport(PreparedStatement.class);
+    public MethodRef insertInto(SqlUnitModel unit, Imports imports) {
+        TypeRef sqlExceptionType = imports.add(SQLException.class);
+        TypeRef typesType = imports.add(Types.class);
+        TypeRef conType = imports.add(Connection.class);
+        TypeRef psType = imports.add(PreparedStatement.class);
         Parameters p = new Parameters();
         ParameterRef con = p.addParameter(conType, "con");
         ParameterRef o = p.addParameter(type, "o");
@@ -142,7 +149,20 @@ public class EntityMirror implements Comparable<EntityMirror> {
             FieldMirror field = fields.get(i);
             if(field instanceof PrimitiveFieldMirror) {
                 TypeHandler typeHandler = generatorConfiguration.typeHandler(field.type);
-                body.add("    stmt." + typeHandler.setter(i + 1, o.name + "." + field.javaName) + ";");
+                String access = o.name + "." + field.javaName;
+                String setter = "    stmt." + typeHandler.setter(i + 1, access) + ";";
+
+                if(field.notNull) {
+                    body.add(setter);
+                }
+                else {
+                    body.add("    " + field.type + " " + field.javaName + " = " + access + ";");
+                    body.add("    if(" + field.javaName + " == null) {");
+                    body.add("        stmt.setNull(" + (i + 1) + ", " + typesType + "." + typeHandler.typeName() + ");");
+                    body.add("    } else {");
+                    body.add("    " + setter);
+                    body.add("    }");
+                }
             } else if (field instanceof ReferenceFieldMirror) {
                 ReferenceFieldMirror ref = (ReferenceFieldMirror) field;
                 EntityMirror referenced = unit.get(ref.type);
@@ -153,14 +173,12 @@ public class EntityMirror implements Comparable<EntityMirror> {
         }
         body.add("    stmt.executeUpdate();");
         body.add("}");
-
-        g.addStaticMethod(body, TypeRef.VOID, "insertInto", p).
-                exception(g.addImport(SQLException.class));
+        return new MethodRef(PUBLIC | STATIC, TypeRef.VOID, "insertInto", p, body).exception(sqlExceptionType);
     }
 
-    public void delete(ClassG g) {
-        TypeRef conType = g.addImport(Connection.class);
-        TypeRef objectType = g.addImport(type);
+    public MethodRef delete(Imports imports) {
+        TypeRef conType = imports.add(Connection.class);
+        TypeRef objectType = imports.add(type);
         Parameters p = new Parameters();
         ParameterRef con = p.addParameter(conType, "con");
         ParameterRef o = p.addParameter(objectType, "o");
@@ -173,13 +191,13 @@ public class EntityMirror implements Comparable<EntityMirror> {
         List<String> body = new ArrayList<>();
         body.add("deleteById(" + join(arguments, ", ") + ");");
 
-        g.addStaticMethod(body, TypeRef.VOID, "delete", p).
-                exception(g.addImport(SQLException.class));
+        return new MethodRef(PUBLIC | STATIC, TypeRef.VOID, "delete", p, body).
+                exception(imports.add(SQLException.class));
     }
 
-    public void deleteById(ClassG g) {
-        TypeRef conType = g.addImport(Connection.class);
-        TypeRef psType = g.addImport(PreparedStatement.class);
+    public MethodRef deleteById(Imports imports) {
+        TypeRef conType = imports.add(Connection.class);
+        TypeRef psType = imports.add(PreparedStatement.class);
         Parameters p = new Parameters();
         ParameterRef con = p.addParameter(conType, "con");
 
@@ -195,30 +213,39 @@ public class EntityMirror implements Comparable<EntityMirror> {
         body.add("    stmt.executeUpdate();");
         body.add("}");
 
-        g.addStaticMethod(body, TypeRef.VOID, "deleteById", p).
-                exception(g.addImport(SQLException.class));
+        return new MethodRef(PUBLIC | STATIC, TypeRef.VOID, "deleteById", p, body).
+                exception(imports.add(SQLException.class));
     }
 
-    public void query(SqlUnitModel sqlUnit, ClassG g) {
-        TypeRef conType = g.addImport(Connection.class);
-        TypeRef abstractQueryType = g.addImport(AbstractTypedQuery.class).args(type);
-        TypeRef typedQueryType = g.addImport(TypedQuery.class).args(type);
+    public ClassG queryType(Imports imports) {
+        TypeRef abstractQueryType = imports.add(AbstractTypedQuery.class).args(type);
+        TypeRef conType = imports.add(Connection.class);
+        TypeRef entityTypedQuery = new TypeRef(type.className + "TypedQuery");
 
         Parameters p = new Parameters();
         ParameterRef c = p.addParameter(conType, "c");
 
-        ClassG typedQuery = g.addInnerClass(PUBLIC | STATIC, new TypeRef(type.className + "TypedQuery")).
+        ClassG typedQuery = new ClassG(PUBLIC | STATIC, entityTypedQuery).
                 extendsType(abstractQueryType);
         typedQuery.addConstructor(p, singletonList("super(" + c.name + ", " + daoType.className + ".desc);"));
-        fromResultSet(sqlUnit, typedQuery);
+        return typedQuery;
+    }
+
+    public MethodRef query(SqlUnitModel sqlUnit, Imports imports) {
+        TypeRef conType = imports.add(Connection.class);
+        TypeRef typedQueryType = imports.add(TypedQuery.class).args(type);
+        TypeRef entityTypedQuery = new TypeRef(type.className + "TypedQuery");
+
+        Parameters p = new Parameters();
+        ParameterRef c = p.addParameter(conType, "c");
 
         List<String> body = new ArrayList<>();
-        body.add("return new " + typedQuery.type + "(" + c.name + ");");
-        g.addStaticMethod(body, typedQueryType, "query", p);
+        body.add("return new " + entityTypedQuery + "(" + c.name + ");");
+        return new MethodRef(PUBLIC | STATIC, typedQueryType, "query", p, body);
     }
 
-    public void fromResultSet(SqlUnitModel unit, ClassG g) {
-        TypeRef rsType = g.addImport(ResultSet.class);
+    public MethodRef fromResultSet(SqlUnitModel unit, Imports g) {
+        TypeRef rsType = g.add(ResultSet.class);
         Parameters p = new Parameters();
         ParameterRef rs = p.addParameter(rsType, "rs");
 
@@ -230,10 +257,10 @@ public class EntityMirror implements Comparable<EntityMirror> {
                 TypeHandler typeHandler = generatorConfiguration.typeHandler(field.type);
                 body.add(field.type + " " + field.javaName + " = " + typeHandler.getter(rs.name, i + 1) + ";");
             } 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);
+//                ReferenceFieldMirror ref = (ReferenceFieldMirror) field;
+//                EntityMirror referenced = unit.get(ref.type);
+//                FieldMirror idField = referenced.getIdField();
+//                TypeHandler typeHandler = generatorConfiguration.typeHandler(idField.type);
 //                body.add(field.type + " " + field.javaName + " = " + typeHandler.getter(rs.name, i + 1) + ";");
                 body.add(field.type + " " + field.javaName + " = null;");
             }
@@ -242,8 +269,8 @@ public class EntityMirror implements Comparable<EntityMirror> {
 
         body.add("return new " + type + "(" + join(names, ", ") + ");");
 
-        g.addMethod(body, type, "fromResultSet", p).
-                exception(g.addImport(SQLException.class));
+        return new MethodRef(PUBLIC, type, "fromResultSet", p, body).
+                exception(g.add(SQLException.class));
     }
 
     @Override
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 0454a96..72941f2 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,13 +1,58 @@
 package io.trygvis.persistence;
 
-public interface TypeHandler {
-    String setter(int i, String expr);
+import io.trygvis.container.compiler.InternalErrorException;
 
-    String getter(String rs, int i);
+import java.lang.reflect.Field;
+import java.sql.Types;
 
-    String sqlType(FieldMirror field);
+import static java.sql.Types.*;
+
+public abstract class TypeHandler {
+    private final String sqlType;
+    private final int type;
+
+    protected TypeHandler(String sqlType, int type) {
+        this.sqlType = sqlType;
+        this.type = type;
+    }
+
+    abstract String setter(int i, String expr);
+
+    abstract String getter(String rs, int i);
+
+    public String sqlType(FieldMirror field) {
+        return sqlType;
+    }
+
+    public int type() {
+        return type;
+    }
+
+    public final String typeName() {
+        try {
+            int type = type();
+            for (Field field : Types.class.getFields()) {
+                Object o = field.get(null);
+                if (o != null && o instanceof Integer && o.equals(type)) {
+                    return field.getName();
+                }
+            }
+        } catch (IllegalAccessException e) {
+            throw new InternalErrorException(e);
+        }
+
+        throw new InternalErrorException("Unknown java.sql.Types: " + type);
+    }
+
+    // -----------------------------------------------------------------------
+    //
+    // -----------------------------------------------------------------------
+
+    public static class IntTypeHandler extends TypeHandler {
+        public IntTypeHandler() {
+            super("INTEGER", INTEGER);
+        }
 
-    public static class IntTypeHandler implements TypeHandler {
         @Override
         public String setter(int i, String expr) {
             return "setInt(" + i + ", " + expr + ")";
@@ -17,14 +62,13 @@ public interface TypeHandler {
         public String getter(String rs, int i) {
             return rs + ".getInt(" + i + ")";
         }
+    }
 
-        @Override
-        public String sqlType(FieldMirror field) {
-            return "INTEGER";
+    public static class LongTypeHandler extends TypeHandler {
+        protected LongTypeHandler() {
+            super("BIGINT", BIGINT);
         }
-    }
 
-    public static class LongTypeHandler implements TypeHandler {
         @Override
         public String setter(int i, String expr) {
             return "setLong(" + i + ", " + expr + ")";
@@ -34,14 +78,13 @@ public interface TypeHandler {
         public String getter(String rs, int i) {
             return rs + ".getLong(" + i + ")";
         }
+    }
 
-        @Override
-        public String sqlType(FieldMirror field) {
-            return "BIGINT";
+    public static class StringTypeHandler extends TypeHandler {
+        protected StringTypeHandler() {
+            super(null, VARCHAR);
         }
-    }
 
-    public static class StringTypeHandler implements TypeHandler {
         @Override
         public String setter(int i, String expr) {
             return "setString(" + i + ", " + expr + ")";
@@ -58,7 +101,11 @@ public interface TypeHandler {
         }
     }
 
-    public static class DateTypeHandler implements TypeHandler {
+    public static class DateTypeHandler extends TypeHandler {
+        protected DateTypeHandler() {
+            super("TIMESTAMP", TIMESTAMP);
+        }
+
         @Override
         public String setter(int i, String expr) {
             return "setTimestamp(" + i + ", new java.sql.Timestamp(" + expr + ".getTime()))";
@@ -68,10 +115,5 @@ public interface TypeHandler {
         public String getter(String rs, int i) {
             return "new java.util.Date(" + rs + ".getTimestamp(" + i + ").getTime())";
         }
-
-        @Override
-        public String sqlType(FieldMirror field) {
-            return "TIMESTAMP";
-        }
     }
 }
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 9d90eef..3f0f345 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,10 +1,15 @@
 package io.trygvis.persistence;
 
 import io.trygvis.container.compiler.SqlUnitModel;
+import io.trygvis.container.compiler.model.Imports;
+import io.trygvis.container.compiler.model.MethodRef;
 import io.trygvis.container.compiler.model.TypeRef;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
+import java.io.CharArrayWriter;
+import java.io.PrintWriter;
+
 import static io.trygvis.persistence.FieldMirror.PrimitiveFieldMirror;
 import static java.lang.String.format;
 import static org.testng.Assert.assertEquals;
@@ -15,6 +20,7 @@ public class EntityMirrorTest {
     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 age = new PrimitiveFieldMirror(new TypeRef(Integer.class), "age", "age", false, true, false);
     static FieldMirror ref = new FieldMirror.ReferenceFieldMirror(entityType, "parent", "parent", false, false);
 
     @DataProvider(name = "insertIntoSql", parallel = true)
@@ -60,11 +66,67 @@ public class EntityMirrorTest {
             myTable.add(field);
         }
 
-        SqlUnitModel unit = new SqlUnitModel();
-        unit.add(myTable);
+        SqlUnitModel unit = new SqlUnitModel().add(myTable);
 
         assertEquals(myTable.insertIntoSql(), insert);
         assertEquals(myTable.deleteFromSql(), delete);
         assertEquals(myTable.createTableSql(unit), create);
     }
+
+    @Test
+    public void testInsertIntoMethod() {
+        eq(insertInto(idString),
+                "try(java.sql.PreparedStatement stmt = con.prepareStatement(insertIntoSql)) {",
+                "    java.lang.String id = o.id;",
+                "    if(id == null) {",
+                "        stmt.setNull(1, java.sql.Types.VARCHAR);",
+                "    } else {",
+                "        stmt.setString(1, o.id);",
+                "    }",
+                "    stmt.executeUpdate();",
+                "}");
+
+        eq(insertInto(age),
+                "try(java.sql.PreparedStatement stmt = con.prepareStatement(insertIntoSql)) {",
+                "    stmt.setInt(1, o.age);",
+                "    stmt.executeUpdate();",
+                "}");
+    }
+
+    @Test
+    public void testFromResultSet() {
+        eq(fromResultSet(age),
+                "java.lang.Integer age = rs.getInt(1);",
+                "return new Wat(age);");
+    }
+
+    private MethodRef insertInto(FieldMirror... fields) {
+        EntityMirror myTable = new EntityMirror(new GeneratorConfiguration(), entityType, "my_table");
+        myTable.add(fields);
+        SqlUnitModel unit = new SqlUnitModel().add(myTable);
+        return myTable.insertInto(unit, new Imports());
+    }
+
+    private MethodRef fromResultSet(FieldMirror... fields) {
+        EntityMirror myTable = new EntityMirror(new GeneratorConfiguration(), entityType, "my_table");
+        myTable.add(fields);
+        SqlUnitModel unit = new SqlUnitModel().add(myTable);
+        return myTable.fromResultSet(unit, new Imports());
+    }
+
+    private void eq(MethodRef m, String... expected) {
+        CharArrayWriter actual = new CharArrayWriter();
+        PrintWriter w = new PrintWriter(actual);
+        for (String s : m.body) {
+            w.println(s);
+        }
+
+        CharArrayWriter e = new CharArrayWriter();
+        w = new PrintWriter(e);
+        for (String s : expected) {
+            w.println(s);
+        }
+
+        assertEquals(actual.toString(), e.toString());
+    }
 }