1 files changed, 7 insertions, 9 deletions
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/SqlSessionGenerator.java b/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/SqlSessionGenerator.java
index 7a46018..ddb4d99 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/SqlSessionGenerator.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/SqlSessionGenerator.java
@@ -20,18 +20,15 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
-import static io.trygvis.container.compiler.Utils.toClassName;
 import static io.trygvis.container.compiler.Utils.toFieldName;
 import static io.trygvis.container.compiler.model.Parameters.ParameterRef;
 import static java.lang.reflect.Modifier.PUBLIC;
 
 public class SqlSessionGenerator {
     public static ClassG generateSqlSession(SqlUnitModel unit) throws IOException {
-        String prefix = unit.getPackageName() + "." + toClassName(unit.getName());
-        TypeRef emType = new TypeRef(prefix + "SqlSession");
         TypeRef sqlSession = new TypeRef(SqlSession.class);
 
-        ClassG g = new ClassG(PUBLIC, emType).
+        ClassG g = new ClassG(PUBLIC, unit.sessionType).
                 extendsType(sqlSession);
 
         Map<EntityMirror, FieldRef> daoFields = new HashMap<>();
@@ -54,20 +51,21 @@ public class SqlSessionGenerator {
 
         for (EntityMirror entity : unit.getEntities().values()) {
             FieldRef f = daoFields.get(entity);
-            body.add("this." + f.name + " = new " + entity.daoType.plainName + "(" + c.name + ");");
+            body.add("this." + f.name + " = new " + entity.daoType.plainName + "(this);");
         }
         return new Constructor(p, body);
     }
 
     public static MethodRef getSqlDao(SqlUnitModel unit, Imports imports, Map<EntityMirror, FieldRef> daoFields) {
-        TypeRef sqlDatoType = imports.add(new TypeRef(SqlDao.class)).args("Id", "T");
+        TypeRef sqlDatoType = imports.add(new TypeRef(SqlDao.class)).
+                args("Id", "T", unit.sessionType.className, "?");
         Parameters p = new Parameters();
-        TypeRef klassType = new TypeRef(Class.class).args("T");
-        ParameterRef klass = p.addParameter(klassType, "klass");
+        p.addParameter(new TypeRef(Class.class).args("T"), "klass");
         List<String> body = new ArrayList<>();
         for (EntityMirror entity : unit.getEntities().values()) {
+            String sqlDao = "SqlDao<Id, T, " + unit.sessionType.className + ", " + entity.rowType.className + ">";
             body.add("if (klass == " + entity.type.plainName + ".class) {");
-            body.add("    return (SqlDao<Id, T>) " + daoFields.get(entity).name + ";");
+            body.add("    return (" + sqlDao + ") " + daoFields.get(entity).name + ";");
             body.add("}");
         }
         body.add("throw new RuntimeException(\"Type is not a part of this persistence unit: \" + klass);");