nl.minbzk.dwr.zoeken.enricher.uploader.SolrResultUploader.java Source code

Java tutorial

Introduction

Here is the source code for nl.minbzk.dwr.zoeken.enricher.uploader.SolrResultUploader.java

Source

/* Copyright (c) 2010 Ministry of the Interior and Kingdom Relations,
 * the Netherlands. All rights reserved.
 * 
 * This file is part of the MinBZK Search Enricher indexing generator.
 * 
 * Search Enricher is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * Search Enricher is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with Search Enricher. If not, see <http://www.gnu.org/licenses/>. */

package nl.minbzk.dwr.zoeken.enricher.uploader;

import java.io.IOException;
import java.io.PrintWriter;
import java.net.MalformedURLException;
import java.util.Arrays;
import java.util.Date;
import java.util.GregorianCalendar;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import nl.minbzk.dwr.zoeken.enricher.GeneratorResult;
import nl.minbzk.dwr.zoeken.enricher.Uploader;
import nl.minbzk.dwr.zoeken.enricher.generator.MultiGeneratorResult;
import nl.minbzk.dwr.zoeken.enricher.notifier.DirectResultNotifier;
import nl.minbzk.dwr.zoeken.enricher.settings.EnricherJob;
import nl.minbzk.dwr.zoeken.enricher.settings.EnricherSettings;
import nl.minbzk.dwr.zoeken.enricher.settings.EnricherSettings.GeneratorType;
import org.apache.solr.common.SolrInputField;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.solr.client.solrj.SolrServer;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.client.solrj.impl.CloudSolrServer;
import org.apache.solr.client.solrj.impl.HttpSolrServer;
import org.apache.solr.client.solrj.request.UpdateRequest;
import org.apache.solr.client.solrj.response.UpdateResponse;
import org.apache.solr.common.SolrInputDocument;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.expression.ExpressionParser;
import org.springframework.expression.spel.standard.SpelExpressionParser;
import org.springframework.expression.spel.support.StandardEvaluationContext;
import org.springframework.util.StringUtils;

/**
 * Solr autn envelope uploader.
 * 
 * @author Jasper van Veghel <j.veghel@rijksoverheid.nl>
 */
public class SolrResultUploader implements Uploader {
    /**
     * Constants.
     */
    private static final String REFERENCE_FIELD = "id";

    /**
     * The logger.
     */
    private static Logger logger = LoggerFactory.getLogger(SolrResultUploader.class);

    /**
     * The Solr servers per database (core).
     */
    private final Map<String, SolrServer> solrServers = new HashMap<String, SolrServer>();

    /**
     * The Solr base URI.
     */
    private final String solrBaseUri;

    /**
     * The SolrCloud URI.
     */
    private final String solrCloudUri;

    /**
     * The direct result notifier.
     */
    @Autowired
    private DirectResultNotifier notifier;

    /**
     * Default constructor.
     * 
     * @param settings
     * @throws MalformedURLException
     */
    @Autowired
    public SolrResultUploader(final EnricherSettings settings) throws MalformedURLException {
        this.solrBaseUri = settings.getSolrUri();
        this.solrCloudUri = settings.getSolrCloudUri();
    }

    /**
     * Get the Solr server based on the job.
     * 
     * @param databaseName
     * @param types
     * @return SolrServer
     * @throws MalformedURLException
     */
    private SolrServer getInstance(final String databaseName, final List<GeneratorType> types)
            throws MalformedURLException {
        if (solrServers.containsKey(databaseName))
            return solrServers.get(databaseName);

        SolrServer instance = null;

        for (GeneratorType type : types)
            switch (type) {
            case Solr:
                instance = new HttpSolrServer(solrBaseUri + databaseName + "/");

                break;
            case SolrCloud:
                instance = new CloudSolrServer(solrCloudUri);

                ((CloudSolrServer) instance).setDefaultCollection(databaseName);

                break;
            default:
                continue;
            }

        if (instance == null)
            throw new IllegalArgumentException(String.format("Unknown generator type(s) '%s'",
                    StringUtils.collectionToCommaDelimitedString(types)));

        solrServers.put(databaseName, instance);

        return instance;
    }

