org.apache.solr.handler.loader.CborLoader Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of solr-core Show documentation
Apache Solr (module: core)
There is a newer version: 9.7.0
/*
 * 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 org.apache.solr.handler.loader;

import com.fasterxml.jackson.core.JsonToken;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.dataformat.cbor.CBORFactory;
import com.fasterxml.jackson.dataformat.cbor.CBORGenerator;
import com.fasterxml.jackson.dataformat.cbor.CBORParser;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;
import org.apache.solr.common.SolrInputDocument;
import org.apache.solr.common.params.SolrParams;
import org.apache.solr.common.params.UpdateParams;
import org.apache.solr.common.util.ContentStream;
import org.apache.solr.request.SolrQueryRequest;
import org.apache.solr.response.SolrQueryResponse;
import org.apache.solr.update.AddUpdateCommand;
import org.apache.solr.update.processor.UpdateRequestProcessor;

/**
 * This class can load a single document or a stream of documents in CBOR format this is equivalent
 * of loading a single json documet or an array of json documents
 */
public class CborLoader {
  final CBORFactory cborFactory;
  private final Consumer sink;

  public CborLoader(CBORFactory cborFactory, Consumer sink) {
    this.cborFactory =
        cborFactory == null
            ? CBORFactory.builder().enable(CBORGenerator.Feature.STRINGREF).build()
            : cborFactory;
    this.sink = sink;
  }

  public void stream(InputStream is) throws IOException {
    ObjectMapper mapper = new ObjectMapper(new CBORFactory());
    try (CBORParser parser = (CBORParser) mapper.getFactory().createParser(is)) {
      JsonToken t;
      while ((t = parser.nextToken()) != null) {

        if (t == JsonToken.START_ARRAY) {
          // this is an array of docs
          t = parser.nextToken();
          if (t == JsonToken.START_OBJECT) {
            handleDoc(parser);
          }

        } else if (t == JsonToken.START_OBJECT) {
          // this is just a single doc
          handleDoc(parser);
        }
      }
    }
  }

  private void handleDoc(CBORParser p) throws IOException {
    SolrInputDocument doc = new SolrInputDocument();
    for (; ; ) {
      JsonToken t = p.nextToken();
      if (t == JsonToken.END_OBJECT) {
        if (!doc.isEmpty()) {
          sink.accept(doc);
        }
        return;
      }
      String name;
      if (t == JsonToken.FIELD_NAME) {
        name = p.getCurrentName();
        t = p.nextToken();
        if (t == JsonToken.START_ARRAY) {

          List