EnumColumnMapper.java

/*
 * Copyright 2018-2019 Providence Authors
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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 net.morimekta.proto.jdbi.v3;

import com.google.protobuf.ProtocolMessageEnum;
import net.morimekta.collect.UnmodifiableMap;
import net.morimekta.proto.ProtoEnum;
import net.morimekta.proto.jdbi.MorimektaJdbiOptions;
import net.morimekta.strings.StringUtil;
import org.jdbi.v3.core.mapper.ColumnMapper;
import org.jdbi.v3.core.result.ResultSetException;
import org.jdbi.v3.core.statement.StatementContext;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Types;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

import static net.morimekta.strings.StringUtil.isNotEmpty;

/**
 * Map a result set to a message based on meta information and the message
 * descriptor.
 *
 * @param <E> The enum value type.
 */
public class EnumColumnMapper<E extends Enum<E> & ProtocolMessageEnum> implements ColumnMapper<E> {
    /**
     * Create a enum value column mapper.
     *
     * @param acceptUnknown If unknown values should be accepted.
     * @param descriptor    Message descriptor.
     */
    public EnumColumnMapper(boolean acceptUnknown, ProtoEnum<E> descriptor) {
        this.acceptUnknown = acceptUnknown;
        this.descriptor = Objects.requireNonNull(descriptor);
        this.valueByName = makeValueByNameMap(descriptor);
    }

    @Override
    public String toString() {
        return "EnumColumnMapper{type=" + descriptor.getTypeName() + "}";
    }

    public Class<E> getType() {
        return descriptor.getEnumClass();
    }

    private final boolean        acceptUnknown;
    private final ProtoEnum<E>   descriptor;
    private final Map<String, E> valueByName;

    @Override
    public E map(ResultSet r, int columnNumber, StatementContext ctx) throws SQLException {
        int columnType = r.getMetaData().getColumnType(columnNumber);
        switch (columnType) {
            case Types.TINYINT:
            case Types.SMALLINT:
            case Types.INTEGER:
            case Types.BIGINT:
            case Types.NUMERIC: {
                int id = r.getInt(columnNumber);
                if (r.wasNull()) {
                    return null;
                }
                E out = descriptor.findByNumber(id);
                if (out != null || acceptUnknown) {
                    return out;
                }
                throw new ResultSetException(
                        "Unknown value " + id +
                        " for enum " +
                        descriptor.getTypeName(),
                        null, ctx);
            }
            case Types.CLOB:
            case Types.NCHAR:
            case Types.VARCHAR:
            case Types.NVARCHAR:
            case Types.LONGVARCHAR:
            case Types.LONGNVARCHAR: {
                String name = r.getString(columnNumber);
                if (StringUtil.isNullOrEmpty(name)) {
                    return null;
                }
                E out = valueByName.get(name);
                if (out != null || acceptUnknown) {
                    return out;
                }
                throw new ResultSetException(
                        "Unknown value " + name +
                        " for enum " +
                        descriptor.getTypeName(),
                        null, ctx);
            }
            default:
                throw new ResultSetException(
                        "Unhandled column type " +
                        r.getMetaData().getColumnTypeName(columnNumber) +
                        "(" +
                        columnType +
                        ") for enum " +
                        descriptor.getTypeName(),
                        null, ctx);
        }
    }

    @SuppressWarnings("unchecked,rawtypes")
    private static <E extends Enum<E> & ProtocolMessageEnum> Map<String, E> makeValueByNameMap(ProtoEnum<E> descriptor) {
        return (Map<String, E>) (Map) valueByNameCache.computeIfAbsent(descriptor.getEnumClass(), c -> {
            LinkedHashMap<String, Object> builder = new LinkedHashMap<>();
            for (var value : descriptor.allValues()) {
                builder.put(value.name(), value);
            }
            for (var value : descriptor.allValues()) {
                var name = value.getValueDescriptor().getOptions().getExtension(MorimektaJdbiOptions.sqlName);
                if (isNotEmpty(name)) {
                    var old = builder.put(name, value);
                    if (old != null && old != value) {
                        // Not allowed to redefine a value for another value
                        // by setting the sql value field to the name of another.
                        throw new IllegalArgumentException(
                                "Enum value " + old + " redefined as " + value +
                                " for enum " + descriptor.getTypeName());
                    }
                }
            }
            return UnmodifiableMap.asMap(builder);
        });
    }

    private static final Map<Class<?>, Map<String, Object>> valueByNameCache = new ConcurrentHashMap<>();
}