org.apache.hive.storage.jdbc.dao.GenericJdbcDatabaseAccessor.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.hive.storage.jdbc.dao.GenericJdbcDatabaseAccessor.java

Source

/*
 *
 * 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 org.apache.hive.storage.jdbc.dao;

import org.apache.commons.dbcp.BasicDataSourceFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.security.Credentials;
import org.apache.hadoop.security.UserGroupInformation;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.hive.storage.jdbc.conf.JdbcStorageConfig;
import org.apache.hive.storage.jdbc.conf.JdbcStorageConfigManager;
import org.apache.hive.storage.jdbc.exception.HiveJdbcDatabaseAccessException;

import javax.sql.DataSource;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Properties;

/**
 * A data accessor that should in theory work with all JDBC compliant database drivers.
 */
public class GenericJdbcDatabaseAccessor implements DatabaseAccessor {

    protected static final String DBCP_CONFIG_PREFIX = JdbcStorageConfigManager.CONFIG_PREFIX + ".dbcp";
    protected static final int DEFAULT_FETCH_SIZE = 1000;
    protected static final Logger LOGGER = LoggerFactory.getLogger(GenericJdbcDatabaseAccessor.class);
    protected DataSource dbcpDataSource = null;
    protected static final Text DBCP_PWD = new Text(DBCP_CONFIG_PREFIX + ".password");

    public GenericJdbcDatabaseAccessor() {
    }

    @Override
    public List<String> getColumnNames(Configuration conf) throws HiveJdbcDatabaseAccessException {
        Connection conn = null;
        PreparedStatement ps = null;
        ResultSet rs = null;

        try {
            initializeDatabaseConnection(conf);
            String sql = JdbcStorageConfigManager.getQueryToExecute(conf);
            String metadataQuery = addLimitToQuery(sql, 1);
            LOGGER.debug("Query to execute is [{}]", metadataQuery);

            conn = dbcpDataSource.getConnection();
            ps = conn.prepareStatement(metadataQuery);
            rs = ps.executeQuery();

            ResultSetMetaData metadata = rs.getMetaData();
            int numColumns = metadata.getColumnCount();
            List<String> columnNames = new ArrayList<String>(numColumns);
            for (int i = 0; i < numColumns; i++) {
                columnNames.add(metadata.getColumnName(i + 1));
            }

            return columnNames;
        } catch (Exception e) {
            LOGGER.error("Error while trying to get column names.", e);
            throw new HiveJdbcDatabaseAccessException("Error while trying to get column names: " + e.getMessage(),
                    e);
        } finally {
            cleanupResources(conn, ps, rs);
        }

    }

    @Override
    public int getTotalNumberOfRecords(Configuration conf) throws HiveJdbcDatabaseAccessException {
        Connection conn = null;
        PreparedStatement ps = null;
        ResultSet rs = null;

        try {
            initializeDatabaseConnection(conf);
            String sql = JdbcStorageConfigManager.getQueryToExecute(conf);
            String countQuery = "SELECT COUNT(*) FROM (" + sql + ") tmptable";
            LOGGER.info("Query to execute is [{}]", countQuery);

            conn = dbcpDataSource.getConnection();
            ps = conn.prepareStatement(countQuery);
            rs = ps.executeQuery();
            if (rs.next()) {
                return rs.getInt(1);
            } else {
                LOGGER.warn("The count query did not return any results.", countQuery);
                throw new HiveJdbcDatabaseAccessException("Count query did not return any results.");
            }
        } catch (HiveJdbcDatabaseAccessException he) {
            throw he;
        } catch (Exception e) {
            LOGGER.error("Caught exception while trying to get the number of records", e);
            throw new HiveJdbcDatabaseAccessException(e);
        } finally {
            cleanupResources(conn, ps, rs);
        }
    }