    /**
     * Determine the database name, based on the composite key, or null if any of the composite key replacements could not be resolved.
     *
     * XXX: We only consider the replacement values from the first document given.
     *
     * @param name
     * @param nameComposition
     * @param namePrerequisitesExpression
     * @param documents
     * @return String
     */
    private String determineAlternateDatabaseName(final String name, final String nameComposition,
            final String namePrerequisitesExpression, final List<SolrInputDocument> documents) {
        GregorianCalendar calendar = new GregorianCalendar();

        calendar.setTime(new Date());

        String result;

        result = nameComposition.replace("{name}", name).trim();

        result = result.replace("{year}", String.format("%04d", calendar.get(calendar.YEAR)));
        result = result.replace("{month}", String.format("%02d", calendar.get(calendar.MONTH)));

        if (documents.size() > 0) {
            final SolrInputDocument document = documents.get(0);

            while (result.contains("{") && result.indexOf("}") > result.indexOf("{")) {
                String fieldName = result.substring(result.indexOf("{") + 1, result.indexOf("}"));

                if (document.containsKey(fieldName))
                    result = result.replace("{" + fieldName + "}", document.getFieldValue(fieldName).toString());
                else {
                    if (logger.isDebugEnabled())
                        logger.debug(String.format(
                                "Field '%s' was missing from document with ID '%s' - will revert back to default collection '%s'",
                                fieldName, document.getFieldValue(REFERENCE_FIELD), name));

                    return null;
                }
            }

            // Also check the pre-requisite expression - only return a composite database name if it's met

            if (StringUtils.hasText(namePrerequisitesExpression)) {
                final ExpressionParser parser = new SpelExpressionParser();

                final Map<String, Object> values = new HashMap<String, Object>();

                for (Map.Entry<String, SolrInputField> entry : document.entrySet()) {
                    // XXX: Always get just the first value

                    /* if (entry.getValue().getValueCount() > 1)
                       values.put(entry.getKey(), entry.getValue().getValues());
                    else */
                    values.put(entry.getKey(), entry.getValue().getFirstValue());
                }

                StandardEvaluationContext context = new StandardEvaluationContext(new Object() {
                    public Map<String, Object> getValues() {
                        return values;
                    }
                });

                if (!parser.parseExpression(namePrerequisitesExpression).getValue(context, Boolean.class)) {
                    if (logger.isDebugEnabled())
                        logger.debug(String.format(
                                "Pre-requisite expression '%s' failed to match against document with ID '%s' - will revert back to default collection '%s'",
                                namePrerequisitesExpression, document.get(REFERENCE_FIELD).toString(), name));

                    return null;
                }
            }
        }

        return result;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void upload(final EnricherJob job, final GeneratorResult result) throws Exception {
        try {
            SolrServer server = getInstance(job.getDatabaseName(), job.getGeneratorTypes());

            // Create a new update request

            UpdateRequest request = createRequest();

            request.add(((MultiGeneratorResult<SolrInputDocument>) result).getDocuments());

            // Set a collection in case of a (resolvable) composite database name

            if (StringUtils.hasText(job.getDatabaseNameComposition())) {
                String compositeDatabaseName = determineAlternateDatabaseName(job.getDatabaseName(),
                        job.getDatabaseNameComposition(), job.getDatabaseNamePrerequisitesExpression(),
                        ((MultiGeneratorResult<SolrInputDocument>) result).getDocuments());

                if (compositeDatabaseName != null) {
                    if (logger.isDebugEnabled())
                        logger.debug(String.format("Composite database name resolved to collection '%s'",
                                compositeDatabaseName));

                    request.setParam("collection", compositeDatabaseName);
                } else {
                    if (logger.isDebugEnabled())
                        logger.debug(String.format(
                                "Composite database name could not be (completely) resolved - will use default collection '%s'",
                                job.getDatabaseName()));
                }
            }

            // Now perform the request

            UpdateResponse response = request.process(server);

            if (response.getStatus() != 0)
                logger.error("Failed to add documents to Solr, status = " + response.getStatus());
            else
                logger.info(String.format("Successfully added document to Solr, using request URL %s",
                        response.getRequestUrl()));

            // Perform direct notification now that the document has been uploaded

            notifier.process(result);
        } catch (SolrServerException e) {
            logger.error("Failed to add documents to Solr", e);
        }
    }

    /**
     * Create a new update request.
     *
     * @return UpdateRequest
     */
    protected UpdateRequest createRequest() {
        return new UpdateRequest();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void commit(final EnricherJob job) throws Exception {
        try {
            SolrServer server = getInstance(job.getDatabaseName(), job.getGeneratorTypes());

            UpdateResponse response = server.commit();

            if (response.getStatus() != 0)
                logger.error("Failed to relay commit message to Solr, status = " + response.getStatus());
            else
                logger.info("Successfully committed documents to Solr");
        } catch (SolrServerException e) {
            logger.error("Failed to relay commit message to Solr", e);
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void deleteByDocId(final EnricherJob job, final String[] documents) throws Exception {
        try {
            StringBuffer query = new StringBuffer();

            // We assume the documents are split by a space

            for (String document : documents)
                query.append(query.length() == 0 ? "" : " ").append("\"").append(document).append("\"");

            SolrServer server = getInstance(job.getDatabaseName(), job.getGeneratorTypes());

            UpdateResponse response = server.deleteByQuery(REFERENCE_FIELD + ":(" + query.toString() + ")");

            if (response.getStatus() != 0)
                logger.error("Failed to delete messages from Solr, status = " + response.getStatus());
            else
                logger.info("Successfully deleted document(s) from Solr");
        } catch (SolrServerException e) {
            logger.error("Failed to delete messages from Solr", e);
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void deleteByReference(final EnricherJob job, final String field, final String[] documents)
            throws Exception {
        try {
            StringBuffer query = new StringBuffer();

            // We assume the documents are split by a space

            for (String document : documents)
                query.append(query.length() == 0 ? "" : " ").append("\"").append(document).append("\"");

            SolrServer server = getInstance(job.getDatabaseName(), job.getGeneratorTypes());

            UpdateResponse response = server.deleteByQuery(field + ":(" + query.toString() + ")");

            if (response.getStatus() != 0)
                logger.error("Failed to delete messages from Solr, status = " + response.getStatus());
            else
                logger.info("Successfully deleted document(s) from Solr");
        } catch (SolrServerException e) {
            logger.error("Failed to delete messages from Solr", e);
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public List<GeneratorType> getType() {
        return Arrays.asList(new GeneratorType[] { GeneratorType.Solr, GeneratorType.SolrCloud });
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void writeOut(final String databaseName, final PrintWriter writer) {
        SolrServer server = solrServers.get(databaseName);

        if (server != null && server instanceof CloudSolrServer)
            writer.write(((CloudSolrServer) server).getZkStateReader().getClusterState().toString());
        else if (server != null && server instanceof HttpSolrServer)
            try {
                writer.write(((HttpSolrServer) server).ping().toString());
            } catch (SolrServerException e) {
                writer.write(String.format("Unable to ping Solr server for database %s: %s", databaseName,
                        e.getMessage()));
            } catch (IOException e) {
                writer.write(String.format("Unable to ping Solr server for database %s: %s", databaseName,
                        e.getMessage()));
            }
        else
            writer.write(String.format("Unknown or empty Solr server for database %s", databaseName));
    }
}