1 files changed, 29 insertions, 9 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 c623af7..19205b1 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
@@ -7,6 +7,7 @@ import io.trygvis.persistence.FieldMirror;
 import io.trygvis.persistence.GeneratorConfiguration;
 import io.trygvis.persistence.SqlEntity;
 import io.trygvis.persistence.TypeHandler;
+import io.trygvis.persistence.sql.SqlEntityDesc;
 
 import javax.annotation.processing.ProcessingEnvironment;
 import javax.lang.model.element.AnnotationMirror;
@@ -28,7 +29,9 @@ import java.util.Set;
 import static io.trygvis.container.compiler.Utils.toJavaString;
 import static io.trygvis.persistence.FieldMirror.PrimitiveFieldMirror;
 import static java.lang.Character.isUpperCase;
+import static java.lang.reflect.Modifier.PUBLIC;
 import static javax.lang.model.util.ElementFilter.fieldsIn;
+import static org.apache.commons.lang.StringUtils.stripEnd;
 
 public class EntityHandler extends AbstractHandler {
     private GeneratorConfiguration generatorConfiguration = new GeneratorConfiguration();
@@ -49,7 +52,7 @@ public class EntityHandler extends AbstractHandler {
 
                         TypeHandler typeHandler = (TypeHandler) typeHandlerClass.newInstance();
                         String type = entity.asType().toString();
-                        generatorConfiguration.addTypeHandler(new TypeRef(type, type), typeHandler);
+                        generatorConfiguration.addTypeHandler(new TypeRef(type), typeHandler);
                         System.out.println("Loaded TypeHandler for " + type + " through " + typeHandlerClass.getCanonicalName());
                         break;
                 }
@@ -105,18 +108,33 @@ public class EntityHandler extends AbstractHandler {
     }
 
     public void generate(EntityMirror entityMirror) throws IOException {
-        ClassG g = new ClassG(entityMirror.daoType);
-        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()));
+        TypeRef stringType = new TypeRef(String.class);
+        TypeRef sqlEntityDescType = new TypeRef(SqlEntityDesc.class);
+
+        ClassG g = new ClassG(PUBLIC, entityMirror.daoType);
+        g.addPublicStaticFinalField(stringType, "createTableSql").
+                value(toJavaString(entityMirror.createTableSql(sqlUnit)));
+        g.addPublicStaticFinalField(stringType, "dropTableSql").
+                value(toJavaString(entityMirror.dropTableSql()));
+        g.addPublicStaticFinalField(stringType, "insertIntoSql").
+                value(toJavaString(entityMirror.insertIntoSql()));
+        g.addPublicStaticFinalField(stringType, "deleteFromSql").
+                value(toJavaString(entityMirror.deleteFromSql()));
+        String desc = "new " + sqlEntityDescType + "(" +
+                toJavaString(entityMirror.tableName) + ", " +
+                toJavaString(entityMirror.defaultFields()) +
+                ")";
+        g.addPublicStaticFinalField(sqlEntityDescType, "desc").value(desc);
         entityMirror.insertInto(sqlUnit, g);
         entityMirror.delete(g);
         entityMirror.deleteById(g);
+        entityMirror.query(sqlUnit, g);
 
         JavaFileObject sourceFile = processingEnv.getFiler().createSourceFile(entityMirror.daoType.fqName, sqlUnit.elementForEntity(entityMirror));
         try (PrintWriter w = new PrintWriter(sourceFile.openWriter())) {
-            g.write(w);
+            for (String s : g.generate()) {
+                w.println(stripEnd(s, " "));
+            }
         }
     }
 
@@ -162,7 +180,7 @@ public class EntityHandler extends AbstractHandler {
         // TODO: Support a name prefix from @SqlEntitySet
         TypeRef type = new TypeRef(p + ".Session");
 
-        ClassG g = new ClassG(type);
+        ClassG g = new ClassG(PUBLIC, type);
 
 /*
         TypeRef conType = g.addImport(Connection.class);
@@ -180,7 +198,9 @@ public class EntityHandler extends AbstractHandler {
 
         JavaFileObject sourceFile = processingEnv.getFiler().createSourceFile(type.fqName, packageElement);
         try (PrintWriter w = new PrintWriter(sourceFile.openWriter())) {
-            g.write(w);
+            for (String s : g.generate()) {
+                w.println(stripEnd(s, " "));
+            }
         }
     }