Immutables 2.0 and JDBI integration

BindValue.java

/*
    Copyright 2014 Immutables Authors and Contributors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 */
package org.immutables.common.jdbi;

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonSerializer;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializerProvider;
import com.fasterxml.jackson.databind.annotation.JsonSerialize;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.databind.util.TokenBuffer;
import com.google.common.annotations.Beta;
import com.google.common.base.Preconditions;
import org.joda.time.DateTime;
import org.joda.time.format.DateTimeFormat;
import org.joda.time.format.DateTimeFormatter;
import org.skife.jdbi.v2.SQLStatement;
import org.skife.jdbi.v2.exceptions.UnableToCreateStatementException;
import org.skife.jdbi.v2.sqlobject.Binder;
import org.skife.jdbi.v2.sqlobject.BinderFactory;
import org.skife.jdbi.v2.sqlobject.BindingAnnotation;

import java.io.IOException;
import java.lang.annotation.*;
import java.util.Date;
import java.util.Map;
import java.util.Map.Entry;

@Beta
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.PARAMETER})
@BindingAnnotation(BindValue.BindValueFactory.class)
public @interface BindValue {

    String value() default "";

    public static class BindValueFactory implements BinderFactory {

        private static final ObjectMapper CODEC = new ObjectMapper();

        static {
            // This magic makes sure that we write out Date and DateTimes in the form MySql understands.
            final DateTimeFormatter JDBC_DATE_FORMATTER = DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss.SSS");
            SimpleModule module = new SimpleModule("ImmutablesJDBIDateSerializer");
            module.addSerializer(DateTime.class, new JsonSerializer<DateTime>() {
                @Override
                public void serialize(DateTime value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
                    jgen.writeString(JDBC_DATE_FORMATTER.print(value));
                }
            });
            module.addSerializer(Date.class, new JsonSerializer<Date>() {
                @Override
                public void serialize(Date value, JsonGenerator jgen, SerializerProvider provider) throws IOException {
                    jgen.writeString(JDBC_DATE_FORMATTER.print(new DateTime(value)));
                }

            });
            
            CODEC.registerModule(new JodaModule());
            CODEC.registerModule(module);
        }


        private static boolean accepts(Class type) {
            if (type.getAnnotation(JsonSerialize.class) != null) {
                return true;
            }

            final Class<?> parent = type.getSuperclass();
            return parent != null && parent.getAnnotation(JsonSerialize.class) != null;
        }

        private static final TypeReference<Map<String, Object>> MAP_OF_OBJECTS =
                new TypeReference<Map<String, Object>>() {};

        @Override
        public Binder build(final Annotation annotation) {
            return new Binder<BindValue, Object>() {
                @Override
                public void bind(SQLStatement<?> q, BindValue bind, Object arg) {
                    Class<?> argumentType = arg.getClass();

                    Preconditions.checkState(accepts(argumentType), "The bound type must be an Immutables class with the @JsonSerialize annotation present");

                    String prefix = prefix(bind);
                    try {
                        @SuppressWarnings("resource")
                        TokenBuffer buffer = new TokenBuffer(CODEC, false);
                        CODEC.writeValue(buffer, arg);
                        Map<String, Object> parameters = buffer.asParser().readValueAs(MAP_OF_OBJECTS);
                        bindStatement(q, prefix, parameters);
                    } catch (Exception exception) {
                        UnableToCreateStatementException statementException = new UnableToCreateStatementException(
                                String.format("Could not bind parameter %s as '%s'", arg, bind.value()),
                                exception,
                                q.getContext());

                        MapperFactory.makeStackTraceUseful(statementException, exception);
                        throw statementException;
                    }
                }

                private void bindStatement(SQLStatement<?> q, String prefix, Map<String, Object> map) {
                    for (Entry<String, Object> entry : map.entrySet()) {
                        Object value = entry.getValue();
                        String name = prefix + entry.getKey();
                        if (value == null) {
                            q.bind(name, (Object) null);
                        } else {
                            q.dynamicBind(value.getClass(), name, value);
                        }
                    }
                }

                private String prefix(BindValue bind) {
                    return !bind.value().isEmpty() ? (bind.value() + ".") : "";
                }

                @Override
                public String toString() {
                    return BindValueFactory.this.toString() + ".build(" + annotation + ")";
                }
            };
        }

        @Override
        public String toString() {
            return BindValue.class.getSimpleName() + "." + BindValueFactory.class.getSimpleName();
        }
    }
}

MapperFactory.java

/*
    Copyright 2014 Immutables Authors and Contributors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 */
package org.immutables.common.jdbi;

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.DeserializationContext;
import com.fasterxml.jackson.databind.JsonDeserializer;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.databind.util.TokenBuffer;
import com.fasterxml.jackson.datatype.joda.deser.DateTimeDeserializer;
import com.google.common.annotations.Beta;
import com.google.common.base.Ascii;
import com.google.common.base.CaseFormat;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.collect.ObjectArrays;
import org.joda.time.DateTime;
import org.skife.jdbi.v2.ResultSetMapperFactory;
import org.skife.jdbi.v2.StatementContext;
import org.skife.jdbi.v2.exceptions.ResultSetException;
import org.skife.jdbi.v2.tweak.ResultSetMapper;

