com.facebook.drift.codec.internal.reflection.ReflectionThriftUnionCodec Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of drift-codec Show documentation
Annotation based encoder and decoder for Thrift
There is a newer version: 1.39
/*
 * Copyright (C) 2012 Facebook, Inc.
 *
 * 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 com.facebook.drift.codec.internal.reflection;

import com.facebook.drift.codec.ThriftCodec;
import com.facebook.drift.codec.ThriftCodecManager;
import com.facebook.drift.codec.internal.ProtocolReader;
import com.facebook.drift.codec.internal.ProtocolWriter;
import com.facebook.drift.codec.metadata.FieldKind;
import com.facebook.drift.codec.metadata.ThriftConstructorInjection;
import com.facebook.drift.codec.metadata.ThriftFieldInjection;
import com.facebook.drift.codec.metadata.ThriftFieldMetadata;
import com.facebook.drift.codec.metadata.ThriftInjection;
import com.facebook.drift.codec.metadata.ThriftMethodInjection;
import com.facebook.drift.codec.metadata.ThriftStructMetadata;
import com.facebook.drift.protocol.TProtocolReader;
import com.facebook.drift.protocol.TProtocolWriter;
import com.google.common.collect.Maps;

import javax.annotation.concurrent.Immutable;

import java.util.Map;

import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.Iterables.getOnlyElement;
import static com.google.common.collect.Maps.uniqueIndex;
import static java.lang.String.format;
import static java.util.Objects.requireNonNull;

@Immutable
public class ReflectionThriftUnionCodec
        extends AbstractReflectionThriftCodec
{
    private final Map metadataMap;
    private final Map.Entry> idField;

    public ReflectionThriftUnionCodec(ThriftCodecManager manager, ThriftStructMetadata metadata)
    {
        super(manager, metadata);

        ThriftFieldMetadata idField = getOnlyElement(metadata.getFields(FieldKind.THRIFT_UNION_ID));

        this.idField = Maps.immutableEntry(idField, manager.getCodec(idField.getThriftType()));
        requireNonNull(this.idField.getValue(), () -> "No codec for ID field found: " + idField);

        this.metadataMap = uniqueIndex(metadata.getFields(), ThriftFieldMetadata::getId);
    }

    @Override
    public T read(TProtocolReader protocol)
            throws Exception
    {
        ProtocolReader reader = new ProtocolReader(protocol);
        reader.readStructBegin();

        Map.Entry data = null;
        Short fieldId = null;
        while (reader.nextField()) {
            checkState(fieldId == null, "Received Union with more than one value (seen id %s, now id %s)", fieldId, reader.getFieldId());

            fieldId = reader.getFieldId();

            // do we have a codec for this field
            ThriftCodec codec = fields.get(fieldId);
            if (codec == null) {
                reader.skipFieldData();
            }
            else {
                // is this field readable
                ThriftFieldMetadata field = metadata.getField(fieldId);
                if (field.isWriteOnly() || field.getType() != FieldKind.THRIFT_FIELD) {
                    reader.skipFieldData();
                    continue;
                }

                // read the value
                Object value = reader.readField(codec);
                if (value == null) {
                    continue;
                }

                data = Maps.immutableEntry(fieldId, value);
            }
        }
        reader.readStructEnd();

        // build the struct
        return constructStruct(data);
    }

    @Override
    public void write(T instance, TProtocolWriter protocol)
            throws Exception
    {
        ProtocolWriter writer = new ProtocolWriter(protocol);

        Short idValue = (Short) getFieldValue(instance, idField.getKey());

        writer.writeStructBegin(metadata.getStructName());

        if (metadataMap.containsKey(idValue)) {
            ThriftFieldMetadata fieldMetadata = metadataMap.get(idValue);

            if (fieldMetadata.isReadOnly() || fieldMetadata.getType() != FieldKind.THRIFT_FIELD) {
                throw new IllegalStateException(format("Field %s is not readable", fieldMetadata.getName()));
            }

            Object fieldValue = getFieldValue(instance, fieldMetadata);

            // write the field
            if (fieldValue != null) {
                @SuppressWarnings("unchecked")
                ThriftCodec