DaliHelper.java

package gaff.dali.jdbi;

import org.skife.jdbi.v2.StatementContext;
import org.skife.jdbi.v2.tweak.ResultColumnMapper;
import org.skife.jdbi.v2.tweak.ResultSetMapper;

import java.io.IOException;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.HashSet;
import java.util.Set;
import java.util.function.Consumer;


/**
 * Static class to store all the helper methods used by Dali mappers
 */
public class DaliHelper {
    private final ResultSet rs;
    private final StatementContext ctx;
    private Set<String> returnedColumns;

    public DaliHelper(ResultSet rs, StatementContext ctx){
        this.rs = rs;
        this.ctx = ctx;
    }

    private boolean hasColumn(String columnName) throws SQLException {
        if(returnedColumns == null){
            returnedColumns = new HashSet<>();
            ResultSetMetaData rsmd = rs.getMetaData();
            int columns = rsmd.getColumnCount();
            for (int x = 1; x <= columns; x++) {
                returnedColumns.add(rsmd.getColumnName(x).toUpperCase());
            }
        }

        return returnedColumns.contains(columnName.toUpperCase());
    }

    public <T> T getEntity(String columnName, boolean isOptional, Class<T> clazz) throws SQLException {
        if(isOptional && !hasColumn(columnName)){
            return null;
        }

        ResultColumnMapper<T> rcm = ctx.columnMapperFor(clazz);
        if(rcm != null ) {
            T ret = rcm.mapColumn(rs, columnName, ctx);
            return ret;
        }

        throw new SQLException("Don't know how to type " + clazz + " (seen in field " + columnName + ")");
        //If you got here and you're not sure why, it's because I retired this. Talk to me (Matthew Shaylor) about it
        //return ObjectMapperFactory.mapper.convertValue(rs.getObject(columnName), new TypeReference<T>() {});
    }

    public <T> void setEntityViaJson(Consumer<T> consumer, String columnName, boolean isOptional, Class<T> clazz) throws SQLException {
        String json = getEntity(columnName, isOptional, String.class);
        if(json == null){
            return;
        }

        try {
            T entity = ObjectMapperFactory.mapper.readValue(json, clazz);
            consumer.accept(entity);
        } catch (IOException e) {
            throw new SQLException(e);
        }
    }

    public <T> void setEmbeddedEntity(Consumer<T> consumer, ResultSetMapper<T> embeddedMapper,
                                      int index, ResultSet r, StatementContext ctx) throws SQLException {
        T entity = embeddedMapper.map(index, r, ctx);
        consumer.accept(entity);
    }

    public <T> void setEntity(Consumer<T> consumer, String columnName, boolean isOptional, Class<T> clazz) throws SQLException {
        T entity = getEntity(columnName, isOptional, clazz);

        //Check if null to fall through to default values within Immutables
        if(entity == null){
            return;
        }

        consumer.accept(entity);

    }

    public void setString(Consumer<String> consumer, String columnName, boolean isOptional, boolean trim) throws SQLException {
        String entity = getEntity(columnName, isOptional, String.class);

        //Check if null to fall through to default values within Immutables
        if(entity == null){
            return;
        }

        if(trim){
            // There are reference data entries in ICTS which the only difference are the leading spaces, so the leading spaces have to be kept
            entity = entity.replaceAll("\\s++$", "");
        }

        consumer.accept(entity);
    }
}

JdbiProcessor.java

package gaff.dali.annotationprocessors;

import gaff.dali.jdbi.*;
import com.squareup.javapoet.*;
import org.immutables.value.Value;
import org.skife.jdbi.v2.StatementContext;
import org.skife.jdbi.v2.tweak.ResultSetMapper;

import javax.annotation.Generated;
import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.Messager;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.TypeMirror;
import javax.tools.Diagnostic;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.lang.annotation.Annotation;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.time.ZonedDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static javax.lang.model.SourceVersion.latestSupported;

/**
 * Created by mshaylor on 08/07/2016.
 */
@SupportedAnnotationTypes({
        "gaff.dali.jdbi.Jdbi"
})
public class JdbiProcessor extends AbstractProcessor {

