rapture.repo.postgres.PostgresDocHandler 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.repo.postgres;

import rapture.common.exception.ExceptionToString;
import rapture.postgres.PostgresException;
import rapture.postgres.PostgresFactory;
import rapture.postgres.PostgresHelper;
import rapture.postgres.TemplateLoader;
import rapture.repo.StoreKeyVisitor;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;

import javax.sql.DataSource;

import org.apache.log4j.Logger;
import org.postgresql.util.PGobject;
import org.springframework.jdbc.core.RowCallbackHandler;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
import org.springframework.jdbc.core.namedparam.SqlParameterSource;

/**
 * This handles a document store - but without the folder handling
 *
 * @author alanmoore
 */
public class PostgresDocHandler {
    private static Logger log = Logger.getLogger(PostgresDocHandler.class);

    private String tableName;
    private NamedParameterJdbcTemplate namedJdbcTemplate;

    public PostgresDocHandler(String instanceName, String tableName) {
        PostgresSanitizer sanitizer = PostgresFactory.getSanitizer(instanceName);
        this.tableName = sanitizer.sanitizeTableName(tableName);
    }

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

    public void initialize() {
        // Check to see if the tables exist. If they don't, create them
        log.debug("Checking that " + tableName + " exists");
        boolean exists = false;
        boolean skip = false;

        log.debug("Checking that " + tableName + " exists");
        try {
            exists = PostgresHelper.tableExists(namedJdbcTemplate, tableName);
        } catch (PostgresException e) {
            log.error(ExceptionToString.format(e));
            skip = true;
        }
        if (!skip && !exists) {
            log.debug("Table " + tableName + " does not exist, recreating");

            String sql = String.format("CREATE TABLE %s (\n"
                    + "    key VARCHAR(1024) NOT NULL,\n"
                    + "    content JSONB NOT NULL,\n"
                    + "    timestamp TIMESTAMPTZ NOT NULL,\n"
                    + "    PRIMARY KEY(key)\n"
                    + ");", tableName);
            namedJdbcTemplate.getJdbcOperations().execute(sql);

            String sqlFunctionTemplate = TemplateLoader.getResourceAsString("/sqltemplates/insertFunc.sql");
            sql = String.format(sqlFunctionTemplate, tableName);
            namedJdbcTemplate.getJdbcOperations().execute(sql);

        }
    }

    /**
     * Put a key/value pair in the database
     *
     * @param key
     * @param value
     */
    public void put(String key, String value) {
        String sql = String.format("select rap_insert_%s(:key, :content);", tableName);
        PGobject valueJson = new PGobject();
        valueJson.setType("jsonb");
        try {
            valueJson.setValue(value);
        } catch (SQLException e) {
            log.error(ExceptionToString.format(e));
            valueJson = null;
        }
        if (valueJson != null) {
            SqlParameterSource paramSource = new MapSqlParameterSource("key", key).addValue("content", valueJson);
            RowMapper rowMapper = new RowMapper() {
                @Override
                public Boolean mapRow(ResultSet rs, int rowNum) throws SQLException {
                    return true;
                }
            };
            namedJdbcTemplate.query(sql, paramSource, rowMapper);
        }
    }

    public String get(String key) {
        String sql = String.format("SELECT content\n"
                + "FROM %s\n"
                + "WHERE key = :key;", tableName);
        List results = namedJdbcTemplate.query(sql, new MapSqlParameterSource("key", key), new RowMapper() {
            @Override
            public String mapRow(ResultSet rs, int rowNum) throws SQLException {
                return rs.getString("content");
            }
        });
        if (results.size() > 0) {
            return results.get(0);
        } else {
            return null;
        }
    }

    public boolean exists(String key) {
        String sql = String.format("SELECT COUNT(key)\n"
                + "FROM %s\n"
                + "WHERE key = :key;", tableName);
        Long count = namedJdbcTemplate.queryForObject(sql, new MapSqlParameterSource("key", key), Long.class);
        return count != 0;
    }

    public long getCount() {
        String sql = String.format("SELECT COUNT(key) "
                + "FROM %s;", tableName);
        return namedJdbcTemplate.queryForObject(sql, new MapSqlParameterSource(), Long.class);
    }

    public boolean delete(String key) {
        String deleteFn = String.format("DELETE\n"
                + "FROM %s\n"
                + "WHERE key=:key;", tableName);
        namedJdbcTemplate.update(deleteFn, new MapSqlParameterSource("key", key));
        return true;
    }

    public boolean drop() {
        String sql = String.format("DROP TABLE %1$s;\n"
                + "DROP FUNCTION IF EXISTS rap_insert_%1$s(character, jsonb);\n", tableName);
        namedJdbcTemplate.getJdbcOperations().execute(sql);
        return true;
    }

    public long getSize() {
        String sql = "SELECT pg_total_relation_size(:table_name);";
        return (long) namedJdbcTemplate.queryForObject(sql, new MapSqlParameterSource("table_name", "public." + tableName), Integer.class);
    }

    public void visitKeys(String prefix, final StoreKeyVisitor iStoreKeyVisitor) {
        String sql = String.format("SELECT key, content\n"
                + "FROM %s\n"
                + "WHERE key LIKE :like_prefix;", tableName);
        RowCallbackHandler callback = new RowCallbackHandler() {
            @Override
            public void processRow(ResultSet rs) throws SQLException {
                String key = rs.getString(1);
                String content = rs.getString(2);
                log.debug("Found key " + key);
                iStoreKeyVisitor.visit(key, content);
            }
        };
        namedJdbcTemplate.query(sql, new MapSqlParameterSource("like_prefix", prefix + "%"), callback);
    }

}