MessageNamedArgumentFinder.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.providence.jdbi.v2;

import net.morimekta.providence.PMessage;
import net.morimekta.providence.PMessageOrBuilder;
import net.morimekta.providence.PType;
import net.morimekta.providence.descriptor.PField;
import net.morimekta.providence.descriptor.PMessageDescriptor;
import net.morimekta.providence.jdbi.v2.util.NullArgument;
import net.morimekta.util.collect.UnmodifiableMap;
import org.skife.jdbi.v2.tweak.Argument;
import org.skife.jdbi.v2.tweak.NamedArgumentFinder;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.Map;

import static net.morimekta.util.Strings.isNullOrEmpty;

/**
 * A {@link NamedArgumentFinder} implementation that uses a message
 * and finds values based on the thrift declared field names. This
 * supports chained calls to any depth as long as each level is a
 * single message field.
 *
 * @param <M> The message type.
 */
public class MessageNamedArgumentFinder<M extends PMessage<M>> implements NamedArgumentFinder {
    private final String                  prefix;
    private final PMessageOrBuilder<M>    message;
    private final Map<PField<M>, Integer> fieldTypes;

    /**
     * Create a named argument finder.
     *
     * @param prefix Optional prefix name. E.g. "x" will make for lookup
     *               tags like ":x.my_field".
     * @param message The message to look up fields in.
     * @param fieldTypes Overriding of default field types. This can contain
     *                   fields for any of the contained message types, and
     *                   will be mapped whenever the field is selected.
     */
    public MessageNamedArgumentFinder(@Nullable String prefix,
                                      @Nonnull PMessageOrBuilder<M> message,
                                      @Nonnull Map<PField<M>, Integer> fieldTypes) {
        this.message = message;
        this.prefix = isNullOrEmpty(prefix) ? "" : prefix + ".";
        this.fieldTypes = UnmodifiableMap.copyOf(fieldTypes);
    }

    @Override
    public String toString() {
        return prefix + "{" + message.descriptor().getQualifiedName() + "}";
    }

    @Override
    @SuppressWarnings("unchecked")
    public Argument find(String name) {
        if (!prefix.isEmpty()) {
            if (name.startsWith(prefix)) {
                name = name.substring(prefix.length());
            } else {
                return null;
            }
        }

        String[]             parts          = name.split("\\.", Byte.MAX_VALUE);
        PMessageOrBuilder<?> leaf           = message;
        PMessageDescriptor   leafDescriptor = message.descriptor();

        for (int i = 0; i < parts.length - 1; ++i) {
            String part = parts[i];
            PField field = leafDescriptor.findFieldByName(part);
            if (field == null) return null;
            if (field.getType() != PType.MESSAGE) {
                throw new IllegalArgumentException("");
            }
            leafDescriptor = (PMessageDescriptor) field.getDescriptor();
            if (leaf != null) {
                leaf = leaf.get(field.getId());
            }
        }
        String leafName = parts[parts.length - 1];
        PField field = leafDescriptor.findFieldByName(leafName);
        if (field != null) {
            if (leaf != null) {
                return new MessageFieldArgument(leaf, field, getColumnType(field));
            }
            return new NullArgument(getColumnType(field));
        }
        return null;
    }

    private int getColumnType(PField field) {
        return fieldTypes.getOrDefault(field, MessageFieldArgument.getDefaultColumnType(field));
    }

}