1 files changed, 23 insertions, 2 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 8e32cfe..90309e4 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
@@ -16,13 +16,17 @@ import io.trygvis.persistence.sql.SqlSession;
 import java.io.IOException;
 import java.sql.Connection;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 import static io.trygvis.container.compiler.Utils.toFieldName;
 import static io.trygvis.container.compiler.model.Parameters.ParameterRef;
+import static java.lang.String.format;
+import static java.lang.reflect.Modifier.FINAL;
 import static java.lang.reflect.Modifier.PUBLIC;
+import static org.apache.commons.lang.StringUtils.join;
 
 public class SqlSessionGenerator {
     public static ClassG generateSqlSession(SqlUnitModel unit) throws IOException {
@@ -41,18 +45,33 @@ public class SqlSessionGenerator {
             daoFields.put(entity, f);
         }
 
-        g.add(constructor(unit, g.imports, daoFields));
+        TypeRef sqlDaoType = new TypeRef(SqlDao.class);
+        TypeRef listOfDaos = g.imports.add(new TypeRef(List.class).args(sqlDaoType));
+        FieldRef daos = new FieldRef(PUBLIC | FINAL, listOfDaos, "daos");
+
+        g.add(daos);
+        g.add(constructor(unit, g.imports, daoFields, daos));
         g.add(getSqlDao(unit, g.imports, daoFields));
+        g.add(getDaos(g.imports, daos));
 
         return g;
     }
 
-    private static Constructor constructor(SqlUnitModel unit, Imports imports, Map<EntityMirror, FieldRef> daoFields) {
+    private static MethodRef getDaos(Imports imports, FieldRef daos) {
+        TypeRef sqlDaoType = new TypeRef(SqlDao.class);
+        TypeRef listOfDaos = imports.add(new TypeRef(List.class).args(sqlDaoType));
+        return new MethodRef(PUBLIC, listOfDaos, "getDaos", "return " + daos.name + ";");
+    }
+
+    private static Constructor constructor(SqlUnitModel unit, Imports imports, Map<EntityMirror, FieldRef> daoFields, FieldRef daos) {
+        TypeRef arraysType = imports.add(Arrays.class);
+
         Parameters p = new Parameters();
         ParameterRef c = p.addParameter(imports.add(Connection.class), "c");
         List<String> body = new ArrayList<>();
         body.add("super(" + c.name + ");");
 
+        List<String> daoNames = new ArrayList<>();
         for (EntityMirror entity : unit.getEntities().values()) {
             if (!entity.concrete) {
                 continue;
@@ -60,7 +79,9 @@ public class SqlSessionGenerator {
 
             FieldRef f = daoFields.get(entity);
             body.add("this." + f.name + " = new " + entity.daoType.plainName + "(this);");
+            daoNames.add(f.name);
         }
+        body.add(format("this." + daos.name + " = " + arraysType.plainName + ".asList(new SqlDao[]{%n" + join(daoNames, ", %n") + "%n});"));
         return new Constructor(p, body);
     }