JDBI + Immutables annotation processor

Jdbi.java

package gaff.jdbi;

import com.google.common.base.CaseFormat;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

@Retention(RetentionPolicy.SOURCE)
public @interface Jdbi {
    CaseFormat dbCaseFormat() default CaseFormat.LOWER_UNDERSCORE;
    CaseFormat javaCaseFormat() default CaseFormat.LOWER_CAMEL;
    String tableName() default "";
}

jdbifactory.moustache

package {{pkg}};

import org.skife.jdbi.v2.tweak.ResultSetMapper;
import org.skife.jdbi.v2.StatementContext;
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;
import gaff.jdbi.JdbiHelper;

public class {{targetSimpleName}}Mapper implements ResultSetMapper<{{target}}> {
    public static final String DB_COLUMNS = "{{#columnList}}{{formattedDbName}}{{^last}}, {{/last}}{{/columnList}}";
    public static final String BIND_COLUMNS = "{{#columnList}}:{{javaName}}{{^last}}, {{/last}}{{/columnList}}";
    {{#tableName}}
    public static final String SQL_INSERT = "insert into {{tableName}} (" + DB_COLUMNS + ") values (" + BIND_COLUMNS + ")";
    {{/tableName}}

    @Override
    public {{target}} map(int index, ResultSet r, StatementContext ctx) throws SQLException {
        Immutable{{targetSimpleName}}.Builder builder = Immutable{{targetSimpleName}}.builder();

        JdbiHelper dh = new JdbiHelper(r);
        {{#members}}
        {{{extractionMethod}}};
        {{/members}}

        return builder.build();
    }
}

JdbiHelper.java

package gaff.jdbi;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;

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;

import static gaff.jdbi.JdbiHelper.ObjectMapperFactory.mapper;

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

    public JdbiHelper(ResultSet rs){
        this.rs = rs;
    }

    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));
            }
        }

        return returnedColumns.contains(columnName);
    }

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

        return mapper.convertValue(rs.getObject(columnName), new TypeReference<T>() {});
    }

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

        //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);

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

        if(trim){
            entity = entity.trim();
        }

        consumer.accept(entity);
    }

    /**
     * This hook allows us to inject custom mappings.
     *
     * Note that the itself mapper is threadsafe - but configuration of it is not. As such one should not
     * register any modules in the mapper once the configuration is released.
     *
     * JDBI creates a new ResultSetMapper each time it performs a query => if there are no queries in flight then
     * nobody is referencing the objectMapper.
     *
     * A Better implementation would use ObjectReader / ObjectWriter, AFAICT these are immutable snapshots of the mapper
     * for a specific type.
     */
    public static class ObjectMapperFactory {
        public static final ObjectMapper mapper = new ObjectMapper().findAndRegisterModules();

        public static ObjectMapper getMapper() {
            return mapper;
        }
    }
}

JdbiOptions.java

package gaff.jdbi;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;

@Retention(RetentionPolicy.SOURCE)
public @interface JdbiOptions {
    String columnName() default "";
    boolean ignore() default false;
    boolean inColumnList() default true;
    boolean optional() default false;
    boolean trim() default false;
}

JdbProcessor.java

package gaff.annotationprocessors;

import com.github.mustachejava.DefaultMustacheFactory;
import com.github.mustachejava.Mustache;
import com.github.mustachejava.MustacheFactory;
import gaff.jdbi.Jdbi;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import javax.lang.model.util.Elements;
import javax.tools.JavaFileObject;
import java.io.IOException;
import java.io.Writer;
import java.util.*;

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

@SupportedAnnotationTypes({
        "gaff.jdbi.Jdbi"
})
public class JdbiProcessor extends AbstractProcessor {
    private final Mustache mustache;
    private Filer filer;
    private Elements elementUtils;

    public JdbiProcessor() {
        MustacheFactory mf = new DefaultMustacheFactory();
        mustache = mf.compile("jdbifactory.mustache");
    }

    @Override public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        //typeUtils = processingEnv.getTypeUtils();
        elementUtils = processingEnv.getElementUtils();
        filer = processingEnv.getFiler();
        //messager = processingEnv.getMessager();
    }


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

        Messager messager = processingEnv.getMessager();

        for(Element elem : env.getElementsAnnotatedWith(Jdbi.class)) {
            try {
                generateCode((TypeElement)elem, filer);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }

        return true;
    }

    private void generateCode(TypeElement elem, Filer filer) throws IOException {
        Model m = new Model(elem, elementUtils);

        JavaFileObject jfo = filer.createSourceFile( m.pkg + "." + m.targetSimpleName + "Mapper");
        Writer writer = jfo.openWriter();
        mustache.execute(writer, m);
        writer.close();
    }

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

Model.java

package gaff.annotationprocessors;

import com.google.common.collect.Lists;
import gaff.jdbi.Jdbi;
import gaff.jdbi.JdbiOptions;

import javax.lang.model.element.*;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.Elements;
import java.lang.annotation.Annotation;
import java.util.ArrayList;
import java.util.List;