import javax.annotation.Nullable;
import java.io.IOException;
import java.math.BigDecimal;
import java.sql.*;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;

@Beta
@SuppressWarnings("unchecked")
public class MapperFactory implements ResultSetMapperFactory {
    private static final String COLUMN_METADATA_ATTRIBUTE = ColumnMetadata.class.getCanonicalName();

    private static final ObjectMapper CODEC = new ObjectMapper();

    private static DateTime deserialiseDate(JsonParser parser, DeserializationContext ctxt) throws IOException {

        // We embed the object below in Mapper.generateTokens()
        Object object = parser.getEmbeddedObject();
        if (object instanceof java.util.Date) {
            return new DateTime(object);
        }

        // Fall back to default JODA deserialisation - e.g. long millis or String.
        return DateTimeDeserializer.forType(DateTime.class).deserialize(parser, ctxt);
    }

    static {
        SimpleModule module = new SimpleModule("ImmutablesJDBIDateDeserializer");
        module.addDeserializer(DateTime.class, new JsonDeserializer <DateTime>() {
           @Override
            public DateTime deserialize(JsonParser jp, DeserializationContext ctxt) throws IOException {
                return deserialiseDate(jp, ctxt);
            }
        });
        module.addDeserializer(Date.class, new JsonDeserializer <Date>(){
            @Override
            public Date deserialize(JsonParser jp, DeserializationContext ctxt) throws IOException {
                return new Date(deserialiseDate(jp, ctxt).getMillis());
            }
        });
        CODEC.registerModule(new JodaModule());
        CODEC.registerModule(module);
    }


    private static Cache <Class, Boolean>  isImmutableType = CacheBuilder.newBuilder().weakKeys().weakValues().build();
    private static boolean isImmutable(final Class type) {

        try {
            return isImmutableType.get(type, new Callable<Boolean>() {
                 @Override
                 public Boolean call() throws Exception {

                     // It's the parent immutable class
                     if (type.getAnnotation(JsonDeserialize.class) != null) {
                         return true;
                     }

                     // It's the Immutable concrete implementation
                     final Class superclass = type.getSuperclass();
                     if (superclass != null && superclass.getAnnotation(JsonDeserialize.class) != null) {
                         return true;
                     }

                     // Check if it's an Immutable class. We still can't use it without the JsonDeserialize but at least we can log a useful message
                     try {
                         if (!type.isPrimitive()) {
                             final Class<?> immutableSubclass = Class.forName(type.getPackage().getName() + ".Immutable" + type.getSimpleName());
                             if (immutableSubclass != null) {
                                 throw new RuntimeException("Found Immutable class for " + type + " in " + MapperFactory.class + " without @JsonDeserialize annotation.\n\tHINT: Add the @JsonDeserialize annotation to the class.");
                             }
                         }
                     } catch (ClassNotFoundException e) {
                         // OK
                     }
                     return false;
                 }
             });
        } catch (ExecutionException e) {
            throw new RuntimeException(e);
        }
    }
    
    @Override
    public final boolean accepts(Class type, StatementContext ctx) {
       return isImmutable(type);
    }

    @SuppressWarnings({"unchecked", "unused"})
    @Override
    public final ResultSetMapper mapperFor(final Class type, StatementContext ctx) {
        return new Mapper<Object>(type);
    }

    /**
     * Override this method in your {@link MapperFactory} subclass to hook into writing of values from
     * result set.
     * @param buffer token buffer
     * @return json generator view of provided buffer
     */
    protected JsonGenerator asGenerator(TokenBuffer buffer) {
        return buffer;
    }

    /**
     * Override this method in your {@link MapperFactory} subclass to hook into reading of values
     * during unmarshaling
     * of result set.
     * @param buffer the buffer
     * @return json parser view of provied buffer
     */
    protected JsonParser asParser(TokenBuffer buffer) {
        return buffer.asParser();
    }

    /**
     * Converts column label into name that should match. While it's overridable, in most cases
     * default conversion might suffice as it does snake-case and camel-case conversion on a best
     * effort basis.
     * @param columnLabel the column label
     * @return immutable object attribute name.
     */
    protected String columnLabelToAttributeName(String columnLabel) {
        boolean hasUnderscoreSeparator = columnLabel.indexOf('_') > 0;
        if (hasUnderscoreSeparator || Ascii.toUpperCase(columnLabel).equals(columnLabel)) {
            columnLabel = Ascii.toLowerCase(columnLabel);
        }
        return hasUnderscoreSeparator
                ? CaseFormat.LOWER_UNDERSCORE.to(CaseFormat.LOWER_CAMEL, columnLabel)
                : CaseFormat.UPPER_CAMEL.to(CaseFormat.LOWER_CAMEL, columnLabel);
    }