    private static final JdbiOptions DEFAULT_OPTIONS = new JdbiOptions() {
        //Would be nice to pull the defaults from the annotation but this is obsecenely difficult
        //as the retention is source level.

        @Override
        public Class<? extends Annotation> annotationType() {
            return JdbiOptions.class;
        }

        @Override
        public String columnName() {
            return "";
        }

        @Override
        public boolean ignore() {
            return false;
        }

        @Override
        public boolean inColumnList() {
            return true;
        }

        @Override
        public boolean optional() {
            return false;
        }

        @Override
        public boolean trim() {
            return false;
        }

        @Override
        public int size() {
            return -1;
        }

        @Override
        public int scale() {
            return -1;
        }
    };

    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment env) {

        Messager messager = processingEnv.getMessager();

        for (Element elem : env.getElementsAnnotatedWith(Jdbi.class)) {
            try {
                generateCode((TypeElement) elem);
            } catch (Exception e) {
                messager.printMessage(Diagnostic.Kind.ERROR, e.getMessage());
                StringWriter stringWriter = new StringWriter();
                e.printStackTrace(new PrintWriter(stringWriter));
                messager.printMessage(Diagnostic.Kind.ERROR, stringWriter.toString());
                throw new JDBIProcessingException(e);
            }
        }

        return true;
    }

    private void generateCode(TypeElement elem) throws IOException {

        PackageElement pkg = processingEnv.getElementUtils().getPackageOf(elem);

        Jdbi jdbiAnnotation = elem.getAnnotation(Jdbi.class);

        AnnotationSpec generated = AnnotationSpec.builder(Generated.class)
                .addMember("value", "$S", JdbiProcessor.class.getCanonicalName())
                .addMember("date", "$S", ZonedDateTime.now())
                .build();

        TypeName typeName = TypeName.get(elem.asType());
        String simpleName = elem.getSimpleName().toString();
        TypeSpec.Builder mapperType = TypeSpec.classBuilder(simpleName + "Mapper")
                .addAnnotation(generated)
                .addModifiers(Modifier.PUBLIC)
                .addSuperinterface(ParameterizedTypeName.get(ClassName.get(ResultSetMapper.class), typeName));

        MethodSpec.Builder mapMethodSpec = MethodSpec.methodBuilder("map")
                .addAnnotation(Override.class)
                .addModifiers(Modifier.PUBLIC)
                .returns(typeName)
                .addParameter(TypeName.INT, "index")
                .addParameter(ResultSet.class, "r")
                .addParameter(StatementContext.class, "ctx")
                .addException(SQLException.class)
                .addStatement("Immutable$L.Builder builder = Immutable$L.builder()", simpleName, simpleName)
                .addStatement("$T dh = new $T(r, ctx)", DaliHelper.class, DaliHelper.class);

        List<JavaDbName> columnList = new ArrayList<>();
        List<ClassName> embeddedMapperClasses = new ArrayList<>();

        for (Element member : elem.getEnclosedElements()) {
            //Check for method annotation
            JdbiOptions jdbiOptions = member.getAnnotation(JdbiOptions.class);
            Value.Derived derivedAnnotation = member.getAnnotation(Value.Derived.class);
            if (jdbiOptions == null) jdbiOptions = DEFAULT_OPTIONS;
            if (jdbiOptions.ignore() || derivedAnnotation != null) {
                continue;
            }
            JdbiStoreAsJson jdbiStoreAsJson = member.getAnnotation(JdbiStoreAsJson.class);
            JdbiEmbedded jdbiEmbedded = member.getAnnotation(JdbiEmbedded.class);

            switch (member.getKind()) {
                case METHOD:
                    ExecutableElement ee = (ExecutableElement) member;
                    //Static stuff should not be included!
                    if (ee.getModifiers().contains(Modifier.STATIC))
                        continue;

                    //Protected methods are assumed not to be JDBI related
                    //Typically these are for @Value.Check() or similar.
                    if (ee.getModifiers().contains(Modifier.PROTECTED) || ee.getModifiers().contains(Modifier.PRIVATE))
                        continue;

                    String javaName = member.getSimpleName().toString();

                    String formattedDbName = DynamicClassBinding.getColumnName(jdbiAnnotation, jdbiOptions, javaName);

                    if (jdbiEmbedded == null && jdbiOptions.inColumnList()) {
                        columnList.add(new JavaDbName(formattedDbName, javaName));
                    }
                    TypeName returnType = getReturnType(ee);
                    ParsingAndRegularType daliType = getDaliType(ee.getReturnType(), jdbiOptions, jdbiStoreAsJson, jdbiEmbedded);
                    switch (daliType.getParsingType()) {
                        case primitive:
                            mapMethodSpec.addStatement("builder.$L(r.get$L($S))", javaName, daliType.getType(), formattedDbName);
                            break;
                        case jdbiobject:
                            mapMethodSpec.addStatement("dh.setEntity(($T) builder::$L, $S, $L, $L.class)",
                                    ParameterizedTypeName.get(ClassName.get(Consumer.class), returnType),
                                    javaName,
                                    formattedDbName,
                                    jdbiOptions.optional(),
                                    returnType);
                            break;
                        case jsonobject:
                            mapMethodSpec.addStatement("dh.setEntityViaJson(($T) builder::$L, $S, $L, $L.class)",
                                    ParameterizedTypeName.get(ClassName.get(Consumer.class), returnType),
                                    javaName,
                                    formattedDbName,
                                    jdbiOptions.optional(),
                                    returnType);
                            break;
                        case embedded:
                            mapMethodSpec.addStatement("dh.setEmbeddedEntity(($T) builder::$L, new $L(), index, r, ctx)",
                                    ParameterizedTypeName.get(ClassName.get(Consumer.class), returnType),
                                    javaName,
                                    getMapperType(returnType));
                            embeddedMapperClasses.add(getMapperType(returnType));
                            break;
                        case string:
                            mapMethodSpec.addStatement("dh.setString(($T) builder::$L, $S, $L, $L)",
                                    ParameterizedTypeName.get(ClassName.get(Consumer.class), returnType),
                                    javaName,
                                    formattedDbName,
                                    jdbiOptions.optional(),
                                    jdbiOptions.trim());
                            break;
                        default:
                            String msg = String.format("Can't handle dali type: %s", returnType);
                            throw new JDBIProcessingException(msg);
                    }
            }
        }
        mapMethodSpec.addStatement("return builder.build()");

        mapperType.addMethod(mapMethodSpec.build());

        addFields(mapperType, columnList, embeddedMapperClasses);

        if (!jdbiAnnotation.tableName().isEmpty()) {
            addFieldsWithTableName(jdbiAnnotation, mapperType);
        }

        JavaFile directClientFile = JavaFile.builder(pkg.toString(), mapperType.build()).build();
        directClientFile.writeTo(processingEnv.getFiler());
    }

    private ClassName getMapperType(TypeName entityType) {
        if (!(entityType instanceof ClassName)) {
            String msg = String.format("@JdbiEmbedded field is not a class: %s", entityType);
            throw new JDBIProcessingException(msg);
        }
        ClassName className = (ClassName) entityType;
        return ClassName.get(className.packageName(), className.simpleName() + "Mapper");
    }

    private void addFieldsWithTableName(Jdbi jdbiAnnotation, TypeSpec.Builder mapperType) {
        String tableName = jdbiAnnotation.tableName();
        mapperType.addField(FieldSpec
                .builder(String.class, "TABLE_NAME", Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
                .initializer("$S", tableName)
                .build());

        mapperType.addField(FieldSpec
                .builder(String.class, "SQL_INSERT", Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
                .initializer("$S + TABLE_NAME + $S + DB_COLUMNS + $S + BIND_COLUMNS + $S", "insert into ", "(", ") values (", ")")
                .build());

        mapperType.addField(FieldSpec
                .builder(String.class, "SQL_UPDATE_PREFIX", Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
                .initializer("$S + TABLE_NAME + $S + SET_COLUMNS", "UPDATE ", " SET ")
                .build());
    }

    private void addFields(TypeSpec.Builder mapperType, List<JavaDbName> columnList, List<ClassName> embeddedMapperClasses) {


        mapperType.addField(FieldSpec
                .builder(String.class, "DB_COLUMNS", Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
                .initializer(
                        buildFieldFormat(embeddedMapperClasses, "DB_COLUMNS"),
                        buildFieldValues(columnList.stream().map(JavaDbName::getFormattedDbName).collect(Collectors.joining(", ")), embeddedMapperClasses))
                .build());

        mapperType.addField(FieldSpec
                .builder(String.class, "BIND_COLUMNS", Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
                .initializer(
                        buildFieldFormat(embeddedMapperClasses, "BIND_COLUMNS"),
                        buildFieldValues(columnList.stream().map(JavaDbName::getJavaName).collect(Collectors.joining(", :", ":", "")), embeddedMapperClasses))
                .build());

        mapperType.addField(FieldSpec
                .builder(String.class, "SET_COLUMNS", Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
                .initializer(
                        buildFieldFormat(embeddedMapperClasses, "SET_COLUMNS"),
                        buildFieldValues(columnList.stream().map(jdbn -> jdbn.getFormattedDbName() + "=:" + jdbn.getJavaName()).collect(Collectors.joining(", ")), embeddedMapperClasses))
                .build());
    }

    private Object[] buildFieldValues(Object item1, List<?> rest) {
        return Stream.concat(Stream.of(item1), rest.stream()).collect(Collectors.toList()).toArray();
    }

    private String buildFieldFormat(List<ClassName> embeddedMapperClasses, String field) {
        String format = "$S";
        for (ClassName className : embeddedMapperClasses) {
            format += " + \", \" + $L." + field;
        }
        return format;
    }

    private TypeName getReturnType(ExecutableElement ee) {
        TypeName typeName = TypeName.get(ee.getReturnType());
        if (typeName instanceof ParameterizedTypeName) {
            if (((ParameterizedTypeName) typeName).rawType.equals(ClassName.get(Optional.class))) {
                return ((ParameterizedTypeName) typeName).typeArguments.get(0);
            }
        }
        return typeName;
    }

    private ParsingAndRegularType getDaliType(TypeMirror retType, JdbiOptions jdbiOptions, JdbiStoreAsJson jdbiStoreAsJson, JdbiEmbedded jdbiEmbedded) {
        ParsingType parsingType;
        String type;
        switch (retType.getKind()) {
            case INT:
                parsingType = ParsingType.primitive;
                type = "Int";
                break;
            case DOUBLE:
                parsingType = ParsingType.primitive;
                type = "Double";
                break;
            case LONG:
                parsingType = ParsingType.primitive;
                type = "Long";
                break;
            case BOOLEAN:
                parsingType = ParsingType.primitive;
                type = "Boolean";
                break;
            case ARRAY:
                if (!retType.toString().equals(byte[].class.getCanonicalName())) {
                    String msg = String.format("Can't handle array type: %s", retType);
                    throw new JDBIProcessingException(msg);
                }
                // byte[] falls through and is treated as Object
            case DECLARED:
                type = retType.toString();
                if (type.startsWith(Optional.class.getCanonicalName())) {
                    type = type.substring(Optional.class.getCanonicalName().length() + 1, type.length() - 1);
                }
                if (jdbiStoreAsJson != null) {
                    parsingType = ParsingType.jsonobject;
                } else if (jdbiEmbedded != null) {
                    parsingType = ParsingType.embedded;
                } else if (String.class.getCanonicalName().equals(type)) {
                    parsingType = ParsingType.string;
                } else {
                    parsingType = ParsingType.jdbiobject;
                }
                break;
            default:
                String msg = String.format("Can't handle return type: %s", retType);
                throw new JDBIProcessingException(msg);
        }

        if (jdbiOptions.trim() && parsingType != ParsingType.string) {
            String msg = String.format("Can't trim type: %s", type);
            throw new JDBIProcessingException(msg);
        }
        return new ParsingAndRegularType(parsingType, type);
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return latestSupported();
    }

    enum ParsingType {
        primitive,
        string,
        jdbiobject,
        jsonobject,
        embedded,
    }
}