    @Override
    public JdbcRecordIterator getRecordIterator(Configuration conf, int limit, int offset)
            throws HiveJdbcDatabaseAccessException {

        Connection conn = null;
        PreparedStatement ps = null;
        ResultSet rs = null;

        try {
            initializeDatabaseConnection(conf);
            String sql = JdbcStorageConfigManager.getQueryToExecute(conf);
            String limitQuery = addLimitAndOffsetToQuery(sql, limit, offset);
            LOGGER.info("Query to execute is [{}]", limitQuery);

            conn = dbcpDataSource.getConnection();
            ps = conn.prepareStatement(limitQuery, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
            ps.setFetchSize(getFetchSize(conf));
            rs = ps.executeQuery();

            return new JdbcRecordIterator(conn, ps, rs);
        } catch (Exception e) {
            LOGGER.error("Caught exception while trying to execute query", e);
            cleanupResources(conn, ps, rs);
            throw new HiveJdbcDatabaseAccessException("Caught exception while trying to execute query", e);
        }
    }

    /**
     * Uses generic JDBC escape functions to add a limit and offset clause to a query string
     *
     * @param sql
     * @param limit
     * @param offset
     * @return
     */
    protected String addLimitAndOffsetToQuery(String sql, int limit, int offset) {
        if (offset == 0) {
            return addLimitToQuery(sql, limit);
        } else {
            return sql + " {LIMIT " + limit + " OFFSET " + offset + "}";
        }
    }

    /*
     * Uses generic JDBC escape functions to add a limit clause to a query string
     */
    protected String addLimitToQuery(String sql, int limit) {
        return sql + " {LIMIT " + limit + "}";
    }

    protected void cleanupResources(Connection conn, PreparedStatement ps, ResultSet rs) {
        try {
            if (rs != null) {
                rs.close();
            }
        } catch (SQLException e) {
            LOGGER.warn("Caught exception during resultset cleanup.", e);
        }

        try {
            if (ps != null) {
                ps.close();
            }
        } catch (SQLException e) {
            LOGGER.warn("Caught exception during statement cleanup.", e);
        }

        try {
            if (conn != null) {
                conn.close();
            }
        } catch (SQLException e) {
            LOGGER.warn("Caught exception during connection cleanup.", e);
        }
    }

    protected void initializeDatabaseConnection(Configuration conf) throws Exception {
        if (dbcpDataSource == null) {
            synchronized (this) {
                if (dbcpDataSource == null) {
                    Properties props = getConnectionPoolProperties(conf);
                    dbcpDataSource = BasicDataSourceFactory.createDataSource(props);
                }
            }
        }
    }

    protected Properties getConnectionPoolProperties(Configuration conf) throws Exception {
        // Create the default properties object
        Properties dbProperties = getDefaultDBCPProperties();

        // override with user defined properties
        Map<String, String> userProperties = conf.getValByRegex(DBCP_CONFIG_PREFIX + "\\.*");
        if ((userProperties != null) && (!userProperties.isEmpty())) {
            for (Entry<String, String> entry : userProperties.entrySet()) {
                dbProperties.put(entry.getKey().replaceFirst(DBCP_CONFIG_PREFIX + "\\.", ""), entry.getValue());
            }
        }

        // handle password
        Credentials credentials = UserGroupInformation.getCurrentUser().getCredentials();
        if (credentials.getSecretKey(DBCP_PWD) != null) {
            LOGGER.info("found token in credentials");
            dbProperties.put(DBCP_PWD, new String(credentials.getSecretKey(DBCP_PWD)));
        }

        // essential properties that shouldn't be overridden by users
        dbProperties.put("url", conf.get(JdbcStorageConfig.JDBC_URL.getPropertyName()));
        dbProperties.put("driverClassName", conf.get(JdbcStorageConfig.JDBC_DRIVER_CLASS.getPropertyName()));
        dbProperties.put("type", "javax.sql.DataSource");
        return dbProperties;
    }

    protected Properties getDefaultDBCPProperties() {
        Properties props = new Properties();
        props.put("initialSize", "1");
        props.put("maxActive", "3");
        props.put("maxIdle", "0");
        props.put("maxWait", "10000");
        props.put("timeBetweenEvictionRunsMillis", "30000");
        return props;
    }

    protected int getFetchSize(Configuration conf) {
        return conf.getInt(JdbcStorageConfig.JDBC_FETCH_SIZE.getPropertyName(), DEFAULT_FETCH_SIZE);
    }
}