    /**
     * Metadata is cached due to possible overhead in attribute name conversions
     * in a hope that hashtable lookup from {@link StatementContext#getAttribute(String)} is less
     * expensive. Also, same string instances will be used for attribute name on each result row.
     */
    private final class ColumnMetadata {
        final int count;
        final int[] types;
        final String[] names;

        ColumnMetadata(ResultSetMetaData metaData) throws SQLException {
            this.count = metaData.getColumnCount();
            this.types = new int[this.count];
            this.names = new String[this.count];
            collectMetadata(metaData);
        }

        private void collectMetadata(ResultSetMetaData metaData) throws SQLException {
            for (int i = 0; i < count; i++) {
                types[i] = metaData.getColumnType(i + 1);
                names[i] = columnLabelToAttributeName(metaData.getColumnLabel(i + 1));
            }
        }
    }

    private ColumnMetadata columnsFrom(ResultSetMetaData metaData, StatementContext ctx) throws SQLException {
        Object attribute = ctx.getAttribute(COLUMN_METADATA_ATTRIBUTE);
        if (attribute instanceof ColumnMetadata) {
            return (ColumnMetadata) attribute;
        }
        ColumnMetadata columnMetadata = new ColumnMetadata(metaData);
        ctx.setAttribute(COLUMN_METADATA_ATTRIBUTE, columnMetadata);
        return columnMetadata;
    }

    /**
     * Stacktrace of wrapper derived exception is almost always irrelevant,
     * while stack of cause is most important. In any case it would be
     * relatively easy to find this place.
     */
    static void makeStackTraceUseful(Exception wrapper, Exception cause) {
        wrapper.setStackTrace(ObjectArrays.concat(
                wrapper.getStackTrace()[0],
                cause.getStackTrace()));
    }

    private class Mapper<T> implements ResultSetMapper<T> {

        private final Class type;

        public Mapper(Class type) {
            this.type = type;
        }

        @SuppressWarnings("resource")
        @Override
        public T map(int index, ResultSet result, StatementContext ctx) throws SQLException {
            try {
                TokenBuffer buffer = new TokenBuffer(null, false);
                ColumnMetadata columns = columnsFrom(result.getMetaData(), ctx);
                generateTokens(result, columns, asGenerator(buffer));
                return (T) CODEC.readValue(asParser(buffer), this.type);
            } catch (IOException ex) {
                ResultSetException resultSetException = new ResultSetException("Unable to map result object", ex, ctx);
                makeStackTraceUseful(resultSetException, ex);
                throw resultSetException;
            }
        }

        private void generateTokens(
                ResultSet result,
                ColumnMetadata columns,
                JsonGenerator generator) throws IOException, SQLException {

            generator.writeStartObject();

            for (int j = 0; j < columns.count; j++) {
                int i = j + 1;
                String name = columns.names[j];
                switch (columns.types[j]) {
                    case Types.VARCHAR://$FALL-THROUGH$
                    case Types.LONGVARCHAR://$FALL-THROUGH$
                    case Types.CHAR: {
                        String v = result.getString(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeString(v);
                        }
                        break;
                    }
                    case Types.NVARCHAR://$FALL-THROUGH$
                    case Types.NCHAR://$FALL-THROUGH$
                    case Types.LONGNVARCHAR: {
                        String v = result.getNString(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeString(v);
                        }
                        break;
                    }
                    case Types.DATE:
                        Date date = result.getDate(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeObject(date);
                        }
                        break;
                    case Types.TIME:
                        Time time = result.getTime(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeObject(time);
                        }
                        break;
                    case Types.TIMESTAMP: {
                        Timestamp timestamp = result.getTimestamp(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeObject(timestamp);
                        }
                        break;
                    }
                    case Types.BOOLEAN://$FALL-THROUGH$
                    case Types.BIT: {
                        boolean v = result.getBoolean(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeBoolean(v);
                        }
                        break;
                    }
                    case Types.INTEGER://$FALL-THROUGH$
                    case Types.TINYINT: {
                        int v = result.getInt(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeNumber(v);
                        }
                        break;
                    }
                    case Types.BIGINT: {
                        long v = result.getLong(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeNumber(v);
                        }
                        break;
                    }
                    case Types.DECIMAL: {
                        BigDecimal v = result.getBigDecimal(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeNumber(v);
                        }
                        break;
                    }
                    case Types.NUMERIC://$FALL-THROUGH$
                    case Types.REAL://$FALL-THROUGH$
                    case Types.FLOAT://$FALL-THROUGH$
                    case Types.DOUBLE: {
                        double v = result.getDouble(i);
                        if (!result.wasNull()) {
                            generator.writeFieldName(name);
                            generator.writeNumber(v);
                        }
                        break;
                    }
                    case Types.NULL: {
                        generator.writeFieldName(name);
                        generator.writeNull();
                        break;
                    }
                    default:
                        @Nullable
                        Object object = result.getObject(i);
                        if (object != null) {
                            generator.writeFieldName(name);
                            generator.writeObject(object);
                        }
                }
            }
            generator.writeEndObject();
        }

        @Override
        public String toString() {
            return MapperFactory.this.getClass().getSimpleName() + ".mapperFor(" + type + ")";
        }
    }
}