nstream.adapter.avro.SwimAvroAssembler Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of nstream-adapter-avro Show documentation
Templates for loading avro data
The newest version!
// Copyright 2015-2024 Nstream, inc.
//
// Licensed under the Redis Source Available License 2.0 (RSALv2) Agreement;
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://redis.com/legal/rsalv2-agreement/
//
// 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 nstream.adapter.avro;

import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import nstream.adapter.common.AdapterUtils;
import nstream.adapter.common.ingress.AssemblyException;
import nstream.adapter.common.ingress.ContentAssembler;
import org.apache.avro.Schema;
import org.apache.avro.generic.GenericDatumReader;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.io.Decoder;
import org.apache.avro.io.DecoderFactory;
import swim.codec.Unicode;
import swim.codec.Utf8;
import swim.structure.Attr;
import swim.structure.Value;

// TODO: replace implementation to use swim-avro libraries
public class SwimAvroAssembler extends ContentAssembler {

  private Schema schema;
  private GenericDatumReader reader;

  public SwimAvroAssembler(Value settings) {
    super(settings);
    if (settings.head() instanceof Attr) {
      final String head = settings.head().key().stringValue(null);
      if ("avroAssembler".equals(head) || "valueAssembler".equals(head)) {
        if (loadFromSchema(settings) || loadFromSchemaConfig(settings)) {
          return;
        }
      }
    }
    throw new IllegalArgumentException("Invalid SwimAvroAssembler structure: " + settings);
  }

  private boolean loadFromSchema(Value settings) {
    final String schemaStr = settings.get("schema").stringValue(null);
    if (schemaStr != null) {
      try {
        this.schema = new Schema.Parser().parse(schemaStr);
        this.reader = new GenericDatumReader<>(this.schema);
        return true;
      } catch (Exception e) {
        // swallow
      }
    }
    return false;
  }

  private boolean loadFromSchemaConfig(Value settings) {
    final String schemaConfig = settings.get("schemaConfig").stringValue(null);
    // final String schemaConfig = settings.get("schemaConfig").stringValue(
    //     settings.get("schema").stringValue(null));
    if (schemaConfig != null) {
      try (InputStream is = AdapterUtils.openFileAsStream(schemaConfig, SwimAvroAssembler.class)) {
        final String schemaStr = Utf8.read(is, Unicode.stringParser());
        if (schemaStr != null) {
          this.schema = new Schema.Parser().parse(schemaStr);
          this.reader = new GenericDatumReader<>(this.schema);
          return true;
        }
      } catch (Exception e) {
        // swallow
      }
    }
    return false;
  }

  @Override
  public String contentType() {
    return "avro";
  }

  @Override
  public Value assembleStream(InputStream decodedStream) throws AssemblyException, IOException {
    return assembleBytes(decodedStream.readAllBytes());
  }

  @Override
  public Value assembleBytes(byte[] bytes) throws AssemblyException {
    final Decoder decoder = DecoderFactory.get().binaryDecoder(bytes, null);
    try {
      return Avro.genericRecordStructureHeadless(this.reader.read(null, decoder));
    } catch (IOException e) {
      throw new RuntimeException("Failed to mold", e);
    }
  }

  @Override
  public Value assemble(String raw) throws AssemblyException {
    return assembleBytes(raw.getBytes(StandardCharsets.UTF_8));
  }

}