Java tutorial
/** * Copyright (C) 2014 Stratio (http://stratio.com) * * 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.stratio.decision.functions; import java.io.IOException; import java.net.URISyntaxException; import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import javax.xml.parsers.ParserConfigurationException; import javax.xml.transform.TransformerException; import org.apache.solr.client.solrj.SolrClient; import org.apache.solr.client.solrj.SolrServerException; import org.apache.solr.client.solrj.impl.CloudSolrClient; import org.apache.solr.client.solrj.impl.HttpSolrClient; import org.apache.solr.common.SolrException; import org.apache.solr.common.SolrInputDocument; import org.apache.spark.api.java.JavaPairRDD; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.xml.sax.SAXException; import com.google.common.collect.Iterables; import com.stratio.decision.commons.constants.StreamAction; import com.stratio.decision.commons.messages.ColumnNameTypeValue; import com.stratio.decision.commons.messages.StratioStreamingMessage; import com.stratio.decision.service.SolrOperationsService; import com.stratio.decision.utils.RetryStrategy; import scala.Tuple2; public class SaveToSolrActionExecutionFunction extends BaseActionExecutionFunction { private static final long serialVersionUID = 3522740757019463301L; private static final Logger log = LoggerFactory.getLogger(SaveToSolrActionExecutionFunction.class); private Map<String, SolrClient> solrClients = new HashMap<>(); private List<String> solrCores = new ArrayList<String>(); private transient SolrOperationsService solrOperationsService; private RetryStrategy retryStrategy; private final String dataDir; private final String solrHost; private final String zkHost; private final Boolean isCloud; private final Integer maxBatchSize; public SaveToSolrActionExecutionFunction(String solrHost, String zkHost, Boolean isCloud, String dataDir, Integer maxBatchSize) { this.solrHost = solrHost; this.zkHost = zkHost; this.dataDir = dataDir; this.isCloud = isCloud; this.retryStrategy = new RetryStrategy(); this.maxBatchSize = maxBatchSize != null ? maxBatchSize : -1; } public SaveToSolrActionExecutionFunction(String solrHost, String zkHost, Boolean isCloud, String dataDir, Integer maxBatchSize, SolrOperationsService solrOperationsService) { this.solrHost = solrHost; this.zkHost = zkHost; this.dataDir = dataDir; this.isCloud = isCloud; this.solrOperationsService = solrOperationsService; this.retryStrategy = new RetryStrategy(); this.maxBatchSize = maxBatchSize != null ? maxBatchSize : -1; } @Override public Boolean check() throws Exception { try { getSolrOperationsService().getCoreList(); return true; } catch (Exception e) { return false; } } @Override public void process(Iterable<StratioStreamingMessage> messages) throws Exception { Integer partitionSize = maxBatchSize; if (partitionSize <= 0) { partitionSize = Iterables.size(messages); } Iterable<List<StratioStreamingMessage>> partitionIterables = Iterables.partition(messages, partitionSize); try { for (List<StratioStreamingMessage> messageList : partitionIterables) { Map<String, Collection<SolrInputDocument>> elemntsToInsert = new HashMap<String, Collection<SolrInputDocument>>(); int count = 0; for (StratioStreamingMessage stratioStreamingMessage : messageList) { count += 1; SolrInputDocument document = new SolrInputDocument(); document.addField("stratio_decision_id", System.nanoTime() + "-" + count); for (ColumnNameTypeValue column : stratioStreamingMessage.getColumns()) { document.addField(column.getColumn(), column.getValue()); } checkCore(stratioStreamingMessage); Collection<SolrInputDocument> collection = elemntsToInsert .get(stratioStreamingMessage.getStreamName()); if (collection == null) { collection = new HashSet<>(); } collection.add(document); elemntsToInsert.put(stratioStreamingMessage.getStreamName(), collection); } while (retryStrategy.shouldRetry()) { try { for (Map.Entry<String, Collection<SolrInputDocument>> elem : elemntsToInsert.entrySet()) { getSolrclient(elem.getKey()).add(elem.getValue()); } break; } catch (SolrException e) { try { log.error("Solr cloud status not yet properly initialized, retrying"); retryStrategy.errorOccured(); } catch (RuntimeException ex) { log.error("Error while initializing Solr Cloud core ", ex.getMessage()); } } } flushClients(); } } catch (Exception ex) { log.error("Error in Solr: " + ex.getMessage()); } } private void checkCore(StratioStreamingMessage message) throws IOException, SolrServerException, ParserConfigurationException, TransformerException, SAXException, URISyntaxException, InterruptedException { String core = message.getStreamName(); //check if core exists if (solrCores.size() == 0) { // Initialize solrcores list solrCores = getSolrOperationsService().getCoreList(); } if (!solrCores.contains(core)) { // Create Core getSolrOperationsService().createCore(message); // Update core list solrCores = getSolrOperationsService().getCoreList(); } } private SolrClient getClient(StratioStreamingMessage message) throws IOException, SolrServerException, URISyntaxException, TransformerException, SAXException, ParserConfigurationException { String core = message.getStreamName(); if (solrClients.containsKey(core)) { //we have a client for this core return solrClients.get(core); } else { SolrClient solrClient = getSolrclient(core); solrClients.put(core, solrClient); return solrClient; } } private void flushClients() throws IOException, SolrServerException, URISyntaxException { //Do commit in all Solrclients for (String core : solrClients.keySet()) { getSolrclient(core).commit(); } } private SolrClient getSolrclient(String core) { SolrClient solrClient; if (solrClients.containsKey(core)) { //we have a client for this core return solrClients.get(core); } else { if (isCloud) { solrClient = new CloudSolrClient(zkHost); ((CloudSolrClient) solrClient).setDefaultCollection(core); } else { solrClient = new HttpSolrClient("http://" + solrHost + "/solr/" + core); } solrClients.put(core, solrClient); } return solrClient; } private SolrOperationsService getSolrOperationsService() { if (solrOperationsService == null) { solrOperationsService = (SolrOperationsService) ActionBaseContext.getInstance().getContext() .getBean("solrOperationsService"); } return solrOperationsService; } }