rapture.table.postgres.PostgresIndexHandler Maven / Gradle / Ivy

Go to download
/**
 * Copyright (C) 2011-2015 Incapture Technologies LLC
 *
 * This is an autogenerated license statement. When copyright notices appear below
 * this one that copyright supercedes this statement.
 *
 * Unless required by applicable law or agreed to in writing, software is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied.
 *
 * Unless explicit permission obtained in writing this software cannot be distributed.
 */
package rapture.table.postgres;

import rapture.common.LockHandle;
import rapture.common.TableQuery;
import rapture.common.TableQueryResult;
import rapture.common.TableRecord;
import rapture.common.model.DocumentMetadata;
import rapture.dsl.idef.FieldDefinition;
import rapture.dsl.idef.IndexDefinition;
import rapture.dsl.idef.IndexFieldType;
import rapture.dsl.iqry.IndexQuery;
import rapture.dsl.iqry.IndexQueryFactory;
import rapture.index.IndexCreationLock;
import rapture.index.IndexHandler;
import rapture.index.IndexProducer;
import rapture.postgres.PostgresFactory;
import rapture.postgres.TemplateLoader;
import rapture.repo.postgres.PostgresSanitizer;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import javax.sql.DataSource;

import org.apache.commons.lang.StringUtils;
import org.apache.log4j.Logger;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;

import rapture.config.MultiValueConfigLoader;

public class PostgresIndexHandler implements IndexHandler {
    private static Logger log = Logger.getLogger(PostgresIndexHandler.class);

    private static final String PREFIX = "prefix";

    private PostgresSanitizer sanitizer;
    private String tableName;
    private IndexProducer indexProducer;
    private NamedParameterJdbcTemplate namedJdbcTemplate;
    private Map fieldToType;
    private String instanceName;

    public PostgresIndexHandler() {
        fieldToType = new HashMap<>();
    }

    @Override
    public void setInstanceName(String instanceName) {
        this.instanceName = instanceName;
        sanitizer = PostgresFactory.getSanitizer(instanceName);
    }

    @Override
    public void setConfig(Map config) {
        String tableName = config.get(PREFIX);
        this.tableName = sanitizer.sanitizeTableName(tableName);
    }

    @Override
    public void setIndexProducer(IndexProducer indexProducer) {
        this.indexProducer = indexProducer;
    }

    public void setDataSource(DataSource dataSource) {
        this.namedJdbcTemplate = new NamedParameterJdbcTemplate(dataSource);
    }

    @Override
    public void initialize() {
        for (IndexDefinition indexDefinition : indexProducer.getIndexDefinitions()) {
            createIndex(indexDefinition, false);
        }
    }

    @Override
    public void ensureIndicesExist() {
        log.info(String.format("About to build indices on table [%s]...", tableName));
        LockHandle lockHandle = null;
        try {
            lockHandle = IndexCreationLock.INSTANCE.grabLock();
            for (IndexDefinition indexDefinition : indexProducer.getIndexDefinitions()) {
                createIndex(indexDefinition, true);
            }
        } finally {
            if (lockHandle != null) {
                IndexCreationLock.INSTANCE.releaseLock(lockHandle);
            }
        }
        log.info(String.format("Done building indices on table [%s]...", tableName));
    }

    private void createIndex(IndexDefinition indexDefinition, boolean force) {
        List pgFields = new LinkedList<>();

        for (int i = 0; i < indexDefinition.getFields().size(); i++) {
            FieldDefinition field = indexDefinition.getFields().get(i);
            IndexFieldType fieldType = field.getType();
            String pgType = FieldTypeMapper.getPgType(fieldType);
            String fieldName = field.getName();
            String pgField = String.format("(content->>'%s')::%s", fieldName, pgType);
            fieldToType.put(fieldName, fieldType);
            pgFields.add(String.format("(%s)", pgField));
        }

        String pgFieldsString = StringUtils.join(pgFields, ", ");
        String pgIndexName = sanitizer.sanitizeIndexName(String.format("%s_%s", tableName, indexDefinition.getIndexName()));
        String sql;
        if (force) {
            String template = TemplateLoader.getResourceAsString("/sqltemplates/ensureIndexForce.sql");
            sql = String.format(template, pgIndexName, tableName, pgFieldsString);
        } else {
            String template = TemplateLoader.getResourceAsString("/sqltemplates/ensureIndex.sql");

            int insertLimit = getInsertLimit();
            sql = String.format(template, pgIndexName, tableName, pgFieldsString, insertLimit);
        }

        namedJdbcTemplate.update(sql, new MapSqlParameterSource());
    }

    private int getInsertLimit() {
        return Integer.parseInt(MultiValueConfigLoader
                .getConfig("POSTGRES-" + instanceName + ".indexCreateLimit", MultiValueConfigLoader.getConfig("POSTGRES-default.indexCreateLimit", "250")));
    }

    @Override
    public TableQueryResult query(String query) {
        log.debug("Performing query " + query);
        final IndexQuery parsedQuery = IndexQueryFactory.parseQuery(query);

        TableQueryResult result = new TableQueryResult();

        List selectFields = formatSelect(parsedQuery.getSelect().getFieldList());
        String selectClause = StringUtils.join(selectFields, ", ");
        PostgresWhereClause whereClause = QueryParser.parseWhereClause(fieldToType, parsedQuery.getWhere());
        String sql = String.format("SELECT %s \n"
                + "FROM %s\n"
                + "WHERE (%s)", selectClause, tableName, whereClause.getSql());

        RowMapper> rowMapper = new RowMapper>() {
            @Override
            public List