gobblin.metrics.reporter.util.SchemaRegistryVersionWriter.java Source code

Introduction

Here is the source code for gobblin.metrics.reporter.util.SchemaRegistryVersionWriter.java

Source

/*
 * 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 gobblin.metrics.reporter.util;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.util.Map;

import org.apache.avro.Schema;
import org.apache.commons.codec.DecoderException;
import org.apache.commons.codec.binary.Hex;

import com.google.common.base.Optional;
import com.google.common.base.Throwables;
import com.google.common.collect.Maps;
import com.typesafe.config.Config;

import gobblin.metrics.kafka.KafkaAvroSchemaRegistry;
import gobblin.metrics.kafka.SchemaRegistryException;
import gobblin.util.ConfigUtils;

/**
 * Implementation of {@link gobblin.metrics.reporter.util.SchemaVersionWriter} that uses a
 * {@link gobblin.metrics.kafka.KafkaAvroSchemaRegistry} to get Schema version identifier and write it to
 * {@link java.io.DataOutputStream}.
 */
public class SchemaRegistryVersionWriter implements SchemaVersionWriter<Schema> {

    private final KafkaAvroSchemaRegistry registry;
    private Map<Schema, String> registrySchemaIds;
    private final Optional<String> overrideName;
    private final Optional<Schema> schema;
    private final Optional<String> schemaId;
    private final int schemaIdLengthBytes;

    public SchemaRegistryVersionWriter(Config config) throws IOException {
        this(new KafkaAvroSchemaRegistry(ConfigUtils.configToProperties(config)), Optional.<String>absent(),
                Optional.<Schema>absent());
    }

    public SchemaRegistryVersionWriter(KafkaAvroSchemaRegistry registry, String overrideName) throws IOException {
        this(registry, overrideName, Optional.<Schema>absent());
    }

    public SchemaRegistryVersionWriter(KafkaAvroSchemaRegistry registry, String overrideName,
            Optional<Schema> singleSchema) throws IOException {
        this(registry, Optional.of(overrideName), singleSchema);
    }

    public SchemaRegistryVersionWriter(KafkaAvroSchemaRegistry registry, Optional<String> overrideName,
            Optional<Schema> singleSchema) throws IOException {
        this.registry = registry;
        this.registrySchemaIds = Maps.newConcurrentMap();
        this.overrideName = overrideName;
        this.schema = singleSchema;
        this.schemaIdLengthBytes = registry.getSchemaIdLengthByte();
        if (this.schema.isPresent()) {
            try {
                this.schemaId = this.overrideName.isPresent()
                        ? Optional.of(this.registry.register(this.schema.get(), this.overrideName.get()))
                        : Optional.of(this.registry.register(this.schema.get()));
            } catch (SchemaRegistryException e) {
                throw Throwables.propagate(e);
            }
        } else {
            this.schemaId = Optional.absent();
        }
    }

    @Override
    public void writeSchemaVersioningInformation(Schema schema, DataOutputStream outputStream) throws IOException {

        String schemaId = this.schemaId.isPresent() ? this.schemaId.get() : this.getIdForSchema(schema);

        outputStream.writeByte(KafkaAvroSchemaRegistry.MAGIC_BYTE);
        try {
            outputStream.write(Hex.decodeHex(schemaId.toCharArray()));
        } catch (DecoderException exception) {
            throw new IOException(exception);
        }
    }

    private String getIdForSchema(Schema schema) {
        if (!this.registrySchemaIds.containsKey(schema)) {
            try {
                String schemaId = this.overrideName.isPresent()
                        ? this.registry.register(schema, this.overrideName.get())
                        : this.registry.register(schema);
                this.registrySchemaIds.put(schema, schemaId);
            } catch (SchemaRegistryException e) {
                throw Throwables.propagate(e);
            }
        }
        return this.registrySchemaIds.get(schema);
    }

    @Override
    public Schema readSchemaVersioningInformation(DataInputStream inputStream) throws IOException {
        if (inputStream.readByte() != KafkaAvroSchemaRegistry.MAGIC_BYTE) {
            throw new IOException("MAGIC_BYTE not found in Avro message.");
        }

        byte[] byteKey = new byte[schemaIdLengthBytes];
        int bytesRead = inputStream.read(byteKey, 0, schemaIdLengthBytes);
        if (bytesRead != schemaIdLengthBytes) {
            throw new IOException(
                    String.format("Could not read enough bytes for schema id. Expected: %d, found: %d.",
                            schemaIdLengthBytes, bytesRead));
        }
        String hexKey = Hex.encodeHexString(byteKey);

        try {
            return this.registry.getSchemaByKey(hexKey);
        } catch (SchemaRegistryException sre) {
            throw new IOException("Failed to retrieve schema for key " + hexKey, sre);
        }
    }
}