public class Model {
    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;
        }
    };


    public final String pkg;
    public final String target;
    public final String targetSimpleName;
    public final String tableName;
    public final List<Member> members;
    public final List<Member> columnList; //Different if you want to skip some fields


    private String determineRelativeName(Element elem) {
        //This method works out the name of the annotated element relative to the package.

        List<String> targetBuilder = new ArrayList<>();
        while(elem.getKind() != ElementKind.PACKAGE) {
            targetBuilder.add(elem.getSimpleName().toString());
            elem = elem.getEnclosingElement();
        }

        return String.join(".", Lists.reverse(targetBuilder));
    }

    public Model(TypeElement te, Elements elementUtils) {

        PackageElement pkg = elementUtils.getPackageOf(te);
        this.pkg = pkg.toString();

        this.target = determineRelativeName(te);
        this.targetSimpleName = te.getSimpleName().toString();

        Jdbi jdbiAnnotation = te.getAnnotation(Jdbi.class);
        if (!jdbiAnnotation.tableName().isEmpty()) {
            this.tableName = jdbiAnnotation.tableName();
        } else {
            this.tableName = null;
        }

        members = new ArrayList<>();
        columnList = new ArrayList<>();

        for (Element member : te.getEnclosedElements()) {
            //Check for method annotation
            JdbiOptions jdbiOptions = member.getAnnotation(JdbiOptions.class);
            if( jdbiOptions == null) jdbiOptions = DEFAULT_OPTIONS;
            if (jdbiOptions.ignore()) {
                continue;
            }

            switch (member.getKind()) {
                case METHOD:
                    Member m = new Member((ExecutableElement) member, jdbiAnnotation, jdbiOptions);
                    members.add(m);
                    if(jdbiOptions.inColumnList()) {
                        columnList.add(m);
                    }
                    break;
                default:
                    //Nothing to do.
            }
        }

        //This is so we can don't have trailing commas.
        columnList.get(columnList.size() - 1).last = true;
    }

    public static class Member {
        public final String javaName;
        public final String formattedDbName;
        public final String extractionMethod;
        public boolean last = false;

        enum JdbiType {
            primitive,
            string,
            object
        }

        public Member(ExecutableElement member, Jdbi jdbiAnnotation, JdbiOptions jdbiOptions) {
            TypeMirror rettype = member.getReturnType();
            this.javaName = member.getSimpleName().toString();

            JdbiType jdbiType;
            String type;
            switch (rettype.getKind()) {
                case INT:
                    jdbiType = JdbiType.primitive;
                    type = "Int";
                    break;
                case DOUBLE:
                    jdbiType = JdbiType.primitive;
                    type = "Double";
                    break;
                case LONG:
                    jdbiType = JdbiType.primitive;
                    type = "Long";
                    break;
                case BOOLEAN:
                    jdbiType = JdbiType.primitive;
                    type = "Boolean";
                    break;
                case DECLARED:
                    type = rettype.toString();
                    if(String.class.getCanonicalName().equals(type)){
                        jdbiType = JdbiType.string;
                    }else {
                        jdbiType = JdbiType.object;
                    }
                    break;
                default:
                    String msg = String.format("Can't handle type: %s", rettype);
                    throw new RuntimeException(msg);
            }

            if(jdbiOptions.trim() && jdbiType != JdbiType.string){
                String msg = String.format("Can't trim type: %s", type);
                throw new RuntimeException(msg);
            }

            String rawDbName;
            if (jdbiOptions.columnName().isEmpty()) {
                rawDbName = javaName;
                formattedDbName = jdbiAnnotation.javaCaseFormat().to(jdbiAnnotation.dbCaseFormat(), rawDbName);
            } else {
                rawDbName = jdbiOptions.columnName();
                formattedDbName = rawDbName;
            }
            String methodRef = String.format("(Consumer<%s>) builder::%s", type, javaName);
            switch (jdbiType){
                case primitive:
                    extractionMethod = String.format("builder.%s(r.get%s(\"%s\"))", javaName, type, formattedDbName);
                    break;
                case object:
                    extractionMethod = String.format("dh.setEntity(%s, \"%s\", %b)", methodRef, formattedDbName, jdbiOptions.optional());
                    break;
                case string:
                    extractionMethod = String.format("dh.setString(%s, \"%s\", %b, %b)", methodRef, formattedDbName, jdbiOptions.optional(), jdbiOptions.trim());
                    break;
                default:
                    String msg = String.format("Can't handle type: %s", rettype);
                    throw new RuntimeException(msg);
            }
        }
    }
}

Here is an example input:

package dummy;

import gaff.jdbi.JdbiOptions;
import org.immutables.value.Value;
import gaff.jdbi.Jdbi;
import java.time.LocalDateTime;
import java.util.UUID;

@Value.Immutable
@Jdbi(tableName = "DummyTable")
public interface Dummy {
    UUID myId();
    int myNumber();
    @JdbiOptions(columnName = "my_string")
    String myString1();

    LocalDateTime myDate();

    @JdbiOptions(ignore = true)
    @Value.Default default String myString2() {
        return "Ignore me!";
    }
}

And the output:

package dummy;
import ...;

public class DummyMapper implements ResultSetMapper<Dummy> {
    public static final String DB_COLUMNS = "my_id, my_number, my_string, my_date";
    public static final String BIND_COLUMNS = ":myId, :myNumber, :myString1, :myDate";
    public static final String SQL_INSERT = "insert into DummyTable (" + DB_COLUMNS + ") values (" + BIND_COLUMNS + ")";

    @Override
    public Dummy map(int index, ResultSet r, StatementContext ctx) throws SQLException {
        ImmutableDummy.Builder builder = ImmutableDummy.builder();

        JdbiHelper dh = new JdbiHelper(r);
        dh.setEntity((Consumer<java.util.UUID>) builder::myId, "my_id", false);
        builder.myNumber(r.getInt("my_number"));
        dh.setString((Consumer<java.lang.String>) builder::myString1, "my_string", false, false);
        dh.setEntity((Consumer<java.time.LocalDateTime>) builder::myDate, "my_date", false);

        return builder.build();
    }
}