AbstractSqlDao.java

package org.wikibrain.core.dao.sql;

import org.jodah.typetools.TypeResolver;
import org.jooq.DSLContext;
import org.jooq.SQLDialect;
import org.jooq.TableField;
import org.wikibrain.core.dao.Dao;
import org.wikibrain.core.dao.DaoException;
import org.wikibrain.core.dao.MetaInfoDao;
import org.wikibrain.core.lang.LanguageSet;

import javax.sql.DataSource;
import java.io.*;
import java.sql.Connection;
import java.sql.SQLException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 *
 * A SQL Dao superclass that contains a few important parameters and utility methods
 * ubiquitous to all SQL Daos.
 *
 * @author Ari Weiland
 * @author Shilad Sen
 *
 */
public abstract class AbstractSqlDao<T> implements Dao<T> {
    public static final Logger LOG = LoggerFactory.getLogger(AbstractSqlDao.class);

    public static final int DEFAULT_FETCH_SIZE = 1000;

    protected final SQLDialect dialect;
    private final String sqlScriptPrefix;
    private final TableField[] fields;
    private final Class<T> klass;
    private final MetaInfoSqlDao metaDao;
    protected WpDataSource wpDs;
    protected SqlCache cache;
    private int fetchSize = DEFAULT_FETCH_SIZE;

    // Used for directly loading csv files for databases that support it.
    FastLoader loader;

    /**
     * @param dataSource Data source for jdbc connections
     * @param fields Ordered list of fields for inserts into the database.
     * @param sqlScriptPrefix The prefix used to find sql scripts in the class path
     *                        (e.g. "/db/raw-page"). This class will append "-create-tables.sql",
     *                        "-create-indexes.sql", "-drop-tables.sql", and "-drop-indexes.sql"
     *                        to the prefix to find sql scripts, and they all must exist.
     *                        The create-tables script must ONLY create the table for the dao
     *                        because it is used by the fast loader.
     * @throws DaoException
     */
    public AbstractSqlDao(WpDataSource dataSource, TableField [] fields, String sqlScriptPrefix) throws DaoException {
        Class<?>[] typeArguments = TypeResolver.resolveRawArguments(AbstractSqlDao.class, getClass());
        this.klass = (Class<T>) typeArguments[0];

        wpDs = dataSource;
        Connection conn = null;
        try {
            conn = wpDs.getConnection();
            this.dialect = JooqUtils.dialect(conn);
        } catch (SQLException e) {
            throw new DaoException("SQL Dao Failed. Check if the table exists / if the desired information has been parsed and stored in the database\n" + e.toString());
        } finally {
            quietlyCloseConn(conn);
        }
        cache = null;

        // TODO: pass this through the constructor
        if (this instanceof MetaInfoDao) {
            this.metaDao = (MetaInfoSqlDao) this;
        } else {
            this.metaDao = new MetaInfoSqlDao(wpDs);
        }
        this.fields = fields;
        this.sqlScriptPrefix = sqlScriptPrefix;
    }

    /**
     * Executes a sql resource on the classpath
     * @param name Resource path - e.g. "/db/local-page.schema.sql"
     * @throws DaoException
     */
    public void executeSqlResource(String name) throws DaoException {
        wpDs.executeSqlResource(name);
    }

    protected DSLContext getJooq() throws DaoException {
        return wpDs.getJooq();
    }

    protected void freeJooq(DSLContext context) {
        wpDs.freeJooq(context);
    }

    @Override
    public LanguageSet getLoadedLanguages() throws DaoException {
        return metaDao.getLoadedLanguages(klass);
    }

    @Override
    public void clear() throws DaoException {
        executeSqlScriptWithSuffix("-drop-indexes.sql");
        executeSqlScriptWithSuffix("-drop-tables.sql");
    }

    @Override
    public void beginLoad() throws  DaoException {
        executeSqlScriptWithSuffix("-drop-indexes.sql");
        executeSqlScriptWithSuffix("-create-tables.sql");
        if (fields != null) {
            loader = new FastLoader(wpDs        , fields);
        }
    }

    /**
     * Inserts values into the database.
     * Call this instead of direct sql inserts because the underlying code may optimize the inserts
     * by creating batch inserts or a cvs that can be directly loaded by the underlying database.
     * @param values
     */
    protected void insert(Object ... values) throws DaoException{
        loader.load(values);
    }


    @Override
    public void endLoad() throws  DaoException {
        if (loader != null) {
            loader.endLoad();
        }
        LOG.info("creating indexes in {}-create-indexes.sql (this can take some time)", sqlScriptPrefix);
        executeSqlScriptWithSuffix("-create-indexes.sql");
        if (fields != null && fields.length > 0) {
            wpDs.optimize(fields[0].getTable());
        }
    }

    /**
     * Executes the appropriate sql script with a particular suffix (.e.g. "-drop-tables.sql").
     * @param suffix
     * @throws DaoException
     */
    protected void executeSqlScriptWithSuffix(String suffix) throws DaoException {
        executeSqlResource(sqlScriptPrefix + suffix);
    }

    public void useCache(File dir) throws DaoException{
        cache = new SqlCache(metaDao, dir);
    }

    /**
     * Close a connection without generating an exception if it fails.
     * @param conn
     */
    public static void quietlyCloseConn(Connection conn) {
        if (conn != null) {
            try {
                conn.close();
            } catch (SQLException e) {
                LOG.warn("Failed to close connection: ", e);
            }
        }
    }

    public int getFetchSize() {
        return fetchSize;
    }

    public void setFetchSize(int fetchSize) {
        this.fetchSize = fetchSize;
    }
}