1 files changed, 13 insertions, 8 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");