1 files changed, 14 insertions, 5 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 ae3443c..74f303f 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
@@ -15,6 +15,7 @@ import javax.annotation.processing.ProcessingEnvironment;
 import javax.lang.model.element.AnnotationMirror;
 import javax.lang.model.element.AnnotationValue;
 import javax.lang.model.element.Element;
+import javax.lang.model.element.ElementKind;
 import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.PackageElement;
@@ -40,14 +41,13 @@ import static io.trygvis.persistence.FieldMirror.AccessorType.FIELD;
 import static io.trygvis.persistence.FieldMirror.AccessorType.METHOD;
 import static io.trygvis.persistence.FieldMirror.FieldType.PRIMITIVE;
 import static io.trygvis.persistence.FieldMirror.FieldType.REFERENCE;
+import static io.trygvis.persistence.TypeHandler.StringEnumTypeHandler;
 import static java.lang.Character.isUpperCase;
 import static java.lang.String.format;
 import static java.lang.reflect.Modifier.PUBLIC;
 import static javax.lang.model.util.ElementFilter.fieldsIn;
 import static javax.lang.model.util.ElementFilter.methodsIn;
-import static org.apache.commons.lang.StringUtils.capitalize;
-import static org.apache.commons.lang.StringUtils.join;
-import static org.apache.commons.lang.StringUtils.stripEnd;
+import static org.apache.commons.lang.StringUtils.*;
 
 public class EntityHandler extends AbstractHandler {
     private GeneratorConfiguration generatorConfiguration = new GeneratorConfiguration();
@@ -119,7 +119,6 @@ public class EntityHandler extends AbstractHandler {
             }
         }
 
-        // Iterate over all fields. Consume relevant getters and setters.
         for (VariableElement f : fieldsIn(elements.getAllMembers(element))) {
             String name = f.getSimpleName().toString();
             ExecutableElement getter = getters.remove(name);
@@ -253,6 +252,7 @@ public class EntityHandler extends AbstractHandler {
         // TODO: check for transient and @Transient
         FieldMirror.AccessorType accessorType;
         TypeRef type;
+        Element element;
         String javaName;
         String sqlName;
         boolean id;
@@ -263,6 +263,7 @@ public class EntityHandler extends AbstractHandler {
             }
             accessorType = FIELD;
             type = new TypeRef(var.asType());
+            element = types.asElement(var.asType());
             javaName = var.getSimpleName().toString();
             id = isId(var);
             processSequenceGenerators(var);
@@ -277,6 +278,7 @@ public class EntityHandler extends AbstractHandler {
             }
             accessorType = METHOD;
             type = new TypeRef(getter.getReturnType());
+            element = types.asElement(getter.getReturnType());
             id = isId(getter);
             processSequenceGenerators(getter);
             // TODO: this might be relaxed, just find the common type and use that.
@@ -288,9 +290,16 @@ public class EntityHandler extends AbstractHandler {
         }
         sqlName = sqlName(javaName);
 
+        // Register type handler if this is an enum
+        if (element != null && element.getKind() == ElementKind.ENUM) {
+            if (!generatorConfiguration.hasTypeHandler(type)) {
+                generatorConfiguration.addTypeHandler(type, new StringEnumTypeHandler(type));
+            }
+        }
+
         boolean notNull = false;
         boolean unique = false;
-        boolean primitive = generatorConfiguration.isPrimitive(type);
+        boolean primitive = generatorConfiguration.hasTypeHandler(type);
         if (id && !primitive) {
             throw new CompilerException(var, "A @Id field has to be a primitive or embedded.");
         }