com.avaje.ebeanservice.elastic.ElasticDocumentStore Maven / Gradle / Ivy

Go to download
package com.avaje.ebeanservice.elastic;

import com.avaje.ebean.DocStoreQueueEntry;
import com.avaje.ebean.DocumentStore;
import com.avaje.ebean.PagedList;
import com.avaje.ebean.PersistenceIOException;
import com.avaje.ebean.Query;
import com.avaje.ebean.QueryEachConsumer;
import com.avaje.ebean.QueryEachWhileConsumer;
import com.avaje.ebean.config.DocStoreConfig;
import com.avaje.ebean.plugin.BeanType;
import com.avaje.ebean.plugin.SpiServer;
import com.avaje.ebeaninternal.api.SpiQuery;
import com.avaje.ebeanservice.docstore.api.DocQueryRequest;
import com.avaje.ebeanservice.docstore.api.DocStoreQueryUpdate;
import com.avaje.ebeanservice.elastic.bulk.BulkUpdate;
import com.avaje.ebeanservice.elastic.index.EIndexService;
import com.avaje.ebeanservice.elastic.query.EQueryService;
import com.avaje.ebeanservice.elastic.support.IndexMessageSender;
import com.fasterxml.jackson.core.JsonFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.List;

/**
 * ElasticSearch based document store.
 */
public class ElasticDocumentStore implements DocumentStore {
  /**
   * Logger that can be used to log Bulk API messages.
   */
  public static final Logger BULK = LoggerFactory.getLogger("org.avaje.ebeanservice.elastic.BULK");

  private final SpiServer server;

  private final ElasticUpdateProcessor updateProcessor;

  private final EQueryService queryService;

  private final EIndexService indexService;

  public ElasticDocumentStore(SpiServer server, ElasticUpdateProcessor updateProcessor, IndexMessageSender sender, JsonFactory jsonFactory) {
    this.server = server;
    this.updateProcessor = updateProcessor;
    this.queryService = new EQueryService(server, jsonFactory, sender);
    this.indexService = new EIndexService(server, jsonFactory, sender);
  }

  @Override
  public long process(List entries) throws IOException {

    BulkUpdate bulk = updateProcessor.createBulkUpdate(0);
    try {
      long count = updateProcessor.processQueue(bulk, entries);
      bulk.flush();
      return count;

    } catch (IOException e) {
      throw new PersistenceIOException(e);
    }
  }

  @Override
  public void dropIndex(String newIndex) {
    try {
      indexService.dropIndex(newIndex);
    } catch (IOException e) {
      throw new PersistenceIOException(e);
    }
  }

  @Override
  public void createIndex(String indexName, String alias) {
    try {
      indexService.createIndex(indexName, alias);
    } catch (IOException e) {
      throw new PersistenceIOException(e);
    }
  }

  @Override
  public long copyIndex(Query query, String newIndex, int bulkBatchSize) {
    try {
      BulkUpdate txn = updateProcessor.createBulkUpdate(bulkBatchSize);
      long count = queryService.copyIndexSince((SpiQuery)query, newIndex, txn);
      txn.flush();

      return count;

    } catch (IOException e) {
      throw new PersistenceIOException(e);
    }
  }

  @Override
  public long copyIndex(Class beanType, String newIndex, long epochMillis) {
    BeanType type = checkMapped(server.getBeanType(beanType));
    try {
      BulkUpdate txn = updateProcessor.createBulkUpdate(0);
      long count = queryService.copyIndexSince(type, newIndex, txn, epochMillis);
      txn.flush();

      return count;

    } catch (IOException e) {
      throw new PersistenceIOException(e);
    }
  }

  @Override
  public long copyIndex(Class beanType, String newIndex) {
    return copyIndex(beanType, newIndex, 0);
  }

  @Override
  public void indexAll(Class beanType) {
    indexByQuery(server.createQuery(beanType));
  }

  @Override
  public  void indexByQuery(Query query) {
    indexByQuery(query, 0);
  }

  @Override
  public  void indexByQuery(Query query, int bulkBatchSize) {

    SpiQuery spiQuery = (SpiQuery) query;
    BeanType desc = checkMapped(spiQuery.getBeanDescriptor());

    try {
      DocStoreQueryUpdate update = updateProcessor.createQueryUpdate(desc, bulkBatchSize);
      indexByQuery(desc, query, update);
      update.flush();

    } catch (IOException e) {
      throw new PersistenceIOException(e);
    }
  }


  private  void indexByQuery(final BeanType desc, Query query, final DocStoreQueryUpdate queryUpdate) throws IOException {

    desc.docStore().applyPath(query);
    query.setLazyLoadBatchSize(100);
    query.findEach(new QueryEachConsumer() {
      @Override
      public void accept(T bean) {
        Object idValue = desc.getBeanId(bean);
        try {
          queryUpdate.store(idValue, bean);
        } catch (Exception e) {
          throw new PersistenceIOException("Error performing query update to doc store", e);
        }
      }
    });
  }

  @Override
  public  void findEach(DocQueryRequest request, QueryEachConsumer consumer) {
    queryService.findEach(request, consumer);
  }

  @Override
  public  void findEachWhile(DocQueryRequest request, QueryEachWhileConsumer consumer) {
    queryService.findEachWhile(request, consumer);
  }

  @Override
  public  List findList(DocQueryRequest request) {
    return queryService.findList(request);
  }

  @Override
  public  PagedList findPagedList(DocQueryRequest request) {
    return queryService.findPagedList(request);
  }

  @Override
  public  T find(DocQueryRequest request) {
    return queryService.findById(request);
  }

  public void onStartup() {

    try {
      DocStoreConfig docStoreConfig = server.getServerConfig().getDocStoreConfig();
      if (docStoreConfig.isActive()) {
        indexService.createIndexesOnStartup();
      }

    } catch (IOException e) {
      throw new PersistenceIOException(e);
    }
  }

  private  BeanType checkMapped(BeanType type) {
    if (type == null) {
      throw new IllegalStateException("No bean type mapping found?");
    }
    if (!type.isDocStoreMapped()) {
      throw new IllegalStateException("No doc store mapping for bean type "+type.getFullName());
    }
    return type;
  }
}