1 files changed, 2 insertions, 2 deletions
diff --git a/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/EntityManagerGenerator.java b/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/EntityManagerGenerator.java
index f621343..4a6df0f 100644
--- a/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/EntityManagerGenerator.java
+++ b/container-compiler-plugin/src/main/java/io/trygvis/persistence/generators/EntityManagerGenerator.java
@@ -60,14 +60,14 @@ public class EntityManagerGenerator {
     }
 
     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");
         List<String> body = new ArrayList<>();
         for (EntityMirror entity : unit.getEntities().values()) {
             body.add("if (klass == " + entity.type.plainName + ".class) {");
-            body.add("    return (SqlDao<Id, T>) " + daoFields.get(entity).name + ";");
+            body.add("    return (SqlDao<Id, T, " +  unit.sessionType.className + ">) " + daoFields.get(entity).name + ";");
             body.add("}");
         }
         body.add("throw new RuntimeException(\"Type is not a part of this persistence unit: \" + klass);");